v19-0004-Add-per-result-relation-extraUpdatedCols-to-Modi.patch

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

In spirit of removing things from RangeTblEntry that are better
carried by Query or PlannedStmt directly.
---
 src/backend/executor/execUtils.c         |  9 ++--
 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     | 17 +++++++
 src/backend/optimizer/prep/preptlist.c   | 49 ++++++++++++++++++
 src/backend/optimizer/util/inherit.c     | 64 +++++++++++++++---------
 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            |  3 ++
 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, 138 insertions(+), 88 deletions(-)

diff --git a/src/backend/executor/execUtils.c b/src/backend/executor/execUtils.c
index 461230b011..ede4c98875 100644
--- a/src/backend/executor/execUtils.c
+++ b/src/backend/executor/execUtils.c
@@ -1378,20 +1378,17 @@ 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);
 
 		if (relinfo->ri_RootToPartitionMap != NULL)
 			return execute_attr_map_cols(relinfo->ri_RootToPartitionMap->attrMap,
-										 rte->extraUpdatedCols);
+										 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 04454ad6e6..e034e3ceee 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -3971,6 +3971,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 d3beb907ea..139d8e095f 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -569,7 +569,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 75bf11c741..bcd380516d 100644
--- a/src/backend/nodes/readfuncs.c
+++ b/src/backend/nodes/readfuncs.c
@@ -532,7 +532,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 f854855951..5c1c7aed2f 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);
@@ -2824,6 +2824,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,
@@ -6980,7 +6981,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)
@@ -7049,6 +7050,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 9576b69f1a..8913200879 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -1746,6 +1746,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;
@@ -1779,15 +1780,24 @@ 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, resultRelation,
+															   extraUpdatedCols,
+															   this_result_rel->top_parent_relids);
+						}
 						updateColnosLists = lappend(updateColnosLists,
 													update_colnos);
+						extraUpdatedColsBitmaps = lappend(extraUpdatedColsBitmaps,
+														  extraUpdatedCols);
 					}
 					if (parse->withCheckOptions)
 					{
@@ -1869,7 +1879,10 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 					 */
 					resultRelations = list_make1_int(parse->resultRelation);
 					if (parse->commandType == CMD_UPDATE)
+					{
 						updateColnosLists = list_make1(root->update_colnos);
+						extraUpdatedColsBitmaps = list_make1(root->extraUpdatedCols);
+					}
 					if (parse->withCheckOptions)
 						withCheckOptionLists = list_make1(parse->withCheckOptions);
 					if (parse->returningList)
@@ -1883,7 +1896,10 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 				/* Single-relation INSERT/UPDATE/DELETE. */
 				resultRelations = list_make1_int(parse->resultRelation);
 				if (parse->commandType == CMD_UPDATE)
+				{
 					updateColnosLists = list_make1(root->update_colnos);
+					extraUpdatedColsBitmaps = list_make1(root->extraUpdatedCols);
+				}
 				if (parse->withCheckOptions)
 					withCheckOptionLists = list_make1(parse->withCheckOptions);
 				if (parse->returningList)
@@ -1922,6 +1938,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 e5c1103316..10a948ed40 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->rtepermlist, 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,43 @@ extract_update_targetlist_colnos(List *tlist)
 	return update_colnos;
 }
 
