Allow whole-row Vars to be used in partitioning expressions.

In the wake of commit 5b931237, there's no particular reason
for this restriction (previously, it was problematic because of
the implied rowtype reference).  A simple constraint on a whole-row
Var probably isn't that useful, but conceivably somebody would want
to pass one to a function that extracts a partitioning key.  Besides
which, we're expending much more code to enforce the restriction than
we save by having it, since the latter quantity is now zero.
So drop the restriction.

Amit Langote

Discussion: https://postgr.es/m/CA+HiwqFUzjfj9HEsJtYWcr1SgQ_=iCAvQ=O2Sx6aQxoDu4OiHw@mail.gmail.com

src/backend/catalog/partition.c

@@ -181,17 +181,14 @@ index_get_partition(Relation partition, Oid indexId)
 }
 
 /*
- * map_partition_varattnos - maps varattno of any Vars in expr from the
- * attno's of 'from_rel' to the attno's of 'to_rel' partition, each of which
- * may be either a leaf partition or a partitioned table, but both of which
- * must be from the same partitioning hierarchy.
+ * map_partition_varattnos - maps varattnos of all Vars in 'expr' (that have
+ * varno 'fromrel_varno') from the attnums of 'from_rel' to the attnums of
+ * 'to_rel', each of which may be either a leaf partition or a partitioned
+ * table, but both of which must be from the same partitioning hierarchy.
  *
- * Even though all of the same column names must be present in all relations
- * in the hierarchy, and they must also have the same types, the attnos may
- * be different.
- *
- * If found_whole_row is not NULL, *found_whole_row returns whether a
- * whole-row variable was found in the input expression.
+ * We need this because even though all of the same column names must be
+ * present in all relations in the hierarchy, and they must also have the
+ * same types, the attnums may be different.
  *
  * Note: this will work on any node tree, so really the argument and result
  * should be declared "Node *".  But a substantial majority of the callers
@@ -199,14 +196,12 @@ index_get_partition(Relation partition, Oid indexId)
  */
 List *
 map_partition_varattnos(List *expr, int fromrel_varno,
-						Relation to_rel, Relation from_rel,
-						bool *found_whole_row)
+						Relation to_rel, Relation from_rel)
 {
-	bool		my_found_whole_row = false;
-
 	if (expr != NIL)
 	{
 		AttrMap    *part_attmap;
+		bool		found_whole_row;
 
 		part_attmap = build_attrmap_by_name(RelationGetDescr(to_rel),
 											RelationGetDescr(from_rel));
@@ -214,12 +209,10 @@ map_partition_varattnos(List *expr, int fromrel_varno,
 											fromrel_varno, 0,
 											part_attmap,
 											RelationGetForm(to_rel)->reltype,
-											&my_found_whole_row);
+											&found_whole_row);
+		/* Since we provided a to_rowtype, we may ignore found_whole_row. */
 	}
 
-	if (found_whole_row)
-		*found_whole_row = my_found_whole_row;
-
 	return expr;
 }
 

src/backend/commands/tablecmds.c

@@ -15168,15 +15168,11 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
 				 */
 
 				/*
-				 * Cannot have expressions containing whole-row references or
-				 * system column references.
+				 * Cannot allow system column references, since that would
+				 * make partition routing impossible: their values won't be
+				 * known yet when we need to do that.
 				 */
 				pull_varattnos(expr, 1, &expr_attrs);
