From c1f8aed91ee1a43871b09d9ba30a6e210c89cb28 Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Thu, 23 Jan 2025 14:49:48 +0900
Subject: [PATCH v60 2/3] Perform runtime initial pruning outside
 ExecInitNode()

This commit builds on the prior change that moved PartitionPruneInfos
out of individual plan nodes into a list in PlannedStmt. It moves the
initialization of PartitionPruneStates and runtime initial pruning
from ExecInitNode() to a new routine, ExecDoInitialPruning().

ExecDoInitialPruning() is called by InitPlan() before calling
ExecInitNode() on the main plan tree and subplans. It performs the
initial pruning and saves the result -- a Bitmapset of indexes for
surviving child subnodes -- in es_part_prune_results, a list in
EState.

PartitionPruneStates created for initial pruning are stored 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, enabling shared
indexing.

PartitionPruneStates initialized in ExecDoInitialPruning() now include
only the PartitionPruneContexts for initial pruning steps. Exec
pruning contexts are initialized later in
ExecInitPartitionExecPruning() when the parent plan node is
initialized, as the exec pruning step expressions depend on the
parent node's PlanState. The existing function
PartitionPruneFixSubPlanMap() has been repurposed for this
initialization to avoid duplicating a similar loop structure for
finding PartitionedRelPruningData to initialize exec contexts for. It
has been renamed to InitExecPruningContexts() to reflect its new
primary responsibility. The original logic to "fix subplan maps"
remains intact but is now encapsulated within the renamed function.

To ensure exec pruning contexts are not accessed before initialization,
a new boolean field, 'initialized', is added to PartitionPruneContext.

This commit removes two obsolete Asserts in partkey_datum_from_expr().
The ExprContext used for pruning expression evaluation is now
independent of the parent PlanState, making these Asserts unnecessary.

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/execMain.c        |  12 +
 src/backend/executor/execPartition.c   | 318 ++++++++++++++++++-------
 src/backend/executor/nodeAppend.c      |  10 +-
 src/backend/executor/nodeMergeAppend.c |  10 +-
 src/backend/partitioning/partprune.c   |   7 +-
 src/include/executor/execPartition.h   |  18 +-
 src/include/nodes/execnodes.h          |   2 +
 src/include/partitioning/partprune.h   |   2 +
 8 files changed, 266 insertions(+), 113 deletions(-)

diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 1d27b840ca9..604cb0625b8 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"
@@ -855,6 +856,17 @@ InitPlan(QueryDesc *queryDesc, int eflags)
 	estate->es_plannedstmt = plannedstmt;
 	estate->es_part_prune_infos = plannedstmt->partPruneInfos;
 
+	/*
+	 * 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.
+	 */
+	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 16aec59d0ec..70f11913ad4 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -181,7 +181,7 @@ 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,
 									  List *pruning_steps,
@@ -189,9 +189,10 @@ static void InitPartitionPruneContext(PartitionPruneContext *context,
 									  PartitionKey partkey,
 									  PlanState *planstate,
 									  ExprContext *econtext);
