v56-0002-Initialize-PartitionPruneContexts-lazily.patch

From 7ba748a1055880ee20f908a2cf2757f2ad82e9ef Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Wed, 18 Sep 2024 11:16:48 +0900
Subject: [PATCH v56 2/5] Initialize PartitionPruneContexts lazily
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This commit moves the initialization of PartitionPruneContexts for
both initial and exec pruning steps from CreatePartitionPruneState()
to find_matching_subplans_recurse(), where they are actually needed.
To track whether the context has been initialized and is ready for
use, a boolean field is_valid has been added to PartitionPruneContext.

The primary motivation is to eliminate the need to perform
CreatePartitionPruneState() during ExecInitNode(), as creating the
exec pruning context requires access to the parent plan node’s
PlanState. By deferring context creation to where it’s needed, this
change enables calling CreatePartitionPruneState() before ExecInitNode().

This will be useful in a future commit, which will move initial
pruning to occur before ExecInitNode().
---
 src/backend/executor/execPartition.c | 150 +++++++++++++++++++--------
 src/include/executor/execPartition.h |  11 ++
 src/include/partitioning/partprune.h |   2 +
 3 files changed, 120 insertions(+), 43 deletions(-)

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index ec730674f2..63c3429fe7 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -181,18 +181,17 @@ 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,
+static PartitionPruneState *CreatePartitionPruneState(EState *estate,
 													  PartitionPruneInfo *pruneinfo);
-static void InitPartitionPruneContext(PartitionPruneContext *context,
+static void InitPartitionPruneContext(PartitionedRelPruningData *pprune,
+									  PartitionPruneContext *context,
 									  List *pruning_steps,
-									  PartitionDesc partdesc,
-									  PartitionKey partkey,
-									  PlanState *planstate,
-									  ExprContext *econtext);
+									  PlanState *planstate);
 static void PartitionPruneFixSubPlanMap(PartitionPruneState *prunestate,
 										Bitmapset *initially_valid_subplans,
 										int n_total_subplans);