-				if (bms_is_member(0 - FirstLowInvalidHeapAttributeNumber,
-								  expr_attrs))
-					ereport(ERROR,
-							(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
-							 errmsg("partition key expressions cannot contain whole-row references")));
 				for (i = FirstLowInvalidHeapAttributeNumber; i < 0; i++)
 				{
 					if (bms_is_member(i - FirstLowInvalidHeapAttributeNumber,
@@ -15196,7 +15192,8 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
 				{
 					AttrNumber	attno = i + FirstLowInvalidHeapAttributeNumber;
 
-					if (TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated)
+					if (attno > 0 &&
+						TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated)
 						ereport(ERROR,
 								(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
 								 errmsg("cannot use generated column in partition key"),
@@ -15451,7 +15448,6 @@ QueuePartitionConstraintValidation(List **wqueue, Relation scanrel,
 		for (i = 0; i < partdesc->nparts; i++)
 		{
 			Relation	part_rel;
-			bool		found_whole_row;
 			List	   *thisPartConstraint;
 
 			/*
@@ -15465,10 +15461,7 @@ QueuePartitionConstraintValidation(List **wqueue, Relation scanrel,
 			 */
 			thisPartConstraint =
 				map_partition_varattnos(partConstraint, 1,
-										part_rel, scanrel, &found_whole_row);
-			/* There can never be a whole-row reference here */
-			if (found_whole_row)
-				elog(ERROR, "unexpected whole-row reference found in partition constraint");
+										part_rel, scanrel);
 
 			QueuePartitionConstraintValidation(wqueue, part_rel,
 											   thisPartConstraint,
@@ -15497,7 +15490,6 @@ ATExecAttachPartition(List **wqueue, Relation rel, PartitionCmd *cmd)
 	TupleDesc	tupleDesc;
 	ObjectAddress address;
 	const char *trigger_name;
-	bool		found_whole_row;
 	Oid			defaultPartOid;
 	List	   *partBoundConstraint;
 
@@ -15714,11 +15706,7 @@ ATExecAttachPartition(List **wqueue, Relation rel, PartitionCmd *cmd)
 		 * numbers.
 		 */
 		partConstraint = map_partition_varattnos(partConstraint, 1, attachrel,
-												 rel, &found_whole_row);
-		/* There can never be a whole-row reference here */
-		if (found_whole_row)
-			elog(ERROR,
-				 "unexpected whole-row reference found in partition key");
+												 rel);
 
 		/* Validate partition constraints against the table being attached. */
 		QueuePartitionConstraintValidation(wqueue, attachrel, partConstraint,
@@ -15750,7 +15738,7 @@ ATExecAttachPartition(List **wqueue, Relation rel, PartitionCmd *cmd)
 		 */
 		defPartConstraint =
 			map_partition_varattnos(defPartConstraint,
-									1, defaultrel, rel, NULL);
+									1, defaultrel, rel);
 		QueuePartitionConstraintValidation(wqueue, defaultrel,
 										   defPartConstraint, true);
 
@@ -16004,19 +15992,11 @@ CloneRowTriggersToPartition(Relation parent, Relation partition)
 							 RelationGetDescr(pg_trigger), &isnull);
 		if (!isnull)
 		{
-			bool		found_whole_row;
-
 			qual = stringToNode(TextDatumGetCString(value));
 			qual = (Node *) map_partition_varattnos((List *) qual, PRS2_OLD_VARNO,
-													partition, parent,
-													&found_whole_row);
-			if (found_whole_row)
-				elog(ERROR, "unexpected whole-row reference found in trigger WHEN clause");
+													partition, parent);
 			qual = (Node *) map_partition_varattnos((List *) qual, PRS2_NEW_VARNO,
-													partition, parent,
-													&found_whole_row);
-			if (found_whole_row)
-				elog(ERROR, "unexpected whole-row reference found in trigger WHEN clause");
+													partition, parent);
 		}
 
 		/*

src/backend/commands/trigger.c

@@ -1144,7 +1144,6 @@ CreateTrigger(CreateTrigStmt *stmt, const char *queryString,
 			CreateTrigStmt *childStmt;
 			Relation	childTbl;
 			Node	   *qual;
-			bool		found_whole_row;
 
 			childTbl = table_open(partdesc->oids[i], ShareRowExclusiveLock);
 
@@ -1177,16 +1176,10 @@ CreateTrigger(CreateTrigStmt *stmt, const char *queryString,
 			qual = copyObject(whenClause);
 			qual = (Node *)
 				map_partition_varattnos((List *) qual, PRS2_OLD_VARNO,
-										childTbl, rel,
-										&found_whole_row);
-			if (found_whole_row)
-				elog(ERROR, "unexpected whole-row reference found in trigger WHEN clause");
+										childTbl, rel);
 			qual = (Node *)
 				map_partition_varattnos((List *) qual, PRS2_NEW_VARNO,
-										childTbl, rel,
-										&found_whole_row);
-			if (found_whole_row)
-				elog(ERROR, "unexpected whole-row reference found in trigger WHEN clause");
+										childTbl, rel);
 
 			CreateTrigger(childStmt, queryString,
 						  partdesc->oids[i], refRelOid,

src/backend/partitioning/partbounds.c

@@ -1247,7 +1247,7 @@ check_default_partition_contents(Relation parent, Relation default_rel,
 	 */
 	def_part_constraints =
 		map_partition_varattnos(def_part_constraints, 1, default_rel,
-								parent, NULL);
+								parent);
 
 	/*
 	 * If the existing constraints on the default partition imply that it will
@@ -1297,7 +1297,7 @@ check_default_partition_contents(Relation parent, Relation default_rel,
 			partition_constraint = make_ands_explicit(def_part_constraints);
 			partition_constraint = (Expr *)
 				map_partition_varattnos((List *) partition_constraint, 1,
-										part_rel, default_rel, NULL);
+										part_rel, default_rel);
 
 			/*
 			 * If the partition constraints on default partition child imply

src/backend/utils/cache/partcache.c

@@ -342,7 +342,6 @@ generate_partition_qual(Relation rel)
 	List	   *my_qual = NIL,
 			   *result = NIL;
 	Relation	parent;
-	bool		found_whole_row;
 
 	/* Guard against stack overflow due to overly deep partition tree */
 	check_stack_depth();
@@ -388,11 +387,7 @@ generate_partition_qual(Relation rel)
 	 * in it to bear this relation's attnos. It's safe to assume varno = 1
 	 * here.
 	 */
-	result = map_partition_varattnos(result, 1, rel, parent,
-									 &found_whole_row);
-	/* There can never be a whole-row reference here */
-	if (found_whole_row)
-		elog(ERROR, "unexpected whole-row reference found in partition key");
+	result = map_partition_varattnos(result, 1, rel, parent);
 
 	/* Assert that we're not leaking any old data during assignments below */
 	Assert(rel->rd_partcheckcxt == NULL);

src/include/catalog/partition.h

@@ -23,8 +23,7 @@ extern Oid	get_partition_parent(Oid relid);
 extern List *get_partition_ancestors(Oid relid);
 extern Oid	index_get_partition(Relation partition, Oid indexId);
 extern List *map_partition_varattnos(List *expr, int fromrel_varno,
-									 Relation to_rel, Relation from_rel,
-									 bool *found_whole_row);
+									 Relation to_rel, Relation from_rel);
 extern bool has_partition_attrs(Relation rel, Bitmapset *attnums,
 								bool *used_in_expr);
 

src/test/regress/expected/create_table.out

@@ -420,11 +420,6 @@ CREATE TABLE partitioned (
 ) PARTITION BY RANGE (immut_func(a));
 ERROR:  functions in partition key expression must be marked IMMUTABLE
 DROP FUNCTION immut_func(int);
--- cannot contain whole-row references
-CREATE TABLE partitioned (
-	a	int
-) PARTITION BY RANGE ((partitioned));
-ERROR:  partition key expressions cannot contain whole-row references
 -- prevent using columns of unsupported types in key (type must have a btree operator class)
 CREATE TABLE partitioned (
 	a point
@@ -527,6 +522,31 @@ select * from partitioned where row(a,b)::partitioned = '(1,2)'::partitioned;
    Filter: (ROW(a, b)::partitioned = '(1,2)'::partitioned)
 (2 rows)
 
+drop table partitioned;
+-- whole-row Var in partition key works too
+create table partitioned (a int, b int)
+  partition by list ((partitioned));
+create table partitioned1
+  partition of partitioned for values in ('(1,2)');
+create table partitioned2
+  partition of partitioned for values in ('(2,4)');
+explain (costs off)
+select * from partitioned where partitioned = '(1,2)'::partitioned;
+                           QUERY PLAN                            
+-----------------------------------------------------------------
+ Seq Scan on partitioned1 partitioned
+   Filter: ((partitioned.*)::partitioned = '(1,2)'::partitioned)
+(2 rows)
+
+\d+ partitioned1
+                               Table "public.partitioned1"
+ Column |  Type   | Collation | Nullable | Default | Storage | Stats target | Description 
+--------+---------+-----------+----------+---------+---------+--------------+-------------
+ a      | integer |           |          |         | plain   |              | 
+ b      | integer |           |          |         | plain   |              | 
+Partition of: partitioned FOR VALUES IN ('(1,2)')
+Partition constraint: (((partitioned1.*)::partitioned IS DISTINCT FROM NULL) AND ((partitioned1.*)::partitioned = '(1,2)'::partitioned))
+
 drop table partitioned;
 -- check that dependencies of partition columns are handled correctly
 create domain intdom1 as int;

src/test/regress/sql/create_table.sql

@@ -401,11 +401,6 @@ CREATE TABLE partitioned (
 ) PARTITION BY RANGE (immut_func(a));
 DROP FUNCTION immut_func(int);
 
--- cannot contain whole-row references
-CREATE TABLE partitioned (
-	a	int
-) PARTITION BY RANGE ((partitioned));
-
 -- prevent using columns of unsupported types in key (type must have a btree operator class)
 CREATE TABLE partitioned (
 	a point
@@ -470,6 +465,18 @@ explain (costs off)
 select * from partitioned where row(a,b)::partitioned = '(1,2)'::partitioned;
 drop table partitioned;
 
+-- whole-row Var in partition key works too
+create table partitioned (a int, b int)
+  partition by list ((partitioned));
+create table partitioned1
+  partition of partitioned for values in ('(1,2)');
+create table partitioned2
+  partition of partitioned for values in ('(2,4)');
+explain (costs off)
+select * from partitioned where partitioned = '(1,2)'::partitioned;
+\d+ partitioned1
+drop table partitioned;
+
 -- check that dependencies of partition columns are handled correctly
 create domain intdom1 as int;