Fix handling of custom SQL-based partitioning functions

erimatnor · erimatnor · commit 0e79df4ac3bc · 2018-02-22T17:38:09.000+01:00
Previously, when the FmgrInfo was initialized for partitioning
functions, the type information was cached in the fn_extra
field. However, the contract for fn_extra is that it is to be set by
the called function and not prior to its invokation. Setting fn_extra
prior to function invokation is an issue for custom partitioning
functions implemented in SQL, as the SQL call manager expects to use
fn_extra for its own cache.

This change avoids setting fn_extra prior to function invokation, and
instead information is cached in fn_extra within our provided
partitioning functions.

The native partitioning functions now also support nested functions,
i.e., the input to the function is the output of another function.
diff --git a/src/partitioning.c b/src/partitioning.c
@@ -12,6 +12,7 @@
 #include <utils/jsonb.h>
 #include <utils/acl.h>
 #include <utils/rangetypes.h>
+#include <utils/memutils.h>
 #include <catalog/namespace.h>
 #include <catalog/pg_type.h>
 #include <access/hash.h>
@@ -158,7 +159,7 @@ partitioning_info_create(const char *schema,
 	pinfo->typcache_entry = lookup_type_cache(columntype, TYPECACHE_HASH_FLAGS);
 
 	if (pinfo->typcache_entry->hash_proc == InvalidOid)
-		elog(ERROR, "No hash function for type %u", columntype);
+		elog(ERROR, "could not find hash function for type %u", columntype);
 
 	partitioning_func_set_func_fmgr(&pinfo->partfunc);
 
@@ -181,20 +182,13 @@ partitioning_info_create(const char *schema,
 
 	fmgr_info_set_expr((Node *) expr, &pinfo->partfunc.func_fmgr);
 
-	/*
-	 * Set the type cache entry in fn_extra to avoid an extry lookup in the
-	 * partition hash function
-	 */
-	pinfo->partfunc.func_fmgr.fn_extra = pinfo->typcache_entry;
-
 	return pinfo;
 }
 
 /*
  * Apply the partitioning function of a hypertable to a value.
  *
- * We support both partitioning functions with the signature int
- * func(anyelement).
+ * We support partitioning functions with the signature (anyelement) -> int.
  */
 int32
 partitioning_func_apply(PartitioningInfo *pinfo, Datum value)
@@ -232,10 +226,10 @@ resolve_function_argtype(FunctionCallInfo fcinfo)
 	fe = (FuncExpr *) fcinfo->flinfo->fn_expr;
 
 	if (NULL == fe || !IsA(fe, FuncExpr))
-		elog(ERROR, "No function expression set when invoking partitioning function");
+		elog(ERROR, "no function expression set when invoking partitioning function");
 
 	if (list_length(fe->args) != 1)
-		elog(ERROR, "Unexpected number of arguments in function expression");
+		elog(ERROR, "unexpected number of arguments in function expression");
 
 	node = linitial(fe->args);
 
@@ -250,13 +244,50 @@ resolve_function_argtype(FunctionCallInfo fcinfo)
 		case T_CoerceViaIO:
 			argtype = ((CoerceViaIO *) node)->resulttype;
 			break;
+		case T_FuncExpr:
+			/* Argument is function, so our input is its result type */
+			argtype = ((FuncExpr *) node)->funcresulttype;
+			break;
 		default:
-			elog(ERROR, "Unsupported expression argument node type %u", nodeTag(node));
+			elog(ERROR, "unsupported expression argument node type %u", nodeTag(node));
 	}
 
 	return argtype;
 }
 
+/*
+ * Partitioning function cache.
+ *
+ * Holds type information to avoid repeated lookups. The cache is allocated on a
+ * child memory context of the context that created the associated FmgrInfo
+ * struct. For partitioning functions invoked on the insert path, this is
+ * typically the Hypertable cache's memory context. Hence, the type cache lives
+ * for the duration of the hypertable cache and can be reused across multiple
+ * invokations of the partitioning function, even across transactions.
+ *
+ * If the partitioning function is invoked outside the insert path, the FmgrInfo
+ * and its memory context has a lifetime corresponding to that invokation.
+ */
+typedef struct PartFuncCache
+{
+	Oid			argtype;
+	Oid			coerce_funcid;
+	TypeCacheEntry *tce;
+} PartFuncCache;
+
+static PartFuncCache *
+part_func_cache_create(Oid argtype, TypeCacheEntry *tce, Oid coerce_funcid, MemoryContext mcxt)
+{
+	PartFuncCache *pfc;
+
+	pfc = MemoryContextAlloc(mcxt, sizeof(PartFuncCache));
+	pfc->argtype = argtype;
+	pfc->tce = tce;
+	pfc->coerce_funcid = coerce_funcid;
+
+	return pfc;
+}
+
 /* _timescaledb_catalog.get_partition_for_key(key anyelement) RETURNS INT */
 PGDLLEXPORT Datum get_partition_for_key(PG_FUNCTION_ARGS);
 
