Hi.

Rebasing and rechecking the code changes again.

typedef struct PartitionKeyData
{
    PartitionStrategy strategy; /* partitioning strategy */
    int16        partnatts;        /* number of columns in the partition key */
    AttrNumber *partattrs;
    List       *partexprs;
....
}

Normally, a partition key is mutually exclusive: it is either a simple column
(partattrs[i] != 0 and partexprs[i] == NULL) or an expression (partattrs[i] == 0
with partexprs[i] != NULL). However, for virtual generated columns, it's
possible for both to exist (partattrs[i] != 0 and partexprs[i] != NULL).

Because of this scenario, I  have double checked all occurrences of
`->partattrs` in the codebase to ensure partition keys with virtual generated
columns are handled properly.



--
jian
https://www.enterprisedb.com/
From 2e4eb97f3a926c2b9cd4f284908b08d3c482adc5 Mon Sep 17 00:00:00 2001
From: jian he <[email protected]>
Date: Wed, 4 Mar 2026 15:56:37 +0800
Subject: [PATCH v4 1/1] virtual generated column as partition key

If the partition key contains a virtual generated column, then the generation
expression for each partition must match with the partitioned table, since it is
used as a partition key. Otherwise, the partition bound would be dynamically
evaluated.

We use a delicate, has_partition_attrs, to handle dependencies on partition key
columns. This function automatically checks operations like altering column
types, altering generation expressions, or dropping columns.  So it's OK to to
register direct dependency between the virtual generated column and the
partition key expression.

A virtual generated column entry in the pg_partitioned_table catalog is flagged
with both non-zero partattrs and non-NULL partexprs, which is abnormal.
normally, partattrs is non-zero or partexprs is NULL, but not both.

Maybe we should mention this in the doc/src/sgml/catalogs.sgml

demo:
CREATE TABLE t(f1 bigint, f2 bigint GENERATED ALWAYS AS (f1 * 2) VIRTUAL) PARTITION BY RANGE (f2);

The partition key cannot be expression on top of virtual generated column. so
the following is not allowed:

CREATE TABLE t(f1 bigint, f2 bigint GENERATED ALWAYS AS (f1 * 2) VIRTUAL) PARTITION BY RANGE ((f2+1));

commitfest: https://commitfest.postgresql.org/patch/5720
discussion: https://postgr.es/m/CACJufxF9RV2_iHBAG0AfCyt5c-1zJkh_Jc7F1tZfk3m3N+E=q...@mail.gmail.com
---
 src/backend/catalog/partition.c               |   3 +
 src/backend/commands/indexcmds.c              |  12 +
 src/backend/commands/tablecmds.c              | 174 +++++++++-
 src/backend/executor/execPartition.c          |   2 +-
 src/backend/optimizer/util/plancat.c          |   2 +-
 src/backend/partitioning/partbounds.c         |  28 +-
 src/backend/utils/adt/ruleutils.c             |   7 +
 src/backend/utils/cache/partcache.c           |   3 +
 src/backend/utils/cache/relcache.c            |  14 +
 src/include/utils/partcache.h                 |  14 +-
 src/include/utils/relcache.h                  |   1 +
 .../regress/expected/generated_stored.out     |  20 +-
 .../regress/expected/generated_virtual.out    | 319 ++++++++++++++++--
 src/test/regress/sql/generated_virtual.sql    | 131 ++++++-
 14 files changed, 665 insertions(+), 65 deletions(-)

