From 84b875b1ca3af89a9242cdaf9bea052223f9530e Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Thu, 12 Sep 2024 15:44:43 +0900
Subject: [PATCH v56 3/5] 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 that it
creates to do the initial pruning to use during exec pruninng 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.
---
 src/backend/executor/execMain.c      | 55 ++++++++++++++++++++++++++++
 src/backend/executor/execPartition.c | 51 ++++++++++++++------------
 src/include/executor/execPartition.h |  2 +
 src/include/nodes/execnodes.h        |  2 +
 4 files changed, 87 insertions(+), 23 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index f263232c67..5222aa9ab3 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"
@@ -828,6 +829,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 = ExecCreatePartitionPruneState(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
@@ -860,7 +909,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 63c3429fe7..40eb74d187 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(EState *estate,
-													  PartitionPruneInfo *pruneinfo);
 static void InitPartitionPruneContext(PartitionedRelPruningData *pprune,
 									  PartitionPruneContext *context,
 									  List *pruning_steps,
@@ -1782,20 +1780,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,
@@ -1820,11 +1824,12 @@ 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(estate, pruneinfo);
+	/*
+	 * ExecDoInitialPruning() must have initialized the PartitionPruneState to
+	 * perform the initial pruning.
+	 */
+	prunestate = list_nth(estate->es_part_prune_states, part_prune_index);
+	Assert(prunestate != NULL);
 
 	/*
 	 * Store PlanState for using it to initialize exec pruning contexts later
@@ -1833,11 +1838,11 @@ ExecInitPartitionPruning(PlanState *planstate,
 	if (prunestate->do_exec_prune)
 		prunestate->parent_plan = planstate;
 
-	/*
-	 * Perform an initial partition prune pass, if required.
-	 */
+	/* 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 */
@@ -1886,8 +1891,8 @@ ExecInitPartitionPruning(PlanState *planstate,
  * (which are stored in each PartitionedRelPruningData) are initialized lazily
  * in find_matching_subplans_recurse() when used for the first time.
  */
-static PartitionPruneState *
-CreatePartitionPruneState(EState *estate, PartitionPruneInfo *pruneinfo)
+PartitionPruneState *
+ExecCreatePartitionPruneState(EState *estate, PartitionPruneInfo *pruneinfo)
 {
 	PartitionPruneState *prunestate;
 	int			n_part_hierarchies;
diff --git a/src/include/executor/execPartition.h b/src/include/executor/execPartition.h
index 41afb522f3..8d85fa990e 100644
--- a/src/include/executor/execPartition.h
+++ b/src/include/executor/execPartition.h
@@ -139,4 +139,6 @@ extern PartitionPruneState *ExecInitPartitionPruning(PlanState *planstate,
 													 Bitmapset **initially_valid_subplans);
 extern Bitmapset *ExecFindMatchingSubPlans(PartitionPruneState *prunestate,
 										   bool initial_prune);
+extern PartitionPruneState *ExecCreatePartitionPruneState(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