-static void PartitionPruneFixSubPlanMap(PartitionPruneState *prunestate,
-										Bitmapset *initially_valid_subplans,
-										int n_total_subplans);
+static void InitExecPartitionPruneContexts(PartitionPruneState *prunstate,
+										   PlanState *parent_plan,
+										   Bitmapset *initially_valid_subplans,
+										   int n_total_subplans);
 static void find_matching_subplans_recurse(PartitionPruningData *prunedata,
 										   PartitionedRelPruningData *pprune,
 										   bool initial_prune,
@@ -1762,48 +1763,106 @@ adjust_partition_colnos_using_map(List *colnos, AttrMap *attrMap)
  *
  * Functions:
  *
- * ExecInitPartitionPruning:
- *		Creates the PartitionPruneState required by ExecFindMatchingSubPlans.
- *		Details stored include how to map the partition index returned by the
- *		partition pruning code into subplan indexes.  Also determines the set
- *		of subplans to initialize considering the result of performing initial
- *		pruning steps if any.  Maps in PartitionPruneState are updated to
+ * ExecDoInitialPruning:
+ *		Perform runtime "initial" pruning, if necessary, to determine the set
+ *		of child subnodes that need to be initialized during ExecInitNode() for
+ *		all plan nodes that contain a PartitionPruneInfo.
+ *
+ * ExecInitPartitionExecPruning:
+ *		Updates the PartitionPruneState found at given part_prune_index in
+ *		EState.es_part_prune_states for use during "exec" pruning if required.
+ *		Also returns the set of subplans to initialize that would be stored at
+ *		part_prune_index in EState.es_part_prune_result by
+ *		ExecDoInitialPruning().  Maps in PartitionPruneState are updated to
  *		account for initial pruning possibly having eliminated some of the
  *		subplans.
  *
  * ExecFindMatchingSubPlans:
  *		Returns indexes of matching subplans after evaluating the expressions
  *		that are safe to evaluate at a given point.  This function is first
- *		called during ExecInitPartitionPruning() to find the initially
- *		matching subplans based on performing the initial pruning steps and
- *		then must be called again each time the value of a Param listed in
+ *		called during ExecDoInitialPruning() to find the initially matching
+ *		subplans based on performing the initial pruning steps and then must be
+ *		called again each time the value of a Param listed in
  *		PartitionPruneState's 'execparamids' changes.
  *-------------------------------------------------------------------------
  */
 
 /*
- * ExecInitPartitionPruning
- *		Initialize data structure needed for run-time partition pruning and
- *		do initial pruning if needed
+ * 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.  ExecInitPartitionExecPruning() accesses
+ * these states through their corresponding indexes in es_part_prune_states and
+ * assign each state to the parent node's PlanState, from where it will be used
+ * for "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 ExecInitPartitionExecPruning() can
+ * use the same index to retrieve the pruning results.
+ */
+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 = CreatePartitionPruneState(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);
+	}
+}
+
+/*
+ * ExecInitPartitionExecPruning
+ *		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.
+ * 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 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.
+ * If subplans were indeed 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,
-						 int n_total_subplans,
-						 int part_prune_index,
-						 Bitmapset *relids,
-						 Bitmapset **initially_valid_subplans)
+ExecInitPartitionExecPruning(PlanState *planstate,
+							 int n_total_subplans,
+							 int part_prune_index,
+							 Bitmapset *relids,
+							 Bitmapset **initially_valid_subplans)
 {
 	PartitionPruneState *prunestate;
 	EState	   *estate = planstate->state;
@@ -1819,17 +1878,19 @@ ExecInitPartitionPruning(PlanState *planstate,
 			 bmsToString(pruneinfo->relids), part_prune_index,
 			 bmsToString(relids));
 
-	/* We may need an expression context to evaluate partition exprs */
-	ExecAssignExprContext(estate, planstate);
-
-	/* Create the working data structure for pruning */
-	prunestate = CreatePartitionPruneState(planstate, pruneinfo);
-
 	/*
-	 * Perform an initial partition prune pass, if required.
+	 * The PartitionPruneState would have been created by
+	 * ExecDoInitialPruning() and stored as the part_prune_index'th element of
+	 * EState.es_part_prune_states.
 	 */
+	prunestate = list_nth(estate->es_part_prune_states, part_prune_index);
+	Assert(prunestate != NULL);
+
+	/* 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 */
@@ -1839,22 +1900,21 @@ 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.
+	 * The exec pruning state must also be initialized, if needed, before it
+	 * can be used for pruning during execution.
+	 *
+	 * This also re-sequences subplan indexes contained in prunestate to
+	 * account for any that were removed due to initial pruning; refer to the
+	 * condition in InitExecPartitionPruneContexts() that is used to determine
+	 * whether to do this.  If no exec pruning needs to be done, we would thus
+	 * leave the maps to be in an invalid invalid state, but that's ok since
+	 * that data won't be consulted again (cf initial Assert in
+	 * ExecFindMatchingSubPlans).
 	 */
-	if (bms_num_members(*initially_valid_subplans) < n_total_subplans)
-	{
-		/*
-		 * We can safely skip this when !do_exec_prune, even though that
-		 * leaves invalid data in prunestate, because that data won't be
-		 * consulted again (cf initial Assert in ExecFindMatchingSubPlans).
-		 */
-		if (prunestate->do_exec_prune)
-			PartitionPruneFixSubPlanMap(prunestate,
-										*initially_valid_subplans,
-										n_total_subplans);
-	}
+	if (prunestate->do_exec_prune)
+		InitExecPartitionPruneContexts(prunestate, planstate,
+									   *initially_valid_subplans,
+									   n_total_subplans);
 
 	return prunestate;
 }
