diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index 42610f50b0b..a2908fdeef9 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -19770,6 +19770,7 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
 			char		partattname[16];
 			Bitmapset  *expr_attrs = NULL;
 			int			i;
+			bool		whole_row = false;
 
 			Assert(expr != NULL);
 			atttype = exprType(expr);
@@ -19799,7 +19800,9 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
 			 * the partitioned table.
 			 */
 			pull_varattnos(expr, 1, &expr_attrs);
-			if (bms_is_member(0 - FirstLowInvalidHeapAttributeNumber, expr_attrs))
+			whole_row = bms_is_member(0 - FirstLowInvalidHeapAttributeNumber, expr_attrs);
+
+			if (whole_row)
 			{
 				expr_attrs = bms_add_range(expr_attrs,
 										   1 - FirstLowInvalidHeapAttributeNumber,
@@ -19839,12 +19842,23 @@ ComputePartitionAttrs(ParseState *pstate, Relation rel, List *partParams, AttrNu
 				 * used in partition keys).  Seems safer to prohibit for now.
 				 */
 				if (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)));
+				{
+					if (!whole_row)
+						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));
+					else
+						ereport(ERROR,
+								errcode(ERRCODE_INVALID_OBJECT_DEFINITION),
+								errmsg("cannot use generated column in partition key"),
+								errdetail("Generated column \"%s\" is part of the partition key of relation \"%s\"",
+										  get_attname(RelationGetRelid(rel), attno, false),
+										  RelationGetRelationName(rel)),
+								parser_errposition(pstate, pelem->location));
+				}
 			}
 
 			if (IsA(expr, Var) &&