From 540079086a25bc18538a78c8990464ab4c77659e Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Wed, 4 Dec 2024 16:16:41 +0900
Subject: [PATCH v58 2/4] Initialize PartitionPruneContexts lazily

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 allow CreatePartitionPruneState() to be
called before ExecInitNode(). Right now, it's coupled with
ExecInitNode() because setting up the exec pruning context requires
access to the parent plan node's PlanState.  By deferring context
creation to where it's actually needed, we break this dependency.

The ExprContext used for both pruning phases is now a standalone
context, independent of the parent PlanState.

This change will be useful in a future commit, which will move initial
pruning to occur outside ExecInitNode(), specifically before it is
called by InitPlan().

Reviewed-by: Robert Haas
Reviewed-by: Tom Lane
Reviewed-by: Tomas Vondra
Discussion: https://postgr.es/m/CA+HiwqFGkMSge6TgC9KQzde0ohpAycLQuV7ooitEEpbKB0O_mg@mail.gmail.com
---
 src/backend/executor/execPartition.c | 151 +++++++++++++++++++--------
 src/backend/partitioning/partprune.c |   7 +-
 src/include/executor/execPartition.h |  12 +++
 src/include/partitioning/partprune.h |   2 +
 4 files changed, 123 insertions(+), 49 deletions(-)

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index 950fa3289c..f4d425cd45 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);
@@ -1823,7 +1822,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.
@@ -1863,8 +1869,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
@@ -1875,16 +1879,24 @@ 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;
+
+	/*
+	 * Expression context that will be used by partkey_datum_from_expr() to
+	 * evaluate expressions for comparison against partition bounds.
+	 */
+	ExprContext *econtext = CreateExprContext(estate);
 
 	/* For data reading, executor always includes detached partitions */
 	if (estate->es_partition_directory == NULL)
@@ -1906,6 +1918,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;
 
 	/*
@@ -1941,16 +1954,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 of partitioned 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);
 
@@ -2061,32 +2083,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.initialized = 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.initialized = 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
@@ -2107,17 +2123,41 @@ 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;
 
+	/*
+	 * Use the ExprContext that CreatePartitionPruneState() should have
+	 * created.
+	 */
+	ExprContext *econtext = pprune->econtext;
+	EState	   *estate = pprune->estate;
+	MemoryContext oldcxt;
+	Relation	partrel = pprune->partrel;
+	PartitionKey partkey;
+	PartitionDesc partdesc;
+
+	Assert(econtext != NULL);
+
+	/* Must allocate the needed stuff in the query lifetime context. */
+	oldcxt = MemoryContextSwitchTo(estate->es_query_cxt);
+
+	/*
+	 * 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);
 
 	context->strategy = partkey->strategy;
@@ -2185,6 +2225,9 @@ InitPartitionPruneContext(PartitionPruneContext *context,
 			}
 		}
 	}
+
+	MemoryContextSwitchTo(oldcxt);
+	context->initialized = true;
 }
 
 /*
@@ -2348,12 +2391,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.initialized)
+		{
+			Assert(pprune->exec_pruning_steps != NIL);
 			ResetExprContext(pprune->exec_context.exprcontext);
+		}
 	}
 
 	/* Add in any subplans that partition pruning didn't account for */
@@ -2376,7 +2423,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)
@@ -2393,11 +2441,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.initialized))
+			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.initialized))
+			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;
 
@@ -2413,7 +2477,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/backend/partitioning/partprune.c b/src/backend/partitioning/partprune.c
index ca5467104d..ae1d69f96c 100644
--- a/src/backend/partitioning/partprune.c
+++ b/src/backend/partitioning/partprune.c
@@ -3783,13 +3783,8 @@ partkey_datum_from_expr(PartitionPruneContext *context,
 		/*
 		 * We should never see a non-Const in a step unless the caller has
 		 * passed a valid ExprContext.
-		 *
-		 * When context->planstate is valid, context->exprcontext is same as
-		 * context->planstate->ps_ExprContext.
 		 */
-		Assert(context->planstate != NULL || context->exprcontext != NULL);
-		Assert(context->planstate == NULL ||
-			   (context->exprcontext == context->planstate->ps_ExprContext));
+		Assert(context->exprcontext != NULL);
 
 		exprstate = context->exprstates[stateidx];
 		ectx = context->exprcontext;
diff --git a/src/include/executor/execPartition.h b/src/include/executor/execPartition.h
index 33d922fe8d..7e470c82f6 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 run-time pruning
+ * partrel						Partitioned table Relation; obtained by
+ * 								ExecGetRangeTableRelation(estate, rti), where
+ *								rti is PartitionedRelPruneInfo.rtindex.
  * 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,8 @@ 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 evaluating partition
+ *								key
  * 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 +64,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 +114,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
+ *						expression contained in "exec" pruning steps.
  * 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 +128,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 6922e04430..0cbcb4fb4e 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.
  *
+ * initialized		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		initialized;
 	char		strategy;
 	int			partnatts;
 	int			nparts;
-- 
2.43.0