+/*
+ * Return the indexes of any generated columns that depend on any columns
+ * mentioned in target_perminfo->updatedCols.
+ */
+Bitmapset *
+get_extraUpdatedCols(Bitmapset *updatedCols, Relation target_relation)
+{
+	TupleDesc	tupdesc = RelationGetDescr(target_relation);
+	TupleConstr *constr = tupdesc->constr;
+	/* These go into ModifyTable.extraUpdatedColsBitmaps, so make as Nodes. */
+	Bitmapset *extraUpdatedCols = makeNode(Bitmapset);
+
+	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;
+}
+
 /*****************************************************************************
  *
  *		TARGETLIST EXPANSION
diff --git a/src/backend/optimizer/util/inherit.c b/src/backend/optimizer/util/inherit.c
index 10c2aa13f6..ea632a94a6 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,
@@ -148,6 +149,7 @@ expand_inherited_rtentry(PlannerInfo *root, RelOptInfo *rel,
 		expand_partitioned_rtentry(root, rel, rte, rti,
 								   oldrelation,
 								   root_perminfo->updatedCols,
+								   root->extraUpdatedCols,
 								   oldrc, lockmode);
 	}
 	else
@@ -313,6 +315,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;
@@ -343,7 +346,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)
@@ -412,14 +415,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);
 		}
 
@@ -456,7 +463,6 @@ expand_single_inheritance_child(PlannerInfo *root, RangeTblEntry *parentrte,
 								Index *childRTindex_p)
 {
 	Query	   *parse = root->parse;
-	Oid			parentOID = RelationGetRelid(parentrel);
 	Oid			childOID = RelationGetRelid(childrel);
 	RangeTblEntry *childrte;
 	Index		childRTindex;
@@ -486,7 +492,7 @@ expand_single_inheritance_child(PlannerInfo *root, RangeTblEntry *parentrte,
 	/* A partitioned child will need to be expanded further. */
 	if (childrte->relkind == RELKIND_PARTITIONED_TABLE)
 	{
-		Assert(childOID != parentOID);
+		Assert(childOID != RelationGetRelid(parentrel));
 		childrte->inh = true;
 	}
 	else
@@ -553,13 +559,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.
@@ -656,7 +655,8 @@ static Bitmapset *
 translate_col_privs(const Bitmapset *parent_privs,
 					List *translated_vars)
 {
-	Bitmapset  *child_privs = NULL;
+	/* These go into ModifyTable.extraUpdatedColsBitmaps, so make as Nodes. */
+	Bitmapset  *child_privs = makeNode(Bitmapset);
 	bool		whole_row;
 	int			attno;
 	ListCell   *lc;
@@ -866,13 +866,16 @@ apply_child_basequals(PlannerInfo *root, RelOptInfo *parentrel,
  * 		'top_parent_cols' to the columns numbers of a descendent relation
  * 		given by 'relid'
  */
-static Bitmapset *
-translate_col_privs_recurse(PlannerInfo *root, Index relid,
-							Bitmapset *top_parent_cols,
-							Relids top_parent_relids)
+Bitmapset *
+translate_col_privs_multilevel(PlannerInfo *root, Index relid,
+							   Bitmapset *top_parent_cols,
+							   Relids top_parent_relids)
 {
 	AppendRelInfo *appinfo;
 
+	if (top_parent_cols == NULL)
+		return NULL;
+
 	Assert(root->append_rel_array != NULL);
 	appinfo = root->append_rel_array[relid];
 	Assert(appinfo != NULL);
@@ -883,9 +886,9 @@ translate_col_privs_recurse(PlannerInfo *root, Index relid,
 	 * and child relation pairs.
 	 */
 	if (bms_singleton_member(top_parent_relids) != appinfo->parent_relid)
-		translate_col_privs_recurse(root, appinfo->parent_relid,
-									top_parent_cols,
-									top_parent_relids);
+		translate_col_privs_multilevel(root, appinfo->parent_relid,
+									   top_parent_cols,
+									   top_parent_relids);
 
 	return translate_col_privs(top_parent_cols, appinfo->translated_vars);
 }
@@ -903,6 +906,8 @@ GetRelAllUpdatedCols(PlannerInfo *root, RelOptInfo *rel)
 	Bitmapset *updatedCols,
 			  *extraUpdatedCols;
 
+	Assert(root->parse->commandType == CMD_UPDATE);
+
 	if (!IS_SIMPLE_REL(rel))
 		return NULL;
 
@@ -918,24 +923,33 @@ GetRelAllUpdatedCols(PlannerInfo *root, RelOptInfo *rel)
 	 * translate the columns found therein to match the given relation.
 	 */
 	if (rel->top_parent_relids != NULL)
+	{
 		rte =  planner_rt_fetch(bms_singleton_member(rel->top_parent_relids),
 								root);
+	}
 	else
+	{
+		Assert(rel->relid == root->parse->resultRelation);
 		rte = planner_rt_fetch(rel->relid, root);
+	}
 
 	Assert(rte->perminfoindex > 0);
 	perminfo = GetRTEPermissionInfo(root->parse->rtepermlist, rte);
 
 	if (rel->top_parent_relids != NULL)
-		updatedCols = translate_col_privs_recurse(root, rel->relid,
-												  perminfo->updatedCols,
-												  rel->top_parent_relids);
+	{
+		updatedCols = translate_col_privs_multilevel(root, rel->relid,
+													 perminfo->updatedCols,
+													 rel->top_parent_relids);
+		extraUpdatedCols = translate_col_privs_multilevel(root, rel->relid,
+														  root->extraUpdatedCols,
+														  rel->top_parent_relids);
+	}
 	else
+	{
 		updatedCols = perminfo->updatedCols;
-
-	/* extraUpdatedCols can be obtained directly from the RTE. */
-	rte = planner_rt_fetch(rel->relid, root);
-	extraUpdatedCols = rte->extraUpdatedCols;
+		extraUpdatedCols = root->extraUpdatedCols;
+	}
 
 	return bms_union(updatedCols, extraUpdatedCols);
 }
diff --git a/src/backend/optimizer/util/pathnode.c b/src/backend/optimizer/util/pathnode.c
index 70f61ae7b1..f70fe2736c 100644
--- a/src/backend/optimizer/util/pathnode.c
+++ b/src/backend/optimizer/util/pathnode.c
@@ -3643,6 +3643,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)
@@ -3658,6 +3660,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)
@@ -3722,6 +3725,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 e834130151..b75455382e 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"
@@ -1791,7 +1792,6 @@ apply_handle_update(StringInfo s)
 	LogicalRepTupleData newtup;
 	bool		has_oldtup;
 	TupleTableSlot *remoteslot;
