From d28e281935e97065d14c70b34f903c385c539a66 Mon Sep 17 00:00:00 2001
From: amitlan <amitlangote09@gmail.com>
Date: Wed, 2 Mar 2022 15:17:55 +0900
Subject: [PATCH v9] Some refactoring of runtime pruning code

* Move the execution pruning initialization steps that are common
between both ExecInitAppend() and ExecInitMergeAppend() into a new
function ExecInitPartitionPruning() defined in execPartition.c.
Those steps include creation of a PartitionPruneState to be used for
all instances of pruning and determining the minimal set of child
subplans that need to be initialized by performing initial pruning if
needed, and finally adjusting the subplan_map arrays in the
PartitionPruneState to reflect the new set of subplans remaining
after initial pruning if it was indeed performed.
ExecCreatePartitionPruneState() is no longer exported out of
execPartition.c and has been renamed to CreatePartitionState()
as a local sub-routine of ExecInitPartitionPruning().

* Likewise, ExecFindInitialMatchingSubPlans() that was in the charge
of performing initial pruning no longer needs to be exported.  In
fact, since it would now have the same body as the more generally
named ExecFindMatchingSubPlans(), except differing in the value of
the initial_prune passed to the common subroutine
find_matching_subplans_recurse(), it seems better to just have
ExecFindMatchingSubPlans() with an initial_prune argument.

