moving extraUpdatedCols out of RangeTblEntry (into ModifyTable)

From aa8a80b3e1e8593e7cfaa5821dc4aad3ebd5b532 Mon Sep 17 00:00:00 2001
From: amitlan <amitlangote09@gmail.com>
Date: Tue, 4 Oct 2022 20:54:03 +0900
Subject: [PATCH v2] Add per-result-relation extraUpdatedCols to ModifyTable

Currently, the rewriter populates the extraUpdatedCols bitmapset
in the main target relation's RTE, which the planner later copies
into the child target relations' RTEs, if needed, with appropriate
attribute number translation.

This commit moves the responsibility of computing even the initial
bitmapset to the the planner, and instead of putting it in the RTEs,
they are now added to a List of Bitmapsets in the ModifyTable node,
containing an element each for all relations in resultRelations.
---
 src/backend/executor/execUtils.c         | 10 ++---
 src/backend/executor/nodeModifyTable.c   |  4 ++
 src/backend/nodes/outfuncs.c             |  1 -
 src/backend/nodes/readfuncs.c            |  1 -
 src/backend/optimizer/plan/createplan.c  |  6 ++-
 src/backend/optimizer/plan/planner.c     | 20 ++++++++++
 src/backend/optimizer/prep/preptlist.c   | 47 ++++++++++++++++++++++++
 src/backend/optimizer/util/inherit.c     | 24 +++++-------
 src/backend/optimizer/util/pathnode.c    |  4 ++
 src/backend/replication/logical/worker.c |  6 +--
 src/backend/rewrite/rewriteHandler.c     | 45 -----------------------
 src/include/nodes/execnodes.h            |  3 ++
 src/include/nodes/parsenodes.h           |  1 -
 src/include/nodes/pathnodes.h            |  7 ++++
 src/include/nodes/plannodes.h            |  1 +
 src/include/optimizer/inherit.h          |  3 ++
 src/include/optimizer/pathnode.h         |  1 +
 src/include/optimizer/prep.h             |  4 ++
 src/include/rewrite/rewriteHandler.h     |  4 --
 19 files changed, 115 insertions(+), 77 deletions(-)