@@ -270,26 +301,35 @@ Datum
 get_partition_for_key(PG_FUNCTION_ARGS)
 {
 	Datum		arg = PG_GETARG_DATUM(0);
+	PartFuncCache *pfc = fcinfo->flinfo->fn_extra;
 	struct varlena *data;
-	Oid			argtype;
 	uint32		hash_u;
 	int32		res;
 
 	if (PG_NARGS() != 1)
-		elog(ERROR, "Unexpected number of arguments to partitioning function");
-
-	argtype = resolve_function_argtype(fcinfo);
+		elog(ERROR, "unexpected number of arguments to partitioning function");
 
-	if (argtype != TEXTOID)
+	if (NULL == pfc)
 	{
-		/* Not TEXT input -> need to convert to text */
-		Oid			funcid = find_text_coercion_func(argtype);
+		Oid			funcid = InvalidOid;
+		Oid			argtype = resolve_function_argtype(fcinfo);
 
-		if (!OidIsValid(funcid))
-			elog(ERROR, "Could not coerce type %u to text",
-				 argtype);
+		if (argtype != TEXTOID)
+		{
+			/* Not TEXT input -> need to convert to text */
+			funcid = find_text_coercion_func(argtype);
+
+			if (!OidIsValid(funcid))
+				elog(ERROR, "could not coerce type %u to text", argtype);
+		}
 
-		arg = OidFunctionCall1(funcid, arg);
+		pfc = part_func_cache_create(argtype, NULL, funcid, fcinfo->flinfo->fn_mcxt);
+		fcinfo->flinfo->fn_extra = pfc;
+	}
+
+	if (pfc->argtype != TEXTOID)
+	{
+		arg = OidFunctionCall1(pfc->coerce_funcid, arg);
 		arg = CStringGetTextDatum(DatumGetCString(arg));
 	}
 
@@ -320,23 +360,26 @@ Datum
 get_partition_hash(PG_FUNCTION_ARGS)
 {
 	Datum		arg = PG_GETARG_DATUM(0);
-	TypeCacheEntry *tce = fcinfo->flinfo->fn_extra;
-	Oid			argtype;
+	PartFuncCache *pfc = fcinfo->flinfo->fn_extra;
 	Datum		hash;
 	int32		res;
 
 	if (PG_NARGS() != 1)
-		elog(ERROR, "Unexpected number of arguments to partitioning function");
+		elog(ERROR, "unexpected number of arguments to partitioning function");
 
-	argtype = resolve_function_argtype(fcinfo);
+	if (NULL == pfc)
+	{
+		Oid			argtype = resolve_function_argtype(fcinfo);
+		TypeCacheEntry *tce = lookup_type_cache(argtype, TYPECACHE_HASH_FLAGS);
 
-	if (tce == NULL)
-		tce = lookup_type_cache(argtype, TYPECACHE_HASH_FLAGS);
+		pfc = part_func_cache_create(argtype, tce, InvalidOid, fcinfo->flinfo->fn_mcxt);
+		fcinfo->flinfo->fn_extra = pfc;
+	}
 
-	if (tce->hash_proc == InvalidOid)
-		elog(ERROR, "No hash function for type %u", argtype);
+	if (pfc->tce->hash_proc == InvalidOid)
+		elog(ERROR, "could not find hash function for type %u", pfc->argtype);
 
-	hash = FunctionCall1(&tce->hash_proc_finfo, arg);
+	hash = FunctionCall1(&pfc->tce->hash_proc_finfo, arg);
 
 	/* Only positive numbers */
 	res = (int32) (DatumGetUInt32(hash) & 0x7fffffff);
diff --git a/test/expected/partitioning.out b/test/expected/partitioning.out
@@ -141,3 +141,56 @@ SELECT * FROM _timescaledb_catalog.dimension;
   9 |             4 | location    | integer                     | f       |          2 | _timescaledb_internal    | get_partition_for_key |                
 (9 rows)
 
+-- Test that we support custom SQL-based partitioning functions and
+-- that our native partitioning function handles function expressions
+-- as argument
+CREATE OR REPLACE FUNCTION custom_partfunc(source anyelement)
+    RETURNS INTEGER LANGUAGE PLPGSQL AS
+$BODY$
+DECLARE
+    retval INTEGER;
+BEGIN
+    retval = _timescaledb_internal.get_partition_hash(substring(source::text FROM '[A-za-z0-9 ]+'));
+    RAISE NOTICE 'hash value for % is %', source, retval;
+    RETURN retval;
+END
+$BODY$;
+CREATE TABLE part_custom_func(time timestamptz, temp float8, device text);
+SELECT create_hypertable('part_custom_func', 'time', 'device', 2, partitioning_func => 'custom_partfunc');
+NOTICE:  adding NOT NULL constraint to column "time"
+ create_hypertable 
+-------------------
+ 
+(1 row)
+
+SELECT _timescaledb_internal.get_partition_hash(substring('dev1' FROM '[A-za-z0-9 ]+'));
+ get_partition_hash 
+--------------------
+         1129986420
+(1 row)
+
+SELECT _timescaledb_internal.get_partition_hash('dev1'::text);
+ get_partition_hash 
+--------------------
+         1129986420
+(1 row)
+
+SELECT _timescaledb_internal.get_partition_hash('dev7'::text);
+ get_partition_hash 
+--------------------
+          449729092
+(1 row)
+
+INSERT INTO part_custom_func VALUES ('2017-03-22T09:18:23', 23.4, 'dev1'),
+                                    ('2017-03-22T09:18:23', 23.4, 'dev7');
+NOTICE:  hash value for dev1 is 1129986420
+NOTICE:  hash value for dev1 is 1129986420
+NOTICE:  hash value for dev7 is 449729092
+NOTICE:  hash value for dev7 is 449729092
+SELECT * FROM test.show_subtables('part_custom_func');
+                 Child                  | Tablespace 
+----------------------------------------+------------
+ _timescaledb_internal._hyper_5_6_chunk | 
+ _timescaledb_internal._hyper_5_7_chunk | 
+(2 rows)
+
diff --git a/test/sql/partitioning.sql b/test/sql/partitioning.sql
@@ -56,3 +56,30 @@ SELECT add_dimension('part_add_dim', 'location', 2, partitioning_func => 'bad_fu
 
 SELECT add_dimension('part_add_dim', 'location', 2, partitioning_func => '_timescaledb_internal.get_partition_for_key');
 SELECT * FROM _timescaledb_catalog.dimension;
+
+-- Test that we support custom SQL-based partitioning functions and
+-- that our native partitioning function handles function expressions
+-- as argument
+CREATE OR REPLACE FUNCTION custom_partfunc(source anyelement)
+    RETURNS INTEGER LANGUAGE PLPGSQL AS
+$BODY$
+DECLARE
+    retval INTEGER;
+BEGIN
+    retval = _timescaledb_internal.get_partition_hash(substring(source::text FROM '[A-za-z0-9 ]+'));
+    RAISE NOTICE 'hash value for % is %', source, retval;
+    RETURN retval;
+END
+$BODY$;
+
+CREATE TABLE part_custom_func(time timestamptz, temp float8, device text);
+SELECT create_hypertable('part_custom_func', 'time', 'device', 2, partitioning_func => 'custom_partfunc');
+
+SELECT _timescaledb_internal.get_partition_hash(substring('dev1' FROM '[A-za-z0-9 ]+'));
+SELECT _timescaledb_internal.get_partition_hash('dev1'::text);
+SELECT _timescaledb_internal.get_partition_hash('dev7'::text);
+
+INSERT INTO part_custom_func VALUES ('2017-03-22T09:18:23', 23.4, 'dev1'),
+                                    ('2017-03-22T09:18:23', 23.4, 'dev7');
+
+SELECT * FROM test.show_subtables('part_custom_func');