@@ -1863,7 +1923,11 @@ ExecInitPartitionPruning(PlanState *planstate,
  * CreatePartitionPruneState
  *		Build the data structure required for calling ExecFindMatchingSubPlans
  *
- * 'planstate' is the parent plan node's execution state.
+ * This includes PartitionPruneContexts (stored in each
+ * PartitionedRelPruningData corresponding to a PartitionedRelPruneInfo),
+ * which hold the ExprStates needed to evaluate pruning expressions, and
+ * mapping arrays to convert partition indexes from the pruning logic
+ * into subplan indexes in the parent plan node's list of child subplans.
  *
  * 'pruneinfo' is a PartitionPruneInfo as generated by
  * make_partition_pruneinfo.  Here we build a PartitionPruneState containing a
@@ -1875,16 +1939,25 @@ 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 only the PartitionPruneContexts for initial pruning are
+ * initialized here. Those required for exec pruning are initialized later in
+ * ExecInitPartitionExecPruning(), as they depend on the availability of the
+ * parent plan node's PlanState.
  */
 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)
@@ -1901,6 +1974,8 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 		palloc(offsetof(PartitionPruneState, partprunedata) +
 			   sizeof(PartitionPruningData *) * n_part_hierarchies);
 
+	/* Save ExprContext for use during InitExecPartitionPruneContexts(). */
+	prunestate->econtext = econtext;
 	prunestate->execparamids = NULL;
 	/* other_subplans can change at runtime, so we need our own copy */
 	prunestate->other_subplans = bms_copy(pruneinfo->other_subplans);
@@ -1950,6 +2025,10 @@ CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 			 * duration of this executor run.
 			 */
 			partrel = ExecGetRangeTableRelation(estate, pinfo->rtindex);
+
+			/* Remember for InitExecPartitionPruneContext(). */
+			pprune->partrel = partrel;
+
 			partkey = RelationGetPartitionKey(partrel);
 			partdesc = PartitionDirectoryLookup(estate->es_partition_directory,
 												partrel);
@@ -2061,29 +2140,31 @@ 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.
+			 * Only initial_context is initialized here.  exec_context is
+			 * initialized during ExecInitPartitionExecPruning() when the
+			 * parent plan's PlanState is available.
+			 *
+			 * Note that we must skip execution-time (both "init" and "exec")
+			 * 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,
+										  pprune->initial_pruning_steps,
+										  partdesc, partkey, NULL,
 										  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;
 			}
@@ -2118,6 +2199,9 @@ InitPartitionPruneContext(PartitionPruneContext *context,
 	int			partnatts;
 	ListCell   *lc;
 
+	/* Never call twice for a context. */
+	Assert(!context->initialized);
+
 	n_steps = list_length(pruning_steps);
 
 	context->strategy = partkey->strategy;
@@ -2185,13 +2269,22 @@ InitPartitionPruneContext(PartitionPruneContext *context,
 			}
 		}
 	}
+
+	context->initialized = true;
 }
 
 /*
- * PartitionPruneFixSubPlanMap
- *		Fix mapping of partition indexes to subplan indexes contained in
- *		prunestate by considering the new list of subplans that survived
- *		initial pruning
+ * InitExecPartitionPruneContexts
+ *		Initialize exec pruning contexts deferred by CreatePartitionPruneState()
+ *
+ * This function finalizes exec pruning setup for a PartitionPruneState by
+ * initializing contexts for pruning steps that require the parent plan's
+ * PlanState. It iterates over PartitionPruningData entries and sets up the
+ * necessary execution contexts for pruning during query execution.
+ *
+ * Also fix the mapping of partition indexes to subplan indexes contained in
+ * prunestate by considering the new list of subplans that survived initial
+ * pruning.
  *
  * Current values of the indexes present in PartitionPruneState count all the
  * subplans that would be present before initial pruning was done.  If initial
@@ -2202,27 +2295,39 @@ InitPartitionPruneContext(PartitionPruneContext *context,
  * subplans in the post-initial-pruning set.
  */
 static void
