From 25986c0aefbbabe022282b7941023fe4fbb9e1dc Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Wed, 23 Apr 2025 19:54:03 +0800
Subject: [PATCH v1 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/
---
 src/backend/commands/tablecmds.c              | 371 ++++++++++++------
 src/backend/executor/execPartition.c          |   3 +-
 src/backend/partitioning/partbounds.c         |  32 +-
 src/backend/utils/cache/partcache.c           |   5 +-
 src/backend/utils/cache/relcache.c            |  13 +
 src/include/utils/relcache.h                  |   1 +
 .../regress/expected/generated_stored.out     |  11 +-
 .../regress/expected/generated_virtual.out    | 113 +++++-
 src/test/regress/sql/generated_stored.sql     |   1 +
 src/test/regress/sql/generated_virtual.sql    |  64 ++-
 10 files changed, 471 insertions(+), 143 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 265b1c397fb..9d05d83b5e3 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -707,6 +707,13 @@ static void RangeVarCallbackForDropRelation(const RangeVar *rel, Oid relOid,
 static void RangeVarCallbackForAlterRelation(const RangeVar *rv, Oid relid,
 											 Oid oldrelid, void *arg);
 static PartitionSpec *transformPartitionSpec(Relation rel, PartitionSpec *partspec);
+static void ComputePartitionExprs(ParseState *pstate,
+								  Relation rel,
+								  int	attn,
+								  AttrNumber *partattrs,
+								  List **partexprs,
+								  PartitionElem *pelem,
+								  Node	   *expr);
 static void ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNumber *partattrs,
 								  List **partexprs, Oid *partopclass, Oid *partcollation,
 								  PartitionStrategy strategy);
@@ -8601,6 +8608,45 @@ 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.
@@ -19700,6 +19746,128 @@ transformPartitionSpec(Relation rel, PartitionSpec *partspec)
 	return newspec;
 }
 
+/*
+ * subroutine of ComputePartitionAttrs.
+*/
+static void ComputePartitionExprs(ParseState *pstate, Relation rel,
+								  int	attn,
+								  AttrNumber *partattrs,
+								  List **partexprs,
+								  PartitionElem *pelem,
+								  Node	   *expr)
+{
+	Oid			atttype;
+	Oid			attcollation;
+	char		partattname[16];
+	Bitmapset  *expr_attrs = NULL;
+	int			i;
+
+	Assert(expr != NULL);
+	atttype = exprType(expr);
+	attcollation = exprCollation(expr);
+
+	/*
+	 * The expression must be of a storable type (e.g., not RECORD).
+	 * The test is the same as for whether a table column is of a safe
+	 * type (which is why we needn't check for the non-expression
+	 * case).
+	*/
+	snprintf(partattname, sizeof(partattname), "%d", attn + 1);
+	CheckAttributeType(partattname,
+					   atttype, attcollation,
+					   NIL, CHKATYPE_IS_PARTKEY);
+
+	partattrs[attn] = 0;	/* marks the column as expression */
+	*partexprs = lappend(*partexprs, expr);
+
+	/*
+	 * transformPartitionSpec() should have already rejected
+	 * subqueries, aggregates, window functions, and SRFs, based
+	 * on the EXPR_KIND_ for partition expressions.
+	*/
+
+	/*
+	 * Cannot allow system column references, since that would
+	 * make partition routing impossible: their values won't be
+	 * known yet when we need to do that.
+	*/
+	pull_varattnos(expr, 1, &expr_attrs);
+	for (i = FirstLowInvalidHeapAttributeNumber; i < 0; i++)
+	{
+		if (bms_is_member(i - FirstLowInvalidHeapAttributeNumber,
+							expr_attrs))
+			ereport(ERROR,
+					errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
+					errmsg("partition key expressions cannot contain system column references"));
+	}
+
+	/*
+	 * 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.
+	*/
+	i = -1;
+	while ((i = bms_next_member(expr_attrs, i)) >= 0)
+	{
+		AttrNumber	attno = i + FirstLowInvalidHeapAttributeNumber;
+
+		if (attno > 0 &&
+			TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated == ATTRIBUTE_GENERATED_STORED)
+			ereport(ERROR,
+					errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
+					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 (attno > 0 &&
+			TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+			ereport(ERROR,
+					errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
+					errmsg("partition key expression cannot use virtual generated column"),
+					errhint("Only plain virtual generated column reference can be used in partition key");
+					parser_errposition(pstate, pelem->location));
+	}
+	/*
+	 * Preprocess the expression before checking for mutability.
+	 * This is essential for the reasons described in
+	 * contain_mutable_functions_after_planning.  However, we call
+	 * expression_planner for ourselves rather than using that
+	 * function, because if constant-folding reduces the
+	 * expression to a constant, we'd like to know that so we can
+	 * complain below.
+	 *
+	 * Like contain_mutable_functions_after_planning, assume that
+	 * expression_planner won't scribble on its input, so this
+	 * won't affect the partexprs entry we saved above.
+	*/
+	expr = (Node *) expression_planner((Expr *) expr);
+
+	/*
+	 * Partition expressions cannot contain mutable functions,
+	 * because a given row must always map to the same partition
+	 * as long as there is no change in the partition boundary
+	 * structure.
+	*/
+	if (contain_mutable_functions(expr))
+		ereport(ERROR,
+				errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
+				errmsg("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"));
+}
+
+
 /*
  * Compute per-partition-column information from a list of PartitionElems.
  * Expressions in the PartitionElems must be parse-analyzed already.
@@ -19746,147 +19914,76 @@ 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 handled below.
 			 */
-			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.",
+						errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
+						errmsg("cannot use stored generated column in partition key"),
+						errdetail("Column \"%s\" is a stored generated column.",
 								   pelem->name),
-						 parser_errposition(pstate, pelem->location)));
+						parser_errposition(pstate, pelem->location));
 