diff --git a/src/backend/executor/execUtils.c b/src/backend/executor/execUtils.c
index e296f44ebd..fcf2fe1aed 100644
--- a/src/backend/executor/execUtils.c
+++ b/src/backend/executor/execUtils.c
@@ -1342,20 +1342,18 @@ ExecGetExtraUpdatedCols(ResultRelInfo *relinfo, EState *estate)
 {
 	if (relinfo->ri_RangeTableIndex != 0)
 	{
-		RangeTblEntry *rte = exec_rt_fetch(relinfo->ri_RangeTableIndex, estate);
-
-		return rte->extraUpdatedCols;
+		return relinfo->ri_extraUpdatedCols;
 	}
 	else if (relinfo->ri_RootResultRelInfo)
 	{
 		ResultRelInfo *rootRelInfo = relinfo->ri_RootResultRelInfo;
-		RangeTblEntry *rte = exec_rt_fetch(rootRelInfo->ri_RangeTableIndex, estate);
 		TupleConversionMap *map = ExecGetRootToChildMap(relinfo, estate);
 
 		if (map != NULL)
-			return execute_attr_map_cols(map->attrMap, rte->extraUpdatedCols);
+			return execute_attr_map_cols(map->attrMap,
+										 rootRelInfo->ri_extraUpdatedCols);
 		else
-			return rte->extraUpdatedCols;
+			return rootRelInfo->ri_extraUpdatedCols;
 	}
 	else
 		return NULL;
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index a3988b1175..f2504901c9 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -4012,6 +4012,10 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	{
 		resultRelInfo = &mtstate->resultRelInfo[i];
 
+		if (node->extraUpdatedColsBitmaps)
+			resultRelInfo->ri_extraUpdatedCols =
+				list_nth_node(Bitmapset, node->extraUpdatedColsBitmaps, i);
+
 		/* Let FDWs init themselves for foreign-table result rels */
 		if (!resultRelInfo->ri_usesFdwDirectModify &&
 			resultRelInfo->ri_FdwRoutine != NULL &&
diff --git a/src/backend/nodes/outfuncs.c b/src/backend/nodes/outfuncs.c
index 59b0fdeb62..2d369e0340 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -561,7 +561,6 @@ _outRangeTblEntry(StringInfo str, const RangeTblEntry *node)
 	WRITE_BOOL_FIELD(lateral);
 	WRITE_BOOL_FIELD(inh);
 	WRITE_BOOL_FIELD(inFromCl);
-	WRITE_BITMAPSET_FIELD(extraUpdatedCols);
 	WRITE_NODE_FIELD(securityQuals);
 }
 
diff --git a/src/backend/nodes/readfuncs.c b/src/backend/nodes/readfuncs.c
index 966b75f5a6..d720aea857 100644
--- a/src/backend/nodes/readfuncs.c
+++ b/src/backend/nodes/readfuncs.c
@@ -537,7 +537,6 @@ _readRangeTblEntry(void)
 	READ_BOOL_FIELD(lateral);
 	READ_BOOL_FIELD(inh);
 	READ_BOOL_FIELD(inFromCl);
-	READ_BITMAPSET_FIELD(extraUpdatedCols);
 	READ_NODE_FIELD(securityQuals);
 
 	READ_DONE();
diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index 66139928e8..5eea7d6996 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -308,7 +308,7 @@ static ModifyTable *make_modifytable(PlannerInfo *root, Plan *subplan,
 									 Index nominalRelation, Index rootRelation,
 									 bool partColsUpdated,
 									 List *resultRelations,
-									 List *updateColnosLists,
+									 List *updateColnosLists, List *extraUpdatedColsBitmaps,
 									 List *withCheckOptionLists, List *returningLists,
 									 List *rowMarks, OnConflictExpr *onconflict,
 									 List *mergeActionLists, int epqParam);
@@ -2826,6 +2826,7 @@ create_modifytable_plan(PlannerInfo *root, ModifyTablePath *best_path)
 							best_path->partColsUpdated,
 							best_path->resultRelations,
 							best_path->updateColnosLists,
+							best_path->extraUpdatedColsBitmaps,
 							best_path->withCheckOptionLists,
 							best_path->returningLists,
 							best_path->rowMarks,
@@ -6982,7 +6983,7 @@ make_modifytable(PlannerInfo *root, Plan *subplan,
 				 Index nominalRelation, Index rootRelation,
 				 bool partColsUpdated,
 				 List *resultRelations,
-				 List *updateColnosLists,
+				 List *updateColnosLists, List *extraUpdatedColsBitmaps,
 				 List *withCheckOptionLists, List *returningLists,
 				 List *rowMarks, OnConflictExpr *onconflict,
 				 List *mergeActionLists, int epqParam)
@@ -7051,6 +7052,7 @@ make_modifytable(PlannerInfo *root, Plan *subplan,
 		node->exclRelTlist = onconflict->exclRelTlist;
 	}
 	node->updateColnosLists = updateColnosLists;
+	node->extraUpdatedColsBitmaps = extraUpdatedColsBitmaps;
 	node->withCheckOptionLists = withCheckOptionLists;
 	node->returningLists = returningLists;
 	node->rowMarks = rowMarks;
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 15aa9c5087..709871a267 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -1747,6 +1747,7 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 			Index		rootRelation;
 			List	   *resultRelations = NIL;
 			List	   *updateColnosLists = NIL;
+			List	   *extraUpdatedColsBitmaps = NIL;
 			List	   *withCheckOptionLists = NIL;
 			List	   *returningLists = NIL;
 			List	   *mergeActionLists = NIL;
@@ -1780,15 +1781,25 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 					if (parse->commandType == CMD_UPDATE)
 					{
 						List	   *update_colnos = root->update_colnos;
+						Bitmapset  *extraUpdatedCols = root->extraUpdatedCols;
 
 						if (this_result_rel != top_result_rel)
+						{
 							update_colnos =
 								adjust_inherited_attnums_multilevel(root,
 																	update_colnos,
 																	this_result_rel->relid,
 																	top_result_rel->relid);
+							extraUpdatedCols =
+								translate_col_privs_multilevel(root, this_result_rel,
+															   top_result_rel,
+															   extraUpdatedCols);
+						}
 						updateColnosLists = lappend(updateColnosLists,
 													update_colnos);
+						if (extraUpdatedCols)
+							extraUpdatedColsBitmaps = lappend(extraUpdatedColsBitmaps,
+															  extraUpdatedCols);
 					}
 					if (parse->withCheckOptions)
 					{
@@ -1870,7 +1881,11 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 					 */
 					resultRelations = list_make1_int(parse->resultRelation);
 					if (parse->commandType == CMD_UPDATE)
+					{
 						updateColnosLists = list_make1(root->update_colnos);
+						if (root->extraUpdatedCols)
+							extraUpdatedColsBitmaps = list_make1(root->extraUpdatedCols);
+					}
 					if (parse->withCheckOptions)
 						withCheckOptionLists = list_make1(parse->withCheckOptions);
 					if (parse->returningList)
@@ -1884,7 +1899,11 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 				/* Single-relation INSERT/UPDATE/DELETE/MERGE. */
 				resultRelations = list_make1_int(parse->resultRelation);
 				if (parse->commandType == CMD_UPDATE)
+				{
 					updateColnosLists = list_make1(root->update_colnos);
+					if (root->extraUpdatedCols)
+						extraUpdatedColsBitmaps = list_make1(root->extraUpdatedCols);
+				}
 				if (parse->withCheckOptions)
 					withCheckOptionLists = list_make1(parse->withCheckOptions);
 				if (parse->returningList)
@@ -1923,6 +1942,7 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 										root->partColsUpdated,
 										resultRelations,
 										updateColnosLists,
+										extraUpdatedColsBitmaps,
 										withCheckOptionLists,
 										returningLists,
 										rowMarks,
diff --git a/src/backend/optimizer/prep/preptlist.c b/src/backend/optimizer/prep/preptlist.c
index 137b28323d..c2a45755ca 100644
--- a/src/backend/optimizer/prep/preptlist.c
+++ b/src/backend/optimizer/prep/preptlist.c
@@ -43,6 +43,7 @@
 #include "optimizer/prep.h"
 #include "optimizer/tlist.h"
 #include "parser/parse_coerce.h"
+#include "parser/parse_relation.h"
 #include "parser/parsetree.h"
 #include "utils/rel.h"
 
@@ -106,6 +107,17 @@ preprocess_targetlist(PlannerInfo *root)
 	else if (command_type == CMD_UPDATE)
 		root->update_colnos = extract_update_targetlist_colnos(tlist);
 
+	/* Also populate extraUpdatedCols (for generated columns) */
+	if (command_type == CMD_UPDATE)
+	{
+		RTEPermissionInfo *target_perminfo =
+			getRTEPermissionInfo(parse->rteperminfos, target_rte);
+
+		root->extraUpdatedCols =
+			get_extraUpdatedCols(target_perminfo->updatedCols,
+								 target_relation);
+	}
+
 	/*
 	 * For non-inherited UPDATE/DELETE/MERGE, register any junk column(s)
 	 * needed to allow the executor to identify the rows to be updated or
@@ -337,6 +349,41 @@ extract_update_targetlist_colnos(List *tlist)
 	return update_colnos;
 }
 
+/*
+ * Return the indexes of any generated columns that depend on any columns
+ * mentioned in updatedCols.
+ */
+Bitmapset *
+get_extraUpdatedCols(Bitmapset *updatedCols, Relation target_relation)
+{
+	TupleDesc	tupdesc = RelationGetDescr(target_relation);
+	TupleConstr *constr = tupdesc->constr;
+	Bitmapset *extraUpdatedCols = NULL;
+
+	if (constr && constr->has_generated_stored)
+	{
+		for (int i = 0; i < constr->num_defval; i++)
+		{
+			AttrDefault *defval = &constr->defval[i];
+			Node	   *expr;
+			Bitmapset  *attrs_used = NULL;
+
+			/* skip if not generated column */
+			if (!TupleDescAttr(tupdesc, defval->adnum - 1)->attgenerated)
+				continue;
+
+			/* identify columns this generated column depends on */
+			expr = stringToNode(defval->adbin);
+			pull_varattnos(expr, 1, &attrs_used);
+
+			if (bms_overlap(updatedCols, attrs_used))
+				extraUpdatedCols = bms_add_member(extraUpdatedCols,
+												  defval->adnum - FirstLowInvalidHeapAttributeNumber);
+		}
+	}
+
+	return extraUpdatedCols;
+}
 
 /*****************************************************************************
  *
diff --git a/src/backend/optimizer/util/inherit.c b/src/backend/optimizer/util/inherit.c
index f51ce45cd3..5d31348457 100644
--- a/src/backend/optimizer/util/inherit.c
+++ b/src/backend/optimizer/util/inherit.c
@@ -40,6 +40,7 @@ static void expand_partitioned_rtentry(PlannerInfo *root, RelOptInfo *relinfo,
 									   RangeTblEntry *parentrte,
 									   Index parentRTindex, Relation parentrel,
 									   Bitmapset *parent_updatedCols,
+									   Bitmapset *parent_extraUpdatedCols,
 									   PlanRowMark *top_parentrc, LOCKMODE lockmode);
 static void expand_single_inheritance_child(PlannerInfo *root,
 											RangeTblEntry *parentrte,
@@ -49,10 +50,6 @@ static void expand_single_inheritance_child(PlannerInfo *root,
 											Index *childRTindex_p);
 static Bitmapset *translate_col_privs(const Bitmapset *parent_privs,
 									  List *translated_vars);
-static Bitmapset *translate_col_privs_multilevel(PlannerInfo *root,
-												 RelOptInfo *rel,
-												 RelOptInfo *top_parent_rel,
-												 Bitmapset *top_parent_cols);
 static void expand_appendrel_subquery(PlannerInfo *root, RelOptInfo *rel,
 									  RangeTblEntry *rte, Index rti);
 
@@ -153,6 +150,7 @@ expand_inherited_rtentry(PlannerInfo *root, RelOptInfo *rel,
 		expand_partitioned_rtentry(root, rel, rte, rti,
 								   oldrelation,
 								   perminfo->updatedCols,
+								   root->extraUpdatedCols,
 								   oldrc, lockmode);
 	}
 	else
@@ -318,6 +316,7 @@ expand_partitioned_rtentry(PlannerInfo *root, RelOptInfo *relinfo,
 						   RangeTblEntry *parentrte,
 						   Index parentRTindex, Relation parentrel,
 						   Bitmapset *parent_updatedCols,
+						   Bitmapset *parent_extraUpdatedCols,
 						   PlanRowMark *top_parentrc, LOCKMODE lockmode)
 {
 	PartitionDesc partdesc;
@@ -348,7 +347,7 @@ expand_partitioned_rtentry(PlannerInfo *root, RelOptInfo *relinfo,
 	/*
 	 * There shouldn't be any generated columns in the partition key.
 	 */
-	Assert(!has_partition_attrs(parentrel, parentrte->extraUpdatedCols, NULL));
+	Assert(!has_partition_attrs(parentrel, parent_extraUpdatedCols, NULL));
 
 	/* Nothing further to do here if there are no partitions. */
 	if (partdesc->nparts == 0)
@@ -417,14 +416,18 @@ expand_partitioned_rtentry(PlannerInfo *root, RelOptInfo *relinfo,
 		{
 			AppendRelInfo *appinfo = root->append_rel_array[childRTindex];
 			Bitmapset  *child_updatedCols;
+			Bitmapset  *child_extraUpdatedCols;
 
 			child_updatedCols = translate_col_privs(parent_updatedCols,
 													appinfo->translated_vars);
+			child_extraUpdatedCols = translate_col_privs(parent_extraUpdatedCols,
+														 appinfo->translated_vars);
 
 			expand_partitioned_rtentry(root, childrelinfo,
 									   childrte, childRTindex,
 									   childrel,
 									   child_updatedCols,
+									   child_extraUpdatedCols,
 									   top_parentrc, lockmode);
 		}
 
@@ -566,13 +569,6 @@ expand_single_inheritance_child(PlannerInfo *root, RangeTblEntry *parentrte,
 	childrte->alias = childrte->eref = makeAlias(parentrte->eref->aliasname,
 												 child_colnames);
 
-	/* Translate the bitmapset of generated columns being updated. */
-	if (childOID != parentOID)
-		childrte->extraUpdatedCols = translate_col_privs(parentrte->extraUpdatedCols,
-														 appinfo->translated_vars);
-	else
-		childrte->extraUpdatedCols = bms_copy(parentrte->extraUpdatedCols);
-
 	/*
 	 * Store the RTE and appinfo in the respective PlannerInfo arrays, which
 	 * the caller must already have allocated space for.
@@ -681,7 +677,7 @@ get_rel_all_updated_cols(PlannerInfo *root, RelOptInfo *rel)
 	perminfo = getRTEPermissionInfo(root->parse->rteperminfos, rte);
 
 	updatedCols = perminfo->updatedCols;
-	extraUpdatedCols = rte->extraUpdatedCols;
+	extraUpdatedCols = root->extraUpdatedCols;
 
 	/*
 	 * For "other" rels, we must look up the root parent relation mentioned in
@@ -927,7 +923,7 @@ apply_child_basequals(PlannerInfo *root, RelOptInfo *parentrel,
  * 		'top_parent_cols' to the columns numbers of a descendent relation
  * 		given by 'relid'
  */
-static Bitmapset *
+Bitmapset *
 translate_col_privs_multilevel(PlannerInfo *root, RelOptInfo *rel,
 							   RelOptInfo *top_parent_rel,
 							   Bitmapset *top_parent_cols)
diff --git a/src/backend/optimizer/util/pathnode.c b/src/backend/optimizer/util/pathnode.c
index 55deee555a..341895d479 100644
--- a/src/backend/optimizer/util/pathnode.c
+++ b/src/backend/optimizer/util/pathnode.c
@@ -3644,6 +3644,8 @@ create_lockrows_path(PlannerInfo *root, RelOptInfo *rel,
  * 'resultRelations' is an integer list of actual RT indexes of target rel(s)
  * 'updateColnosLists' is a list of UPDATE target column number lists
  *		(one sublist per rel); or NIL if not an UPDATE
+ * 'extraUpdatedColsBitmaps' is a list of generated column attribute number
+ *		bitmapsets (one bitmapset per rel); or NIL if not an UPDATE
  * 'withCheckOptionLists' is a list of WCO lists (one per rel)
  * 'returningLists' is a list of RETURNING tlists (one per rel)
  * 'rowMarks' is a list of PlanRowMarks (non-locking only)
@@ -3659,6 +3661,7 @@ create_modifytable_path(PlannerInfo *root, RelOptInfo *rel,
 						bool partColsUpdated,
 						List *resultRelations,
 						List *updateColnosLists,
+						List *extraUpdatedColsBitmaps,
 						List *withCheckOptionLists, List *returningLists,
 						List *rowMarks, OnConflictExpr *onconflict,
 						List *mergeActionLists, int epqParam)
@@ -3723,6 +3726,7 @@ create_modifytable_path(PlannerInfo *root, RelOptInfo *rel,
 	pathnode->partColsUpdated = partColsUpdated;
 	pathnode->resultRelations = resultRelations;
 	pathnode->updateColnosLists = updateColnosLists;
+	pathnode->extraUpdatedColsBitmaps = extraUpdatedColsBitmaps;
 	pathnode->withCheckOptionLists = withCheckOptionLists;
 	pathnode->returningLists = returningLists;
 	pathnode->rowMarks = rowMarks;
diff --git a/src/backend/replication/logical/worker.c b/src/backend/replication/logical/worker.c
index 96772e4d73..109a62635d 100644
--- a/src/backend/replication/logical/worker.c
+++ b/src/backend/replication/logical/worker.c
@@ -156,6 +156,7 @@
 #include "miscadmin.h"
 #include "nodes/makefuncs.h"
 #include "optimizer/optimizer.h"
+#include "optimizer/prep.h"
 #include "parser/parse_relation.h"
 #include "pgstat.h"
 #include "postmaster/bgworker.h"
@@ -1815,7 +1816,6 @@ apply_handle_update(StringInfo s)
 	LogicalRepTupleData newtup;
 	bool		has_oldtup;
 	TupleTableSlot *remoteslot;
-	RangeTblEntry *target_rte;
 	RTEPermissionInfo *target_perminfo;
 	MemoryContext oldctx;
 
@@ -1864,7 +1864,6 @@ apply_handle_update(StringInfo s)
 	 * information.  But it would for example exclude columns that only exist
 	 * on the subscriber, since we are not touching those.
 	 */
-	target_rte = list_nth(estate->es_range_table, 0);
 	target_perminfo = list_nth(estate->es_rteperminfos, 0);
 	for (int i = 0; i < remoteslot->tts_tupleDescriptor->natts; i++)
 	{
@@ -1882,7 +1881,8 @@ apply_handle_update(StringInfo s)
 	}
 
 	/* Also populate extraUpdatedCols, in case we have generated columns */
-	fill_extraUpdatedCols(target_rte, target_perminfo, rel->localrel);
+	edata->targetRelInfo->ri_extraUpdatedCols =
+		get_extraUpdatedCols(target_perminfo->updatedCols, rel->localrel);
 
 	/* Build the search tuple. */
 	oldctx = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index 7cf0ceacc3..5d70bd2742 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -1633,46 +1633,6 @@ rewriteValuesRTEToNulls(Query *parsetree, RangeTblEntry *rte)
 }
 
 
-/*
- * Record in target_rte->extraUpdatedCols the indexes of any generated columns
- * columns that depend on any columns mentioned in
- * target_perminfo->updatedCols.
- */
-void
-fill_extraUpdatedCols(RangeTblEntry *target_rte,
-					  RTEPermissionInfo *target_perminfo,
-					  Relation target_relation)
-{
-	TupleDesc	tupdesc = RelationGetDescr(target_relation);
-	TupleConstr *constr = tupdesc->constr;
-
-	target_rte->extraUpdatedCols = NULL;
-
-	if (constr && constr->has_generated_stored)
-	{
-		for (int i = 0; i < constr->num_defval; i++)
-		{
-			AttrDefault *defval = &constr->defval[i];
-			Node	   *expr;
-			Bitmapset  *attrs_used = NULL;
-
-			/* skip if not generated column */
-			if (!TupleDescAttr(tupdesc, defval->adnum - 1)->attgenerated)
-				continue;
-
-			/* identify columns this generated column depends on */
-			expr = stringToNode(defval->adbin);
-			pull_varattnos(expr, 1, &attrs_used);
-
-			if (bms_overlap(target_perminfo->updatedCols, attrs_used))
-				target_rte->extraUpdatedCols =
-					bms_add_member(target_rte->extraUpdatedCols,
-								   defval->adnum - FirstLowInvalidHeapAttributeNumber);
-		}
-	}
-}
-
-
 /*
  * matchLocks -
  *	  match the list of locks and returns the matching rules
@@ -3738,7 +3698,6 @@ RewriteQuery(Query *parsetree, List *rewrite_events, int orig_rt_length)
 	{
 		int			result_relation;
 		RangeTblEntry *rt_entry;
-		RTEPermissionInfo *rt_perminfo;
 		Relation	rt_entry_relation;
 		List	   *locks;
 		int			product_orig_rt_length;
@@ -3751,7 +3710,6 @@ RewriteQuery(Query *parsetree, List *rewrite_events, int orig_rt_length)
 		Assert(result_relation != 0);
 		rt_entry = rt_fetch(result_relation, parsetree->rtable);
 		Assert(rt_entry->rtekind == RTE_RELATION);
-		rt_perminfo = getRTEPermissionInfo(parsetree->rteperminfos, rt_entry);
 
 		/*
 		 * We can use NoLock here since either the parser or
@@ -3843,9 +3801,6 @@ RewriteQuery(Query *parsetree, List *rewrite_events, int orig_rt_length)
 									parsetree->override,
 									rt_entry_relation,
 									NULL, 0, NULL);
-
-			/* Also populate extraUpdatedCols (for generated columns) */
-			fill_extraUpdatedCols(rt_entry, rt_perminfo, rt_entry_relation);
 		}
 		else if (event == CMD_MERGE)
 		{
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index 08ad9d6be3..2c052bc2d3 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -471,6 +471,9 @@ typedef struct ResultRelInfo
 	/* Have the projection and the slots above been initialized? */
 	bool		ri_projectNewInfoValid;
 
+	/* generated column attribute numbers */
+	Bitmapset   *ri_extraUpdatedCols;
+
 	/* triggers to be fired, if any */
 	TriggerDesc *ri_TrigDesc;
 
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 6a6d3293e4..f3ac316171 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -1154,7 +1154,6 @@ typedef struct RangeTblEntry
 	bool		lateral;		/* subquery, function, or values is LATERAL? */
 	bool		inh;			/* inheritance requested? */
 	bool		inFromCl;		/* present in FROM clause? */
-	Bitmapset  *extraUpdatedCols;	/* generated columns being updated */
 	List	   *securityQuals;	/* security barrier quals to apply, if any */
 } RangeTblEntry;
 
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index 654dba61aa..e383cf2557 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -425,6 +425,12 @@ struct PlannerInfo
 	 */
 	List	   *update_colnos;
 
+	/*
+	 * For UPDATE, generated column attribute numbers of the target relation,
+	 * if any.
+	 */
+	Bitmapset  *extraUpdatedCols;
+
 	/*
 	 * Fields filled during create_plan() for use in setrefs.c
 	 */
@@ -2257,6 +2263,7 @@ typedef struct ModifyTablePath
 	bool		partColsUpdated;	/* some part key in hierarchy updated? */
 	List	   *resultRelations;	/* integer list of RT indexes */
 	List	   *updateColnosLists;	/* per-target-table update_colnos lists */
+	List	   *extraUpdatedColsBitmaps; /* per-target-table extraUpdatedCols bitmaps */
 	List	   *withCheckOptionLists;	/* per-target-table WCO lists */
 	List	   *returningLists; /* per-target-table RETURNING tlists */
 	List	   *rowMarks;		/* PlanRowMarks (non-locking only) */
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index bddfe86191..136688b04d 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -239,6 +239,7 @@ typedef struct ModifyTable
 	bool		partColsUpdated;	/* some part key in hierarchy updated? */
 	List	   *resultRelations;	/* integer list of RT indexes */
 	List	   *updateColnosLists;	/* per-target-table update_colnos lists */
+	List	   *extraUpdatedColsBitmaps; /* per-target-table extraUpdatedCols bitmaps */
 	List	   *withCheckOptionLists;	/* per-target-table WCO lists */
 	List	   *returningLists; /* per-target-table RETURNING tlists */
 	List	   *fdwPrivLists;	/* per-target-table FDW private data lists */
diff --git a/src/include/optimizer/inherit.h b/src/include/optimizer/inherit.h
index 8ebd42b757..1fc5c37ade 100644
--- a/src/include/optimizer/inherit.h
+++ b/src/include/optimizer/inherit.h
@@ -25,5 +25,8 @@ extern Bitmapset *get_rel_all_updated_cols(PlannerInfo *root, RelOptInfo *rel);
 extern bool apply_child_basequals(PlannerInfo *root, RelOptInfo *parentrel,
 								  RelOptInfo *childrel, RangeTblEntry *childRTE,
 								  AppendRelInfo *appinfo);
+extern Bitmapset *translate_col_privs_multilevel(PlannerInfo *root, RelOptInfo *rel,
+							   RelOptInfo *top_parent_rel,
+							   Bitmapset *top_parent_cols);
 
 #endif							/* INHERIT_H */
diff --git a/src/include/optimizer/pathnode.h b/src/include/optimizer/pathnode.h
index 050f00e79a..fd16d94916 100644
--- a/src/include/optimizer/pathnode.h
+++ b/src/include/optimizer/pathnode.h
@@ -278,6 +278,7 @@ extern ModifyTablePath *create_modifytable_path(PlannerInfo *root,
 												bool partColsUpdated,
 												List *resultRelations,
 												List *updateColnosLists,
+												List *extraUpdatedColsBitmaps,
 												List *withCheckOptionLists, List *returningLists,
 												List *rowMarks, OnConflictExpr *onconflict,
 												List *mergeActionLists, int epqParam);
diff --git a/src/include/optimizer/prep.h b/src/include/optimizer/prep.h
index 5b4f350b33..92753c9670 100644
--- a/src/include/optimizer/prep.h
+++ b/src/include/optimizer/prep.h
@@ -16,6 +16,7 @@
 
 #include "nodes/pathnodes.h"
 #include "nodes/plannodes.h"
+#include "utils/relcache.h"
 
 
 /*
@@ -39,6 +40,9 @@ extern void preprocess_targetlist(PlannerInfo *root);
 
 extern List *extract_update_targetlist_colnos(List *tlist);
 
+extern Bitmapset *get_extraUpdatedCols(Bitmapset *updatedCols,
+									   Relation target_relation);
+
 extern PlanRowMark *get_plan_rowmark(List *rowmarks, Index rtindex);
 
 /*
diff --git a/src/include/rewrite/rewriteHandler.h b/src/include/rewrite/rewriteHandler.h
index 05c3680cd6..b4f96f298b 100644
--- a/src/include/rewrite/rewriteHandler.h
+++ b/src/include/rewrite/rewriteHandler.h
@@ -24,10 +24,6 @@ extern void AcquireRewriteLocks(Query *parsetree,
 
 extern Node *build_column_default(Relation rel, int attrno);
 
-extern void fill_extraUpdatedCols(RangeTblEntry *target_rte,
-								  RTEPermissionInfo *target_perminfo,
-								  Relation target_relation);
-
 extern Query *get_view_query(Relation view);
 extern const char *view_query_is_auto_updatable(Query *viewquery,
 												bool check_cols);
-- 
2.35.3

diff --git a/src/backend/executor/execUtils.c b/src/backend/executor/execUtils.c
index 8bcf4784eb..8c82c71675 100644
--- a/src/backend/executor/execUtils.c
+++ b/src/backend/executor/execUtils.c
@@ -1342,25 +1342,16 @@ ExecGetUpdatedCols(ResultRelInfo *relinfo, EState *estate)
 Bitmapset *
 ExecGetExtraUpdatedCols(ResultRelInfo *relinfo, EState *estate)
 {
-	if (relinfo->ri_RangeTableIndex != 0)
-	{
-		RangeTblEntry *rte = exec_rt_fetch(relinfo->ri_RangeTableIndex, estate);
+#ifdef USE_ASSERT_CHECKING
+	/* Verify that ExecInitStoredGenerated has been called if needed. */
+	Relation	rel = relinfo->ri_RelationDesc;
+	TupleDesc	tupdesc = RelationGetDescr(rel);
 
-		return rte->extraUpdatedCols;
-	}
-	else if (relinfo->ri_RootResultRelInfo)
-	{
-		ResultRelInfo *rootRelInfo = relinfo->ri_RootResultRelInfo;
-		RangeTblEntry *rte = exec_rt_fetch(rootRelInfo->ri_RangeTableIndex, estate);
-		TupleConversionMap *map = ExecGetRootToChildMap(relinfo, estate);
+	if (tupdesc->constr && tupdesc->constr->has_generated_stored)
+		Assert(relinfo->ri_GeneratedExprs != NULL);
+#endif
 
-		if (map != NULL)
-			return execute_attr_map_cols(map->attrMap, rte->extraUpdatedCols);
-		else
-			return rte->extraUpdatedCols;
-	}
-	else
-		return NULL;
+	return relinfo->ri_extraUpdatedCols;
 }
 
 /* Return columns being updated, including generated columns */
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index 56398c399c..687a5422ea 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -54,6 +54,7 @@
 #include "foreign/fdwapi.h"
 #include "miscadmin.h"
 #include "nodes/nodeFuncs.h"
+#include "optimizer/optimizer.h"
 #include "rewrite/rewriteHandler.h"
 #include "storage/bufmgr.h"
 #include "storage/lmgr.h"
@@ -352,6 +353,93 @@ ExecCheckTIDVisible(EState *estate,
 	ExecClearTuple(tempSlot);
 }
 
+/*
+ * Initialize to compute stored generated columns for a tuple
+ *
+ * This fills the resultRelInfo's ri_GeneratedExprs and ri_extraUpdatedCols
+ * fields.  (Currently, ri_extraUpdatedCols is consulted only in UPDATE,
+ * but we might as well fill it for INSERT too.)
+ */
+static void
+ExecInitStoredGenerated(ResultRelInfo *resultRelInfo,
+						EState *estate,
+						CmdType cmdtype)
+{
+	Relation	rel = resultRelInfo->ri_RelationDesc;
+	TupleDesc	tupdesc = RelationGetDescr(rel);
+	int			natts = tupdesc->natts;
+	Bitmapset  *updatedCols;
+	MemoryContext oldContext;
+
+	/* Don't call twice */
+	Assert(resultRelInfo->ri_GeneratedExprs == NULL);
+
+	/* Nothing to do if no generated columns */
+	if (!(tupdesc->constr && tupdesc->constr->has_generated_stored))
+		return;
+
+	/*
+	 * In an UPDATE, we can skip computing any generated columns that do not
+	 * depend on any UPDATE target column.  But if there is a BEFORE ROW
+	 * UPDATE trigger, we cannot skip because the trigger might change more
+	 * columns.
+	 */
+	if (cmdtype == CMD_UPDATE &&
+		!(rel->trigdesc && rel->trigdesc->trig_update_before_row))
+		updatedCols = ExecGetUpdatedCols(resultRelInfo, estate);
+	else
+		updatedCols = NULL;
+
+	/*
+	 * Make sure these data structures are built in the per-query memory
+	 * context so they'll survive throughout the query.
+	 */
+	oldContext = MemoryContextSwitchTo(estate->es_query_cxt);
+
+	resultRelInfo->ri_GeneratedExprs =
+		(ExprState **) palloc0(natts * sizeof(ExprState *));
+	resultRelInfo->ri_NumGeneratedNeeded = 0;
+
+	for (int i = 0; i < natts; i++)
+	{
+		if (TupleDescAttr(tupdesc, i)->attgenerated == ATTRIBUTE_GENERATED_STORED)
+		{
+			Expr	   *expr;
+
+			/* Fetch the GENERATED AS expression tree */
+			expr = (Expr *) build_column_default(rel, i + 1);
+			if (expr == NULL)
+				elog(ERROR, "no generation expression found for column number %d of table \"%s\"",
+					 i + 1, RelationGetRelationName(rel));
+
+			/*
+			 * If it's an update with a known set of update target columns,
+			 * see if we can skip the computation.
+			 */
+			if (updatedCols)
+			{
+				Bitmapset  *attrs_used = NULL;
+
+				pull_varattnos((Node *) expr, 1, &attrs_used);
+
+				if (!bms_overlap(updatedCols, attrs_used))
+					continue;	/* need not update this column */
+			}
+
+			/* No luck, so prepare the expression for execution */
+			resultRelInfo->ri_GeneratedExprs[i] = ExecPrepareExpr(expr, estate);
+			resultRelInfo->ri_NumGeneratedNeeded++;
+
+			/* And mark this column in resultRelInfo->ri_extraUpdatedCols */
+			resultRelInfo->ri_extraUpdatedCols =
+				bms_add_member(resultRelInfo->ri_extraUpdatedCols,
+							   i + 1 - FirstLowInvalidHeapAttributeNumber);
+		}
+	}
+
+	MemoryContextSwitchTo(oldContext);
+}
+
 /*
  * Compute stored generated columns for a tuple
  */
@@ -363,58 +451,22 @@ ExecComputeStoredGenerated(ResultRelInfo *resultRelInfo,
 	Relation	rel = resultRelInfo->ri_RelationDesc;
 	TupleDesc	tupdesc = RelationGetDescr(rel);
 	int			natts = tupdesc->natts;
+	ExprContext *econtext = GetPerTupleExprContext(estate);
 	MemoryContext oldContext;
 	Datum	   *values;
 	bool	   *nulls;
 
+	/* We should not be called unless this is true */
 	Assert(tupdesc->constr && tupdesc->constr->has_generated_stored);
 
 	/*
-	 * If first time through for this result relation, build expression
-	 * nodetrees for rel's stored generation expressions.  Keep them in the
-	 * per-query memory context so they'll survive throughout the query.
+	 * For relations named directly in the query, ExecInitStoredGenerated
+	 * should have been called already; but this might not have happened yet
+	 * for a partition child rel.  Also, it's convenient for outside callers
+	 * to not have to call ExecInitStoredGenerated explicitly.
 	 */
 	if (resultRelInfo->ri_GeneratedExprs == NULL)
-	{
-		oldContext = MemoryContextSwitchTo(estate->es_query_cxt);
-
-		resultRelInfo->ri_GeneratedExprs =
-			(ExprState **) palloc(natts * sizeof(ExprState *));
-		resultRelInfo->ri_NumGeneratedNeeded = 0;
-
-		for (int i = 0; i < natts; i++)
-		{
-			if (TupleDescAttr(tupdesc, i)->attgenerated == ATTRIBUTE_GENERATED_STORED)
-			{
-				Expr	   *expr;
-
-				/*
-				 * If it's an update and the current column was not marked as
-				 * being updated, then we can skip the computation.  But if
-				 * there is a BEFORE ROW UPDATE trigger, we cannot skip
-				 * because the trigger might affect additional columns.
-				 */
-				if (cmdtype == CMD_UPDATE &&
-					!(rel->trigdesc && rel->trigdesc->trig_update_before_row) &&
-					!bms_is_member(i + 1 - FirstLowInvalidHeapAttributeNumber,
-								   ExecGetExtraUpdatedCols(resultRelInfo, estate)))
-				{
-					resultRelInfo->ri_GeneratedExprs[i] = NULL;
-					continue;
-				}
-
-				expr = (Expr *) build_column_default(rel, i + 1);
-				if (expr == NULL)
-					elog(ERROR, "no generation expression found for column number %d of table \"%s\"",
-						 i + 1, RelationGetRelationName(rel));
-
-				resultRelInfo->ri_GeneratedExprs[i] = ExecPrepareExpr(expr, estate);
-				resultRelInfo->ri_NumGeneratedNeeded++;
-			}
-		}
-
-		MemoryContextSwitchTo(oldContext);
-	}
+		ExecInitStoredGenerated(resultRelInfo, estate, cmdtype);
 
 	/*
 	 * If no generated columns have been affected by this change, then skip
@@ -435,14 +487,13 @@ ExecComputeStoredGenerated(ResultRelInfo *resultRelInfo,
 	{
 		Form_pg_attribute attr = TupleDescAttr(tupdesc, i);
 
-		if (attr->attgenerated == ATTRIBUTE_GENERATED_STORED &&
-			resultRelInfo->ri_GeneratedExprs[i])
+		if (resultRelInfo->ri_GeneratedExprs[i])
 		{
-			ExprContext *econtext;
 			Datum		val;
 			bool		isnull;
 
-			econtext = GetPerTupleExprContext(estate);
+			Assert(attr->attgenerated == ATTRIBUTE_GENERATED_STORED);
+
 			econtext->ecxt_scantuple = slot;
 
 			val = ExecEvalExpr(resultRelInfo->ri_GeneratedExprs[i], econtext, &isnull);
@@ -4088,6 +4139,15 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 					elog(ERROR, "could not find junk wholerow column");
 			}
 		}
+
+		/*
+		 * For INSERT and UPDATE, prepare to evaluate any generated columns.
+		 * We must do this now, even if we never insert or update any rows,
+		 * because we have to fill resultRelInfo->ri_extraUpdatedCols for
+		 * possible use by the trigger machinery.
+		 */
+		if (operation == CMD_INSERT || operation == CMD_UPDATE)
+			ExecInitStoredGenerated(resultRelInfo, estate, operation);
 	}
 
 	/*
diff --git a/src/backend/nodes/outfuncs.c b/src/backend/nodes/outfuncs.c
index 0cb235bfda..69324d5a9a 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -561,7 +561,6 @@ _outRangeTblEntry(StringInfo str, const RangeTblEntry *node)
 	WRITE_BOOL_FIELD(lateral);
 	WRITE_BOOL_FIELD(inh);
 	WRITE_BOOL_FIELD(inFromCl);
-	WRITE_BITMAPSET_FIELD(extraUpdatedCols);
 	WRITE_NODE_FIELD(securityQuals);
 }
 
diff --git a/src/backend/nodes/readfuncs.c b/src/backend/nodes/readfuncs.c
index c284b96a54..30cd7a0da6 100644
--- a/src/backend/nodes/readfuncs.c
+++ b/src/backend/nodes/readfuncs.c
@@ -537,7 +537,6 @@ _readRangeTblEntry(void)
 	READ_BOOL_FIELD(lateral);
 	READ_BOOL_FIELD(inh);
 	READ_BOOL_FIELD(inFromCl);
-	READ_BITMAPSET_FIELD(extraUpdatedCols);
 	READ_NODE_FIELD(securityQuals);
 
 	READ_DONE();
diff --git a/src/backend/optimizer/util/inherit.c b/src/backend/optimizer/util/inherit.c
index 6c93c22477..07a8818d6c 100644
--- a/src/backend/optimizer/util/inherit.c
+++ b/src/backend/optimizer/util/inherit.c
@@ -25,6 +25,7 @@
 #include "optimizer/inherit.h"
 #include "optimizer/optimizer.h"
 #include "optimizer/pathnode.h"
+#include "optimizer/plancat.h"
 #include "optimizer/planmain.h"
 #include "optimizer/planner.h"
 #include "optimizer/prep.h"
@@ -345,11 +346,6 @@ expand_partitioned_rtentry(PlannerInfo *root, RelOptInfo *relinfo,
 		root->partColsUpdated =
 			has_partition_attrs(parentrel, parent_updatedCols, NULL);
 
-	/*
-	 * There shouldn't be any generated columns in the partition key.
-	 */
-	Assert(!has_partition_attrs(parentrel, parentrte->extraUpdatedCols, NULL));
-
 	/* Nothing further to do here if there are no partitions. */
 	if (partdesc->nparts == 0)
 		return;
@@ -566,13 +562,6 @@ expand_single_inheritance_child(PlannerInfo *root, RangeTblEntry *parentrte,
 	childrte->alias = childrte->eref = makeAlias(parentrte->eref->aliasname,
 												 child_colnames);
 
-	/* Translate the bitmapset of generated columns being updated. */
-	if (childOID != parentOID)
-		childrte->extraUpdatedCols = translate_col_privs(parentrte->extraUpdatedCols,
-														 appinfo->translated_vars);
-	else
-		childrte->extraUpdatedCols = bms_copy(parentrte->extraUpdatedCols);
-
 	/*
 	 * Store the RTE and appinfo in the respective PlannerInfo arrays, which
 	 * the caller must already have allocated space for.
@@ -672,21 +661,16 @@ get_rel_all_updated_cols(PlannerInfo *root, RelOptInfo *rel)
 	Assert(IS_SIMPLE_REL(rel));
 
 	/*
-	 * We obtain updatedCols and extraUpdatedCols for the query's result
-	 * relation.  Then, if necessary, we map it to the column numbers of the
-	 * relation for which they were requested.
+	 * We obtain updatedCols for the query's result relation.  Then, if
+	 * necessary, we map it to the column numbers of the relation for which
+	 * they were requested.
 	 */
 	relid = root->parse->resultRelation;
 	rte = planner_rt_fetch(relid, root);
 	perminfo = getRTEPermissionInfo(root->parse->rteperminfos, rte);
 
 	updatedCols = perminfo->updatedCols;
-	extraUpdatedCols = rte->extraUpdatedCols;
 
-	/*
-	 * For "other" rels, we must look up the root parent relation mentioned in
-	 * the query, and translate the column numbers.
-	 */
 	if (rel->relid != relid)
 	{
 		RelOptInfo *top_parent_rel = find_base_rel(root, relid);
@@ -695,10 +679,15 @@ get_rel_all_updated_cols(PlannerInfo *root, RelOptInfo *rel)
 
 		updatedCols = translate_col_privs_multilevel(root, rel, top_parent_rel,
 													 updatedCols);
-		extraUpdatedCols = translate_col_privs_multilevel(root, rel, top_parent_rel,
-														  extraUpdatedCols);
 	}
 
+	/*
+	 * Now we must check to see if there are any generated columns that depend
+	 * on the updatedCols, and add them to the result.
+	 */
+	extraUpdatedCols = get_dependent_generated_columns(root, rel->relid,
+													   updatedCols);
+
 	return bms_union(updatedCols, extraUpdatedCols);
 }
 
diff --git a/src/backend/optimizer/util/plancat.c b/src/backend/optimizer/util/plancat.c
index c24bdae717..9f158f2421 100644
--- a/src/backend/optimizer/util/plancat.c
+++ b/src/backend/optimizer/util/plancat.c
@@ -2199,6 +2199,11 @@ has_row_triggers(PlannerInfo *root, Index rti, CmdType event)
 	return result;
 }
 
+/*
+ * has_stored_generated_columns
+ *
+ * Does table identified by RTI have any STORED GENERATED columns?
+ */
 bool
 has_stored_generated_columns(PlannerInfo *root, Index rti)
 {
@@ -2218,6 +2223,57 @@ has_stored_generated_columns(PlannerInfo *root, Index rti)
 	return result;
 }
 
+/*
+ * get_dependent_generated_columns
+ *
+ * Get the column numbers of any STORED GENERATED columns of the relation
+ * that depend on any column listed in target_cols.  Both the input and
+ * result bitmapsets contain column numbers offset by
+ * FirstLowInvalidHeapAttributeNumber.
+ */
+Bitmapset *
+get_dependent_generated_columns(PlannerInfo *root, Index rti,
+								Bitmapset *target_cols)
+{
+	Bitmapset  *dependentCols = NULL;
+	RangeTblEntry *rte = planner_rt_fetch(rti, root);
+	Relation	relation;
+	TupleDesc	tupdesc;
+	TupleConstr *constr;
+
+	/* Assume we already have adequate lock */
+	relation = table_open(rte->relid, NoLock);
+
+	tupdesc = RelationGetDescr(relation);
+	constr = tupdesc->constr;
+
+	if (constr && constr->has_generated_stored)
+	{
+		for (int i = 0; i < constr->num_defval; i++)
+		{
+			AttrDefault *defval = &constr->defval[i];
+			Node	   *expr;
+			Bitmapset  *attrs_used = NULL;
+
+			/* skip if not generated column */
+			if (!TupleDescAttr(tupdesc, defval->adnum - 1)->attgenerated)
+				continue;
+
+			/* identify columns this generated column depends on */
+			expr = stringToNode(defval->adbin);
+			pull_varattnos(expr, 1, &attrs_used);
+
+			if (bms_overlap(target_cols, attrs_used))
+				dependentCols = bms_add_member(dependentCols,
+											   defval->adnum - FirstLowInvalidHeapAttributeNumber);
+		}
+	}
+
+	table_close(relation, NoLock);
+
+	return dependentCols;
+}
+
 /*
  * set_relation_partition_info
  *
diff --git a/src/backend/replication/logical/worker.c b/src/backend/replication/logical/worker.c
index 446f84fa97..f8e8cf71eb 100644
--- a/src/backend/replication/logical/worker.c
+++ b/src/backend/replication/logical/worker.c
@@ -1815,7 +1815,6 @@ apply_handle_update(StringInfo s)
 	LogicalRepTupleData newtup;
 	bool		has_oldtup;
 	TupleTableSlot *remoteslot;
-	RangeTblEntry *target_rte;
 	RTEPermissionInfo *target_perminfo;
 	MemoryContext oldctx;
 
@@ -1864,7 +1863,6 @@ apply_handle_update(StringInfo s)
 	 * information.  But it would for example exclude columns that only exist
 	 * on the subscriber, since we are not touching those.
 	 */
-	target_rte = list_nth(estate->es_range_table, 0);
 	target_perminfo = list_nth(estate->es_rteperminfos, 0);
 	for (int i = 0; i < remoteslot->tts_tupleDescriptor->natts; i++)
 	{
@@ -1881,9 +1879,6 @@ apply_handle_update(StringInfo s)
 		}
 	}
 
-	/* Also populate extraUpdatedCols, in case we have generated columns */
-	fill_extraUpdatedCols(target_rte, target_perminfo, rel->localrel);
-
 	/* Build the search tuple. */
 	oldctx = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
 	slot_store_data(remoteslot, rel,
diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index ca8b543bc1..1960dad701 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -1633,46 +1633,6 @@ rewriteValuesRTEToNulls(Query *parsetree, RangeTblEntry *rte)
 }
 
 
-/*
- * Record in target_rte->extraUpdatedCols the indexes of any generated columns
- * columns that depend on any columns mentioned in
- * target_perminfo->updatedCols.
- */
-void
-fill_extraUpdatedCols(RangeTblEntry *target_rte,
-					  RTEPermissionInfo *target_perminfo,
-					  Relation target_relation)
-{
-	TupleDesc	tupdesc = RelationGetDescr(target_relation);
-	TupleConstr *constr = tupdesc->constr;
-
-	target_rte->extraUpdatedCols = NULL;
-
-	if (constr && constr->has_generated_stored)
-	{
-		for (int i = 0; i < constr->num_defval; i++)
-		{
-			AttrDefault *defval = &constr->defval[i];
-			Node	   *expr;
-			Bitmapset  *attrs_used = NULL;
-
-			/* skip if not generated column */
-			if (!TupleDescAttr(tupdesc, defval->adnum - 1)->attgenerated)
-				continue;
-
-			/* identify columns this generated column depends on */
-			expr = stringToNode(defval->adbin);
-			pull_varattnos(expr, 1, &attrs_used);
-
-			if (bms_overlap(target_perminfo->updatedCols, attrs_used))
-				target_rte->extraUpdatedCols =
-					bms_add_member(target_rte->extraUpdatedCols,
-								   defval->adnum - FirstLowInvalidHeapAttributeNumber);
-		}
-	}
-}
-
-
 /*
  * matchLocks -
  *	  match the list of locks and returns the matching rules
@@ -3738,7 +3698,6 @@ RewriteQuery(Query *parsetree, List *rewrite_events, int orig_rt_length)
 	{
 		int			result_relation;
 		RangeTblEntry *rt_entry;
-		RTEPermissionInfo *rt_perminfo;
 		Relation	rt_entry_relation;
 		List	   *locks;
 		int			product_orig_rt_length;
@@ -3751,7 +3710,6 @@ RewriteQuery(Query *parsetree, List *rewrite_events, int orig_rt_length)
 		Assert(result_relation != 0);
 		rt_entry = rt_fetch(result_relation, parsetree->rtable);
 		Assert(rt_entry->rtekind == RTE_RELATION);
-		rt_perminfo = getRTEPermissionInfo(parsetree->rteperminfos, rt_entry);
 
 		/*
 		 * We can use NoLock here since either the parser or
@@ -3843,9 +3801,6 @@ RewriteQuery(Query *parsetree, List *rewrite_events, int orig_rt_length)
 									parsetree->override,
 									rt_entry_relation,
 									NULL, 0, NULL);
-
-			/* Also populate extraUpdatedCols (for generated columns) */
-			fill_extraUpdatedCols(rt_entry, rt_perminfo, rt_entry_relation);
 		}
 		else if (event == CMD_MERGE)
 		{
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index 2cd0a4f472..2e7c30b63e 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -462,6 +462,9 @@ typedef struct ResultRelInfo
 	 */
 	AttrNumber	ri_RowIdAttNo;
 
+	/* If UPDATE, attribute numbers of generated columns to be updated */
+	Bitmapset  *ri_extraUpdatedCols;
+
 	/* Projection to generate new tuple in an INSERT/UPDATE */
 	ProjectionInfo *ri_projectNew;
 	/* Slot to hold that tuple */
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 6300945734..cfeca96d53 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -1153,7 +1153,6 @@ typedef struct RangeTblEntry
 	bool		lateral;		/* subquery, function, or values is LATERAL? */
 	bool		inh;			/* inheritance requested? */
 	bool		inFromCl;		/* present in FROM clause? */
-	Bitmapset  *extraUpdatedCols;	/* generated columns being updated */
 	List	   *securityQuals;	/* security barrier quals to apply, if any */
 } RangeTblEntry;
 
