From 1deb363d5d0b7573d116198798a3e550be9a320f Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Thu, 12 Sep 2024 15:44:43 +0900
Subject: [PATCH v54 2/4] Perform runtime initial pruning outside
 ExecInitNode()

This commit follows up on the previous change that moved
PartitionPruneInfos out of individual plan nodes into a list in
PlannedStmt. It moves the initialization of PartitionPruneStates
and runtime initial pruning out of ExecInitNode() and into a new
routine, ExecDoInitialPruning(), which is called by InitPlan()
before ExecInitNode() is invoked on the main plan tree and subplans.

ExecDoInitialPruning() stores the PartitionPruneStates in a list
matching the length of es_part_prune_infos (which holds the
PartitionPruneInfos from PlannedStmt), allowing both lists to share
the same index. It also saves the initial pruning result -- a
bitmapset of indexes for surviving child subnodes -- in a similarly
indexed list.

While the initial pruning is done earlier, the execution pruning
context information (needed for runtime pruning) is initialized
later during ExecInitNode() for the parent plan node, as it requires
access to the parent node's PlanState struct.
---
 src/backend/executor/execMain.c      |  55 ++++++++
 src/backend/executor/execPartition.c | 179 +++++++++++++++++++++------
 src/include/executor/execPartition.h |   6 +
 src/include/nodes/execnodes.h        |   2 +
 4 files changed, 202 insertions(+), 40 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index e6197c165e..8fab8dbccd 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -46,6 +46,7 @@
 #include "commands/matview.h"
 #include "commands/trigger.h"
 #include "executor/executor.h"
+#include "executor/execPartition.h"
 #include "executor/nodeSubplan.h"
 #include "foreign/fdwapi.h"
 #include "mb/pg_wchar.h"
@@ -818,6 +819,54 @@ ExecCheckXactReadOnly(PlannedStmt *plannedstmt)
 		PreventCommandIfParallelMode(CreateCommandName((Node *) plannedstmt));
 }
 
+/*
+ * ExecDoInitialPruning
+ *		Perform runtime "initial" pruning, if necessary, to determine the set
+ *		of child subnodes that need to be initialized during ExecInitNode()
+ *		for plan nodes that support partition pruning.
+ *
+ * For each PartitionPruneInfo in estate->es_part_prune_infos, this function
+ * creates a PartitionPruneState (even if no initial pruning is done) and adds
+ * it to es_part_prune_states. For PartitionPruneInfo entries that include
+ * initial pruning steps, the result of those steps is saved as a bitmapset
+ * of indexes representing child subnodes that are "valid" and should be
+ * initialized for execution.
+ */
+static void
+ExecDoInitialPruning(EState *estate)
+{
+	ListCell *lc;
+
+	foreach(lc, estate->es_part_prune_infos)
+	{
+		PartitionPruneInfo *pruneinfo = lfirst_node(PartitionPruneInfo, lc);
+		PartitionPruneState *prunestate;
+		Bitmapset *validsubplans = NULL;
+
+		/*
+		 * Create the working data structure for pruning, and save it for use
+		 * later in ExecInitPartitionPruning(), which will be called by the
+		 * parent plan node's ExecInit* function.
+		 */
+		prunestate = CreatePartitionPruneState(estate, pruneinfo);
+		estate->es_part_prune_states = lappend(estate->es_part_prune_states,
+											   prunestate);
+
+		/*
+		 * Perform an initial partition pruning pass, if necessary, and save
+		 * the bitmapset of valid subplans for use in
+		 * ExecInitPartitionPruning().  If no initial pruning is performed, we
+		 * still store a NULL to ensure that es_part_prune_results is the same
+		 * length as es_part_prune_infos.  This ensures that
+		 * ExecInitPartitionPruning() can use the same index to locate the
+		 * result.
+		 */
+		if (prunestate->do_initial_prune)
+			validsubplans = ExecFindMatchingSubPlans(prunestate, true);
+		estate->es_part_prune_results = lappend(estate->es_part_prune_results,
+												validsubplans);
+	}
+}
 
 /* ----------------------------------------------------------------
  *		InitPlan
@@ -850,7 +899,13 @@ InitPlan(QueryDesc *queryDesc, int eflags)
 	ExecInitRangeTable(estate, rangeTable, plannedstmt->permInfos);
 
 	estate->es_plannedstmt = plannedstmt;
+
+	/*
+	 * Perform runtime "initial" pruning to determine the plan nodes that will
+	 * not be executed.
+	 */
 	estate->es_part_prune_infos = plannedstmt->partPruneInfos;