-			partattrs[attn] = attform->attnum;
 			atttype = attform->atttypid;
 			attcollation = attform->attcollation;
+
+			if (attform->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+			{
+				Node	*expr;
+
+				expr = build_generation_expression(rel, attform->attnum);
+
+				ComputePartitionExprs(pstate, rel, attn, partattrs, partexprs, pelem, expr);
+			}
+
+			partattrs[attn] = attform->attnum;
 			ReleaseSysCache(atttuple);
 		}
 		else
 		{
-			/* Expression */
-			Node	   *expr = pelem->expr;
-			char		partattname[16];
+			Node	*expr;
 
-			Assert(expr != NULL);
+			expr = pelem->expr;
 			atttype = exprType(expr);
 			attcollation = exprCollation(expr);
 
-			/*
-			 * The expression must be of a storable type (e.g., not RECORD).
-			 * The test is the same as for whether a table column is of a safe
-			 * type (which is why we needn't check for the non-expression
-			 * case).
-			 */
-			snprintf(partattname, sizeof(partattname), "%d", attn + 1);
-			CheckAttributeType(partattname,
-							   atttype, attcollation,
-							   NIL, CHKATYPE_IS_PARTKEY);
-
 			/*
 			 * Strip any top-level COLLATE clause.  This ensures that we treat
 			 * "x COLLATE y" and "(x COLLATE y)" alike.
-			 */
+			*/
 			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_STORED)
+					ereport(ERROR,
+							errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
+							errmsg("cannot use stored generated column in partition key"),
+							errdetail("Column \"%s\" is a generated column.",
+									get_attname(RelationGetRelid(rel), var->varattno, false)),
+							parser_errposition(pstate, pelem->location));
+
+				if (TupleDescAttr(RelationGetDescr(rel), var->varattno - 1)->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+				{
+					Node	*defexpr;
+
+					defexpr = build_generation_expression(rel, var->varattno);
+
+					ComputePartitionExprs(pstate, rel, attn, partattrs, partexprs, pelem, defexpr);
+				}
+
 				/*
 				 * User wrote "(column)" or "(column COLLATE something)".
 				 * Treat it like simple attribute anyway.
-				 */
-				partattrs[attn] = ((Var *) expr)->varattno;
+				*/
+				partattrs[attn] = var->varattno;
 			}
 			else