@@ -1189,15 +1188,6 @@ typedef struct RangeTblEntry
  * updatedCols is also used in some other places, for example, to determine
  * which triggers to fire and in FDWs to know which changed columns they need
  * to ship off.
- *
- * Generated columns that are caused to be updated by an update to a base
- * column are listed in extraUpdatedCols.  This is not considered for
- * permission checking, but it is useful in those places that want to know the
- * full set of columns being updated as opposed to only the ones the user
- * explicitly mentioned in the query.  (There is currently no need for an
- * extraInsertedCols, but it could exist.)  Note that extraUpdatedCols is
- * populated during query rewrite, NOT in the parser, since generated columns
- * could be added after a rule has been parsed and stored.
  */
 typedef struct RTEPermissionInfo
 {
diff --git a/src/include/optimizer/plancat.h b/src/include/optimizer/plancat.h
index 21edac04ea..eb1c3ccc4b 100644
--- a/src/include/optimizer/plancat.h
+++ b/src/include/optimizer/plancat.h
@@ -74,4 +74,7 @@ extern bool has_row_triggers(PlannerInfo *root, Index rti, CmdType event);
 
 extern bool has_stored_generated_columns(PlannerInfo *root, Index rti);
 
+extern Bitmapset *get_dependent_generated_columns(PlannerInfo *root, Index rti,
+												  Bitmapset *target_cols);
+
 #endif							/* PLANCAT_H */
diff --git a/src/include/rewrite/rewriteHandler.h b/src/include/rewrite/rewriteHandler.h
index 62bca7cfdf..b71e20b087 100644
--- a/src/include/rewrite/rewriteHandler.h
+++ b/src/include/rewrite/rewriteHandler.h
@@ -24,10 +24,6 @@ extern void AcquireRewriteLocks(Query *parsetree,
 
 extern Node *build_column_default(Relation rel, int attrno);
 
-extern void fill_extraUpdatedCols(RangeTblEntry *target_rte,
-								  RTEPermissionInfo *target_perminfo,
-								  Relation target_relation);
-
 extern Query *get_view_query(Relation view);
 extern const char *view_query_is_auto_updatable(Query *viewquery,
 												bool check_cols);
diff --git a/src/test/regress/expected/generated.out b/src/test/regress/expected/generated.out
index bb4190340e..1db5f9ed47 100644
--- a/src/test/regress/expected/generated.out
+++ b/src/test/regress/expected/generated.out
@@ -337,6 +337,25 @@ CREATE TABLE gtest1_2 () INHERITS (gtest0, gtesty);  -- error
 NOTICE:  merging multiple inherited definitions of column "b"
 ERROR:  inherited column "b" has a generation conflict
 DROP TABLE gtesty;
+-- test correct handling of GENERATED column that's only in child
+CREATE TABLE gtestp (f1 int);
+CREATE TABLE gtestc (f2 int GENERATED ALWAYS AS (f1+1) STORED) INHERITS(gtestp);
+INSERT INTO gtestc values(42);
+TABLE gtestc;
+ f1 | f2 
+----+----
+ 42 | 43
+(1 row)
+
+UPDATE gtestp SET f1 = f1 * 10;
+TABLE gtestc;
+ f1  | f2  
+-----+-----
+ 420 | 421
+(1 row)
+
+DROP TABLE gtestp CASCADE;
+NOTICE:  drop cascades to table gtestc
 -- test stored update
 CREATE TABLE gtest3 (a int, b int GENERATED ALWAYS AS (a * 3) STORED);
 INSERT INTO gtest3 (a) VALUES (1), (2), (3), (NULL);
diff --git a/src/test/regress/sql/generated.sql b/src/test/regress/sql/generated.sql
index 378297e6ea..39eec40bce 100644
--- a/src/test/regress/sql/generated.sql
+++ b/src/test/regress/sql/generated.sql
@@ -149,6 +149,15 @@ CREATE TABLE gtesty (x int, b int DEFAULT 55);
 CREATE TABLE gtest1_2 () INHERITS (gtest0, gtesty);  -- error
 DROP TABLE gtesty;
 
+-- test correct handling of GENERATED column that's only in child
+CREATE TABLE gtestp (f1 int);
+CREATE TABLE gtestc (f2 int GENERATED ALWAYS AS (f1+1) STORED) INHERITS(gtestp);
+INSERT INTO gtestc values(42);
+TABLE gtestc;
+UPDATE gtestp SET f1 = f1 * 10;
+TABLE gtestc;
+DROP TABLE gtestp CASCADE;
+
 -- test stored update
 CREATE TABLE gtest3 (a int, b int GENERATED ALWAYS AS (a * 3) STORED);
 INSERT INTO gtest3 (a) VALUES (1), (2), (3), (NULL);