virtual generated column as partition key

From 266121c64eb53a3248d33bc0c2a8070662e80ba4 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Mon, 17 Nov 2025 15:58:25 +0800
Subject: [PATCH v2 1/1] virtual generated column as partition key

demo:
CREATE TABLE t(f1 bigint, f2 bigint GENERATED ALWAYS AS (f1 * 2) VIRTUAL) PARTITION BY RANGE (f2);
but partition key can not 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));

The virtual generated column 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.
cross partition update tests added.

A virtual generated column entry in the pg_partitioned_table catalog is marked
as non-zero partattrs and a non-null partexprs, which is abnormal. Normally,
either partattrs is non-zero or partexprs is null.
we should mention this in the doc/src/sgml/catalogs.sgml

discussion: https://postgr.es/m/CACJufxF9RV2_iHBAG0AfCyt5c-1zJkh_Jc7F1tZfk3m3N+E=qA@mail.gmail.com
---
 src/backend/catalog/partition.c               |   3 +
 src/backend/commands/indexcmds.c              |   3 +-
 src/backend/commands/tablecmds.c              | 179 +++++++++-
 src/backend/executor/execPartition.c          |   2 +-
 src/backend/optimizer/util/plancat.c          |   3 +-
 src/backend/partitioning/partbounds.c         |  30 +-
 src/backend/utils/adt/ruleutils.c             |   8 +
 src/backend/utils/cache/partcache.c           |   2 +-
 src/backend/utils/cache/relcache.c            |  11 +
 src/include/catalog/pg_partitioned_table.h    |   3 +
 src/include/utils/relcache.h                  |   1 +
 .../regress/expected/generated_stored.out     |  20 +-
 .../regress/expected/generated_virtual.out    | 310 ++++++++++++++++--
 src/test/regress/sql/generated_virtual.sql    | 129 +++++++-
 14 files changed, 640 insertions(+), 64 deletions(-)