-	RangeTblEntry *target_rte;
 	RTEPermissionInfo *target_perminfo;
 	MemoryContext oldctx;
 
@@ -1840,7 +1840,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_rtepermlist, 0);
 	for (int i = 0; i < remoteslot->tts_tupleDescriptor->natts; i++)
 	{
@@ -1858,7 +1857,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 156c033bda..d62d457fc0 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -1592,46 +1592,6 @@ rewriteValuesRTE(Query *parsetree, RangeTblEntry *rte, int rti,
 }
 
 
-/*
- * 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
@@ -3670,7 +3630,6 @@ RewriteQuery(Query *parsetree, List *rewrite_events)
 	{
 		int			result_relation;
 		RangeTblEntry *rt_entry;
-		RTEPermissionInfo *rt_perminfo;
 		Relation	rt_entry_relation;
 		List	   *locks;
 		List	   *product_queries;
@@ -3682,7 +3641,6 @@ RewriteQuery(Query *parsetree, List *rewrite_events)
 		Assert(result_relation != 0);
 		rt_entry = rt_fetch(result_relation, parsetree->rtable);
 		Assert(rt_entry->rtekind == RTE_RELATION);
-		rt_perminfo = GetRTEPermissionInfo(parsetree->rtepermlist, rt_entry);
 
 		/*
 		 * We can use NoLock here since either the parser or
@@ -3767,9 +3725,6 @@ RewriteQuery(Query *parsetree, List *rewrite_events)
 									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 c32834a9e8..a85570b1de 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 080680ecd0..118a150d3c 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -1152,7 +1152,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 99c8d4f611..04c7403897 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -422,6 +422,8 @@ struct PlannerInfo
 	 */
 	List	   *update_colnos;
 
+	Bitmapset  *extraUpdatedCols;
+
 	/*
 	 * Fields filled during create_plan() for use in setrefs.c
 	 */
@@ -2250,6 +2252,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 aaff24256e..2f8c9f65cc 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -237,6 +237,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 9a4f86920c..3cc54ef7e3 100644
--- a/src/include/optimizer/inherit.h
+++ b/src/include/optimizer/inherit.h
@@ -24,5 +24,8 @@ extern bool apply_child_basequals(PlannerInfo *root, RelOptInfo *parentrel,
 								  RelOptInfo *childrel, RangeTblEntry *childRTE,
 								  AppendRelInfo *appinfo);
 extern Bitmapset *GetRelAllUpdatedCols(PlannerInfo *root, RelOptInfo *rel);
+extern Bitmapset *translate_col_privs_multilevel(PlannerInfo *root, Index relid,
+							   Bitmapset *top_parent_cols,
+							   Relids top_parent_relids);
 
 #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