From 14a1198bdaad007b1dc835f24caa42d3667c7048 Mon Sep 17 00:00:00 2001
From: amitlan <amitlangote09@gmail.com>
Date: Tue, 13 Dec 2022 11:58:07 +0900
Subject: [PATCH v29 1/2] Preparatory refactoring before reworking CachedPlan
 locking

Remember the RT indexes of RTEs that AcquireExecutorLocks() must
look at to consider locking in a bitmapset, so that nstead of looping
over the range table to find those RTEs, it can look them up using
the RT indexes set in the bitmapset.

This also adds some extra information related to execution-time
pruning to the relevant plan nodes.
---
 src/backend/executor/execParallel.c  |  1 +
 src/backend/executor/execPartition.c |  6 ++++
 src/backend/nodes/readfuncs.c        |  8 ++++--
 src/backend/optimizer/plan/planner.c |  2 ++
 src/backend/optimizer/plan/setrefs.c | 12 ++++++++
 src/backend/partitioning/partprune.c | 42 ++++++++++++++++++++++++++--
 src/backend/utils/cache/plancache.c  | 10 +++++--
 src/include/executor/execPartition.h |  2 ++
 src/include/nodes/nodes.h            |  1 +
 src/include/nodes/pathnodes.h        | 11 ++++++++
 src/include/nodes/plannodes.h        | 19 +++++++++++++
 11 files changed, 106 insertions(+), 8 deletions(-)

diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c
index a5b8e43ec5..65c4b63bbd 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -182,6 +182,7 @@ ExecSerializePlan(Plan *plan, EState *estate)
 	pstmt->transientPlan = false;
 	pstmt->dependsOnRole = false;
 	pstmt->parallelModeNeeded = false;
+	pstmt->containsInitialPruning = false;	/* workers need not know! */
 	pstmt->planTree = plan;
 	pstmt->partPruneInfos = estate->es_part_prune_infos;
 	pstmt->rtable = estate->es_range_table;
diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index 76d79b9741..5b62157712 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -1956,6 +1956,7 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 			Assert(partdesc->nparts >= pinfo->nparts);
 			pprune->nparts = partdesc->nparts;
 			pprune->subplan_map = palloc(sizeof(int) * partdesc->nparts);
+			pprune->rti_map = palloc(sizeof(Index) * partdesc->nparts);
 			if (partdesc->nparts == pinfo->nparts)
 			{
 				/*
@@ -1966,6 +1967,8 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 				pprune->subpart_map = pinfo->subpart_map;
 				memcpy(pprune->subplan_map, pinfo->subplan_map,
 					   sizeof(int) * pinfo->nparts);
+				memcpy(pprune->rti_map, pinfo->rti_map,
+					   sizeof(int) * pinfo->nparts);
 
 				/*
 				 * Double-check that the list of unpruned relations has not
@@ -2016,6 +2019,8 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 							pinfo->subplan_map[pd_idx];
 						pprune->subpart_map[pp_idx] =
 							pinfo->subpart_map[pd_idx];
+						pprune->rti_map[pp_idx] =
+							pinfo->rti_map[pd_idx];
 						pd_idx++;
 					}
 					else
@@ -2023,6 +2028,7 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 						/* this partdesc entry is not in the plan */
 						pprune->subplan_map[pp_idx] = -1;
 						pprune->subpart_map[pp_idx] = -1;
+						pprune->rti_map[pp_idx] = 0;
 					}
 				}
 
diff --git a/src/backend/nodes/readfuncs.c b/src/backend/nodes/readfuncs.c
index 966b75f5a6..1161671fa4 100644
--- a/src/backend/nodes/readfuncs.c
+++ b/src/backend/nodes/readfuncs.c
@@ -158,6 +158,11 @@
 	token = pg_strtok(&length);		/* skip :fldname */ \
 	local_node->fldname = readIntCols(len)
 