-PartitionPruneFixSubPlanMap(PartitionPruneState *prunestate,
-							Bitmapset *initially_valid_subplans,
-							int n_total_subplans)
+InitExecPartitionPruneContexts(PartitionPruneState *prunestate,
+							   PlanState *parent_plan,
+							   Bitmapset *initially_valid_subplans,
+							   int n_total_subplans)
 {
-	int		   *new_subplan_indexes;
+	EState	   *estate;
+	int		   *new_subplan_indexes = NULL;
 	Bitmapset  *new_other_subplans;
 	int			i;
 	int			newidx;
+	bool		fix_subplan_map = false;
 
-	/*
-	 * First we must build a temporary array which maps old subplan indexes to
-	 * new ones.  For convenience of initialization, we use 1-based indexes in
-	 * this array and leave pruned items as 0.
-	 */
-	new_subplan_indexes = (int *) palloc0(sizeof(int) * n_total_subplans);
-	newidx = 1;
-	i = -1;
-	while ((i = bms_next_member(initially_valid_subplans, i)) >= 0)
+	Assert(prunestate->do_exec_prune);
+	Assert(parent_plan != NULL);
+	estate = parent_plan->state;
+
+	if (bms_num_members(initially_valid_subplans) < n_total_subplans)
 	{
-		Assert(i < n_total_subplans);
-		new_subplan_indexes[i] = newidx++;
+		fix_subplan_map = true;
+
+		/*
+		 * First we must build a temporary array which maps old subplan
+		 * indexes to new ones.  For convenience of initialization, we use
+		 * 1-based indexes in this array and leave pruned items as 0.
+		 */
+		new_subplan_indexes = (int *) palloc0(sizeof(int) * n_total_subplans);
+		newidx = 1;
+		i = -1;
+		while ((i = bms_next_member(initially_valid_subplans, i)) >= 0)
+		{
+			Assert(i < n_total_subplans);
+			new_subplan_indexes[i] = newidx++;
+		}
 	}
 
 	/*
@@ -2247,6 +2352,29 @@ PartitionPruneFixSubPlanMap(PartitionPruneState *prunestate,
 			int			nparts = pprune->nparts;
 			int			k;
 
+			/* Initialize PartitionPruneContext for exec pruning, if needed. */
+			if (pprune->exec_pruning_steps != NIL)
+			{
+				PartitionKey partkey;
+				PartitionDesc partdesc;
+
+				/*
+				 * See the comment in CreatePartitionPruneState() regarding
+				 * the usage of partdesc and partkey.
+				 */
+				partkey = RelationGetPartitionKey(pprune->partrel);
+				partdesc = PartitionDirectoryLookup(estate->es_partition_directory,
+													pprune->partrel);
+
+				InitPartitionPruneContext(&pprune->exec_context,
+										  pprune->exec_pruning_steps,
+										  partdesc, partkey, parent_plan,
+										  prunestate->econtext);
+			}
+
+			if (!fix_subplan_map)
+				continue;
+
 			/* We just rebuild present_parts from scratch */
 			bms_free(pprune->present_parts);
 			pprune->present_parts = NULL;
@@ -2288,19 +2416,22 @@ PartitionPruneFixSubPlanMap(PartitionPruneState *prunestate,
 	}
 
 	/*
-	 * We must also recompute the other_subplans set, since indexes in it may
-	 * change.
+	 * If we fixed subplan maps, we must also recompute the other_subplans
+	 * set, since indexes in it may change.
 	 */
-	new_other_subplans = NULL;
-	i = -1;
-	while ((i = bms_next_member(prunestate->other_subplans, i)) >= 0)
-		new_other_subplans = bms_add_member(new_other_subplans,
-											new_subplan_indexes[i] - 1);
+	if (fix_subplan_map)
+	{
+		new_other_subplans = NULL;
+		i = -1;
+		while ((i = bms_next_member(prunestate->other_subplans, i)) >= 0)
+			new_other_subplans = bms_add_member(new_other_subplans,
+												new_subplan_indexes[i] - 1);
 
-	bms_free(prunestate->other_subplans);
-	prunestate->other_subplans = new_other_subplans;
+		bms_free(prunestate->other_subplans);
+		prunestate->other_subplans = new_other_subplans;
 
-	pfree(new_subplan_indexes);
+		pfree(new_subplan_indexes);
+	}
 }
 
 /*
@@ -2352,8 +2483,11 @@ ExecFindMatchingSubPlans(PartitionPruneState *prunestate,
 									   &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);
 			ResetExprContext(pprune->exec_context.exprcontext);
+		}
 	}
 
 	/* Add in any subplans that partition pruning didn't account for */
diff --git a/src/backend/executor/nodeAppend.c b/src/backend/executor/nodeAppend.c
index 986ef34030a..2397e5e17b0 100644
--- a/src/backend/executor/nodeAppend.c
+++ b/src/backend/executor/nodeAppend.c
@@ -144,11 +144,11 @@ ExecInitAppend(Append *node, EState *estate, int eflags)
 		 * subplans to initialize (validsubplans) by taking into account the
 		 * result of performing initial pruning if any.
 		 */
