[CF 5720] v1 - virtual generated column as partition key

This branch was automatically generated by a robot using patches from an
email thread registered at:

https://commitfest.postgresql.org/patch/5720

The branch will be overwritten each time a new patch version is posted to
the thread, and also periodically to check for bitrot caused by changes
on the master branch.

Patch(es): https://www.postgresql.org/message-id/CACJufxF9RV2_iHBAG0AfCyt5c-1zJkh_Jc7F1tZfk3m3N+E=qA@mail.gmail.com
Author(s): Jian He

Author: Commitfest Bot

src/backend/commands/tablecmds.c

@@ -1325,7 +1325,8 @@ FormPartitionKeyDatum(PartitionDispatch pd,
 		Datum		datum;
 		bool		isNull;
 
-		if (keycol != 0)
+		if (keycol != 0 &&
+			RelationGetColumnGenerated(pd->reldesc, keycol) != ATTRIBUTE_GENERATED_VIRTUAL)
 		{
 			/* Plain column; get the value directly from the heap tuple */
 			datum = slot_getattr(slot, keycol, &isNull);

src/backend/executor/execPartition.c

@@ -232,13 +232,13 @@ static List *get_qual_for_hash(Relation parent, PartitionBoundSpec *spec);
 static List *get_qual_for_list(Relation parent, PartitionBoundSpec *spec);
 static List *get_qual_for_range(Relation parent, PartitionBoundSpec *spec,
 								bool for_default);
-static void get_range_key_properties(PartitionKey key, int keynum,
+static void get_range_key_properties(Relation rel, PartitionKey key, int keynum,
 									 PartitionRangeDatum *ldatum,
 									 PartitionRangeDatum *udatum,
 									 ListCell **partexprs_item,
 									 Expr **keyCol,
 									 Const **lower_val, Const **upper_val);
-static List *get_range_nulltest(PartitionKey key);
+static List *get_range_nulltest(Relation rel, PartitionKey key);
 
 /*
  * get_qual_from_partbound
@@ -4026,7 +4026,8 @@ get_qual_for_hash(Relation parent, PartitionBoundSpec *spec)
 		Node	   *keyCol;
 
 		/* Left operand */
-		if (key->partattrs[i] != 0)
+		if (key->partattrs[i] != 0 &&
+			RelationGetColumnGenerated(parent, key->partattrs[i]) != ATTRIBUTE_GENERATED_VIRTUAL)
 		{
 			keyCol = (Node *) makeVar(1,
 									  key->partattrs[i],
@@ -4082,7 +4083,8 @@ get_qual_for_list(Relation parent, PartitionBoundSpec *spec)
 	Assert(key->partnatts == 1);
 
 	/* Construct Var or expression representing the partition column */
-	if (key->partattrs[0] != 0)
+	if (key->partattrs[0] != 0 &&
+		RelationGetColumnGenerated(parent, key->partattrs[0]) != ATTRIBUTE_GENERATED_VIRTUAL)
 		keyCol = (Expr *) makeVar(1,
 								  key->partattrs[0],
 								  key->parttypid[0],
@@ -4353,7 +4355,7 @@ get_qual_for_range(Relation parent, PartitionBoundSpec *spec,
 			 */
 			other_parts_constr =
 				makeBoolExpr(AND_EXPR,
-							 lappend(get_range_nulltest(key),
+							 lappend(get_range_nulltest(parent, key),
 									 list_length(or_expr_args) > 1
 									 ? makeBoolExpr(OR_EXPR, or_expr_args,
 													-1)
@@ -4376,7 +4378,7 @@ get_qual_for_range(Relation parent, PartitionBoundSpec *spec,
 	 * to avoid accumulating the NullTest on the same keys for each partition.
 	 */
 	if (!for_default)
-		result = get_range_nulltest(key);
+		result = get_range_nulltest(parent, key);
 
 	/*
 	 * Iterate over the key columns and check if the corresponding lower and
@@ -4408,7 +4410,7 @@ get_qual_for_range(Relation parent, PartitionBoundSpec *spec,
 		 */
 		partexprs_item_saved = partexprs_item;
 
-		get_range_key_properties(key, i, ldatum, udatum,
+		get_range_key_properties(parent, key, i, ldatum, udatum,
 								 &partexprs_item,
 								 &keyCol,
 								 &lower_val, &upper_val);
@@ -4488,7 +4490,7 @@ get_qual_for_range(Relation parent, PartitionBoundSpec *spec,
 			if (lnext(spec->upperdatums, cell2))
 				udatum_next = castNode(PartitionRangeDatum,
 									   lfirst(lnext(spec->upperdatums, cell2)));
-			get_range_key_properties(key, j, ldatum, udatum,
+			get_range_key_properties(parent, key, j, ldatum, udatum,
 									 &partexprs_item,
 									 &keyCol,
 									 &lower_val, &upper_val);
@@ -4608,7 +4610,7 @@ get_qual_for_range(Relation parent, PartitionBoundSpec *spec,
 	 */
 	if (result == NIL)
 		result = for_default
-			? get_range_nulltest(key)
+			? get_range_nulltest(parent, key)
 			: list_make1(makeBoolConst(true, false));
 
 	return result;
@@ -4630,15 +4632,16 @@ get_qual_for_range(Relation parent, PartitionBoundSpec *spec,
  * the key->partexprs list, or NULL.  It may be advanced upon return.
  */
 static void
-get_range_key_properties(PartitionKey key, int keynum,
+get_range_key_properties(Relation rel, PartitionKey key, int keynum,
 						 PartitionRangeDatum *ldatum,
 						 PartitionRangeDatum *udatum,
 						 ListCell **partexprs_item,
 						 Expr **keyCol,
 						 Const **lower_val, Const **upper_val)
 {
 	/* Get partition key expression for this column */
-	if (key->partattrs[keynum] != 0)
+	if (key->partattrs[keynum] != 0 &&
+		RelationGetColumnGenerated(rel, key->partattrs[keynum]) != ATTRIBUTE_GENERATED_VIRTUAL)
 	{
 		*keyCol = (Expr *) makeVar(1,
 								   key->partattrs[keynum],
@@ -4674,7 +4677,7 @@ get_range_key_properties(PartitionKey key, int keynum,
  * keys to be null, so emit an IS NOT NULL expression for each key column.
  */
 static List *
-get_range_nulltest(PartitionKey key)
+get_range_nulltest(Relation rel, PartitionKey key)
 {
 	List	   *result = NIL;
 	NullTest   *nulltest;
@@ -4685,8 +4688,11 @@ get_range_nulltest(PartitionKey key)
 	for (i = 0; i < key->partnatts; i++)
 	{
 		Expr	   *keyCol;
+		bool	key_virtual;
 
-		if (key->partattrs[i] != 0)
+		key_virtual = (RelationGetColumnGenerated(rel, key->partattrs[i]) == ATTRIBUTE_GENERATED_VIRTUAL);
+
+		if (key->partattrs[i] != 0 && !key_virtual)
 		{
 			keyCol = (Expr *) makeVar(1,
 									  key->partattrs[i],

src/backend/partitioning/partbounds.c

@@ -194,6 +194,7 @@ RelationBuildPartitionKey(Relation relation)
 		HeapTuple	opclasstup;
 		Form_pg_opclass opclassform;
 		Oid			funcid;
+		bool		key_virtual;
 
 		/* Collect opfamily information */
 		opclasstup = SearchSysCache1(CLAOID,
@@ -225,8 +226,10 @@ RelationBuildPartitionKey(Relation relation)
 		/* Collation */
 		key->partcollation[i] = collation->values[i];
 
+		key_virtual = (RelationGetColumnGenerated(relation, attno) == ATTRIBUTE_GENERATED_VIRTUAL);
+
 		/* Collect type information */
-		if (attno != 0)
+		if (attno != 0 && !key_virtual)
 		{
 			Form_pg_attribute att = TupleDescAttr(relation->rd_att, attno - 1);
 

src/backend/utils/cache/partcache.c

@@ -6033,6 +6033,19 @@ RelationGetIndexAttOptions(Relation relation, bool copy)
 	return relation->rd_opcoptions;
 }
 
+char
+RelationGetColumnGenerated(Relation rel, int attnum)
+{
+	TupleDesc tupdesc;
+	Form_pg_attribute attr;
+
+	tupdesc = RelationGetDescr(rel);
+
+	attr = TupleDescAttr(tupdesc, attnum -1);
+
+	return attr->attgenerated;
+}
+
 /*
  * Routines to support ereport() reports of relation-related errors
  *

src/backend/utils/cache/relcache.c

@@ -60,6 +60,7 @@ extern List *RelationGetIndexExpressions(Relation relation);
 extern List *RelationGetDummyIndexExpressions(Relation relation);
 extern List *RelationGetIndexPredicate(Relation relation);
 extern bytea **RelationGetIndexAttOptions(Relation relation, bool copy);
+extern char RelationGetColumnGenerated(Relation relation, int attnum);
 
 /*
  * Which set of columns to return by RelationGetIndexAttrBitmap.

src/include/utils/relcache.h

@@ -1070,15 +1070,20 @@ SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
 -- we leave these tables around for purposes of testing dump/reload/upgrade
 -- generated columns in partition key (not allowed)
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE (f3);
-ERROR:  cannot use generated column in partition key
+ERROR:  cannot use stored generated column in partition key
 LINE 1: ...ENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE (f3);
                                                                    ^
+DETAIL:  Column "f3" is a stored generated column.
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
+ERROR:  cannot use stored generated column in partition key
+LINE 1: ...ERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
+                                                                 ^
 DETAIL:  Column "f3" is a generated column.
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
-ERROR:  cannot use generated column in partition key
+ERROR:  cannot use stored generated column in partition key
 LINE 1: ...ED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
                                                              ^
-DETAIL:  Column "f3" is a generated column.
+DETAIL:  Column "f3" is a stored generated column.
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);
 INSERT INTO gtest25 VALUES (3), (4);

src/test/regress/expected/generated_stored.out

@@ -1030,17 +1030,112 @@ SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
 (3 rows)
 
 -- we leave these tables around for purposes of testing dump/reload/upgrade
--- generated columns in partition key (not allowed)
-CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3);
-ERROR:  cannot use generated column in partition key
-LINE 1: ...NERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3);
-                                                                   ^
-DETAIL:  Column "f3" is a generated column.
-CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
-ERROR:  cannot use generated column in partition key
+-- tests for virtual generated columns in partition key
+CREATE TABLE gtest_part_key1 (f1 date, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3)); --error
+ERROR:  partition key expression cannot use virtual generated column
 LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
                                                              ^
-DETAIL:  Column "f3" is a generated column.
+HINT:  Only plain virtual generated column reference can be used in partition key
+CREATE TABLE gtest_part_key1 (f1 date, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3)); --ok
+ALTER TABLE gtest_part_key1 ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2); --error
+ERROR:  cannot alter column "f3" because it is part of the partition key of relation "gtest_part_key1"
+ALTER TABLE gtest_part_key1 ALTER COLUMN f2 set data type int; --error
+ERROR:  cannot alter type of a column used by a generated column
+DETAIL:  Column "f2" is used by generated column "f3".
+ALTER TABLE gtest_part_key1 ALTER COLUMN f3 set data type int; --error
+ERROR:  cannot alter column "f3" because it is part of the partition key of relation "gtest_part_key1"
+LINE 1: ALTER TABLE gtest_part_key1 ALTER COLUMN f3 set data type in...
+                                                 ^
+CREATE TABLE gtest_part_key1_0(f2 bigint, f1 date, f3 bigint GENERATED ALWAYS AS (f2 * 3) VIRTUAL);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); --error
+ERROR:  cannot attach table "gtest_part_key1_0" as a partition because it has with different generation expression
+ALTER TABLE gtest_part_key1_0 ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); --now ok
+CREATE TABLE gtest_part_key1_1 PARTITION OF gtest_part_key1 FOR VALUES FROM (30) TO (50);
+CREATE TABLE gtest_part_key1_2 PARTITION OF gtest_part_key1 FOR VALUES FROM (50) TO (100);
+\d+ gtest_part_key1
+                         Partitioned table "generated_virtual_tests.gtest_part_key1"
+ Column |  Type  | Collation | Nullable |           Default            | Storage | Stats target | Description 
+--------+--------+-----------+----------+------------------------------+---------+--------------+-------------
+ f1     | date   |           |          |                              | plain   |              | 
+ f2     | bigint |           |          |                              | plain   |              | 
+ f3     | bigint |           |          | generated always as (f2 * 2) | plain   |              | 
+Partition key: RANGE (f3)
+Partitions: gtest_part_key1_0 FOR VALUES FROM ('20') TO ('30'),
+            gtest_part_key1_1 FOR VALUES FROM ('30') TO ('50'),
+            gtest_part_key1_2 FOR VALUES FROM ('50') TO ('100')
+
+CREATE OR REPLACE FUNCTION gtest_trigger_info() RETURNS trigger
+  LANGUAGE plpgsql
+AS $$
+BEGIN
+	RAISE INFO 'TG_WHEN: % TG_RELNAME: % trigger name: % tg_op: %', TG_WHEN, TG_relname, TG_NAME, tg_op;
+  IF tg_op IN ('DELETE') THEN
+    RAISE INFO 'old = %', OLD;
+    RETURN OLD;
+  ELSIF tg_op IN ('INSERT') THEN
+    RAISE INFO 'new = %', NEW;
+    RETURN NEW;
+  ELSIF tg_op IN ('UPDATE') THEN
+    RAISE INFO 'old = %d; new = %', OLD, NEW;
+    RETURN NEW;
+  ELSE
+    RETURN NEW;
+  END IF;
+END
+$$;
+CREATE TRIGGER gkey1_0 BEFORE INSERT OR UPDATE ON gtest_part_key1_0
+FOR EACH ROW
+EXECUTE PROCEDURE gtest_trigger_info();
+CREATE TRIGGER gkey1_1 BEFORE INSERT OR UPDATE ON gtest_part_key1_1
+FOR EACH ROW
+EXECUTE PROCEDURE gtest_trigger_info();
+CREATE TRIGGER gkey1_2 BEFORE INSERT OR UPDATE ON gtest_part_key1_2
+FOR EACH ROW
+EXECUTE PROCEDURE gtest_trigger_info();
+INSERT INTO gtest_part_key1(f2, f3) VALUES (9, default); --error
+ERROR:  no partition of relation "gtest_part_key1" found for row
+DETAIL:  Partition key of the failing row contains (f3) = (18).
+INSERT INTO gtest_part_key1(f2, f3) VALUES (10, default) returning tableoid::regclass, *; --ok
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gkey1_0 tg_op: INSERT
+INFO:  new = (10,,)
+     tableoid      | f1 | f2 | f3 
+-------------------+----+----+----
+ gtest_part_key1_0 |    | 10 | 20
+(1 row)
+
+INSERT INTO gtest_part_key1_0(f2, f3) VALUES (12, default) returning tableoid::regclass, *; --ok
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gkey1_0 tg_op: INSERT
+INFO:  new = (12,,)
+     tableoid      | f2 | f1 | f3 
+-------------------+----+----+----
+ gtest_part_key1_0 | 12 |    | 24
+(1 row)
+
+MERGE INTO gtest_part_key1
+  USING (VALUES (10, 100), (12, 25), (14, 30)) AS s(sid, delta)
+ON gtest_part_key1.f2 = s.sid
+WHEN MATCHED AND f2 = 12 THEN UPDATE SET f2 = 20
+WHEN MATCHED AND f2 = 10 THEN UPDATE SET f2 = 30
+WHEN NOT MATCHED THEN INSERT(f2) VALUES (s.sid)
+RETURNING merge_action(), tableoid::regclass, old.f2, old.f3, new.f2, new.f3;
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gkey1_0 tg_op: UPDATE
+INFO:  old = (10,,)d; new = (30,,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_2 trigger name: gkey1_2 tg_op: INSERT
+INFO:  new = (,30,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gkey1_0 tg_op: UPDATE
+INFO:  old = (12,,)d; new = (20,,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_1 trigger name: gkey1_1 tg_op: INSERT
+INFO:  new = (,20,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gkey1_0 tg_op: INSERT
+INFO:  new = (14,,)
+ merge_action |     tableoid      | f2 | f3 | f2 | f3 
+--------------+-------------------+----+----+----+----
+ UPDATE       | gtest_part_key1_2 | 10 | 20 | 30 | 60
+ UPDATE       | gtest_part_key1_1 | 12 | 24 | 20 | 40
+ INSERT       | gtest_part_key1_0 |    |    | 14 | 28
+(3 rows)
+
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);
 INSERT INTO gtest25 VALUES (3), (4);

src/test/regress/expected/generated_virtual.out

@@ -500,6 +500,7 @@ SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
 
 -- generated columns in partition key (not allowed)
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE (f3);
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
 
 -- ALTER TABLE ... ADD COLUMN