+/* Read an Index array */
+#define READ_INDEX_ARRAY(fldname, len) \
+	token = pg_strtok(&length);		/* skip :fldname */ \
+	local_node->fldname = readIndexCols(len)
+
 /* Read a bool array */
 #define READ_BOOL_ARRAY(fldname, len) \
 	token = pg_strtok(&length);		/* skip :fldname */ \
@@ -796,7 +801,6 @@ fnname(int numCols) \
  */
 READ_SCALAR_ARRAY(readAttrNumberCols, int16, atoi)
 READ_SCALAR_ARRAY(readOidCols, Oid, atooid)
-/* outfuncs.c has writeIndexCols, but we don't yet need that here */
-/* READ_SCALAR_ARRAY(readIndexCols, Index, atoui) */
+READ_SCALAR_ARRAY(readIndexCols, Index, atoui)
 READ_SCALAR_ARRAY(readIntCols, int, atoi)
 READ_SCALAR_ARRAY(readBoolCols, bool, strtobool)
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 5dd4f92720..620b163ef9 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -523,8 +523,10 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 	result->parallelModeNeeded = glob->parallelModeNeeded;
 	result->planTree = top_plan;
 	result->partPruneInfos = glob->partPruneInfos;
+	result->containsInitialPruning = glob->containsInitialPruning;
 	result->rtable = glob->finalrtable;
 	result->permInfos = glob->finalrteperminfos;
+	result->minLockRelids = glob->minLockRelids;
 	result->resultRelations = glob->resultRelations;
 	result->appendRelations = glob->appendRelations;
 	result->subplans = glob->subplans;
diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index 596f1fbc8e..ed43d5936d 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -279,6 +279,16 @@ set_plan_references(PlannerInfo *root, Plan *plan)
 	 */
 	add_rtes_to_flat_rtable(root, false);
 
+	/*
+	 * Add the query's adjusted range of RT indexes to glob->minLockRelids.
+	 * The adjusted RT indexes of prunable relations will be deleted from the
+	 * set below where PartitionPruneInfos are processed.
+	 */
+	glob->minLockRelids =
+		bms_add_range(glob->minLockRelids,
+					  rtoffset + 1,
+					  rtoffset + list_length(root->parse->rtable));
+
 	/*
 	 * Adjust RT indexes of PlanRowMarks and add to final rowmarks list
 	 */
@@ -377,9 +387,11 @@ set_plan_references(PlannerInfo *root, Plan *plan)
 				/* RT index of the table to which the pinfo belongs. */
 				pinfo->rtindex += rtoffset;
 			}
+
 		}
 
 		glob->partPruneInfos = lappend(glob->partPruneInfos, pruneinfo);
+		glob->containsInitialPruning |= pruneinfo->needs_init_pruning;
 	}
 
 	return result;
diff --git a/src/backend/partitioning/partprune.c b/src/backend/partitioning/partprune.c
index d48f6784c1..56270d7670 100644
--- a/src/backend/partitioning/partprune.c
+++ b/src/backend/partitioning/partprune.c
@@ -144,7 +144,9 @@ static List *make_partitionedrel_pruneinfo(PlannerInfo *root,
 										   List *prunequal,
 										   Bitmapset *partrelids,
 										   int *relid_subplan_map,
-										   Bitmapset **matchedsubplans);
+										   Bitmapset **matchedsubplans,
+										   bool *needs_init_pruning,
+										   bool *needs_exec_pruning);
 static void gen_partprune_steps(RelOptInfo *rel, List *clauses,
 								PartClauseTarget target,
 								GeneratePruningStepsContext *context);
@@ -234,6 +236,8 @@ make_partition_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 	int		   *relid_subplan_map;
 	ListCell   *lc;
 	int			i;