+	ExecDoInitialPruning(estate);
 
 	/*
 	 * Next, build the ExecRowMark array from the PlanRowMark(s), if any.
diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index ec730674f2..d205e64e84 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -181,8 +181,6 @@ static char *ExecBuildSlotPartitionKeyDescription(Relation rel,
 												  int maxfieldlen);
 static List *adjust_partition_colnos(List *colnos, ResultRelInfo *leaf_part_rri);
 static List *adjust_partition_colnos_using_map(List *colnos, AttrMap *attrMap);
-static PartitionPruneState *CreatePartitionPruneState(PlanState *planstate,
-													  PartitionPruneInfo *pruneinfo);
 static void InitPartitionPruneContext(PartitionPruneContext *context,
 									  List *pruning_steps,
 									  PartitionDesc partdesc,
@@ -192,6 +190,9 @@ static void InitPartitionPruneContext(PartitionPruneContext *context,
 static void PartitionPruneFixSubPlanMap(PartitionPruneState *prunestate,
 										Bitmapset *initially_valid_subplans,
 										int n_total_subplans);
+static void PartitionPruneInitExecPruning(PartitionPruneInfo *pruneinfo,
+										  PartitionPruneState *prunestate,
+										  PlanState *planstate);
 static void find_matching_subplans_recurse(PartitionPruningData *prunedata,
 										   PartitionedRelPruningData *pprune,
 										   bool initial_prune,
@@ -1783,20 +1784,26 @@ adjust_partition_colnos_using_map(List *colnos, AttrMap *attrMap)
 
 /*
  * ExecInitPartitionPruning
- *		Initialize data structure needed for run-time partition pruning and
- *		do initial pruning if needed
+ *		Initialize the data structures needed for runtime "exec" partition
+ *		pruning and return the result of initial pruning, if available.
  *
  * 'root_parent_relids' identifies the relation to which both the parent plan
- * and the PartitionPruneInfo given by 'part_prune_index' belong.
+ * and the PartitionPruneInfo associated with 'part_prune_index' belong.
  *
- * On return, *initially_valid_subplans is assigned the set of indexes of
- * child subplans that must be initialized along with the parent plan node.
- * Initial pruning is performed here if needed and in that case only the
- * surviving subplans' indexes are added.
+ * The PartitionPruneState would have been created by ExecDoInitialPruning()
+ * and stored as the part_prune_index'th element of EState.es_part_prune_states.
+ * Here, we initialize only the PartitionPruneContext necessary for execution
+ * pruning.
  *
- * If subplans are indeed pruned, subplan_map arrays contained in the returned
- * PartitionPruneState are re-sequenced to not count those, though only if the
- * maps will be needed for subsequent execution pruning passes.
+ * On return, *initially_valid_subplans is assigned the set of indexes of child
+ * subplans that must be initialized alongside the parent plan node. Initial
+ * pruning would have been performed by ExecDoInitialPruning() if necessary,
+ * and the bitmapset of surviving subplans' indexes would have been stored as
+ * the part_prune_index'th element of EState.es_part_prune_results.
+ *
+ * If subplans are pruned, the subplan_map arrays in the returned
+ * PartitionPruneState are re-sequenced to exclude those subplans, but only if
+ * the maps will be needed for subsequent execution pruning passes.
  */
 PartitionPruneState *
 ExecInitPartitionPruning(PlanState *planstate,
@@ -1821,17 +1828,21 @@ ExecInitPartitionPruning(PlanState *planstate,
 								   bmsToString(root_parent_relids),
 								   bmsToString(pruneinfo->root_parent_relids)));
 
-	/* We may need an expression context to evaluate partition exprs */
-	ExecAssignExprContext(estate, planstate);
-
-	/* Create the working data structure for pruning */
-	prunestate = CreatePartitionPruneState(planstate, pruneinfo);
-
 	/*
-	 * Perform an initial partition prune pass, if required.
+	 * ExecDoInitialPruning() must have initialized the PartitionPruneState to
+	 * perform the initial pruning.  Now we simply need to initialize the
+	 * context information for exec pruning.
 	 */
+	prunestate = list_nth(estate->es_part_prune_states, part_prune_index);
+	Assert(prunestate != NULL);
+	if (prunestate->do_exec_prune)
+		PartitionPruneInitExecPruning(pruneinfo, prunestate, planstate);
+
+	/* Use the result of initial pruning done by ExecDoInitialPruning(). */
 	if (prunestate->do_initial_prune)