diff --git a/src/backend/catalog/partition.c b/src/backend/catalog/partition.c
index 28f3cade6ff..9a95a2358cf 100644
--- a/src/backend/catalog/partition.c
+++ b/src/backend/catalog/partition.c
@@ -274,6 +274,9 @@ has_partition_attrs(Relation rel, Bitmapset *attnums, bool *used_in_expr)
 
 		if (partattno != 0)
 		{
+			if (attrIsVirtualGenerated(rel, partattno))
+				partexprs_item = lnext(partexprs, partexprs_item);
+
 			if (bms_is_member(partattno - FirstLowInvalidHeapAttributeNumber,
 							  attnums))
 			{
diff --git a/src/backend/commands/indexcmds.c b/src/backend/commands/indexcmds.c
index 635679cc1f2..39808d0690a 100644
--- a/src/backend/commands/indexcmds.c
+++ b/src/backend/commands/indexcmds.c
@@ -1024,6 +1024,18 @@ DefineIndex(ParseState *pstate,
 						 errdetail("%s constraints cannot be used when partition keys include expressions.",
 								   constraint_type)));
 
+			/*
+			 * Since we do not support indexes on virtual generated columns,
+			 * UNIQUE constraints on these columns are also unsupported
+			 */
+			if (attrIsVirtualGenerated(rel, key->partattrs[i]))
+				ereport(ERROR,
+						errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						errmsg("unsupported %s constraint with partition key definition",
+							   constraint_type),
+						errdetail("%s constraints cannot be used when partition keys include virtual generated column.",
+								  constraint_type));
+
 			/* Search the index column(s) for a match */
 			for (j = 0; j < indexInfo->ii_NumIndexKeyAttrs; j++)
 			{
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index b04b0dbd2a0..e416696b1a9 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -8678,6 +8678,51 @@ ATExecSetExpression(AlteredTableInfo *tab, Relation rel, const char *colName,
 	if (attgenerated == ATTRIBUTE_GENERATED_VIRTUAL && attTup->attnotnull)
 		tab->verify_new_notnull = true;
 
+	/*
+	 * We cannot alter the generation expression of a virtual generated column
+	 * if it's used in the partition key. Note that stored generated columns
+	 * are already rejected as partition keys in ComputePartitionAttrs.
+	 */
+	if (attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+	{
+		if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE &&
+			has_partition_attrs(rel,
+								bms_make_singleton(attnum - FirstLowInvalidHeapAttributeNumber),
+								NULL))
+			ereport(ERROR,
+					errcode(ERRCODE_INVALID_COLUMN_DEFINITION),
+					errmsg("cannot alter column \"%s\" because it is part of the partition key of relation \"%s\"",
+						   colName, RelationGetRelationName(rel)));
+
+		if (rel->rd_rel->relispartition)
+		{
+			AttrNumber	parent_attnum;
+
+			Oid			parentId = get_partition_parent(RelationGetRelid(rel),
+														false);
+			Relation	parent = table_open(parentId, AccessShareLock);
+
+			AttrMap    *map = build_attrmap_by_name_if_req(RelationGetDescr(parent),
+														   RelationGetDescr(rel),
+														   false);
+
+			if (map != NULL)
+				parent_attnum = map->attnums[attnum - 1];
+			else
+				parent_attnum = attnum;
+
+			if (has_partition_attrs(parent,
+									bms_make_singleton(parent_attnum - FirstLowInvalidHeapAttributeNumber),
+									NULL))
+				ereport(ERROR,
+						errcode(ERRCODE_INVALID_COLUMN_DEFINITION),
+						errmsg("cannot alter column \"%s\" because it is part of the partition key of relation \"%s\"",
+							   colName, RelationGetRelationName(parent)));
+
+			table_close(parent, AccessShareLock);
+		}
+	}
+
 	/*
 	 * We need to prevent this because a change of expression could affect a
 	 * row filter and inject expressions that are not permitted in a row
@@ -19821,6 +19866,7 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
 		PartitionElem *pelem = lfirst_node(PartitionElem, lc);
 		Oid			atttype;
 		Oid			attcollation;
+		AttrNumber	virtualattnum = InvalidAttrNumber;
 
 		if (pelem->name != NULL)
 		{
@@ -19848,24 +19894,70 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
 			/*
 			 * Stored generated columns cannot work: They are computed after
 			 * BEFORE triggers, but partition routing is done before all
-			 * triggers.  Maybe virtual generated columns could be made to
-			 * work, but then they would need to be handled as an expression
-			 * below.
+			 * triggers. However virtual generated columns is supported.
 			 */
-			if (attform->attgenerated)
+			if (attform->attgenerated == ATTRIBUTE_GENERATED_STORED)
 				ereport(ERROR,
 						(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
-						 errmsg("cannot use generated column in partition key"),
-						 errdetail("Column \"%s\" is a generated column.",
+						 errmsg("cannot use stored generated column in partition key"),
+						 errdetail("Column \"%s\" is a stored generated column.",
 								   pelem->name),
 						 parser_errposition(pstate, pelem->location)));
 
+			if (attform->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+				virtualattnum = attform->attnum;
+
 			partattrs[attn] = attform->attnum;
 			atttype = attform->atttypid;
 			attcollation = attform->attcollation;
 			ReleaseSysCache(atttuple);
 		}
 		else
+		{
+			/* Expression */
+			Node	   *expr = pelem->expr;
+
+			atttype = exprType(expr);
+			attcollation = exprCollation(expr);
+
+			while (expr && IsA(expr, CollateExpr))
+				expr = (Node *) ((CollateExpr *) expr)->arg;
+
+			if (IsA(expr, Var) && ((Var *) expr)->varattno > 0)
+			{
+				Var		   *var = (Var *) expr;
+
+				if (TupleDescAttr(RelationGetDescr(rel), var->varattno - 1)->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+					virtualattnum = var->varattno;
+			}
+		}
+
+		if (AttributeNumberIsValid(virtualattnum))
+		{
+			Node	   *expr = build_generation_expression(rel, virtualattnum);
+
+			expr = (Node *) expression_planner((Expr *) expr);
+
+			/*
+			 * Generation expression expected to be IMMUTABLE, So this is
+			 * unlikely to happen.
+			 */
+			if (contain_mutable_functions(expr))
+				elog(ERROR, "functions in partition key expression must be marked IMMUTABLE");
+
+			/*
+			 * While it is not exactly *wrong* for a partition expression to
+			 * be a constant, it seems better to reject such keys.
+			 */
+			if (IsA(expr, Const))
+				ereport(ERROR,
+						errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
+						errmsg("cannot use constant expression as partition key"));
+
+			partattrs[attn] = virtualattnum;
+			*partexprs = lappend(*partexprs, expr);
+		}
+		else if (pelem->expr != NULL)
 		{
 			/* Expression */
 			Node	   *expr = pelem->expr;
@@ -19929,18 +20021,25 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
 				/*
 				 * Stored generated columns cannot work: They are computed
 				 * after BEFORE triggers, but partition routing is done before
-				 * all triggers.  Virtual generated columns could probably
-				 * work, but it would require more work elsewhere (for example
-				 * SET EXPRESSION would need to check whether the column is
-				 * used in partition keys).  Seems safer to prohibit for now.
+				 * all triggers.
+				 *
+				 * Virtual generated columns are supported, but expression
+				 * over virtual generated column is not supported.
 				 */
-				if (TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated)
+				if (TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated == ATTRIBUTE_GENERATED_STORED)
 					ereport(ERROR,
 							(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
-							 errmsg("cannot use generated column in partition key"),
-							 errdetail("Column \"%s\" is a generated column.",
+							 errmsg("cannot use stored generated column in partition key"),
+							 errdetail("Column \"%s\" is a stored generated column.",
 									   get_attname(RelationGetRelid(rel), attno, false)),
 							 parser_errposition(pstate, pelem->location)));
+
+				if (TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+					ereport(ERROR,
+							errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+							errmsg("cannot use expression over virtual generated column in partition key"),
+							errdetail("Partition key expression over virtual generated column is not supported"),
+							parser_errposition(pstate, pelem->location));
 			}
 
 			if (IsA(expr, Var) &&
@@ -20480,6 +20579,55 @@ ATExecAttachPartition(List **wqueue, Relation rel, PartitionCmd *cmd,
 					 errdetail("The new partition may contain only the columns present in parent.")));
 	}
 
+	/*
+	 * If the partition key contains virtual generated columns, then the
+	 * generated expression in partition key must match that of the
+	 * partitioned table.
+	 */
+	if (tupleDesc->constr && tupleDesc->constr->has_generated_virtual)
+	{
+		AttrMap    *map = build_attrmap_by_name_if_req(RelationGetDescr(rel),
+													   tupleDesc, false);
+
+		PartitionKey key = RelationGetPartitionKey(rel);
+
+		for (int i = 0; i < key->partnatts; i++)
+		{
+			if (AttributeNumberIsValid(key->partattrs[i]) &&
+				attrIsVirtualGenerated(rel, key->partattrs[i]))
+			{
+				Node	   *attachrel_defval = NULL;
+				bool		found_whole_row = false;
+
+				Node	   *rel_defval = build_generation_expression(rel,
+																	 key->partattrs[i]);
+
+				if (map)
+				{
+					attachrel_defval = build_generation_expression(attachrel,
+																   map->attnums[key->partattrs[i] - 1]);
+					attachrel_defval = map_variable_attnos(attachrel_defval,
+														   1, 0,
+														   map,
+														   InvalidOid,
+														   &found_whole_row);
+				}
+				else
+					attachrel_defval = build_generation_expression(attachrel,
+																   key->partattrs[i]);
+
+				if (found_whole_row)
+					elog(ERROR, "cannot use whole-row variable in column generation expression");
+
+				if (!equal(rel_defval, attachrel_defval))
+					ereport(ERROR,
+							errcode(ERRCODE_WRONG_OBJECT_TYPE),
+							errmsg("cannot attach table \"%s\" as a partition because it has with different generation expression",
+								   RelationGetRelationName(attachrel)));
+			}
+		}
+	}
+
 	/*
 	 * If child_rel has row-level triggers with transition tables, we
 	 * currently don't allow it to become a partition.  See also prohibitions
diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index bab294f5e91..7c8975fcfbb 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -1495,7 +1495,7 @@ FormPartitionKeyDatum(PartitionDispatch pd,
 		Datum		datum;
 		bool		isNull;
 
-		if (keycol != 0)
+		if (keycol != 0 && !attrIsVirtualGenerated(pd->reldesc, keycol))
 		{
 			/* Plain column; get the value directly from the heap tuple */
 			datum = slot_getattr(slot, keycol, &isNull);
diff --git a/src/backend/optimizer/util/plancat.c b/src/backend/optimizer/util/plancat.c
index d63e7390be7..e6c41bd73c7 100644
--- a/src/backend/optimizer/util/plancat.c
+++ b/src/backend/optimizer/util/plancat.c
@@ -2856,7 +2856,7 @@ set_baserel_partition_key_exprs(Relation relation,
 		Expr	   *partexpr;
 		AttrNumber	attno = partkey->partattrs[cnt];
 
-		if (attno != InvalidAttrNumber)
+		if (attno != InvalidAttrNumber && !attrIsVirtualGenerated(relation, attno))
 		{
 			/* Single column partition key is stored as a Var node. */
 			Assert(attno > 0);
diff --git a/src/backend/partitioning/partbounds.c b/src/backend/partitioning/partbounds.c
index 0ca312ac27d..abec13cb356 100644
--- a/src/backend/partitioning/partbounds.c
+++ b/src/backend/partitioning/partbounds.c
@@ -233,13 +233,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
@@ -4017,7 +4017,7 @@ get_qual_for_hash(Relation parent, PartitionBoundSpec *spec)
 		Node	   *keyCol;
 
 		/* Left operand */
-		if (key->partattrs[i] != 0)
+		if (key->partattrs[i] != 0 && !attrIsVirtualGenerated(parent, key->partattrs[i]))
 		{
 			keyCol = (Node *) makeVar(1,
 									  key->partattrs[i],
@@ -4073,7 +4073,7 @@ 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 && !attrIsVirtualGenerated(parent, key->partattrs[0]))
 		keyCol = (Expr *) makeVar(1,
 								  key->partattrs[0],
 								  key->parttypid[0],
@@ -4344,7 +4344,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)
@@ -4367,7 +4367,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
@@ -4399,7 +4399,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);
@@ -4479,7 +4479,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);
@@ -4599,7 +4599,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;
@@ -4621,7 +4621,7 @@ 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,
@@ -4629,7 +4629,8 @@ get_range_key_properties(PartitionKey key, int keynum,
 						 Const **lower_val, Const **upper_val)
 {
 	/* Get partition key expression for this column */
-	if (key->partattrs[keynum] != 0)
+	if (key->partattrs[keynum] != 0 &&
+		!attrIsVirtualGenerated(rel, key->partattrs[keynum]))
 	{
 		*keyCol = (Expr *) makeVar(1,
 								   key->partattrs[keynum],
@@ -4665,7 +4666,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;
@@ -4677,7 +4678,8 @@ get_range_nulltest(PartitionKey key)
 	{
 		Expr	   *keyCol;
 
-		if (key->partattrs[i] != 0)
+		if (key->partattrs[i] != 0 &&
+			!attrIsVirtualGenerated(rel, key->partattrs[i]))
 		{
 			keyCol = (Expr *) makeVar(1,
 									  key->partattrs[i],
diff --git a/src/backend/utils/adt/ruleutils.c b/src/backend/utils/adt/ruleutils.c
index f16f1535785..886df60c868 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -2043,6 +2043,13 @@ pg_get_partkeydef_worker(Oid relid, int prettyFlags,
 			get_atttypetypmodcoll(relid, attnum,
 								  &keycoltype, &keycoltypmod,
 								  &keycolcollation);
+
+			/*
+			 * Advance partexprs if the partition key is a virtual generated
+			 * column, as it holds the generation expression
+			 */
+			if (get_attgenerated(relid, attnum) == ATTRIBUTE_GENERATED_VIRTUAL)
+				partexpr_item = lnext(partexprs, partexpr_item);
 		}
 		else
 		{
diff --git a/src/backend/utils/cache/partcache.c b/src/backend/utils/cache/partcache.c
index 3107075c9ad..83f051b3623 100644
--- a/src/backend/utils/cache/partcache.c
+++ b/src/backend/utils/cache/partcache.c
@@ -233,6 +233,9 @@ RelationBuildPartitionKey(Relation relation)
 			key->parttypid[i] = att->atttypid;
 			key->parttypmod[i] = att->atttypmod;
 			key->parttypcoll[i] = att->attcollation;
+
+			if (att->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+				partexprs_item = lnext(key->partexprs, partexprs_item);
 		}
 		else
 		{
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 6b634c9fff1..7a35034876f 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -6028,6 +6028,20 @@ RelationGetIndexAttOptions(Relation relation, bool copy)
 	return relation->rd_opcoptions;
 }
 
+bool
+attrIsVirtualGenerated(Relation rel, AttrNumber attnum)
+{
+	Form_pg_attribute attr;
+
+	TupleDesc	tupdesc = RelationGetDescr(rel);
+
+	Assert(attnum > 0);
+
+	attr = TupleDescAttr(tupdesc, attnum - 1);
+
+	return (attr->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL);
+}
+
 /*
  * Routines to support ereport() reports of relation-related errors
  *
diff --git a/src/include/utils/partcache.h b/src/include/utils/partcache.h
index b1608c985ee..213639dfa84 100644
--- a/src/include/utils/partcache.h
+++ b/src/include/utils/partcache.h
@@ -27,9 +27,17 @@ typedef struct PartitionKeyData
 	PartitionStrategy strategy; /* partitioning strategy */
 	int16		partnatts;		/* number of columns in the partition key */
 	AttrNumber *partattrs;		/* attribute numbers of columns in the
-								 * partition key or 0 if it's an expr */
-	List	   *partexprs;		/* list of expressions in the partitioning
-								 * key, one for each zero-valued partattrs */
+								 * partition key or 0 if it's an expr. Note:
+								 * These can include attribute numbers for
+								 * virtual generated columns. */
+
+	/*
+	 * list of expressions in the partitioning key, one for each zero-valued
+	 * partattrs. For a virtual generated column, this is its generation
+	 * expression, and partattrs stores the attribute number of the virtual
+	 * generated column.
+	 */
+	List	   *partexprs;
 
 	Oid		   *partopfamily;	/* OIDs of operator families */
 	Oid		   *partopcintype;	/* OIDs of opclass declared input data types */
diff --git a/src/include/utils/relcache.h b/src/include/utils/relcache.h
index 2700224939a..d0c6f48a3e4 100644
--- a/src/include/utils/relcache.h
+++ b/src/include/utils/relcache.h
@@ -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 bool attrIsVirtualGenerated(Relation rel, AttrNumber attnum);
 
 /*
  * Which set of columns to return by RelationGetIndexAttrBitmap.
diff --git a/src/test/regress/expected/generated_stored.out b/src/test/regress/expected/generated_stored.out
index 8b7a71d8f0c..eb8fc19e1a8 100644
--- a/src/test/regress/expected/generated_stored.out
+++ b/src/test/regress/expected/generated_stored.out
@@ -1076,30 +1076,30 @@ 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 generated column.
+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 generated column in partition key
+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.
+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 * 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.
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((gtest_part_key));
-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 ((gtest_par...
                                                              ^
-DETAIL:  Column "f3" is a generated column.
+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 ((gtest_part_key is not null));
-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 ((gtest_par...
                                                              ^
-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);
diff --git a/src/test/regress/expected/generated_virtual.out b/src/test/regress/expected/generated_virtual.out
index 6dab60c937b..3591c97b17f 100644
--- a/src/test/regress/expected/generated_virtual.out
+++ b/src/test/regress/expected/generated_virtual.out
@@ -1063,32 +1063,311 @@ 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));
-ERROR:  cannot use generated column in partition key
-LINE 1: ...RATED 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
-LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
-                                                             ^
-DETAIL:  Column "f3" is a generated column.
+-- generated columns in partition key
+-- wholerow with virtual generated column is not supported
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_part_key));
-ERROR:  cannot use generated column in partition key
+ERROR:  cannot use expression over virtual generated column in partition key
 LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_par...
                                                              ^
-DETAIL:  Column "f3" is a generated column.
+DETAIL:  Partition key expression over virtual generated column is not supported
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_part_key is not null));
-ERROR:  cannot use generated column in partition key
+ERROR:  cannot use expression over virtual generated column in partition key
 LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_par...
                                                              ^
-DETAIL:  Column "f3" is a generated column.
+DETAIL:  Partition key expression over virtual generated column is not supported
+-- unique/primary key interact with partition key will fail
+CREATE TABLE gtest_idxpart (a int PRIMARY KEY, b int GENERATED ALWAYS AS (a) ) PARTITION BY RANGE ((b));
+ERROR:  unsupported PRIMARY KEY constraint with partition key definition
+DETAIL:  PRIMARY KEY constraints cannot be used when partition keys include virtual generated column.
+CREATE TABLE gtest_idxpart (a int UNIQUE, b int GENERATED ALWAYS AS (a) ) PARTITION BY RANGE ((b));
+ERROR:  unsupported UNIQUE constraint with partition key definition
+DETAIL:  UNIQUE constraints cannot be used when partition keys include virtual generated column.
+CREATE TABLE gtest_idxpart (a int UNIQUE, b int GENERATED ALWAYS AS (a) ) PARTITION BY RANGE (a, b);
+ERROR:  unsupported UNIQUE constraint with partition key definition
+DETAIL:  UNIQUE constraints cannot be used when partition keys include virtual generated column.
+CREATE TABLE gtest_idxpart (a int UNIQUE, b int GENERATED ALWAYS AS (a) ) PARTITION BY RANGE (a); -- ok
+-- partition key expression over virtual generated column is not supported
+CREATE TABLE gtest_part_key (f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
+ERROR:  cannot use expression over virtual generated column in partition key
+LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
+                                                             ^
+DETAIL:  Partition key expression over virtual generated column is not supported
+-- constant virtual generated column expression is not supported
+CREATE TABLE gtest_part_key (f3 bigint GENERATED ALWAYS AS (2)) PARTITION BY RANGE (f3); -- error
+ERROR:  cannot use constant expression as partition key
+CREATE TABLE gtest_part_keyxx (f2 text, f3 TEXT COLLATE "POSIX" GENERATED ALWAYS AS ( (upper(f2)) COLLATE "C")  VIRTUAL) PARTITION BY RANGE (f3, f3, f2, f3, (upper(f2)));
+SELECT pg_get_partkeydef('gtest_part_keyxx'::regclass);
+         pg_get_partkeydef         
+-----------------------------------
+ RANGE (f3, f3, f2, f3, upper(f2))
+(1 row)
+
+SELECT  partrelid::regclass, partnatts, partattrs,  partcollation[0]::regcollation
+FROM    pg_partitioned_table
+WHERE   partrelid = 'gtest_part_keyxx'::regclass;
+    partrelid     | partnatts | partattrs | partcollation 
+------------------+-----------+-----------+---------------
+ gtest_part_keyxx |         5 | 2 2 1 2 0 | "POSIX"
+(1 row)
+
+DROP TABLE gtest_part_keyxx;
+-- tests for virtual generated columns in partition key
+CREATE TABLE gtest_part_key  (f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3); -- ok
+CREATE TABLE gtest_part_key1 (f1 date, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3)); -- ok
+-- Once a virtual generated column is used as a partition key, its data type and
+-- generation expression can no longer be changed. All below should result
+-- error.
+ALTER TABLE gtest_part_key1 ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2);
+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 f3 SET DATA TYPE INT;
+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...
+                                                 ^
+ALTER TABLE gtest_part_key1 ALTER COLUMN f3 DROP EXPRESSION;
+ERROR:  ALTER TABLE / DROP EXPRESSION is not supported for virtual generated columns
+DETAIL:  Column "f3" of relation "gtest_part_key1" is a virtual generated column.
+CREATE TABLE gtest_part_key1_0(f2 bigint, f1 date, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED);
+-- error: The partition's generation kind must match the parent partitioned table
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30);
+ERROR:  column "f3" inherits from generated column of different kind
+DETAIL:  Parent column is VIRTUAL, child column is STORED.
+DROP TABLE gtest_part_key1_0;
+CREATE TABLE gtest_part_key1_2(f2 bigint, f1 date, f3 bigint GENERATED ALWAYS AS (f2 * 3) VIRTUAL);
+-- error: The partition's generation expression must match the parent partitioned table
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_2 FOR VALUES FROM (20) TO (30);
+ERROR:  cannot attach table "gtest_part_key1_2" as a partition because it has with different generation expression
+DROP TABLE gtest_part_key1_2;
+CREATE TABLE gtest_part_key1_0(f2 bigint, f1 date, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); -- ok
+ALTER TABLE gtest_part_key1_0 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"
+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')
+
+INSERT INTO gtest_part_key1(f2) VALUES (9);     -- 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_2(f2) VALUES (50);  -- error
+ERROR:  new row for relation "gtest_part_key1_2" violates partition constraint
+DETAIL:  Failing row contains (null, 50, virtual).
+INSERT INTO gtest_part_key1(f2) VALUES (10), (12), (25), (30), (20) RETURNING tableoid::regclass, *; -- ok
+     tableoid      | f1 | f2 | f3 
+-------------------+----+----+----
+ gtest_part_key1_0 |    | 10 | 20
+ gtest_part_key1_0 |    | 12 | 24
+ gtest_part_key1_2 |    | 25 | 50
+ gtest_part_key1_2 |    | 30 | 60
+ gtest_part_key1_1 |    | 20 | 40
+(5 rows)
+
+UPDATE gtest_part_key1 SET f2 = 50 WHERE f2 = 30; -- error
+ERROR:  no partition of relation "gtest_part_key1" found for row
+DETAIL:  Partition key of the failing row contains (f3) = (100).
+UPDATE gtest_part_key1 SET f2 = 13 WHERE f2 = 20 RETURNING new.tableoid::regclass, old.tableoid::regclass, OLD.*, NEW.*;
+     tableoid      |     tableoid      | f1 | f2 | f3 | f1 | f2 | f3 
+-------------------+-------------------+----+----+----+----+----+----
+ gtest_part_key1_0 | gtest_part_key1_1 |    | 20 | 40 |    | 13 | 26
+(1 row)
+
+SET max_parallel_workers_per_gather TO 0;
+SET enable_incremental_sort TO off;
+SET enable_partition_pruning TO true;
+EXPLAIN(COSTS OFF) SELECT * FROM gtest_part_key1 WHERE f3 < 50;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Append
+   ->  Seq Scan on gtest_part_key1_0 gtest_part_key1_1
+         Filter: ((f2 * 2) < 50)
+   ->  Seq Scan on gtest_part_key1_1 gtest_part_key1_2
+         Filter: ((f2 * 2) < 50)
+(5 rows)
+
+SET enable_partition_pruning TO false;
+EXPLAIN(COSTS OFF) SELECT * FROM gtest_part_key1 WHERE f3 < 50;
+                      QUERY PLAN                       
+-------------------------------------------------------
+ Append
+   ->  Seq Scan on gtest_part_key1_0 gtest_part_key1_1
+         Filter: ((f2 * 2) < 50)
+   ->  Seq Scan on gtest_part_key1_1 gtest_part_key1_2
+         Filter: ((f2 * 2) < 50)
+   ->  Seq Scan on gtest_part_key1_2 gtest_part_key1_3
+         Filter: ((f2 * 2) < 50)
+(7 rows)
+
+SET enable_partitionwise_aggregate TO true;
+EXPLAIN (COSTS OFF) SELECT f3, count(*) FROM gtest_part_key1 GROUP BY f3 ORDER BY 1;
+                           QUERY PLAN                            
+-----------------------------------------------------------------
+ Sort
+   Sort Key: ((gtest_part_key1.f2 * 2))
+   ->  Append
+         ->  HashAggregate
+               Group Key: (gtest_part_key1.f2 * 2)
+               ->  Seq Scan on gtest_part_key1_0 gtest_part_key1
+         ->  HashAggregate
+               Group Key: (gtest_part_key1_1.f2 * 2)
+               ->  Seq Scan on gtest_part_key1_1
+         ->  HashAggregate
+               Group Key: (gtest_part_key1_2.f2 * 2)
+               ->  Seq Scan on gtest_part_key1_2
+(12 rows)
+
+SET enable_partitionwise_aggregate TO false;
+EXPLAIN (COSTS OFF) SELECT f3, count(*) FROM gtest_part_key1 GROUP BY f3 ORDER BY 1;
+                            QUERY PLAN                             
+-------------------------------------------------------------------
+ Sort
+   Sort Key: ((gtest_part_key1.f2 * 2))
+   ->  HashAggregate
+         Group Key: ((gtest_part_key1.f2 * 2))
+         ->  Append
+               ->  Seq Scan on gtest_part_key1_0 gtest_part_key1_1
+               ->  Seq Scan on gtest_part_key1_1 gtest_part_key1_2
+               ->  Seq Scan on gtest_part_key1_2 gtest_part_key1_3
+(8 rows)
+
+SET enable_partitionwise_join TO true;
+EXPLAIN (COSTS OFF)
+SELECT t1.f3, count(t2.f3) FROM gtest_part_key1 t1 JOIN gtest_part_key1 t2 ON t1.f3 = t2.f3 GROUP BY 1 ORDER BY 1;
+                             QUERY PLAN                              
+---------------------------------------------------------------------
+ Sort
+   Sort Key: ((t1.f2 * 2))
+   ->  HashAggregate
+         Group Key: ((t1.f2 * 2))
+         ->  Append
+               ->  Merge Join
+                     Merge Cond: (((t1_1.f2 * 2)) = ((t2_1.f2 * 2)))
+                     ->  Sort
+                           Sort Key: ((t1_1.f2 * 2))
+                           ->  Seq Scan on gtest_part_key1_0 t1_1
+                     ->  Sort
+                           Sort Key: ((t2_1.f2 * 2))
+                           ->  Seq Scan on gtest_part_key1_0 t2_1
+               ->  Merge Join
+                     Merge Cond: (((t1_2.f2 * 2)) = ((t2_2.f2 * 2)))
+                     ->  Sort
+                           Sort Key: ((t1_2.f2 * 2))
+                           ->  Seq Scan on gtest_part_key1_1 t1_2
+                     ->  Sort
+                           Sort Key: ((t2_2.f2 * 2))
+                           ->  Seq Scan on gtest_part_key1_1 t2_2
+               ->  Merge Join
+                     Merge Cond: (((t1_3.f2 * 2)) = ((t2_3.f2 * 2)))
+                     ->  Sort
+                           Sort Key: ((t1_3.f2 * 2))
+                           ->  Seq Scan on gtest_part_key1_2 t1_3
+                     ->  Sort
+                           Sort Key: ((t2_3.f2 * 2))
+                           ->  Seq Scan on gtest_part_key1_2 t2_3
+(29 rows)
+
+SET enable_partitionwise_join TO false;
+EXPLAIN (COSTS OFF)
+SELECT t1.f3, count(t2.f3) FROM gtest_part_key1 t1 JOIN gtest_part_key1 t2 ON t1.f3 = t2.f3 GROUP BY 1 ORDER BY 1;
+                            QUERY PLAN                            
+------------------------------------------------------------------
+ GroupAggregate
+   Group Key: ((t1.f2 * 2))
+   ->  Merge Join
+         Merge Cond: (((t1.f2 * 2)) = ((t2.f2 * 2)))
+         ->  Sort
+               Sort Key: ((t1.f2 * 2))
+               ->  Result
+                     ->  Append
+                           ->  Seq Scan on gtest_part_key1_0 t1_1
+                           ->  Seq Scan on gtest_part_key1_1 t1_2
+                           ->  Seq Scan on gtest_part_key1_2 t1_3
+         ->  Sort
+               Sort Key: ((t2.f2 * 2))
+               ->  Result
+                     ->  Append
+                           ->  Seq Scan on gtest_part_key1_0 t2_1
+                           ->  Seq Scan on gtest_part_key1_1 t2_2
+                           ->  Seq Scan on gtest_part_key1_2 t2_3
+(18 rows)
+
+RESET enable_partitionwise_join;
+RESET enable_partition_pruning;
+RESET enable_partitionwise_aggregate;
+RESET max_parallel_workers_per_gather;
+RESET enable_incremental_sort;
+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 = %; new = %', OLD, NEW;
+    RETURN NEW;
+  ELSE
+    RETURN NEW;
+  END IF;
+END
+$$;
+CREATE TRIGGER gtest_part_key1_trigger BEFORE INSERT OR UPDATE ON gtest_part_key1 FOR EACH ROW EXECUTE PROCEDURE gtest_trigger_info();
+SELECT * FROM gtest_part_key1 ORDER BY f2;
+ f1 | f2 | f3 
+----+----+----
+    | 10 | 20
+    | 12 | 24
+    | 13 | 26
+    | 25 | 50
+    | 30 | 60
+(5 rows)
+
+-- error
+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 = 14 THEN UPDATE SET f2 = 30
+  WHEN MATCHED AND f2 = 12 THEN UPDATE SET f2 = 50;
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gtest_part_key1_trigger tg_op: UPDATE
+INFO:  old = (12,,); new = (50,,)
+ERROR:  no partition of relation "gtest_part_key1" found for row
+DETAIL:  Partition key of the failing row contains (f3) = (100).
+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: gtest_part_key1_trigger tg_op: UPDATE
+INFO:  old = (10,,); new = (30,,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_2 trigger name: gtest_part_key1_trigger tg_op: INSERT
+INFO:  new = (,30,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gtest_part_key1_trigger tg_op: UPDATE
+INFO:  old = (12,,); new = (20,,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_1 trigger name: gtest_part_key1_trigger tg_op: INSERT
+INFO:  new = (,20,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gtest_part_key1_trigger 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);
diff --git a/src/test/regress/sql/generated_virtual.sql b/src/test/regress/sql/generated_virtual.sql
index e750866d2d8..f2fa974e7d6 100644
--- a/src/test/regress/sql/generated_virtual.sql
+++ b/src/test/regress/sql/generated_virtual.sql
@@ -557,13 +557,136 @@ ALTER TABLE gtest_parent ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2);
 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) VIRTUAL) PARTITION BY RANGE (f3);
-CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3));
-CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
+-- generated columns in partition key
+-- wholerow with virtual generated column is not supported
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_part_key));
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_part_key is not null));
 
+-- unique/primary key interact with partition key will fail
+CREATE TABLE gtest_idxpart (a int PRIMARY KEY, b int GENERATED ALWAYS AS (a) ) PARTITION BY RANGE ((b));
+CREATE TABLE gtest_idxpart (a int UNIQUE, b int GENERATED ALWAYS AS (a) ) PARTITION BY RANGE ((b));
+CREATE TABLE gtest_idxpart (a int UNIQUE, b int GENERATED ALWAYS AS (a) ) PARTITION BY RANGE (a, b);
+CREATE TABLE gtest_idxpart (a int UNIQUE, b int GENERATED ALWAYS AS (a) ) PARTITION BY RANGE (a); -- ok
+
+-- partition key expression over virtual generated column is not supported
+CREATE TABLE gtest_part_key (f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
+-- constant virtual generated column expression is not supported
+CREATE TABLE gtest_part_key (f3 bigint GENERATED ALWAYS AS (2)) PARTITION BY RANGE (f3); -- error
+
+CREATE TABLE gtest_part_keyxx (f2 text, f3 TEXT COLLATE "POSIX" GENERATED ALWAYS AS ( (upper(f2)) COLLATE "C")  VIRTUAL) PARTITION BY RANGE (f3, f3, f2, f3, (upper(f2)));
+
+SELECT pg_get_partkeydef('gtest_part_keyxx'::regclass);
+SELECT  partrelid::regclass, partnatts, partattrs,  partcollation[0]::regcollation
+FROM    pg_partitioned_table
+WHERE   partrelid = 'gtest_part_keyxx'::regclass;
+DROP TABLE gtest_part_keyxx;
+
+-- tests for virtual generated columns in partition key
+CREATE TABLE gtest_part_key  (f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3); -- ok
+CREATE TABLE gtest_part_key1 (f1 date, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3)); -- ok
+
+-- Once a virtual generated column is used as a partition key, its data type and
+-- generation expression can no longer be changed. All below should result
+-- error.
+ALTER TABLE gtest_part_key1 ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2);
+ALTER TABLE gtest_part_key1 ALTER COLUMN f3 SET DATA TYPE INT;
+ALTER TABLE gtest_part_key1 ALTER COLUMN f3 DROP EXPRESSION;
+
+CREATE TABLE gtest_part_key1_0(f2 bigint, f1 date, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED);
+-- error: The partition's generation kind must match the parent partitioned table
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30);
+DROP TABLE gtest_part_key1_0;
+
+CREATE TABLE gtest_part_key1_2(f2 bigint, f1 date, f3 bigint GENERATED ALWAYS AS (f2 * 3) VIRTUAL);
+-- error: The partition's generation expression must match the parent partitioned table
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_2 FOR VALUES FROM (20) TO (30);
+DROP TABLE gtest_part_key1_2;
+
+CREATE TABLE gtest_part_key1_0(f2 bigint, f1 date, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); -- ok
+
+ALTER TABLE gtest_part_key1_0 ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2); -- error
+
+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
+
+INSERT INTO gtest_part_key1(f2) VALUES (9);     -- error
+INSERT INTO gtest_part_key1_2(f2) VALUES (50);  -- error
+INSERT INTO gtest_part_key1(f2) VALUES (10), (12), (25), (30), (20) RETURNING tableoid::regclass, *; -- ok
+
+UPDATE gtest_part_key1 SET f2 = 50 WHERE f2 = 30; -- error
+UPDATE gtest_part_key1 SET f2 = 13 WHERE f2 = 20 RETURNING new.tableoid::regclass, old.tableoid::regclass, OLD.*, NEW.*;
+
+SET max_parallel_workers_per_gather TO 0;
+SET enable_incremental_sort TO off;
+
+SET enable_partition_pruning TO true;
+EXPLAIN(COSTS OFF) SELECT * FROM gtest_part_key1 WHERE f3 < 50;
+
+SET enable_partition_pruning TO false;
+EXPLAIN(COSTS OFF) SELECT * FROM gtest_part_key1 WHERE f3 < 50;
+
+SET enable_partitionwise_aggregate TO true;
+EXPLAIN (COSTS OFF) SELECT f3, count(*) FROM gtest_part_key1 GROUP BY f3 ORDER BY 1;
+
+SET enable_partitionwise_aggregate TO false;
+EXPLAIN (COSTS OFF) SELECT f3, count(*) FROM gtest_part_key1 GROUP BY f3 ORDER BY 1;
+
+SET enable_partitionwise_join TO true;
+EXPLAIN (COSTS OFF)
+SELECT t1.f3, count(t2.f3) FROM gtest_part_key1 t1 JOIN gtest_part_key1 t2 ON t1.f3 = t2.f3 GROUP BY 1 ORDER BY 1;
+
+SET enable_partitionwise_join TO false;
+EXPLAIN (COSTS OFF)
+SELECT t1.f3, count(t2.f3) FROM gtest_part_key1 t1 JOIN gtest_part_key1 t2 ON t1.f3 = t2.f3 GROUP BY 1 ORDER BY 1;
+
+RESET enable_partitionwise_join;
+RESET enable_partition_pruning;
+RESET enable_partitionwise_aggregate;
+RESET max_parallel_workers_per_gather;
+RESET enable_incremental_sort;
+
+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 = %; new = %', OLD, NEW;
+    RETURN NEW;
+  ELSE
+    RETURN NEW;
+  END IF;
+END
+$$;
+
+CREATE TRIGGER gtest_part_key1_trigger BEFORE INSERT OR UPDATE ON gtest_part_key1 FOR EACH ROW EXECUTE PROCEDURE gtest_trigger_info();
+
+SELECT * FROM gtest_part_key1 ORDER BY f2;
+
+-- error
+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 = 14 THEN UPDATE SET f2 = 30
+  WHEN MATCHED AND f2 = 12 THEN UPDATE SET f2 = 50;
+
+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;
+
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);
 INSERT INTO gtest25 VALUES (3), (4);
-- 
2.34.1

Reply via email to