+	bool		needs_init_pruning = false;
+	bool		needs_exec_pruning = false;
 
 	/*
 	 * Scan the subpaths to see which ones are scans of partition child
@@ -313,12 +317,16 @@ make_partition_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 		Bitmapset  *partrelids = (Bitmapset *) lfirst(lc);
 		List	   *pinfolist;
 		Bitmapset  *matchedsubplans = NULL;
+		bool		partrel_needs_init_pruning;
+		bool		partrel_needs_exec_pruning;
 
 		pinfolist = make_partitionedrel_pruneinfo(root, parentrel,
 												  prunequal,
 												  partrelids,
 												  relid_subplan_map,
-												  &matchedsubplans);
+												  &matchedsubplans,
+												  &partrel_needs_init_pruning,
+												  &partrel_needs_exec_pruning);
 
 		/* When pruning is possible, record the matched subplans */
 		if (pinfolist != NIL)
@@ -327,6 +335,9 @@ make_partition_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 			allmatchedsubplans = bms_join(matchedsubplans,
 										  allmatchedsubplans);
 		}
+
+		needs_init_pruning |= partrel_needs_init_pruning;
+		needs_exec_pruning |= partrel_needs_exec_pruning;
 	}
 
 	pfree(relid_subplan_map);
@@ -342,6 +353,8 @@ make_partition_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 	pruneinfo = makeNode(PartitionPruneInfo);
 	pruneinfo->root_parent_relids = parentrel->relids;
 	pruneinfo->prune_infos = prunerelinfos;
+	pruneinfo->needs_init_pruning = needs_init_pruning;
+	pruneinfo->needs_exec_pruning = needs_exec_pruning;
 
 	/*
 	 * Some subplans may not belong to any of the identified partitioned rels.
@@ -442,13 +455,19 @@ add_part_relids(List *allpartrelids, Bitmapset *partrelids)
  * If we cannot find any useful run-time pruning steps, return NIL.
  * However, on success, each rel identified in partrelids will have
  * an element in the result list, even if some of them are useless.
+ * *needs_init_pruning and *needs_exec_pruning are set to indicate whether
+ * the pruning steps contained in the returned PartitionedRelPruneInfos
+ * can be performed during executor startup and during execution,
+ * respectively.
  */
 static List *
 make_partitionedrel_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 							  List *prunequal,
 							  Bitmapset *partrelids,
 							  int *relid_subplan_map,
-							  Bitmapset **matchedsubplans)
+							  Bitmapset **matchedsubplans,
+							  bool *needs_init_pruning,
+							  bool *needs_exec_pruning)
 {
 	RelOptInfo *targetpart = NULL;
 	List	   *pinfolist = NIL;
@@ -459,6 +478,10 @@ make_partitionedrel_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 	int			rti;
 	int			i;
 
+	/* Will find out below. */
+	*needs_init_pruning = false;
+	*needs_exec_pruning = false;
+
 	/*
 	 * Examine each partitioned rel, constructing a temporary array to map
 	 * from planner relids to index of the partitioned rel, and building a
@@ -546,6 +569,9 @@ make_partitionedrel_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 		 * executor per-scan pruning steps.  This first pass creates startup
 		 * pruning steps and detects whether there's any possibly-useful quals
 		 * that would require per-scan pruning.
+		 *
+		 * In the first pass, we note whether the 2nd pass is necessary by
+		 * noting the presence of EXEC parameters.
 		 */
 		gen_partprune_steps(subpart, partprunequal, PARTTARGET_INITIAL,
 							&context);
@@ -620,6 +646,12 @@ make_partitionedrel_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 		pinfo->execparamids = execparamids;
 		/* Remaining fields will be filled in the next loop */
 
+		/* record which types of pruning steps we've seen so far */
+		if (initial_pruning_steps != NIL)
+			*needs_init_pruning = true;
+		if (exec_pruning_steps != NIL)
+			*needs_exec_pruning = true;
+
 		pinfolist = lappend(pinfolist, pinfo);
 	}
 