-		*initially_valid_subplans = ExecFindMatchingSubPlans(prunestate, true);
+		*initially_valid_subplans = list_nth_node(Bitmapset,
+												  estate->es_part_prune_results,
+												  part_prune_index);
 	else
 	{
 		/* No pruning, so we'll need to initialize all subplans */
@@ -1877,16 +1888,23 @@ ExecInitPartitionPruning(PlanState *planstate,
  * stored in each PartitionedRelPruningData can be re-used each time we
  * re-evaluate which partitions match the pruning steps provided in each
  * PartitionedRelPruneInfo.
+ *
+ * Note that we only initialize the PartitionPruneContext (which is placed into
+ * each PartitionedRelPruningData) for initial pruning here. Execution pruning
+ * requires access to the parent plan node's PlanState, which is not available
+ * when this function is called from ExecDoInitialPruning(), so it is
+ * initialized later during ExecInitPartitionPruning() by calling
+ * PartitionPruneInitExecPruning().
  */
-static PartitionPruneState *
-CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
+PartitionPruneState *
+CreatePartitionPruneState(EState *estate, PartitionPruneInfo *pruneinfo)
 {
-	EState	   *estate = planstate->state;
 	PartitionPruneState *prunestate;
 	int			n_part_hierarchies;
 	ListCell   *lc;
 	int			i;
-	ExprContext *econtext = planstate->ps_ExprContext;
+	/* We may need an expression context to evaluate partition exprs */
+	ExprContext *econtext = CreateExprContext(estate);
 
 	/* For data reading, executor always includes detached partitions */
 	if (estate->es_partition_directory == NULL)
@@ -1974,6 +1992,7 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 			 * set to -1, as if they were pruned.  By construction, both
 			 * arrays are in partition bounds order.
 			 */
+			pprune->partrel = partrel;
 			pprune->nparts = partdesc->nparts;
 			pprune->subplan_map = palloc(sizeof(int) * partdesc->nparts);
 
@@ -2073,29 +2092,31 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 			{
 				InitPartitionPruneContext(&pprune->initial_context,
 										  pinfo->initial_pruning_steps,
-										  partdesc, partkey, planstate,
+										  partdesc, partkey, NULL,
 										  econtext);
 				/* Record whether initial pruning is needed at any level */
 				prunestate->do_initial_prune = true;
 			}
-			pprune->exec_pruning_steps = pinfo->exec_pruning_steps;
-			if (pinfo->exec_pruning_steps &&
-				!(econtext->ecxt_estate->es_top_eflags & EXEC_FLAG_EXPLAIN_GENERIC))
-			{
-				InitPartitionPruneContext(&pprune->exec_context,
-										  pinfo->exec_pruning_steps,
-										  partdesc, partkey, planstate,
-										  econtext);
-				/* Record whether exec pruning is needed at any level */
-				prunestate->do_exec_prune = true;
-			}
 
 			/*
-			 * Accumulate the IDs of all PARAM_EXEC Params affecting the
-			 * partitioning decisions at this plan node.
+			 * The exec pruning context will be initialized in
+			 * ExecInitPartitionPruning() when called during the initialization
+			 * of the parent plan node.
+			 *
+			 * pprune->exec_pruning_steps is set to NIL to prevent
+			 * ExecFindMatchingSubPlans() from accessing an uninitialized
+			 * pprune->exec_context during the initial pruning by
+			 * ExecDoInitialPruning().
+			 *
+			 * prunestate->do_exec_prune is set to indicate whether
+			 * PartitionPruneInitExecPruning() needs to be called by
+			 * ExecInitPartitionPruning(). This optimization avoids
+			 * unnecessary cycles when only initial pruning is required.
 			 */
-			prunestate->execparamids = bms_add_members(prunestate->execparamids,
-													   pinfo->execparamids);
+			pprune->exec_pruning_steps = NIL;
+			if (pinfo->exec_pruning_steps &&
+				 !(econtext->ecxt_estate->es_top_eflags & EXEC_FLAG_EXPLAIN_GENERIC))
+				prunestate->do_exec_prune = true;
 
 			j++;
 		}
@@ -2305,6 +2326,84 @@ PartitionPruneFixSubPlanMap(PartitionPruneState *prunestate,
 	pfree(new_subplan_indexes);
 }
 
+/*
+ * PartitionPruneInitExecPruning
+ *		Initialize PartitionPruneState for exec pruning.
+ */
+static void
+PartitionPruneInitExecPruning(PartitionPruneInfo *pruneinfo,
+							  PartitionPruneState *prunestate,
+							  PlanState *planstate)
+{
+	EState *estate = planstate->state;
+	int		i;
+	ExprContext *econtext;
+
+	/* CreatePartitionPruneState() must have initialized. */
+	Assert(estate->es_partition_directory != NULL);
+
+	/* CreatePartitionPruneState() must have set this. */
+	Assert(prunestate->do_exec_prune);
+
+	/*
+	 * Create ExprContext if not already done for the planstate.  We may need
+	 * an expression context to evaluate partition exprs.
+	 */
+	ExecAssignExprContext(estate, planstate);
+	econtext = planstate->ps_ExprContext;
+	for (i = 0; i < prunestate->num_partprunedata; i++)
+	{
+		List *partrel_pruneinfos =
+			list_nth_node(List, pruneinfo->prune_infos, i);
+		PartitionPruningData *prunedata = prunestate->partprunedata[i];
+		int		j;
+
+		for (j = 0; j < prunedata->num_partrelprunedata; j++)
+		{
+			PartitionedRelPruneInfo *pinfo =
+				list_nth_node(PartitionedRelPruneInfo, partrel_pruneinfos, j);
+			PartitionedRelPruningData *pprune = &prunedata->partrelprunedata[j];
+			Relation	partrel = pprune->partrel;
+			PartitionDesc partdesc;
+			PartitionKey partkey;
+
+			/*
+			 * Nothing to do if there are no exec pruning steps, but do set
+			 * pprune->exec_pruning_steps, becasue
+			 * find_matching_subplans_recurse() looks at it.
+			 *
+			 * Also skip if doing EXPLAIN (GENERIC_PLAN), since parameter
+			 * values may be missing.
+			 */
+			pprune->exec_pruning_steps = pinfo->exec_pruning_steps;
+			if (pprune->exec_pruning_steps == NIL ||
+				(econtext->ecxt_estate->es_top_eflags & EXEC_FLAG_EXPLAIN_GENERIC))
+				continue;
+
+			/*
+			 * We can rely on the copies of the partitioned table's partition
+			 * key and partition descriptor appearing in its relcache entry,
+			 * because that entry will be held open and locked for the
+			 * duration of this executor run.
+			 */
+			partkey = RelationGetPartitionKey(partrel);
+			partdesc = PartitionDirectoryLookup(estate->es_partition_directory,
+												partrel);
+			InitPartitionPruneContext(&pprune->exec_context,
+									  pprune->exec_pruning_steps,
+									  partdesc, partkey, planstate,
+									  econtext);
+
+			/*
+			 * Accumulate the IDs of all PARAM_EXEC Params affecting the
+			 * partitioning decisions at this plan node.
+			 */
+			prunestate->execparamids = bms_add_members(prunestate->execparamids,
+													   pinfo->execparamids);
+		}
+	}
+}
+
 /*
  * ExecFindMatchingSubPlans
  *		Determine which subplans match the pruning steps detailed in
diff --git a/src/include/executor/execPartition.h b/src/include/executor/execPartition.h
index 12aacc84ff..c0ba23097f 100644
--- a/src/include/executor/execPartition.h
+++ b/src/include/executor/execPartition.h
@@ -42,6 +42,9 @@ extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
  * PartitionedRelPruneInfo (see plannodes.h); though note that here,
  * subpart_map contains indexes into PartitionPruningData.partrelprunedata[].
  *
+ * partrel						Partitioned table; points to
+ *								EState.es_relations[rti-1], where rti is the
+ *								table's RT index
  * 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.
@@ -58,6 +61,7 @@ extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
  */
 typedef struct PartitionedRelPruningData
 {
+	Relation	partrel;
 	int			nparts;
 	int		   *subplan_map;
 	int		   *subpart_map;
@@ -128,4 +132,6 @@ extern PartitionPruneState *ExecInitPartitionPruning(PlanState *planstate,
 													 Bitmapset **initially_valid_subplans);
 extern Bitmapset *ExecFindMatchingSubPlans(PartitionPruneState *prunestate,
 										   bool initial_prune);
+extern PartitionPruneState *CreatePartitionPruneState(EState *estate,
+													  PartitionPruneInfo *pruneinfo);
 #endif							/* EXECPARTITION_H */
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index 22b928e085..518a9fcd15 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -637,6 +637,8 @@ typedef struct EState
 	List	   *es_rteperminfos;	/* List of RTEPermissionInfo */
 	PlannedStmt *es_plannedstmt;	/* link to top of plan tree */
 	List	   *es_part_prune_infos;	/* PlannedStmt.partPruneInfos */
+	List	   *es_part_prune_states;	/* List of PartitionPruneState */
+	List	   *es_part_prune_results;	/* List of Bitmapset */
 	const char *es_sourceText;	/* Source text from QueryDesc */
 
 	JunkFilter *es_junkFilter;	/* top-level junk filter, if any */
-- 
2.43.0

