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

In spirit of removing things from RangeTblEntry that are better
carried by Query or PlannedStmt directly.

Instead of the rewriter populating extraUpdatedCols bitmapset in the
target relation's RTE, it is now planner's responsibility to compute
the bitmapset and pass it to the executor using the ModifyTable node.
In the inherited UPDATE's case, a copy is added for each child result
relation adjusted to account for any attribute number differences.
---
 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..25df4eb8e4 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->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 ea56ff79c8..8a8088e9a1 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
@@ -3737,7 +3697,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;
@@ -3750,7 +3709,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
@@ -3842,9 +3800,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