* Add an ExprContext field to PartitionPruneContext to remove the
implicit assumption in the runtime pruning code that the ExprContext
to use to compute pruning expressions that need one can always rely
on the PlanState providing it.  A future patch will allow runtime
pruning (at least the initial pruning steps) to be performed without
the corresponding PlanState yet having been created, so this will
help.
---
 src/backend/executor/execPartition.c   | 396 +++++++++++++------------
 src/backend/executor/nodeAppend.c      |  41 +--
 src/backend/executor/nodeMergeAppend.c |  34 +--
 src/backend/partitioning/partprune.c   |  20 +-
 src/include/executor/execPartition.h   |  12 +-
 src/include/partitioning/partprune.h   |   2 +
 6 files changed, 260 insertions(+), 245 deletions(-)

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index aca42ca5b8..83451cf654 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -184,11 +184,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,
+							  PartitionPruneInfo *partitionpruneinfo);
 static void ExecInitPruningContext(PartitionPruneContext *context,
 								   List *pruning_steps,
 								   PartitionDesc partdesc,
 								   PartitionKey partkey,
-								   PlanState *planstate);
+								   PlanState *planstate,
+								   ExprContext *econtext);
+static void PartitionPruneStateFixSubPlanMap(PartitionPruneState *prunestate,
+								 Bitmapset *initially_valid_subplans,
+								 int n_total_subplans);
 static void find_matching_subplans_recurse(PartitionPruningData *prunedata,
 										   PartitionedRelPruningData *pprune,
 										   bool initial_prune,
@@ -1590,34 +1596,91 @@ adjust_partition_colnos_using_map(List *colnos, AttrMap *attrMap)
  *
  * Functions:
  *
- * ExecCreatePartitionPruneState:
- *		Creates the PartitionPruneState required by each of the two pruning
- *		functions.  Details stored include how to map the partition index
- *		returned by the partition pruning code into subplan indexes.
- *
- * ExecFindInitialMatchingSubPlans:
- *		Returns indexes of matching subplans.  Partition pruning is attempted
- *		without any evaluation of expressions containing PARAM_EXEC Params.
- *		This function must be called during executor startup for the parent
- *		plan before the subplans themselves are initialized.  Subplans which
- *		are found not to match by this function must be removed from the
- *		plan's list of subplans during execution, as this function performs a
- *		remap of the partition index to subplan index map and the newly
- *		created map provides indexes only for subplans which remain after
- *		calling this function.
+ * 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 initially valid subplans by performing initial pruning steps, only
+ *		which need be initialized by the caller such as ExecInitAppend.  Maps
+ *		in PartitionPruneState are updated to account for initial pruning
+ *		having eliminated some of the subplans, if any.
  *
  * ExecFindMatchingSubPlans:
- *		Returns indexes of matching subplans after evaluating all available
- *		expressions.  This function can only be called during execution and
- *		must be called again each time the value of a Param listed in
+ *		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
  *		PartitionPruneState's 'execparamids' changes.
  *-------------------------------------------------------------------------
  */
 
 /*
- * ExecCreatePartitionPruneState
- *		Build the data structure required for calling
- *		ExecFindInitialMatchingSubPlans and ExecFindMatchingSubPlans.
+ * ExecInitPartitionPruning
+ *		Initialize data structure needed for run-time partition pruning and
+ *		do initial pruning if needed
+ *
+ * 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.
+ *
+ * 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.
+ */
+PartitionPruneState *
+ExecInitPartitionPruning(PlanState *planstate,
+						 int n_total_subplans,
+						 PartitionPruneInfo *pruneinfo,
+						 Bitmapset **initially_valid_subplans)
+{
+	PartitionPruneState *prunestate;
+	EState *estate = planstate->state;
+
+	/* 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.
+	 */
+	if (prunestate->do_initial_prune)
+	{
+		*initially_valid_subplans = ExecFindMatchingSubPlans(prunestate, true);
+	}
+	else
+	{
+		/* No pruning, so we'll need to initialize all subplans */
+		Assert(n_total_subplans > 0);
+		*initially_valid_subplans = bms_add_range(NULL, 0,
+												  n_total_subplans - 1);
+	}
+
+	/*
+	 * Re-sequence subplan indexes contained in prunestate to account for any
+	 * that were removed above due to initial pruning.
+	 *
+	 * 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 &&
+		bms_num_members(*initially_valid_subplans) < n_total_subplans)
+		PartitionPruneStateFixSubPlanMap(prunestate,
+										 *initially_valid_subplans,
+										 n_total_subplans);
+
+	return prunestate;
+}
+
+/*
+ * CreatePartitionPruneState
+ *		Build the data structure required for calling ExecFindMatchingSubPlans
  *
  * 'planstate' is the parent plan node's execution state.
  *
@@ -1632,8 +1695,8 @@ adjust_partition_colnos_using_map(List *colnos, AttrMap *attrMap)
  * re-used each time we re-evaluate which partitions match the pruning steps
  * provided in each PartitionedRelPruneInfo.
  */
-PartitionPruneState *
-ExecCreatePartitionPruneState(PlanState *planstate,
+static PartitionPruneState *
+CreatePartitionPruneState(PlanState *planstate,
 							  PartitionPruneInfo *partitionpruneinfo)
 {
 	EState	   *estate = planstate->state;
@@ -1641,6 +1704,7 @@ ExecCreatePartitionPruneState(PlanState *planstate,
 	int			n_part_hierarchies;
 	ListCell   *lc;
 	int			i;
+	ExprContext	*econtext = planstate->ps_ExprContext;
 
 	/* For data reading, executor always omits detached partitions */
 	if (estate->es_partition_directory == NULL)
@@ -1814,7 +1878,8 @@ ExecCreatePartitionPruneState(PlanState *planstate,
 			{
 				ExecInitPruningContext(&pprune->initial_context,
 									   pinfo->initial_pruning_steps,
-									   partdesc, partkey, planstate);
+									   partdesc, partkey, planstate,
+									   econtext);
 				/* Record whether initial pruning is needed at any level */
 				prunestate->do_initial_prune = true;
 			}
@@ -1823,7 +1888,8 @@ ExecCreatePartitionPruneState(PlanState *planstate,
 			{
 				ExecInitPruningContext(&pprune->exec_context,
 									   pinfo->exec_pruning_steps,
-									   partdesc, partkey, planstate);
+									   partdesc, partkey, planstate,
+									   econtext);
 				/* Record whether exec pruning is needed at any level */
 				prunestate->do_exec_prune = true;
 			}
@@ -1851,7 +1917,8 @@ ExecInitPruningContext(PartitionPruneContext *context,
 					   List *pruning_steps,
 					   PartitionDesc partdesc,
 					   PartitionKey partkey,
-					   PlanState *planstate)
+					   PlanState *planstate,
+					   ExprContext *econtext)
 {
 	int			n_steps;
 	int			partnatts;
@@ -1872,6 +1939,7 @@ ExecInitPruningContext(PartitionPruneContext *context,
 
 	context->ppccontext = CurrentMemoryContext;
 	context->planstate = planstate;
+	context->exprcontext = econtext;
 
 	/* Initialize expression state for each expression we need */
 	context->exprstates = (ExprState **)
@@ -1900,8 +1968,20 @@ ExecInitPruningContext(PartitionPruneContext *context,
 														step->step.step_id,
 														keyno);
 
-				context->exprstates[stateidx] =
-					ExecInitExpr(expr, context->planstate);
+				/*
+				 * When planstate is NULL, pruning_steps is known not to
+				 * contain any expressions that depend on the parent plan.
+				 * Information of any available EXTERN parameters must be
+				 * passed explicitly in that case, which the caller must
+				 * have made available via econtext.
+				 */
+				if (planstate == NULL)
+					context->exprstates[stateidx] =
+						ExecInitExprWithParams(expr,
+											   econtext->ecxt_param_list_info);
+				else
+					context->exprstates[stateidx] =
+						ExecInitExpr(expr, context->planstate);
 			}
 			keyno++;
 		}
@@ -1909,179 +1989,121 @@ ExecInitPruningContext(PartitionPruneContext *context,
 }
 
 /*
- * ExecFindInitialMatchingSubPlans
- *		Identify the set of subplans that cannot be eliminated by initial
- *		pruning, disregarding any pruning constraints involving PARAM_EXEC
- *		Params.
- *
- * If additional pruning passes will be required (because of PARAM_EXEC
- * Params), we must also update the translation data that allows conversion
- * of partition indexes into subplan indexes to account for the unneeded
- * subplans having been removed.
- *
- * Must only be called once per 'prunestate', and only if initial pruning
- * is required.
+ * PartitionPruneStateFixSubPlanMap
+ *		Fix mapping of partition indexes to subplan indexes contained in
+ *		prunestate by considering the new list of subplans that survived
+ *		initial pruning
  *
- * 'nsubplans' must be passed as the total number of unpruned subplans.
+ * Current values of the indexes present in PartitionPruneState count all the
+ * subplans that would be present before initial pruning was done.  If initial
+ * pruning got rid of some of the subplans, any subsequent pruning passes will
+ * will be looking at a different set of target subplans to choose from than
+ * those in the pre-initial-pruning set, so the maps in PartitionPruneState
+ * containing those indexes must be updated to reflect the new indexes of
+ * subplans in the post-initial-pruning set.
  */
-Bitmapset *
-ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubplans)
+static void
+PartitionPruneStateFixSubPlanMap(PartitionPruneState *prunestate,
+								 Bitmapset *initially_valid_subplans,
+								 int n_total_subplans)
 {
-	Bitmapset  *result = NULL;
-	MemoryContext oldcontext;
+	int		   *new_subplan_indexes;
+	Bitmapset  *new_other_subplans;
 	int			i;
-
-	/* Caller error if we get here without do_initial_prune */
-	Assert(prunestate->do_initial_prune);
+	int			newidx;
 
 	/*
-	 * Switch to a temp context to avoid leaking memory in the executor's
-	 * query-lifespan memory context.
+	 * 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.
 	 */
-	oldcontext = MemoryContextSwitchTo(prunestate->prune_context);
-
-	/*
-	 * For each hierarchy, do the pruning tests, and add nondeletable
-	 * subplans' indexes to "result".
-	 */
-	for (i = 0; i < prunestate->num_partprunedata; i++)
+	new_subplan_indexes = (int *) palloc0(sizeof(int) * n_total_subplans);
+	newidx = 1;
+	i = -1;
+	while ((i = bms_next_member(initially_valid_subplans, i)) >= 0)
 	{
-		PartitionPruningData *prunedata;
-		PartitionedRelPruningData *pprune;
-
-		prunedata = prunestate->partprunedata[i];
-		pprune = &prunedata->partrelprunedata[0];
-
-		/* Perform pruning without using PARAM_EXEC Params */
-		find_matching_subplans_recurse(prunedata, pprune, true, &result);
-
-		/* Expression eval may have used space in node's ps_ExprContext too */
-		if (pprune->initial_pruning_steps)
-			ResetExprContext(pprune->initial_context.planstate->ps_ExprContext);
+		Assert(i < n_total_subplans);
+		new_subplan_indexes[i] = newidx++;
 	}
 
-	/* Add in any subplans that partition pruning didn't account for */
-	result = bms_add_members(result, prunestate->other_subplans);
-
-	MemoryContextSwitchTo(oldcontext);
-
-	/* Copy result out of the temp context before we reset it */
-	result = bms_copy(result);
-
-	MemoryContextReset(prunestate->prune_context);
-
 	/*
-	 * If exec-time pruning is required and we pruned subplans above, then we
-	 * must re-sequence the subplan indexes so that ExecFindMatchingSubPlans
-	 * properly returns the indexes from the subplans which will remain after
-	 * execution of this function.
-	 *
-	 * 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).
+	 * Now we can update each PartitionedRelPruneInfo's subplan_map with
+	 * new subplan indexes.  We must also recompute its present_parts
+	 * bitmap.
 	 */
-	if (prunestate->do_exec_prune && bms_num_members(result) < nsubplans)
+	for (i = 0; i < prunestate->num_partprunedata; i++)
 	{
-		int		   *new_subplan_indexes;
-		Bitmapset  *new_other_subplans;
-		int			i;
-		int			newidx;
+		PartitionPruningData *prunedata = prunestate->partprunedata[i];
+		int			j;
 
 		/*
-		 * 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.
+		 * Within each hierarchy, we perform this loop in back-to-front
+		 * order so that we determine present_parts for the lowest-level
+		 * partitioned tables first.  This way we can tell whether a
+		 * sub-partitioned table's partitions were entirely pruned so we
+		 * can exclude it from the current level's present_parts.
 		 */
-		new_subplan_indexes = (int *) palloc0(sizeof(int) * nsubplans);
-		newidx = 1;
-		i = -1;
-		while ((i = bms_next_member(result, i)) >= 0)
+		for (j = prunedata->num_partrelprunedata - 1; j >= 0; j--)
 		{
-			Assert(i < nsubplans);
-			new_subplan_indexes[i] = newidx++;
-		}
+			PartitionedRelPruningData *pprune = &prunedata->partrelprunedata[j];
+			int			nparts = pprune->nparts;
+			int			k;
 
-		/*
-		 * Now we can update each PartitionedRelPruneInfo's subplan_map with
-		 * new subplan indexes.  We must also recompute its present_parts
-		 * bitmap.
-		 */
-		for (i = 0; i < prunestate->num_partprunedata; i++)
-		{
-			PartitionPruningData *prunedata = prunestate->partprunedata[i];
-			int			j;
+			/* We just rebuild present_parts from scratch */
+			bms_free(pprune->present_parts);
+			pprune->present_parts = NULL;
 
-			/*
-			 * Within each hierarchy, we perform this loop in back-to-front
-			 * order so that we determine present_parts for the lowest-level
-			 * partitioned tables first.  This way we can tell whether a
-			 * sub-partitioned table's partitions were entirely pruned so we
-			 * can exclude it from the current level's present_parts.
-			 */
-			for (j = prunedata->num_partrelprunedata - 1; j >= 0; j--)
+			for (k = 0; k < nparts; k++)
 			{
-				PartitionedRelPruningData *pprune = &prunedata->partrelprunedata[j];
-				int			nparts = pprune->nparts;
-				int			k;
-
-				/* We just rebuild present_parts from scratch */
-				bms_free(pprune->present_parts);
-				pprune->present_parts = NULL;
+				int			oldidx = pprune->subplan_map[k];
+				int			subidx;
 
-				for (k = 0; k < nparts; k++)
+				/*
+				 * If this partition existed as a subplan then change the
+				 * old subplan index to the new subplan index.  The new
+				 * index may become -1 if the partition was pruned above,
+				 * or it may just come earlier in the subplan list due to
+				 * some subplans being removed earlier in the list.  If
+				 * it's a subpartition, add it to present_parts unless
+				 * it's entirely pruned.
+				 */
+				if (oldidx >= 0)
 				{
-					int			oldidx = pprune->subplan_map[k];
-					int			subidx;
+					Assert(oldidx < n_total_subplans);
+					pprune->subplan_map[k] = new_subplan_indexes[oldidx] - 1;
 
-					/*
-					 * If this partition existed as a subplan then change the
-					 * old subplan index to the new subplan index.  The new
-					 * index may become -1 if the partition was pruned above,
-					 * or it may just come earlier in the subplan list due to
-					 * some subplans being removed earlier in the list.  If
-					 * it's a subpartition, add it to present_parts unless
-					 * it's entirely pruned.
-					 */
-					if (oldidx >= 0)
-					{
-						Assert(oldidx < nsubplans);
-						pprune->subplan_map[k] = new_subplan_indexes[oldidx] - 1;
-
-						if (new_subplan_indexes[oldidx] > 0)
-							pprune->present_parts =
-								bms_add_member(pprune->present_parts, k);
-					}
-					else if ((subidx = pprune->subpart_map[k]) >= 0)
-					{
-						PartitionedRelPruningData *subprune;
+					if (new_subplan_indexes[oldidx] > 0)
+						pprune->present_parts =
+							bms_add_member(pprune->present_parts, k);
+				}
+				else if ((subidx = pprune->subpart_map[k]) >= 0)
+				{
+					PartitionedRelPruningData *subprune;
 
-						subprune = &prunedata->partrelprunedata[subidx];
+					subprune = &prunedata->partrelprunedata[subidx];
 
-						if (!bms_is_empty(subprune->present_parts))
-							pprune->present_parts =
-								bms_add_member(pprune->present_parts, k);
-					}
+					if (!bms_is_empty(subprune->present_parts))
+						pprune->present_parts =
+							bms_add_member(pprune->present_parts, k);
 				}
 			}
 		}
+	}
 
-		/*
-		 * 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);
-
-		bms_free(prunestate->other_subplans);
-		prunestate->other_subplans = new_other_subplans;
+	/*
+	 * 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);
 
-		pfree(new_subplan_indexes);
-	}
+	bms_free(prunestate->other_subplans);
+	prunestate->other_subplans = new_other_subplans;
 
-	return result;
+	pfree(new_subplan_indexes);
 }
 
 /*
@@ -2089,21 +2111,26 @@ ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubplans)
  *		Determine which subplans match the pruning steps detailed in
  *		'prunestate' for the current comparison expression values.
  *
- * Here we assume we may evaluate PARAM_EXEC Params.
+ * For example, if initial_prune is true, the caller is telling us that only
+ * those pruning steps that are known to not contain any expressions involving
+ * PARAM_EXEC Params are safe to evaluate at this point.  Whereas when it's
+ * false, it is telling that PARAM_EXEC Params can be safely evaluated and so
+ * also the pruning steps that contain them.
  */
 Bitmapset *
-ExecFindMatchingSubPlans(PartitionPruneState *prunestate)
+ExecFindMatchingSubPlans(PartitionPruneState *prunestate,
+						 bool initial_prune)
 {
 	Bitmapset  *result = NULL;
 	MemoryContext oldcontext;
 	int			i;
 
 	/*
-	 * If !do_exec_prune, we've got problems because
-	 * ExecFindInitialMatchingSubPlans will not have bothered to update
-	 * prunestate for whatever pruning it did.
+	 * Only get here if prunestate->do_exec_prune, because otherwise
+	 * ExecInitPartitionPruning() would not have bothered to update prunestate
+	 * to account for the subplans removed by initial pruning.
 	 */
-	Assert(prunestate->do_exec_prune);
+	Assert(prunestate->do_exec_prune || initial_prune);
 
 	/*
 	 * Switch to a temp context to avoid leaking memory in the executor's
@@ -2123,11 +2150,17 @@ ExecFindMatchingSubPlans(PartitionPruneState *prunestate)
 		prunedata = prunestate->partprunedata[i];
 		pprune = &prunedata->partrelprunedata[0];
 
-		find_matching_subplans_recurse(prunedata, pprune, false, &result);
+		/*
+		 * We pass the 1st item belonging to the root table of the hierarchy
+		 * and find_matching_subplans_recurse() takes care of recursing to
+		 * other (lower-level) parents as needed.
+		 */
+		find_matching_subplans_recurse(prunedata, pprune, initial_prune,
+									   &result);
 
-		/* Expression eval may have used space in node's ps_ExprContext too */
+		/* Expression eval may have used space in ExprContext too */
 		if (pprune->exec_pruning_steps)
-			ResetExprContext(pprune->exec_context.planstate->ps_ExprContext);
+			ResetExprContext(pprune->exec_context.exprcontext);
 	}
 
 	/* Add in any subplans that partition pruning didn't account for */
@@ -2145,8 +2178,7 @@ ExecFindMatchingSubPlans(PartitionPruneState *prunestate)
 
 /*
  * find_matching_subplans_recurse
- *		Recursive worker function for ExecFindMatchingSubPlans and
- *		ExecFindInitialMatchingSubPlans
+ *		Recursive worker function for ExecFindMatchingSubPlans
  *
  * Adds valid (non-prunable) subplan IDs to *validsubplans
  */
diff --git a/src/backend/executor/nodeAppend.c b/src/backend/executor/nodeAppend.c
index 7937f1c88f..357e10a1d7 100644
--- a/src/backend/executor/nodeAppend.c
+++ b/src/backend/executor/nodeAppend.c
@@ -138,30 +138,17 @@ ExecInitAppend(Append *node, EState *estate, int eflags)
 	{
 		PartitionPruneState *prunestate;
 
-		/* We may need an expression context to evaluate partition exprs */
-		ExecAssignExprContext(estate, &appendstate->ps);
-
-		/* Create the working data structure for pruning. */
-		prunestate = ExecCreatePartitionPruneState(&appendstate->ps,
-												   node->part_prune_info);
+		/*
+		 * Set up pruning data structure.  This also initializes the set of
+		 * 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_info,
+											  &validsubplans);
 		appendstate->as_prune_state = prunestate;
-
-		/* Perform an initial partition prune, if required. */
-		if (prunestate->do_initial_prune)
-		{
-			/* Determine which subplans survive initial pruning */
-			validsubplans = ExecFindInitialMatchingSubPlans(prunestate,
-															list_length(node->appendplans));
-
-			nplans = bms_num_members(validsubplans);
-		}
-		else
-		{
-			/* We'll need to initialize all subplans */
-			nplans = list_length(node->appendplans);
-			Assert(nplans > 0);
-			validsubplans = bms_add_range(NULL, 0, nplans - 1);
-		}
+		nplans = bms_num_members(validsubplans);
 
 		/*
 		 * When no run-time pruning is required and there's at least one
@@ -590,7 +577,7 @@ choose_next_subplan_locally(AppendState *node)
 		}
 		else if (node->as_valid_subplans == NULL)
 			node->as_valid_subplans =
-				ExecFindMatchingSubPlans(node->as_prune_state);
+				ExecFindMatchingSubPlans(node->as_prune_state, false);
 
 		whichplan = -1;
 	}
@@ -655,7 +642,7 @@ choose_next_subplan_for_leader(AppendState *node)
 		if (node->as_valid_subplans == NULL)
 		{
 			node->as_valid_subplans =
-				ExecFindMatchingSubPlans(node->as_prune_state);
+				ExecFindMatchingSubPlans(node->as_prune_state, false);
 
 			/*
 			 * Mark each invalid plan as finished to allow the loop below to
@@ -730,7 +717,7 @@ choose_next_subplan_for_worker(AppendState *node)
 	else if (node->as_valid_subplans == NULL)
 	{
 		node->as_valid_subplans =
-			ExecFindMatchingSubPlans(node->as_prune_state);
+			ExecFindMatchingSubPlans(node->as_prune_state, false);
 		mark_invalid_subplans_as_finished(node);
 	}
 
@@ -881,7 +868,7 @@ ExecAppendAsyncBegin(AppendState *node)
 	if (node->as_valid_subplans == NULL)
 	{
 		node->as_valid_subplans =
-			ExecFindMatchingSubPlans(node->as_prune_state);
+			ExecFindMatchingSubPlans(node->as_prune_state, false);
 
 		classify_matching_subplans(node);
 	}
diff --git a/src/backend/executor/nodeMergeAppend.c b/src/backend/executor/nodeMergeAppend.c
index 418f89dea8..ecf9052e03 100644
--- a/src/backend/executor/nodeMergeAppend.c
+++ b/src/backend/executor/nodeMergeAppend.c
@@ -86,29 +86,17 @@ ExecInitMergeAppend(MergeAppend *node, EState *estate, int eflags)
 	{
 		PartitionPruneState *prunestate;
 
-		/* We may need an expression context to evaluate partition exprs */
-		ExecAssignExprContext(estate, &mergestate->ps);
-
-		prunestate = ExecCreatePartitionPruneState(&mergestate->ps,
-												   node->part_prune_info);
+		/*
+		 * Set up pruning data structure.  This also initializes the set of
+		 * 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_info,
+											  &validsubplans);
 		mergestate->ms_prune_state = prunestate;
-
-		/* Perform an initial partition prune, if required. */
-		if (prunestate->do_initial_prune)
-		{
-			/* Determine which subplans survive initial pruning */
-			validsubplans = ExecFindInitialMatchingSubPlans(prunestate,
-															list_length(node->mergeplans));
-
-			nplans = bms_num_members(validsubplans);
-		}
-		else
-		{
-			/* We'll need to initialize all subplans */
-			nplans = list_length(node->mergeplans);
-			Assert(nplans > 0);
-			validsubplans = bms_add_range(NULL, 0, nplans - 1);
-		}
+		nplans = bms_num_members(validsubplans);
 
 		/*
 		 * When no run-time pruning is required and there's at least one
@@ -230,7 +218,7 @@ ExecMergeAppend(PlanState *pstate)
 		 */
 		if (node->ms_valid_subplans == NULL)
 			node->ms_valid_subplans =
-				ExecFindMatchingSubPlans(node->ms_prune_state);
+				ExecFindMatchingSubPlans(node->ms_prune_state, false);
 
 		/*
 		 * First time through: pull the first tuple from each valid subplan,
diff --git a/src/backend/partitioning/partprune.c b/src/backend/partitioning/partprune.c
index 1bc00826c1..7080cb25d9 100644
--- a/src/backend/partitioning/partprune.c
+++ b/src/backend/partitioning/partprune.c
@@ -798,6 +798,7 @@ prune_append_rel_partitions(RelOptInfo *rel)
 
 	/* These are not valid when being called from the planner */
 	context.planstate = NULL;
+	context.exprcontext = NULL;
 	context.exprstates = NULL;
 
 	/* Actual pruning happens here. */
@@ -808,8 +809,8 @@ prune_append_rel_partitions(RelOptInfo *rel)
  * get_matching_partitions
  *		Determine partitions that survive partition pruning
  *
- * Note: context->planstate must be set to a valid PlanState when the
- * pruning_steps were generated with a target other than PARTTARGET_PLANNER.
+ * Note: context->exprcontext must be valid when the pruning_steps were
+ * generated with a target other than PARTTARGET_PLANNER.
  *
  * Returns a Bitmapset of the RelOptInfo->part_rels indexes of the surviving
  * partitions.
@@ -3654,7 +3655,7 @@ match_boolean_partition_clause(Oid partopfamily, Expr *clause, Expr *partkey,
  * exprstate array.
  *
  * Note that the evaluated result may be in the per-tuple memory context of
- * context->planstate->ps_ExprContext, and we may have leaked other memory
+ * context->exprcontext, and we may have leaked other memory
  * there too.  This memory must be recovered by resetting that ExprContext
  * after we're done with the pruning operation (see execPartition.c).
  */
@@ -3677,13 +3678,18 @@ partkey_datum_from_expr(PartitionPruneContext *context,
 		ExprContext *ectx;
 
 		/*
-		 * We should never see a non-Const in a step unless we're running in
-		 * the executor.
+		 * 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);
+		Assert(context->planstate != NULL || context->exprcontext != NULL);
+		Assert(context->planstate == NULL ||
+			   (context->exprcontext == context->planstate->ps_ExprContext));
 
 		exprstate = context->exprstates[stateidx];
-		ectx = context->planstate->ps_ExprContext;
+		ectx = context->exprcontext;
 		*value = ExecEvalExprSwitchContext(exprstate, ectx, isnull);
 	}
 }
diff --git a/src/include/executor/execPartition.h b/src/include/executor/execPartition.h
index 603d8becc4..4c706c11b9 100644
--- a/src/include/executor/execPartition.h
+++ b/src/include/executor/execPartition.h
@@ -119,10 +119,10 @@ extern ResultRelInfo *ExecFindPartition(ModifyTableState *mtstate,
 										EState *estate);
 extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
 									PartitionTupleRouting *proute);
-extern PartitionPruneState *ExecCreatePartitionPruneState(PlanState *planstate,
-														  PartitionPruneInfo *partitionpruneinfo);
-extern Bitmapset *ExecFindMatchingSubPlans(PartitionPruneState *prunestate);
-extern Bitmapset *ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate,
-												  int nsubplans);
-
+extern PartitionPruneState *ExecInitPartitionPruning(PlanState *planstate,
+						 int n_total_subplans,
+						 PartitionPruneInfo *pruneinfo,
+						 Bitmapset **initially_valid_subplans);
+extern Bitmapset *ExecFindMatchingSubPlans(PartitionPruneState *prunestate,
+						 bool initial_prune);
 #endif							/* EXECPARTITION_H */
diff --git a/src/include/partitioning/partprune.h b/src/include/partitioning/partprune.h
index ee11b6feae..90684efa25 100644
--- a/src/include/partitioning/partprune.h
+++ b/src/include/partitioning/partprune.h
@@ -41,6 +41,7 @@ struct RelOptInfo;
  *					subsidiary data, such as the FmgrInfos.
  * planstate		Points to the parent plan node's PlanState when called
  *					during execution; NULL when called from the planner.
+ * exprcontext		ExprContext to use when evaluating pruning expressions
  * exprstates		Array of ExprStates, indexed as per PruneCxtStateIdx; one
  *					for each partition key in each pruning step.  Allocated if
  *					planstate is non-NULL, otherwise NULL.
@@ -56,6 +57,7 @@ typedef struct PartitionPruneContext
 	FmgrInfo   *stepcmpfuncs;
 	MemoryContext ppccontext;
 	PlanState  *planstate;
+	ExprContext *exprcontext;
 	ExprState **exprstates;
 } PartitionPruneContext;
 
-- 
2.24.1