diff --git a/src/backend/catalog/partition.c b/src/backend/catalog/partition.c
index 93d72157a46..710cfb6ec46 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 (ColumnIsVirtualGenerated(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 5712fac3697..4c84db34a41 100644
--- a/src/backend/commands/indexcmds.c
+++ b/src/backend/commands/indexcmds.c
@@ -1012,7 +1012,8 @@ DefineIndex(Oid tableId,
 			 * It may be possible to support UNIQUE constraints when partition
 			 * keys are expressions, but is it worth it?  Give up for now.
 			 */
-			if (key->partattrs[i] == 0)
+			if (key->partattrs[i] == 0 ||
+				ColumnIsVirtualGenerated(rel, key->partattrs[i]))
 				ereport(ERROR,
 						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						 errmsg("unsupported %s constraint with partition key definition",
diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 23ebaa3f230..fb07c80134c 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -8628,6 +8628,47 @@ ATExecSetExpression(AlteredTableInfo *tab, Relation rel, const char *colName,
 				 errmsg("column \"%s\" of relation \"%s\" is not a generated column",
 						colName, RelationGetRelationName(rel))));
 
+	if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
+	{
+		if (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;
+		Relation	parent;
+		AttrMap		*map	= NULL;
+
+		parentId = get_partition_parent(RelationGetRelid(rel), false);
+
+		parent = table_open(parentId, AccessShareLock);
+		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);
+	}
+
 	/*
 	 * TODO: This could be done, just need to recheck any constraints
 	 * afterwards.
@@ -19786,6 +19827,7 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
 		PartitionElem *pelem = lfirst_node(PartitionElem, lc);
 		Oid			atttype;
 		Oid			attcollation;
+		AttrNumber	virtual_attr = InvalidAttrNumber;
 
 		if (pelem->name != NULL)
 		{
@@ -19813,24 +19855,71 @@ 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. 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)
+				virtual_attr = attform->attnum;
+
 			partattrs[attn] = attform->attnum;
 			atttype = attform->atttypid;
 			attcollation = attform->attcollation;
 			ReleaseSysCache(atttuple);
 		}
 		else
+		{
+			Node	*expr = pelem->expr;
+
+			atttype = exprType(expr);
+			attcollation = exprCollation(expr);
+
+			while (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)
+					virtual_attr = var->varattno;
+			}
+		}
+
+		if (AttributeNumberIsValid(virtual_attr))
+		{
+			Node	   *expr;
+
+			expr = build_generation_expression(rel, virtual_attr);
+
+			expr = (Node *) expression_planner((Expr *) expr);
+
+			/*
+			 * generated expression expect to be IMMUTABLE, 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] = virtual_attr;
+			*partexprs = lappend(*partexprs, expr);
+		}
+		else if (pelem->expr != NULL)
 		{
 			/* Expression */
 			Node	   *expr = pelem->expr;
@@ -19894,18 +19983,26 @@ 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 partition key
+				 * cannot use expression that reference them.
 				 */
-				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 virtual generated column in partition key expression"),
+							errdetail("Column \"%s\" is a virtual generated column.",
+									  get_attname(RelationGetRelid(rel), attno, false)),
+							parser_errposition(pstate, pelem->location));
 			}
 
 			if (IsA(expr, Var) &&
@@ -20411,6 +20508,62 @@ 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, the generated
+	 * expression in partition must match that of the partitioned table.
+	 */
+	if (tupleDesc->constr && tupleDesc->constr->has_generated_virtual)
+	{
+		Node	*rel_defval = NULL;
+		Node	*attachrel_defval = NULL;
+		bool	found_whole_row = false;
+		AttrMap	*map	= NULL;
+
+		TupleDesc rootdesc = RelationGetDescr(rel);
+		PartitionKey key = RelationGetPartitionKey(rel);
+
+		map = build_attrmap_by_name_if_req(rootdesc, tupleDesc, false);
+
+		/* Add an argument for each key column. */
+		for (int i = 0; i < key->partnatts; i++)
+		{
+			if (AttributeNumberIsValid(key->partattrs[i]))
+			{
+				Form_pg_attribute attr = TupleDescAttr(rootdesc, key->partattrs[i] - 1);
+
+				if (attr->attgenerated  == ATTRIBUTE_GENERATED_VIRTUAL)
+				{
+					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 aa12e9ad2ea..204d4f88905 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -1325,7 +1325,7 @@ FormPartitionKeyDatum(PartitionDispatch pd,
 		Datum		datum;
 		bool		isNull;
 
-		if (keycol != 0)
+		if (keycol != 0 && !ColumnIsVirtualGenerated(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 d950bd93002..5b940c3e474 100644
--- a/src/backend/optimizer/util/plancat.c
+++ b/src/backend/optimizer/util/plancat.c
@@ -2699,7 +2699,8 @@ set_baserel_partition_key_exprs(Relation relation,
 		Expr	   *partexpr;
 		AttrNumber	attno = partkey->partattrs[cnt];
 
-		if (attno != InvalidAttrNumber)
+		if (attno != InvalidAttrNumber &&
+			!ColumnIsVirtualGenerated(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 8ba038c5ef4..9fd95f9c11a 100644
--- a/src/backend/partitioning/partbounds.c
+++ b/src/backend/partitioning/partbounds.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 &&
+			!ColumnIsVirtualGenerated(parent, key->partattrs[i]))
 		{
 			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 &&
+		!ColumnIsVirtualGenerated(parent, key->partattrs[0]))
 		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,7 +4632,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,
@@ -4638,7 +4640,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 &&
+		!ColumnIsVirtualGenerated(rel, key->partattrs[keynum]))
 	{
 		*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;
@@ -4686,7 +4689,8 @@ get_range_nulltest(PartitionKey key)
 	{
 		Expr	   *keyCol;
 
-		if (key->partattrs[i] != 0)
+		if (key->partattrs[i] != 0 &&
+			!ColumnIsVirtualGenerated(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 556ab057e5a..fe9bdce35b2 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -2042,6 +2042,14 @@ pg_get_partkeydef_worker(Oid relid, int prettyFlags,
 			get_atttypetypmodcoll(relid, attnum,
 								  &keycoltype, &keycoltypmod,
 								  &keycolcollation);
+
+			/*
+			 * When the partition key is a virtual generated column, partexprs
+			 * contains its generation expression. In that case, we need to
+			 * advance partexprs.
+			 */
+			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 f5d7d70def0..ba5b627218e 100644
--- a/src/backend/utils/cache/partcache.c
+++ b/src/backend/utils/cache/partcache.c
@@ -226,7 +226,7 @@ RelationBuildPartitionKey(Relation relation)
 		key->partcollation[i] = collation->values[i];
 
 		/* Collect type information */
-		if (attno != 0)
+		if (attno != 0 && !ColumnIsVirtualGenerated(relation, attno))
 		{
 			Form_pg_attribute att = TupleDescAttr(relation->rd_att, attno - 1);
 
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 915d0bc9084..6acb1321537 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -6033,6 +6033,17 @@ RelationGetIndexAttOptions(Relation relation, bool copy)
 	return relation->rd_opcoptions;
 }
 
+bool
+ColumnIsVirtualGenerated(Relation rel, AttrNumber attnum)
+{
+	TupleDesc tupdesc = RelationGetDescr(rel);
+
+	Form_pg_attribute 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/catalog/pg_partitioned_table.h b/src/include/catalog/pg_partitioned_table.h
index 0527f347690..45d052a504a 100644
--- a/src/include/catalog/pg_partitioned_table.h
+++ b/src/include/catalog/pg_partitioned_table.h
@@ -54,6 +54,9 @@ CATALOG(pg_partitioned_table,3350,PartitionedRelationId)
 																				 * collation for keys */
 	pg_node_tree partexprs;		/* list of expressions in the partition key;
 								 * one item for each zero entry in partattrs[] */
+	pg_node_tree partvirtualexprs;		/* list of expressions in the partition key;
+								 * one item for each zero entry in partattrs[] */
+
 #endif
 } FormData_pg_partitioned_table;
 
diff --git a/src/include/utils/relcache.h b/src/include/utils/relcache.h
index 3561c6bef0b..529878b9fc0 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 ColumnIsVirtualGenerated(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 dde325e46c6..05a5add5faf 100644
--- a/src/test/regress/expected/generated_virtual.out
+++ b/src/test/regress/expected/generated_virtual.out
@@ -1036,32 +1036,302 @@ 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 virtual generated column in partition key expression
 LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_par...
                                                              ^
-DETAIL:  Column "f3" is a generated column.
+DETAIL:  Column "f3" is a virtual generated column.
 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 virtual generated column in partition key expression
 LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((gtest_par...
                                                              ^
-DETAIL:  Column "f3" is a generated column.
+DETAIL:  Column "f3" is a virtual generated column.
+--expression over virtual generated column is not supported now
+CREATE TABLE gtest_part_key (f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
+ERROR:  cannot use virtual generated column in partition key expression
+LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
+                                                             ^
+DETAIL:  Column "f3" is a virtual generated column.
+--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 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3, f3, f2, f3, (f2+1));
+SELECT pg_get_partkeydef('gtest_part_keyxx'::regclass);
+         pg_get_partkeydef          
+------------------------------------
+ RANGE (f3, f3, f2, f3, ((f2 + 1)))
+(1 row)
+
+SELECT partrelid::regclass, partnatts, partattrs FROM pg_partitioned_table WHERE partrelid = ('gtest_part_keyxx'::regclass);
+    partrelid     | partnatts | partattrs 
+------------------+-----------+-----------
+ gtest_part_keyxx |         5 | 2 2 1 2 0
+(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);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); --error
+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);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_2 FOR VALUES FROM (20) TO (30); --error
+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) 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
+(4 rows)
+
+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 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();
+SELECT * FROM gtest_part_key1;
+ f1 | f2 | f3 
+----+----+----
+    | 10 | 20
+    | 12 | 24
+    | 25 | 50
+    | 30 | 60
+(4 rows)
+
+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;
+---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: gkey1_0 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: gkey1_0 tg_op: UPDATE
+INFO:  old = (10,,); 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,,); 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)
+
+SELECT * FROM gtest_part_key1;
+ f1 | f2 | f3 
+----+----+----
+    | 14 | 28
+    | 20 | 40
+    | 25 | 50
+    | 30 | 60
+    | 30 | 60
+(5 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 2911439776c..c23c59b91ab 100644
--- a/src/test/regress/sql/generated_virtual.sql
+++ b/src/test/regress/sql/generated_virtual.sql
@@ -545,13 +545,134 @@ 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));
 
+--expression over virtual generated column is not supported now
+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 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3, f3, f2, f3, (f2+1));
+SELECT pg_get_partkeydef('gtest_part_keyxx'::regclass);
+SELECT partrelid::regclass, partnatts, partattrs 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);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); --error
+DROP TABLE gtest_part_key1_0;
+
+CREATE TABLE gtest_part_key1_2(f2 bigint, f1 date, f3 bigint GENERATED ALWAYS AS (f2 * 3) VIRTUAL);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_2 FOR VALUES FROM (20) TO (30); --error
+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) RETURNING tableoid::regclass, *; --ok
+
+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 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();
+
+SELECT * FROM gtest_part_key1;
+
+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;
+
+---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;
+
+SELECT * FROM gtest_part_key1;
+
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);
 INSERT INTO gtest25 VALUES (3), (4);
-- 
2.34.1

From cb14769865b29dc5824daa82d0a2b07e0da15d63 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Mon, 5 Jan 2026 11:23:48 +0800
Subject: [PATCH v3 1/1] virtual generated column as partition key

If the partition key contains a virtual generated column, 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.

cross partition update regress tests is added.

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);
but partition key can not 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=qA@mail.gmail.com
---
 src/backend/catalog/partition.c               |   3 +
 src/backend/commands/indexcmds.c              |  13 +
 src/backend/commands/tablecmds.c              | 181 +++++++++-
 src/backend/executor/execPartition.c          |   2 +-
 src/backend/optimizer/util/plancat.c          |   3 +-
 src/backend/partitioning/partbounds.c         |  30 +-
 src/backend/utils/adt/ruleutils.c             |   8 +
 src/backend/utils/cache/partcache.c           |   3 +
 src/backend/utils/cache/relcache.c            |  14 +
 src/include/utils/partcache.h                 |  10 +-
 src/include/utils/relcache.h                  |   1 +
 .../regress/expected/generated_stored.out     |  20 +-
 .../regress/expected/generated_virtual.out    | 336 ++++++++++++++++--
 src/test/regress/sql/generated_virtual.sql    | 144 +++++++-
 14 files changed, 704 insertions(+), 64 deletions(-)

diff --git a/src/backend/catalog/partition.c b/src/backend/catalog/partition.c
index 93d72157a46..57478e2943e 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 d9cccb6ac18..8b236b2b7b0 100644
--- a/src/backend/commands/indexcmds.c
+++ b/src/backend/commands/indexcmds.c
@@ -1020,6 +1020,19 @@ DefineIndex(Oid tableId,
 						 errdetail("%s constraints cannot be used when partition keys include expressions.",
 								   constraint_type)));
 
+			/*
+			 * It may be possible to support UNIQUE constraints when partition
+			 * keys are virtual generated column, but is it worth it?  Give up
+			 * for now.
+			 */
+			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 1c9ef53be20..c08f7618f8a 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -8628,6 +8628,54 @@ ATExecSetExpression(AlteredTableInfo *tab, Relation rel, const char *colName,
 				 errmsg("column \"%s\" of relation \"%s\" is not a generated column",
 						colName, RelationGetRelationName(rel))));
 
+	/*
+	 * If a virtual generated column is used in the partition key, then we can
+	 * not change the generation expression of it.  We already reject stored
+	 * generated column as partition key 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;
+			Relation	parent;
+			AttrMap    *map = NULL;
+
+			Oid			parentId = get_partition_parent(RelationGetRelid(rel),
+														false);
+
+			parent = table_open(parentId, AccessShareLock);
+
+			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);
+		}
+	}
+
 	/*
 	 * TODO: This could be done, just need to recheck any constraints
 	 * afterwards.
@@ -19786,6 +19834,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)
 		{
@@ -19813,24 +19862,71 @@ 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
+		{
+			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;
+
+			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;
@@ -19894,18 +19990,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) &&
@@ -20411,6 +20514,58 @@ 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 must match that of the partitioned
+	 * table.
+	 */
+	if (tupleDesc->constr && tupleDesc->constr->has_generated_virtual)
+	{
+		Node	   *rel_defval = NULL;
+		Node	   *attachrel_defval = NULL;
+		bool		found_whole_row = false;
+		AttrMap    *map = NULL;
+
+		PartitionKey key = RelationGetPartitionKey(rel);
+
+		map = build_attrmap_by_name_if_req(RelationGetDescr(rel), tupleDesc, false);
+
+		/* Add an argument for each key column. */
+		for (int i = 0; i < key->partnatts; i++)
+		{
+			if (AttributeNumberIsValid(key->partattrs[i]) &&
+				attrIsVirtualGenerated(rel, key->partattrs[i]))
+			{
+				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 e30db12113b..bc6f5e9cb68 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -1454,7 +1454,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 bf45c355b77..4690c6c94bb 100644
--- a/src/backend/optimizer/util/plancat.c
+++ b/src/backend/optimizer/util/plancat.c
@@ -2844,7 +2844,8 @@ 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 40ac700d529..b5905b68671 100644
--- a/src/backend/partitioning/partbounds.c
+++ b/src/backend/partitioning/partbounds.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
@@ -4016,7 +4016,8 @@ 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],
@@ -4072,7 +4073,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 &&
+		!attrIsVirtualGenerated(parent, key->partattrs[0]))
 		keyCol = (Expr *) makeVar(1,
 								  key->partattrs[0],
 								  key->parttypid[0],
@@ -4343,7 +4345,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)
@@ -4366,7 +4368,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
@@ -4398,7 +4400,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);
@@ -4478,7 +4480,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);
@@ -4598,7 +4600,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;
@@ -4620,7 +4622,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,
@@ -4628,7 +4630,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],
@@ -4664,7 +4667,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;
@@ -4676,7 +4679,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 9f85eb86da1..dc49970f594 100644
--- a/src/backend/utils/adt/ruleutils.c
+++ b/src/backend/utils/adt/ruleutils.c
@@ -2042,6 +2042,14 @@ pg_get_partkeydef_worker(Oid relid, int prettyFlags,
 			get_atttypetypmodcoll(relid, attnum,
 								  &keycoltype, &keycoltypmod,
 								  &keycolcollation);
+
+			/*
+			 * When the partition key is a virtual generated column, partexprs
+			 * contains its generation expression. In that case, we need to
+			 * advance partexprs.
+			 */
+			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 67e88440038..2dcaa2aba35 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 2d0cb7bcfd4..7b929b15c4e 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -6033,6 +6033,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 0fb6fc1c808..75b03801915 100644
--- a/src/include/utils/partcache.h
+++ b/src/include/utils/partcache.h
@@ -28,8 +28,14 @@ typedef struct PartitionKeyData
 	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 */
+
+	/*
+	 * 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 3561c6bef0b..1a8e8964894 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 dde325e46c6..2f83e43bfe5 100644
--- a/src/test/regress/expected/generated_virtual.out
+++ b/src/test/regress/expected/generated_virtual.out
@@ -1036,32 +1036,328 @@ 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
+---error when unique/primary key not included in the partition key
+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.
+--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 bigint,
+  f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL)
+PARTITION BY RANGE (f3, f3, f2, f3, (f2+1));
+SELECT pg_get_partkeydef('gtest_part_keyxx'::regclass);
+         pg_get_partkeydef          
+------------------------------------
+ RANGE (f3, f3, f2, f3, ((f2 + 1)))
+(1 row)
+
+SELECT  partrelid::regclass, partnatts, partattrs
+FROM    pg_partitioned_table
+WHERE   partrelid = ('gtest_part_keyxx'::regclass);
+    partrelid     | partnatts | partattrs 
+------------------+-----------+-----------
+ gtest_part_keyxx |         5 | 2 2 1 2 0
+(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);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); --error
+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);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_2 FOR VALUES FROM (20) TO (30); --error
+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)
+
+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 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();
+SELECT * FROM gtest_part_key1;
+ f1 | f2 | f3 
+----+----+----
+    | 10 | 20
+    | 12 | 24
+    | 13 | 26
+    | 25 | 50
+    | 30 | 60
+(5 rows)
+
+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;
+---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: gkey1_0 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: gkey1_0 tg_op: UPDATE
+INFO:  old = (10,,); 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,,); 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)
+
+SELECT * FROM gtest_part_key1;
+ f1 | f2 | f3 
+----+----+----
+    | 13 | 26
+    | 14 | 28
+    | 20 | 40
+    | 25 | 50
+    | 30 | 60
+    | 30 | 60
+(6 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 2911439776c..999e471bdbc 100644
--- a/src/test/regress/sql/generated_virtual.sql
+++ b/src/test/regress/sql/generated_virtual.sql
@@ -545,13 +545,149 @@ 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));
 
+---error when unique/primary key not included in the partition key
+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));
+
+--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 bigint,
+  f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL)
+PARTITION BY RANGE (f3, f3, f2, f3, (f2+1));
+
+SELECT pg_get_partkeydef('gtest_part_keyxx'::regclass);
+SELECT  partrelid::regclass, partnatts, partattrs
+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);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); --error
+DROP TABLE gtest_part_key1_0;
+
+CREATE TABLE gtest_part_key1_2(f2 bigint, f1 date, f3 bigint GENERATED ALWAYS AS (f2 * 3) VIRTUAL);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_2 FOR VALUES FROM (20) TO (30); --error
+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.*;
+
+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 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();
+
+SELECT * FROM gtest_part_key1;
+
+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;
+
+---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;
+
+SELECT * FROM gtest_part_key1;
+
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);
 INSERT INTO gtest25 VALUES (3), (4);
-- 
2.34.1