-		prunestate = ExecInitPartitionPruning(&appendstate->ps,
-											  list_length(node->appendplans),
-											  node->part_prune_index,
-											  node->apprelids,
-											  &validsubplans);
+		prunestate = ExecInitPartitionExecPruning(&appendstate->ps,
+												  list_length(node->appendplans),
+												  node->part_prune_index,
+												  node->apprelids,
+												  &validsubplans);
 		appendstate->as_prune_state = prunestate;
 		nplans = bms_num_members(validsubplans);
 
diff --git a/src/backend/executor/nodeMergeAppend.c b/src/backend/executor/nodeMergeAppend.c
index 1468f942388..b2dc6626c99 100644
--- a/src/backend/executor/nodeMergeAppend.c
+++ b/src/backend/executor/nodeMergeAppend.c
@@ -92,11 +92,11 @@ ExecInitMergeAppend(MergeAppend *node, EState *estate, int eflags)
 		 * subplans to initialize (validsubplans) by taking into account the
 		 * result of performing initial pruning if any.
 		 */
-		prunestate = ExecInitPartitionPruning(&mergestate->ps,
-											  list_length(node->mergeplans),
-											  node->part_prune_index,
-											  node->apprelids,
-											  &validsubplans);
+		prunestate = ExecInitPartitionExecPruning(&mergestate->ps,
+												  list_length(node->mergeplans),
+												  node->part_prune_index,
+												  node->apprelids,
+												  &validsubplans);
 		mergestate->ms_prune_state = prunestate;
 		nplans = bms_num_members(validsubplans);
 
diff --git a/src/backend/partitioning/partprune.c b/src/backend/partitioning/partprune.c
index d3f60cc87c9..4693eef0c58 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 13177831d9f..855fed4fea5 100644
--- a/src/include/executor/execPartition.h
+++ b/src/include/executor/execPartition.h
@@ -42,6 +42,9 @@ extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
  * PartitionedRelPruneInfo (see plannodes.h); though note that here,
  * subpart_map contains indexes into PartitionPruningData.partrelprunedata[].
  *
+ * 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.
@@ -58,6 +61,7 @@ extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
  */
 typedef struct PartitionedRelPruningData
 {
+	Relation	partrel;
 	int			nparts;
 	int		   *subplan_map;
 	int		   *subpart_map;
@@ -90,6 +94,8 @@ typedef struct PartitionPruningData
  * the clauses being unable to match to any tuple that the subplan could
  * possibly produce.
  *
+ * econtext				Standalone ExprContext to evaluate expressions in
+ *						the pruning steps
  * execparamids			Contains paramids of PARAM_EXEC Params found within
  *						any of the partprunedata structs.  Pruning must be
  *						done again each time the value of one of these
@@ -112,6 +118,7 @@ typedef struct PartitionPruningData
  */
 typedef struct PartitionPruneState
 {
+	ExprContext *econtext;
 	Bitmapset  *execparamids;
 	Bitmapset  *other_subplans;
 	MemoryContext prune_context;
@@ -121,11 +128,12 @@ typedef struct PartitionPruneState
 	PartitionPruningData *partprunedata[FLEXIBLE_ARRAY_MEMBER];
 } PartitionPruneState;
 
-extern PartitionPruneState *ExecInitPartitionPruning(PlanState *planstate,
-													 int n_total_subplans,
-													 int part_prune_index,
-													 Bitmapset *relids,
-													 Bitmapset **initially_valid_subplans);
+extern void ExecDoInitialPruning(EState *estate);
+extern PartitionPruneState *ExecInitPartitionExecPruning(PlanState *planstate,
+														 int n_total_subplans,
+														 int part_prune_index,
+														 Bitmapset *relids,
+														 Bitmapset **initially_valid_subplans);
 extern Bitmapset *ExecFindMatchingSubPlans(PartitionPruneState *prunestate,
 										   bool initial_prune);
 
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index 8ce4430af04..aca15f771a2 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -656,6 +656,8 @@ typedef struct EState
 	List	   *es_rteperminfos;	/* List of RTEPermissionInfo */
 	PlannedStmt *es_plannedstmt;	/* link to top of plan tree */
 	List	   *es_part_prune_infos;	/* List of PartitionPruneInfo */
+	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 */
diff --git a/src/include/partitioning/partprune.h b/src/include/partitioning/partprune.h
index c413734789a..0ed39c89c3d 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

