v55-0003-Initialize-PartitionPruneContext-for-exec-prunin.patch

From 92d87cdbb3ad675ac6ffa2767f1d7d5876bd5369 Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Wed, 18 Sep 2024 11:16:48 +0900
Subject: [PATCH v55 3/5] Initialize PartitionPruneContext for exec pruning
 lazily

Currently, ExecInitPartitionPruning() iterates over PartitionPruningDatas
and nested PartitionedRelPruningDatas in a PartitionPruneState solely
to initialize the exec_context of the PartitionedRelPruningData.

This commit moves the initialization to find_matching_subplans_recurse(),
where the exec_context is actually needed, eliminating the need for
the above iteration. To track whether the context has been initialized
and is ready for use, a boolean field is_valid is added to
PartitionPruneContext.
---
 src/backend/executor/execPartition.c | 166 ++++++++++-----------------
 src/include/executor/execPartition.h |   1 +
 src/include/partitioning/partprune.h |   2 +
 3 files changed, 65 insertions(+), 104 deletions(-)

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index 3c7c631867..d9fa593785 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -190,10 +190,8 @@ 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,
+static void find_matching_subplans_recurse(PlanState *parent_plan,
+										   PartitionPruningData *prunedata,
 										   PartitionedRelPruningData *pprune,
 										   bool initial_prune,
 										   Bitmapset **validsubplans);
@@ -1830,13 +1828,14 @@ ExecInitPartitionPruning(PlanState *planstate,
 
 	/*
 	 * ExecDoInitialPruning() must have initialized the PartitionPruneState to
-	 * perform the initial pruning.  Now we simply need to initialize the
-	 * context information for exec pruning.
+	 * perform the initial pruning.  Store PlanState so that the exec_context
+	 * can be initialized using it later when find_matching_subplans_recurse()
+	 * needs it.
 	 */
 	prunestate = list_nth(estate->es_part_prune_states, part_prune_index);
 	Assert(prunestate != NULL);
 	if (prunestate->do_exec_prune)
-		PartitionPruneInitExecPruning(pruneinfo, prunestate, planstate);
+		prunestate->parent_plan = planstate;
 
 	/* Use the result of initial pruning done by ExecDoInitialPruning(). */
 	if (prunestate->do_initial_prune)
@@ -1893,8 +1892,7 @@ ExecInitPartitionPruning(PlanState *planstate,
  * 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().
+ * initialized lazily during find_matching_subplans_recurse().
  */
 PartitionPruneState *
 ExecCreatePartitionPruneState(EState *estate, PartitionPruneInfo *pruneinfo)
@@ -2099,25 +2097,30 @@ ExecCreatePartitionPruneState(EState *estate, PartitionPruneInfo *pruneinfo)
 			}
 
 			/*
-			 * The exec pruning context will be initialized in
-			 * ExecInitPartitionPruning() when called during the initialization
-			 * of the parent plan node.
+			 * The exec pruning context will be initialized lazily when it
+			 * will be used for the first time in
+			 * find_matching_subplans_recurse().
 			 *
-			 * 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->do_exec_prune is set to indicate whether we're
+			 * actually going to perform exec pruning to inform
+			 * ExecInitPartitionPruning() whether it should fix the
+			 * subplan_map array based on the result of initial pruning
+			 * and also the parent node's code to allow it set up its
+			 * data structure accordingly.
 			 */
-			pprune->exec_pruning_steps = NIL;
+			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))
 				prunestate->do_exec_prune = true;
 
+			/*
+			 * 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);
+
 			j++;
 		}
 		i++;
@@ -2208,6 +2211,8 @@ InitPartitionPruneContext(PartitionPruneContext *context,
 			}
 		}
 	}
+
+	context->is_valid = true;
 }
 
 /*
@@ -2326,84 +2331,6 @@ 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
@@ -2449,12 +2376,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 */
@@ -2477,7 +2408,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)
@@ -2497,8 +2429,33 @@ find_matching_subplans_recurse(PartitionPruningData *prunedata,
 		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))
+		{
+			ExprContext	   *econtext;
+			EState		   *estate = parent_plan->state;
+			/* Must allocate the needed stuff in the query lifetime context. */
+			MemoryContext 	oldcxt = MemoryContextSwitchTo(estate->es_query_cxt);
+			Relation		partrel = pprune->partrel;
+			PartitionKey	partkey = RelationGetPartitionKey(partrel);
+			PartitionDesc	partdesc = PartitionDirectoryLookup(estate->es_partition_directory,
+																partrel);
+
+			if (parent_plan->ps_ExprContext == NULL)
+				ExecAssignExprContext(estate, parent_plan);
+			econtext = parent_plan->ps_ExprContext;
+
+			InitPartitionPruneContext(&pprune->exec_context,
+									  pprune->exec_pruning_steps,
+									  partdesc, partkey, parent_plan,
+									  econtext);
+
+			MemoryContextSwitchTo(oldcxt);
+		}
 		partset = get_matching_partitions(&pprune->exec_context,
 										  pprune->exec_pruning_steps);
+	}
 	else
 		partset = pprune->present_parts;
 
@@ -2514,7 +2471,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 2f45ac1cc8..ef6d8b2d48 100644
--- a/src/include/executor/execPartition.h
+++ b/src/include/executor/execPartition.h
@@ -122,6 +122,7 @@ typedef struct PartitionPruneState
 	bool		do_initial_prune;
 	bool		do_exec_prune;
 	int			num_partprunedata;
+	PlanState  *parent_plan;
 	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