From 0e18a7dd1e5b1e9aaea442b140e9591ccb25644a Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Wed, 23 Oct 2024 17:02:33 +0900
Subject: [PATCH v57 3/5] Perform runtime initial pruning outside
 ExecInitNode()
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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() performs the initial pruning and saves the
result—a bitmapset of indexes for the surviving child subnodes—in
es_part_prune_results, a list in EState. The PartitionPruneStates
created for initial pruning are also saved in es_part_prune_states,
another list in EState, for later use during exec pruning. Both lists
are parallel to es_part_prune_infos (which holds the
PartitionPruneInfos from PlannedStmt), allowing them to share the
same index.

Reviewed-by: Robert Haas
---
 src/backend/executor/execMain.c      | 59 ++++++++++++++++++++++++++++
 src/backend/executor/execPartition.c | 51 +++++++++++++-----------
 src/include/executor/execPartition.h |  2 +
 src/include/nodes/execnodes.h        |  2 +
 4 files changed, 90 insertions(+), 24 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index c460c6aa32..2fcec32dcb 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"
@@ -819,6 +820,53 @@ 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.
+ *
+ * This function iterates over each PartitionPruneInfo entry in
+ * estate->es_part_prune_infos. For each entry, it creates a PartitionPruneState
+ * and adds it to es_part_prune_states, where ExecInitPartitionPruning() can
+ * access it for use during "exec" pruning.
+ *
+ * If initial pruning steps exist for a PartitionPruneInfo entry, this function
+ * executes those pruning steps and stores the result as a bitmapset of valid
+ * child subplans, identifying which subplans should be initialized for
+ * execution.  The results are saved in estate->es_part_prune_results.
+ *
+ * If no initial pruning is performed for a given PartitionPruneInfo, a NULL
+ * entry  is still added to es_part_prune_results to maintain alignment with
+ * es_part_prune_infos. This ensures that ExecInitPartitionPruning() can use
+ * the same index to retrieve the pruning results.
+ */
+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 and save the PartitionPruneState. */
+		prunestate = ExecCreatePartitionPruneState(estate, pruneinfo);
+		estate->es_part_prune_states = lappend(estate->es_part_prune_states,
+											   prunestate);
+
+		/*
+		 * Perform initial pruning steps, if any, and save the result
+		 * bitmapset or NULL as described in the header comment.
+		 */
+		if (prunestate->do_initial_prune)
+			validsubplans = ExecFindMatchingSubPlans(prunestate, true);
+		estate->es_part_prune_results = lappend(estate->es_part_prune_results,
+												validsubplans);
+	}
+}
 
 /* ----------------------------------------------------------------
  *		InitPlan
@@ -851,7 +899,18 @@ InitPlan(QueryDesc *queryDesc, int eflags)
 	ExecInitRangeTable(estate, rangeTable, plannedstmt->permInfos);
 
 	estate->es_plannedstmt = plannedstmt;
+
+	/*
+	 * Perform runtime "initial" pruning to identify which child subplans,
+	 * corresponding to the children of plan nodes that contain
+	 * PartitionPruneInfo such as Append, will not be executed. The results,
+	 * which are bitmapsets of indexes of the child subplans that will be
+	 * executed, are saved in es_part_prune_results.  These results correspond
+	 * to each PartitionPruneInfo entry, and the es_part_prune_results list is
+	 * parallel to es_part_prune_infos.
+	 */
 	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 38311d2991..83d1b61101 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,24 @@ 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.
  *
  * 'relids' identifies the relation to which both the parent plan and the
  * PartitionPruneInfo given by '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.
  *
- * 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 were pruned during initial pruning, 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 +1822,12 @@ ExecInitPartitionPruning(PlanState *planstate,
 								   bmsToString(relids),
 								   bmsToString(pruneinfo->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 +1836,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 */
@@ -1848,8 +1851,8 @@ ExecInitPartitionPruning(PlanState *planstate,
 
 	/*
 	 * Re-sequence subplan indexes contained in prunestate to account for any
-	 * that were removed above due to initial pruning.  No need to do this if
-	 * no steps were removed.
+	 * that were removed due to initial pruning.  No need to do this if no
+	 * partitions were removed.
 	 */
 	if (bms_num_members(*initially_valid_subplans) < n_total_subplans)
 	{
@@ -1886,8 +1889,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 5178c27743..1497aed533 100644
--- a/src/include/executor/execPartition.h
+++ b/src/include/executor/execPartition.h
@@ -140,4 +140,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 5deed9232a..b0ceb1ab05 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -640,6 +640,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