-static void find_matching_subplans_recurse(PartitionPruningData *prunedata,
+static void find_matching_subplans_recurse(PlanState *parent_plan,
+										   PartitionPruningData *prunedata,
 										   PartitionedRelPruningData *pprune,
 										   bool initial_prune,
 										   Bitmapset **validsubplans);
@@ -1825,7 +1824,14 @@ ExecInitPartitionPruning(PlanState *planstate,
 	ExecAssignExprContext(estate, planstate);
 
 	/* Create the working data structure for pruning */
-	prunestate = CreatePartitionPruneState(planstate, pruneinfo);
+	prunestate = CreatePartitionPruneState(estate, pruneinfo);
+
+	/*
+	 * Store PlanState for using it to initialize exec pruning contexts later
+	 * in find_matching_subplans_recurse() where they are needed.
+	 */
+	if (prunestate->do_exec_prune)
+		prunestate->parent_plan = planstate;
 
 	/*
 	 * Perform an initial partition prune pass, if required.
@@ -1865,8 +1871,6 @@ ExecInitPartitionPruning(PlanState *planstate,
  * CreatePartitionPruneState
  *		Build the data structure required for calling ExecFindMatchingSubPlans
  *
- * 'planstate' is the parent plan node's execution state.
- *
  * 'pruneinfo' is a PartitionPruneInfo as generated by
  * make_partition_pruneinfo.  Here we build a PartitionPruneState containing a
  * PartitionPruningData for each partitioning hierarchy (i.e., each sublist of
@@ -1877,16 +1881,20 @@ 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 the PartitionPruneContexts for both initial and exec pruning
+ * (which are stored in each PartitionedRelPruningData) are initialized lazily
+ * in find_matching_subplans_recurse() when used for the first time.
  */
 static PartitionPruneState *
-CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
+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)
@@ -1908,6 +1916,7 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 	prunestate->other_subplans = bms_copy(pruneinfo->other_subplans);
 	prunestate->do_initial_prune = false;	/* may be set below */
 	prunestate->do_exec_prune = false;	/* may be set below */
+	prunestate->parent_plan = NULL;
 	prunestate->num_partprunedata = n_part_hierarchies;
 
 	/*
@@ -1943,16 +1952,25 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 			PartitionedRelPruningData *pprune = &prunedata->partrelprunedata[j];
 			Relation	partrel;
 			PartitionDesc partdesc;
-			PartitionKey partkey;
 
 			/*
-			 * 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.
+			 * Used for initializing the expressions in initial pruning steps.
+			 * For exec pruning steps, the parent plan node's PlanState's
+			 * ps_ExprContext will be used.
 			 */
+			pprune->estate = estate;
+			pprune->econtext = econtext;
+
+			/* Remember Relation for use in InitPartitionPruneContext. */
 			partrel = ExecGetRangeTableRelation(estate, pinfo->rtindex);
-			partkey = RelationGetPartitionKey(partrel);
+			pprune->partrel = partrel;
+
+			/*
+			 * We can rely on the copy partrtitioned table's partition
+			 * descriptor appearing in its relcache entry, because that entry
+			 * will be held open and locked for the duration of this executor
+			 * run.
+			 */
 			partdesc = PartitionDirectoryLookup(estate->es_partition_directory,
 												partrel);
 
@@ -2063,32 +2081,26 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 			pprune->present_parts = bms_copy(pinfo->present_parts);
 
 			/*
-			 * Initialize pruning contexts as needed.  Note that we must skip
-			 * execution-time partition pruning in EXPLAIN (GENERIC_PLAN),
-			 * since parameter values may be missing.
+			 * Pruning contexts (initial_context and exec_context) are
+			 * initialized lazily in find_matching_subplans_recurse() when used
+			 * for the first time.
+			 *
+			 * Note that we must skip execution-time partition pruning in
+			 * EXPLAIN (GENERIC_PLAN), since parameter values may be missing.
 			 */
 			pprune->initial_pruning_steps = pinfo->initial_pruning_steps;
+			pprune->initial_context.is_valid = false;
 			if (pinfo->initial_pruning_steps &&
 				!(econtext->ecxt_estate->es_top_eflags & EXEC_FLAG_EXPLAIN_GENERIC))
-			{
-				InitPartitionPruneContext(&pprune->initial_context,
-										  pinfo->initial_pruning_steps,
-										  partdesc, partkey, planstate,
-										  econtext);
 				/* Record whether initial pruning is needed at any level */
 				prunestate->do_initial_prune = true;
-			}
+
 			pprune->exec_pruning_steps = pinfo->exec_pruning_steps;
+			pprune->exec_context.is_valid = false;
 			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
@@ -2109,16 +2121,43 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
  * Initialize a PartitionPruneContext for the given list of pruning steps.
  */
 static void
-InitPartitionPruneContext(PartitionPruneContext *context,
+InitPartitionPruneContext(PartitionedRelPruningData *pprune,
+						  PartitionPruneContext *context,
 						  List *pruning_steps,
-						  PartitionDesc partdesc,
-						  PartitionKey partkey,
-						  PlanState *planstate,
-						  ExprContext *econtext)
+						  PlanState *planstate)
 {
 	int			n_steps;
 	int			partnatts;
 	ListCell   *lc;
+	ExprContext *econtext;
+	EState	   *estate = pprune->estate;
+	MemoryContext oldcxt;
+	Relation	partrel = pprune->partrel;
+	PartitionKey partkey;
+	PartitionDesc partdesc;
+
+	/* Must allocate the needed stuff in the query lifetime context. */
+	oldcxt = MemoryContextSwitchTo(estate->es_query_cxt);
+
+	/* Use parent_plan's ExprContext when available. */
+	if (planstate)
+	{
+		if (planstate->ps_ExprContext == NULL)
+			ExecAssignExprContext(estate, planstate);
+		econtext = planstate->ps_ExprContext;
+	}
+	else
+		econtext = pprune->econtext;
+
+	/*
+	 * 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);
 
 	n_steps = list_length(pruning_steps);
 
@@ -2187,6 +2226,9 @@ InitPartitionPruneContext(PartitionPruneContext *context,
 			}
 		}
 	}
+
+	MemoryContextSwitchTo(oldcxt);
+	context->is_valid = true;
 }
 
 /*
@@ -2350,12 +2392,16 @@ ExecFindMatchingSubPlans(PartitionPruneState *prunestate,
 		 * recursing to other (lower-level) parents as needed.
 		 */
 		pprune = &prunedata->partrelprunedata[0];
-		find_matching_subplans_recurse(prunedata, pprune, initial_prune,
+		find_matching_subplans_recurse(prunestate->parent_plan,
+									   prunedata, pprune, initial_prune,
 									   &result);
 
 		/* Expression eval may have used space in ExprContext too */
-		if (pprune->exec_pruning_steps)
+		if (pprune->exec_context.is_valid)
+		{
+			Assert(pprune->exec_pruning_steps != NIL);
 			ResetExprContext(pprune->exec_context.exprcontext);
+		}
 	}
 
 	/* Add in any subplans that partition pruning didn't account for */
@@ -2378,7 +2424,8 @@ ExecFindMatchingSubPlans(PartitionPruneState *prunestate,
  * Adds valid (non-prunable) subplan IDs to *validsubplans
  */
 static void
-find_matching_subplans_recurse(PartitionPruningData *prunedata,
+find_matching_subplans_recurse(PlanState *parent_plan,
+							   PartitionPruningData *prunedata,
 							   PartitionedRelPruningData *pprune,
 							   bool initial_prune,
 							   Bitmapset **validsubplans)
@@ -2395,11 +2442,27 @@ find_matching_subplans_recurse(PartitionPruningData *prunedata,
 	 * level.
 	 */
 	if (initial_prune && pprune->initial_pruning_steps)
+	{
+		/* Initialize initial_context if not already done. */
+		if (unlikely(!pprune->initial_context.is_valid))
+			InitPartitionPruneContext(pprune,
+									  &pprune->initial_context,
+									  pprune->initial_pruning_steps,
+									  parent_plan);
 		partset = get_matching_partitions(&pprune->initial_context,
 										  pprune->initial_pruning_steps);
+	}
 	else if (!initial_prune && pprune->exec_pruning_steps)
+	{
+		/* Initialize exec_context if not already done. */
+		if (unlikely(!pprune->exec_context.is_valid))
+			InitPartitionPruneContext(pprune,
+									  &pprune->exec_context,
+									  pprune->exec_pruning_steps,
+									  parent_plan);
 		partset = get_matching_partitions(&pprune->exec_context,
 										  pprune->exec_pruning_steps);
+	}
 	else
 		partset = pprune->present_parts;
 
@@ -2415,7 +2478,8 @@ find_matching_subplans_recurse(PartitionPruningData *prunedata,
 			int			partidx = pprune->subpart_map[i];
 
 			if (partidx >= 0)
-				find_matching_subplans_recurse(prunedata,
+				find_matching_subplans_recurse(parent_plan,
+											   prunedata,
 											   &prunedata->partrelprunedata[partidx],
 											   initial_prune, validsubplans);
 			else
diff --git a/src/include/executor/execPartition.h b/src/include/executor/execPartition.h
index 12aacc84ff..41afb522f3 100644
--- a/src/include/executor/execPartition.h
+++ b/src/include/executor/execPartition.h
@@ -42,6 +42,10 @@ extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
  * PartitionedRelPruneInfo (see plannodes.h); though note that here,
  * subpart_map contains indexes into PartitionPruningData.partrelprunedata[].
  *
+ * estate						The EState for the query doing runtime pruning
+ * partrel						Partitioned table Relation; 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.
@@ -51,6 +55,7 @@ extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
  *								perform executor startup pruning.
  * exec_pruning_steps			List of PartitionPruneSteps used to
  *								perform per-scan pruning.
+ * econtext						ExprContext to use for initial pruning steps
  * initial_context				If initial_pruning_steps isn't NIL, contains
  *								the details needed to execute those steps.
  * exec_context					If exec_pruning_steps isn't NIL, contains
@@ -58,12 +63,15 @@ extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
  */
 typedef struct PartitionedRelPruningData
 {
+	EState	   *estate;
+	Relation	partrel;
 	int			nparts;
 	int		   *subplan_map;
 	int		   *subpart_map;
 	Bitmapset  *present_parts;
 	List	   *initial_pruning_steps;
 	List	   *exec_pruning_steps;
+	ExprContext *econtext;
 	PartitionPruneContext initial_context;
 	PartitionPruneContext exec_context;
 } PartitionedRelPruningData;
@@ -105,6 +113,8 @@ typedef struct PartitionPruningData
  *						startup (at any hierarchy level).
  * do_exec_prune		true if pruning should be performed during
  *						executor run (at any hierarchy level).
+ * parent_plan			Parent plan node's PlanState used to initialize exec
+ *						pruning contexts
  * num_partprunedata	Number of items in "partprunedata" array.
  * partprunedata		Array of PartitionPruningData pointers for the plan's
  *						partitioned relation(s), one for each partitioning
@@ -117,6 +127,7 @@ typedef struct PartitionPruneState
 	MemoryContext prune_context;
 	bool		do_initial_prune;
 	bool		do_exec_prune;
+	PlanState  *parent_plan;
 	int			num_partprunedata;
 	PartitionPruningData *partprunedata[FLEXIBLE_ARRAY_MEMBER];
 } PartitionPruneState;
diff --git a/src/include/partitioning/partprune.h b/src/include/partitioning/partprune.h
index c536a1fe19..b7f48eefcc 100644
--- a/src/include/partitioning/partprune.h
+++ b/src/include/partitioning/partprune.h
@@ -26,6 +26,7 @@ struct RelOptInfo;
  *		Stores information needed at runtime for pruning computations
  *		related to a single partitioned table.
  *
+ * is_valid			Has the information in this struct been initialized?
  * strategy			Partition strategy, e.g. LIST, RANGE, HASH.
  * partnatts		Number of columns in the partition key.
  * nparts			Number of partitions in this partitioned table.
@@ -48,6 +49,7 @@ struct RelOptInfo;
  */
 typedef struct PartitionPruneContext
 {
+	bool		is_valid;
 	char		strategy;
 	int			partnatts;
 	int			nparts;
-- 
2.43.0