@@ -647,6 +679,7 @@ make_partitionedrel_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 		int		   *subplan_map;
 		int		   *subpart_map;
 		Oid		   *relid_map;
+		Index	   *rti_map;
 
 		/*
 		 * Construct the subplan and subpart maps for this partitioning level.
@@ -659,6 +692,7 @@ make_partitionedrel_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 		subpart_map = (int *) palloc(nparts * sizeof(int));
 		memset(subpart_map, -1, nparts * sizeof(int));
 		relid_map = (Oid *) palloc0(nparts * sizeof(Oid));
+		rti_map = (Index *) palloc0(nparts * sizeof(Index));
 		present_parts = NULL;
 
 		i = -1;
@@ -673,6 +707,7 @@ make_partitionedrel_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 			subplan_map[i] = subplanidx = relid_subplan_map[partrel->relid] - 1;
 			subpart_map[i] = subpartidx = relid_subpart_map[partrel->relid] - 1;
 			relid_map[i] = planner_rt_fetch(partrel->relid, root)->relid;
+			rti_map[i] = partrel->relid;
 			if (subplanidx >= 0)
 			{
 				present_parts = bms_add_member(present_parts, i);
@@ -697,6 +732,7 @@ make_partitionedrel_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 		pinfo->subplan_map = subplan_map;
 		pinfo->subpart_map = subpart_map;
 		pinfo->relid_map = relid_map;
+		pinfo->rti_map = rti_map;
 	}
 
 	pfree(relid_subpart_map);
diff --git a/src/backend/utils/cache/plancache.c b/src/backend/utils/cache/plancache.c
index cc943205d3..339bb603f7 100644
--- a/src/backend/utils/cache/plancache.c
+++ b/src/backend/utils/cache/plancache.c
@@ -1747,7 +1747,8 @@ AcquireExecutorLocks(List *stmt_list, bool acquire)
 	foreach(lc1, stmt_list)
 	{
 		PlannedStmt *plannedstmt = lfirst_node(PlannedStmt, lc1);
-		ListCell   *lc2;
+		Bitmapset  *allLockRelids;
+		int			rti;
 
 		if (plannedstmt->commandType == CMD_UTILITY)
 		{
@@ -1760,14 +1761,17 @@ AcquireExecutorLocks(List *stmt_list, bool acquire)
 			 */
 			Query	   *query = UtilityContainsQuery(plannedstmt->utilityStmt);
 
+			Assert(plannedstmt->minLockRelids == NULL);
 			if (query)
 				ScanQueryForLocks(query, acquire);
 			continue;
 		}
 