-			{
-				Bitmapset  *expr_attrs = NULL;
-				int			i;
-
-				partattrs[attn] = 0;	/* marks the column as expression */
-				*partexprs = lappend(*partexprs, expr);
-
-				/*
-				 * transformPartitionSpec() should have already rejected
-				 * subqueries, aggregates, window functions, and SRFs, based
-				 * on the EXPR_KIND_ for partition expressions.
-				 */
-
-				/*
-				 * Cannot allow system column references, since that would
-				 * make partition routing impossible: their values won't be
-				 * known yet when we need to do that.
-				 */
-				pull_varattnos(expr, 1, &expr_attrs);
-				for (i = FirstLowInvalidHeapAttributeNumber; i < 0; i++)
-				{
-					if (bms_is_member(i - FirstLowInvalidHeapAttributeNumber,
-									  expr_attrs))
-						ereport(ERROR,
-								(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
-								 errmsg("partition key expressions cannot contain system column references")));
-				}
-
-				/*
-				 * 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.
-				 */
-				i = -1;
-				while ((i = bms_next_member(expr_attrs, i)) >= 0)
-				{
-					AttrNumber	attno = i + FirstLowInvalidHeapAttributeNumber;
-
-					if (attno > 0 &&
-						TupleDescAttr(RelationGetDescr(rel), attno - 1)->attgenerated)
-						ereport(ERROR,
-								(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
-								 errmsg("cannot use generated column in partition key"),
-								 errdetail("Column \"%s\" is a generated column.",
-										   get_attname(RelationGetRelid(rel), attno, false)),
-								 parser_errposition(pstate, pelem->location)));
-				}
-
-				/*
-				 * Preprocess the expression before checking for mutability.
-				 * This is essential for the reasons described in
-				 * contain_mutable_functions_after_planning.  However, we call
-				 * expression_planner for ourselves rather than using that
-				 * function, because if constant-folding reduces the
-				 * expression to a constant, we'd like to know that so we can
-				 * complain below.
-				 *
-				 * Like contain_mutable_functions_after_planning, assume that
-				 * expression_planner won't scribble on its input, so this
-				 * won't affect the partexprs entry we saved above.
-				 */
-				expr = (Node *) expression_planner((Expr *) expr);
-
-				/*
-				 * Partition expressions cannot contain mutable functions,
-				 * because a given row must always map to the same partition
-				 * as long as there is no change in the partition boundary
-				 * structure.
-				 */
-				if (contain_mutable_functions(expr))
-					ereport(ERROR,
-							(errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
-							 errmsg("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")));
-			}
+				ComputePartitionExprs(pstate, rel, attn, partattrs, partexprs, pelem, pelem->expr);
 		}
 
 		/*
@@ -20336,6 +20433,54 @@ ATExecAttachPartition(List **wqueue, Relation rel, PartitionCmd *cmd,
 					 errdetail("The new partition may contain only the columns present in parent.")));
 	}
 
+
+	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 (key->partattrs[i] != 0)
+			{
+				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");
+
+					/* XXX the error message need polish */
+					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 3f8a4cb5244..921df429288 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -1321,7 +1321,8 @@ FormPartitionKeyDatum(PartitionDispatch pd,
 		Datum		datum;
 		bool		isNull;
 
-		if (keycol != 0)
+		if (keycol != 0 &&
+			RelationGetColumnGenerated(pd->reldesc, keycol) != ATTRIBUTE_GENERATED_VIRTUAL)
 		{
 			/* Plain column; get the value directly from the heap tuple */
 			datum = slot_getattr(slot, keycol, &isNull);
diff --git a/src/backend/partitioning/partbounds.c b/src/backend/partitioning/partbounds.c
index 4bdc2941efb..405bfda01ec 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
@@ -4025,7 +4025,8 @@ get_qual_for_hash(Relation parent, PartitionBoundSpec *spec)
 		Node	   *keyCol;
 
 		/* Left operand */
-		if (key->partattrs[i] != 0)
+		if (key->partattrs[i] != 0 &&
+			RelationGetColumnGenerated(parent, key->partattrs[i]) != ATTRIBUTE_GENERATED_VIRTUAL)
 		{
 			keyCol = (Node *) makeVar(1,
 									  key->partattrs[i],
@@ -4081,7 +4082,8 @@ get_qual_for_list(Relation parent, PartitionBoundSpec *spec)
 	Assert(key->partnatts == 1);
 
 	/* Construct Var or expression representing the partition column */
-	if (key->partattrs[0] != 0)
+	if (key->partattrs[0] != 0 &&
+		RelationGetColumnGenerated(parent, key->partattrs[0]) != ATTRIBUTE_GENERATED_VIRTUAL)
 		keyCol = (Expr *) makeVar(1,
 								  key->partattrs[0],
 								  key->parttypid[0],
@@ -4352,7 +4354,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)
@@ -4375,7 +4377,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
@@ -4407,7 +4409,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);
@@ -4487,7 +4489,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);
@@ -4607,7 +4609,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;
@@ -4629,7 +4631,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,
@@ -4637,7 +4639,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 &&
+		RelationGetColumnGenerated(rel, key->partattrs[keynum]) != ATTRIBUTE_GENERATED_VIRTUAL)
 	{
 		*keyCol = (Expr *) makeVar(1,
 								   key->partattrs[keynum],
@@ -4673,7 +4676,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;
@@ -4684,8 +4687,11 @@ get_range_nulltest(PartitionKey key)
 	for (i = 0; i < key->partnatts; i++)
 	{
 		Expr	   *keyCol;
+		bool	key_virtual;
 
-		if (key->partattrs[i] != 0)
+		key_virtual = (RelationGetColumnGenerated(rel, key->partattrs[i]) == ATTRIBUTE_GENERATED_VIRTUAL);
+
+		if (key->partattrs[i] != 0 && !key_virtual)
 		{
 			keyCol = (Expr *) makeVar(1,
 									  key->partattrs[i],
diff --git a/src/backend/utils/cache/partcache.c b/src/backend/utils/cache/partcache.c
index f5d7d70def0..5c6c2a96b8c 100644
--- a/src/backend/utils/cache/partcache.c
+++ b/src/backend/utils/cache/partcache.c
@@ -194,6 +194,7 @@ RelationBuildPartitionKey(Relation relation)
 		HeapTuple	opclasstup;
 		Form_pg_opclass opclassform;
 		Oid			funcid;
+		bool		key_virtual;
 
 		/* Collect opfamily information */
 		opclasstup = SearchSysCache1(CLAOID,
@@ -225,8 +226,10 @@ RelationBuildPartitionKey(Relation relation)
 		/* Collation */
 		key->partcollation[i] = collation->values[i];
 
+		key_virtual = (RelationGetColumnGenerated(relation, attno) == ATTRIBUTE_GENERATED_VIRTUAL);
+
 		/* Collect type information */
-		if (attno != 0)
+		if (attno != 0 && !key_virtual)
 		{
 			Form_pg_attribute att = TupleDescAttr(relation->rd_att, attno - 1);
 
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 68ff67de549..434d2ced696 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -6030,6 +6030,19 @@ RelationGetIndexAttOptions(Relation relation, bool copy)
 	return relation->rd_opcoptions;
 }
 
+char
+RelationGetColumnGenerated(Relation rel, int attnum)
+{
+	TupleDesc tupdesc;
+	Form_pg_attribute attr;
+
+	tupdesc = RelationGetDescr(rel);
+
+	attr = TupleDescAttr(tupdesc, attnum -1);
+
+	return attr->attgenerated;
+}
+
 /*
  * Routines to support ereport() reports of relation-related errors
  *
diff --git a/src/include/utils/relcache.h b/src/include/utils/relcache.h
index 3561c6bef0b..506922d13d4 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 char RelationGetColumnGenerated(Relation relation, int 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 16de30ab191..3a013b4cd49 100644
--- a/src/test/regress/expected/generated_stored.out
+++ b/src/test/regress/expected/generated_stored.out
@@ -1070,15 +1070,20 @@ SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
 -- we leave these tables around for purposes of testing dump/reload/upgrade
 -- generated columns in partition key (not allowed)
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE (f3);
-ERROR:  cannot use generated column in partition key
+ERROR:  cannot use stored generated column in partition key
 LINE 1: ...ENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE (f3);
                                                                    ^
+DETAIL:  Column "f3" is a stored generated column.
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
+ERROR:  cannot use stored generated column in partition key
+LINE 1: ...ERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
+                                                                 ^
 DETAIL:  Column "f3" is a generated column.
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
-ERROR:  cannot use generated column in partition key
+ERROR:  cannot use stored generated column in partition key
 LINE 1: ...ED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
                                                              ^
-DETAIL:  Column "f3" is a generated column.
+DETAIL:  Column "f3" is a stored generated column.
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);
 INSERT INTO gtest25 VALUES (3), (4);
diff --git a/src/test/regress/expected/generated_virtual.out b/src/test/regress/expected/generated_virtual.out
index 6300e7c1d96..625f810611f 100644
--- a/src/test/regress/expected/generated_virtual.out
+++ b/src/test/regress/expected/generated_virtual.out
@@ -1021,17 +1021,112 @@ SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
 (3 rows)
 
 -- we leave these tables around for purposes of testing dump/reload/upgrade
--- generated columns in partition key (not allowed)
-CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3);
-ERROR:  cannot use generated column in partition key
-LINE 1: ...NERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE (f3);
-                                                                   ^
-DETAIL:  Column "f3" is a generated column.
-CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
-ERROR:  cannot use generated column in partition key
+-- tests for virtual generated columns in partition key
+CREATE TABLE gtest_part_key1 (f1 date, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3)); --error
+ERROR:  partition key expression cannot use virtual generated column
 LINE 1: ...D ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3));
                                                              ^
-DETAIL:  Column "f3" is a generated column.
+HINT:  Only plain virtual generated column reference can be used in partition key
+CREATE TABLE gtest_part_key1 (f1 date, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3)); --ok
+ALTER TABLE gtest_part_key1 ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2); --error
+ERROR:  cannot alter column "f3" because it is part of the partition key of relation "gtest_part_key1"
+ALTER TABLE gtest_part_key1 ALTER COLUMN f2 set data type int; --error
+ERROR:  cannot alter type of a column used by a generated column
+DETAIL:  Column "f2" is used by generated column "f3".
+ALTER TABLE gtest_part_key1 ALTER COLUMN f3 set data type int; --error
+ERROR:  cannot alter column "f3" because it is part of the partition key of relation "gtest_part_key1"
+LINE 1: ALTER TABLE gtest_part_key1 ALTER COLUMN f3 set data type in...
+                                                 ^
+CREATE TABLE gtest_part_key1_0(f2 bigint, f1 date, f3 bigint GENERATED ALWAYS AS (f2 * 3) VIRTUAL);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); --error
+ERROR:  cannot attach table "gtest_part_key1_0" as a partition because it has with different generation expression
+ALTER TABLE gtest_part_key1_0 ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); --now ok
+CREATE TABLE gtest_part_key1_1 PARTITION OF gtest_part_key1 FOR VALUES FROM (30) TO (50);
+CREATE TABLE gtest_part_key1_2 PARTITION OF gtest_part_key1 FOR VALUES FROM (50) TO (100);
+\d+ gtest_part_key1
+                         Partitioned table "generated_virtual_tests.gtest_part_key1"
+ Column |  Type  | Collation | Nullable |           Default            | Storage | Stats target | Description 
+--------+--------+-----------+----------+------------------------------+---------+--------------+-------------
+ f1     | date   |           |          |                              | plain   |              | 
+ f2     | bigint |           |          |                              | plain   |              | 
+ f3     | bigint |           |          | generated always as (f2 * 2) | plain   |              | 
+Partition key: RANGE (f3)
+Partitions: gtest_part_key1_0 FOR VALUES FROM ('20') TO ('30'),
+            gtest_part_key1_1 FOR VALUES FROM ('30') TO ('50'),
+            gtest_part_key1_2 FOR VALUES FROM ('50') TO ('100')
+
+CREATE OR REPLACE FUNCTION gtest_trigger_info() RETURNS trigger
+  LANGUAGE plpgsql
+AS $$
+BEGIN
+	RAISE INFO 'TG_WHEN: % TG_RELNAME: % trigger name: % tg_op: %', TG_WHEN, TG_relname, TG_NAME, tg_op;
+  IF tg_op IN ('DELETE') THEN
+    RAISE INFO 'old = %', OLD;
+    RETURN OLD;
+  ELSIF tg_op IN ('INSERT') THEN
+    RAISE INFO 'new = %', NEW;
+    RETURN NEW;
+  ELSIF tg_op IN ('UPDATE') THEN
+    RAISE INFO 'old = %d; new = %', OLD, NEW;
+    RETURN NEW;
+  ELSE
+    RETURN NEW;
+  END IF;
+END
+$$;
+CREATE TRIGGER gkey1_0 BEFORE INSERT OR UPDATE ON gtest_part_key1_0
+FOR EACH ROW
+EXECUTE PROCEDURE gtest_trigger_info();
+CREATE TRIGGER gkey1_1 BEFORE INSERT OR UPDATE ON gtest_part_key1_1
+FOR EACH ROW
+EXECUTE PROCEDURE gtest_trigger_info();
+CREATE TRIGGER gkey1_2 BEFORE INSERT OR UPDATE ON gtest_part_key1_2
+FOR EACH ROW
+EXECUTE PROCEDURE gtest_trigger_info();
+INSERT INTO gtest_part_key1(f2, f3) VALUES (9, default); --error
+ERROR:  no partition of relation "gtest_part_key1" found for row
+DETAIL:  Partition key of the failing row contains (f3) = (18).
+INSERT INTO gtest_part_key1(f2, f3) VALUES (10, default) returning tableoid::regclass, *; --ok
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gkey1_0 tg_op: INSERT
+INFO:  new = (10,,)
+     tableoid      | f1 | f2 | f3 
+-------------------+----+----+----
+ gtest_part_key1_0 |    | 10 | 20
+(1 row)
+
+INSERT INTO gtest_part_key1_0(f2, f3) VALUES (12, default) returning tableoid::regclass, *; --ok
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gkey1_0 tg_op: INSERT
+INFO:  new = (12,,)
+     tableoid      | f2 | f1 | f3 
+-------------------+----+----+----
+ gtest_part_key1_0 | 12 |    | 24
+(1 row)
+
+MERGE INTO gtest_part_key1
+  USING (VALUES (10, 100), (12, 25), (14, 30)) AS s(sid, delta)
+ON gtest_part_key1.f2 = s.sid
+WHEN MATCHED AND f2 = 12 THEN UPDATE SET f2 = 20
+WHEN MATCHED AND f2 = 10 THEN UPDATE SET f2 = 30
+WHEN NOT MATCHED THEN INSERT(f2) VALUES (s.sid)
+RETURNING merge_action(), tableoid::regclass, old.f2, old.f3, new.f2, new.f3;
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gkey1_0 tg_op: UPDATE
+INFO:  old = (10,,)d; new = (30,,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_2 trigger name: gkey1_2 tg_op: INSERT
+INFO:  new = (,30,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gkey1_0 tg_op: UPDATE
+INFO:  old = (12,,)d; new = (20,,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_1 trigger name: gkey1_1 tg_op: INSERT
+INFO:  new = (,20,)
+INFO:  TG_WHEN: BEFORE TG_RELNAME: gtest_part_key1_0 trigger name: gkey1_0 tg_op: INSERT
+INFO:  new = (14,,)
+ merge_action |     tableoid      | f2 | f3 | f2 | f3 
+--------------+-------------------+----+----+----+----
+ UPDATE       | gtest_part_key1_2 | 10 | 20 | 30 | 60
+ UPDATE       | gtest_part_key1_1 | 12 | 24 | 20 | 40
+ INSERT       | gtest_part_key1_0 |    |    | 14 | 28
+(3 rows)
+
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);
 INSERT INTO gtest25 VALUES (3), (4);
diff --git a/src/test/regress/sql/generated_stored.sql b/src/test/regress/sql/generated_stored.sql
index 4ec155f2da9..c9427c2d4d8 100644
--- a/src/test/regress/sql/generated_stored.sql
+++ b/src/test/regress/sql/generated_stored.sql
@@ -500,6 +500,7 @@ SELECT tableoid::regclass, * FROM gtest_parent ORDER BY 1, 2, 3;
 
 -- generated columns in partition key (not allowed)
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE (f3);
+CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3));
 CREATE TABLE gtest_part_key (f1 date NOT NULL, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) STORED) PARTITION BY RANGE ((f3 * 3));
 
 -- ALTER TABLE ... ADD COLUMN
diff --git a/src/test/regress/sql/generated_virtual.sql b/src/test/regress/sql/generated_virtual.sql
index b4eedeee2fb..f1d6b1c325f 100644
--- a/src/test/regress/sql/generated_virtual.sql
+++ b/src/test/regress/sql/generated_virtual.sql
@@ -532,9 +532,67 @@ 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 * 3));
+-- tests for virtual generated columns in partition key
+CREATE TABLE gtest_part_key1 (f1 date, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3 * 3)); --error
+CREATE TABLE gtest_part_key1 (f1 date, f2 bigint, f3 bigint GENERATED ALWAYS AS (f2 * 2) VIRTUAL) PARTITION BY RANGE ((f3)); --ok
+ALTER TABLE gtest_part_key1 ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2); --error
+ALTER TABLE gtest_part_key1 ALTER COLUMN f2 set data type int; --error
+ALTER TABLE gtest_part_key1 ALTER COLUMN f3 set data type int; --error
+
+CREATE TABLE gtest_part_key1_0(f2 bigint, f1 date, f3 bigint GENERATED ALWAYS AS (f2 * 3) VIRTUAL);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); --error
+
+ALTER TABLE gtest_part_key1_0 ALTER COLUMN f3 SET EXPRESSION AS (f2 * 2);
+ALTER TABLE gtest_part_key1 ATTACH PARTITION gtest_part_key1_0 FOR VALUES FROM (20) TO (30); --now ok
+
+CREATE TABLE gtest_part_key1_1 PARTITION OF gtest_part_key1 FOR VALUES FROM (30) TO (50);
+CREATE TABLE gtest_part_key1_2 PARTITION OF gtest_part_key1 FOR VALUES FROM (50) TO (100);
+
+\d+ gtest_part_key1
+
+CREATE OR REPLACE FUNCTION gtest_trigger_info() RETURNS trigger
+  LANGUAGE plpgsql
+AS $$
+BEGIN
+	RAISE INFO 'TG_WHEN: % TG_RELNAME: % trigger name: % tg_op: %', TG_WHEN, TG_relname, TG_NAME, tg_op;
+  IF tg_op IN ('DELETE') THEN
+    RAISE INFO 'old = %', OLD;
+    RETURN OLD;
+  ELSIF tg_op IN ('INSERT') THEN
+    RAISE INFO 'new = %', NEW;
+    RETURN NEW;
+  ELSIF tg_op IN ('UPDATE') THEN
+    RAISE INFO 'old = %d; new = %', OLD, NEW;
+    RETURN NEW;
+  ELSE
+    RETURN NEW;
+  END IF;
+END
+$$;
+
+CREATE TRIGGER gkey1_0 BEFORE INSERT OR UPDATE ON gtest_part_key1_0
+FOR EACH ROW
+EXECUTE PROCEDURE gtest_trigger_info();
+
+CREATE TRIGGER gkey1_1 BEFORE INSERT OR UPDATE ON gtest_part_key1_1
+FOR EACH ROW
+EXECUTE PROCEDURE gtest_trigger_info();
+
+CREATE TRIGGER gkey1_2 BEFORE INSERT OR UPDATE ON gtest_part_key1_2
+FOR EACH ROW
+EXECUTE PROCEDURE gtest_trigger_info();
+
+INSERT INTO gtest_part_key1(f2, f3) VALUES (9, default); --error
+INSERT INTO gtest_part_key1(f2, f3) VALUES (10, default) returning tableoid::regclass, *; --ok
+INSERT INTO gtest_part_key1_0(f2, f3) VALUES (12, default) returning tableoid::regclass, *; --ok
+
+MERGE INTO gtest_part_key1
+  USING (VALUES (10, 100), (12, 25), (14, 30)) AS s(sid, delta)
+ON gtest_part_key1.f2 = s.sid
+WHEN MATCHED AND f2 = 12 THEN UPDATE SET f2 = 20
+WHEN MATCHED AND f2 = 10 THEN UPDATE SET f2 = 30
+WHEN NOT MATCHED THEN INSERT(f2) VALUES (s.sid)
+RETURNING merge_action(), tableoid::regclass, old.f2, old.f3, new.f2, new.f3;
 
 -- ALTER TABLE ... ADD COLUMN
 CREATE TABLE gtest25 (a int PRIMARY KEY);
-- 
2.34.1