-		foreach(lc2, plannedstmt->rtable)
+		allLockRelids = plannedstmt->minLockRelids;
+		rti = -1;
+		while ((rti = bms_next_member(allLockRelids, rti)) > 0)
 		{
-			RangeTblEntry *rte = (RangeTblEntry *) lfirst(lc2);
+			RangeTblEntry *rte = rt_fetch(rti, plannedstmt->rtable);
 
 			if (rte->rtekind != RTE_RELATION)
 				continue;
diff --git a/src/include/executor/execPartition.h b/src/include/executor/execPartition.h
index 17fabc18c9..aeeaeb7884 100644
--- a/src/include/executor/execPartition.h
+++ b/src/include/executor/execPartition.h
@@ -45,6 +45,7 @@ extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
  * nparts						Length of subplan_map[] and subpart_map[].
  * subplan_map					Subplan index by partition index, or -1.
  * subpart_map					Subpart index by partition index, or -1.
+ * rti_map						Range table index by partition index, or 0.
  * present_parts				A Bitmapset of the partition indexes that we
  *								have subplans or subparts for.
  * initial_pruning_steps		List of PartitionPruneSteps used to
@@ -61,6 +62,7 @@ typedef struct PartitionedRelPruningData
 	int			nparts;
 	int		   *subplan_map;
 	int		   *subpart_map;
+	Index	   *rti_map;
 	Bitmapset  *present_parts;
 	List	   *initial_pruning_steps;
 	List	   *exec_pruning_steps;
diff --git a/src/include/nodes/nodes.h b/src/include/nodes/nodes.h
index 1f33902947..c2f2544df5 100644
--- a/src/include/nodes/nodes.h
+++ b/src/include/nodes/nodes.h
@@ -218,6 +218,7 @@ extern struct Bitmapset *readBitmapset(void);
 extern uintptr_t readDatum(bool typbyval);
 extern bool *readBoolCols(int numCols);
 extern int *readIntCols(int numCols);
+extern Index *readIndexCols(int numCols);
 extern Oid *readOidCols(int numCols);
 extern int16 *readAttrNumberCols(int numCols);
 
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index 654dba61aa..4337e7aa34 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -128,6 +128,17 @@ typedef struct PlannerGlobal
 	/* List of PartitionPruneInfo contained in the plan */
 	List	   *partPruneInfos;
 
+	/*
+	 * Do any of those PartitionPruneInfos have initial pruning steps in them?
+	 */
+	bool		containsInitialPruning;
+
+	/*
+	 * Indexes of all range table entries; for AcquireExecutorLocks()'s
+	 * perusal.
+	 */
+	Bitmapset  *minLockRelids;
+
 	/* OIDs of relations the plan depends on */
 	List	   *relationOids;
 
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index bddfe86191..eb0a007946 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -73,11 +73,18 @@ typedef struct PlannedStmt
 	List	   *partPruneInfos; /* List of PartitionPruneInfo contained in the
 								 * plan */
 
+	bool		containsInitialPruning;	/* Do any of those PartitionPruneInfos
+										 * have initial pruning steps in them?
+										 */
+
 	List	   *rtable;			/* list of RangeTblEntry nodes */
 
 	List	   *permInfos;		/* list of RTEPermissionInfo nodes for rtable
 								 * entries needing one */
 
+	Bitmapset  *minLockRelids;	/* Indexes of all range table entries; for
+								 * AcquireExecutorLocks()'s perusal */
+
 	/* rtable indexes of target relations for INSERT/UPDATE/DELETE/MERGE */
 	List	   *resultRelations;	/* integer list of RT indexes, or NIL */
 
@@ -1417,6 +1424,13 @@ typedef struct PlanRowMark
  * prune_infos			List of Lists containing PartitionedRelPruneInfo nodes,
  *						one sublist per run-time-prunable partition hierarchy
  *						appearing in the parent plan node's subplans.
+ *
+ * needs_init_pruning	Does any of the PartitionedRelPruneInfos in
+ *						prune_infos have its initial_pruning_steps set?
+ *
+ * needs_exec_pruning	Does any of the PartitionedRelPruneInfos in
+ *						prune_infos have its exec_pruning_steps set?
+ *
  * other_subplans		Indexes of any subplans that are not accounted for
  *						by any of the PartitionedRelPruneInfo nodes in
  *						"prune_infos".  These subplans must not be pruned.
@@ -1428,6 +1442,8 @@ typedef struct PartitionPruneInfo
 	NodeTag		type;
 	Bitmapset  *root_parent_relids;
 	List	   *prune_infos;
+	bool		needs_init_pruning;
+	bool		needs_exec_pruning;
 	Bitmapset  *other_subplans;
 } PartitionPruneInfo;
 
@@ -1472,6 +1488,9 @@ typedef struct PartitionedRelPruneInfo
 	/* relation OID by partition index, or 0 */
 	Oid		   *relid_map pg_node_attr(array_size(nparts));
 
+	/* Range table index by partition index, or 0. */
+	Index	   *rti_map pg_node_attr(array_size(nparts));
+
 	/*
 	 * initial_pruning_steps shows how to prune during executor startup (i.e.,
 	 * without use of any PARAM_EXEC Params); it is NIL if no startup pruning
-- 
2.35.3

