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

This does two things mainly:

* Move the execution pruning initialization steps that are common
between both ExecInitAppend() and ExecInitMergeAppend() into a new
function ExecInitPartitionPruning() defined in execPartition.c.
Thus, ExecFindInitialMatchingSubPlans() need not be exported.

* 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   | 340 ++++++++++++++++---------
 src/backend/executor/nodeAppend.c      |  33 +--
 src/backend/executor/nodeMergeAppend.c |  32 +--
 src/backend/partitioning/partprune.c   |  20 +-
 src/include/executor/execPartition.h   |   9 +-
 src/include/partitioning/partprune.h   |   2 +
 6 files changed, 255 insertions(+), 181 deletions(-)

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index 90ed1485d1..21953f253b 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -182,11 +182,18 @@ static char *ExecBuildSlotPartitionKeyDescription(Relation rel,
 												  bool *isnull,
 												  int maxfieldlen);
 static List *adjust_partition_colnos(List *colnos, ResultRelInfo *leaf_part_rri);
+static PartitionPruneState *ExecCreatePartitionPruneState(PlanState *planstate,
+							  PartitionPruneInfo *partitionpruneinfo);
+static Bitmapset *ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate);
 static void ExecInitPruningContext(PartitionPruneContext *context,
 								   List *pruning_steps,
 								   PartitionDesc partdesc,
 								   PartitionKey partkey,
-								   PlanState *planstate);
+								   PlanState *planstate,
+								   ExprContext *econtext);
+static void ExecPartitionPruneFixSubPlanIndexes(PartitionPruneState *prunestate,
+									Bitmapset *initially_valid_subplans,
+									int n_total_subplans);
 static void find_matching_subplans_recurse(PartitionPruningData *prunedata,
 										   PartitionedRelPruningData *pprune,
 										   bool initial_prune,
@@ -1485,30 +1492,87 @@ adjust_partition_colnos(List *colnos, ResultRelInfo *leaf_part_rri)
  *
  * 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:
+ *		Sets up run-time pruning data structure (PartitionPruneState) that is
+ *		needed by each of the two pruning functions.  Also determines the set
+ *		of initially valid subplans by performing initial pruning steps,
+ *		telling the caller (such as ExecInitAppend) to initialize only those
+ *		for execution.  Maps in PartitionPruneState that are used to map the
+ *		partition indexes returned by partprune.c functions into the indexes
+ *		of partition's subplans in the parent node (such as Append) 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
- *		PartitionPruneState's 'execparamids' changes.
+ *		expressions, that is, using execution pruning steps.  This function can
+ *		can only be called during execution and 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
+ *
+ * Initial pruning can be done immediately, so it is done here if needed and
+ * the set of surviving partition subplans' indexes are added to the output
+ * parameter *initially_valid_subplans.  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 = ExecCreatePartitionPruneState(planstate, pruneinfo);
+
+	/*
+	 * Perform an initial partition prune, if required.
+	 */
+	if (prunestate->do_initial_prune)
+	{
+		/* Determine which subplans survive initial pruning */
+		*initially_valid_subplans = ExecFindInitialMatchingSubPlans(prunestate);
+	}
+	else
+	{
+		/* 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)
+		ExecPartitionPruneFixSubPlanIndexes(prunestate,
+											*initially_valid_subplans,
+											n_total_subplans);
+
+	return prunestate;
+}
+
 /*
  * ExecCreatePartitionPruneState
  *		Build the data structure required for calling
@@ -1527,7 +1591,7 @@ adjust_partition_colnos(List *colnos, ResultRelInfo *leaf_part_rri)
  * re-used each time we re-evaluate which partitions match the pruning steps
  * provided in each PartitionedRelPruneInfo.
  */
-PartitionPruneState *
+static PartitionPruneState *
 ExecCreatePartitionPruneState(PlanState *planstate,
 							  PartitionPruneInfo *partitionpruneinfo)
 {
@@ -1536,6 +1600,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)
@@ -1709,7 +1774,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;
 			}
@@ -1718,7 +1784,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;
 			}
@@ -1746,7 +1813,8 @@ ExecInitPruningContext(PartitionPruneContext *context,
 					   List *pruning_steps,
 					   PartitionDesc partdesc,
 					   PartitionKey partkey,
-					   PlanState *planstate)
+					   PlanState *planstate,
+					   ExprContext *econtext)
 {
 	int			n_steps;
 	int			partnatts;
@@ -1767,6 +1835,7 @@ ExecInitPruningContext(PartitionPruneContext *context,
 
 	context->ppccontext = CurrentMemoryContext;
 	context->planstate = planstate;
+	context->exprcontext = econtext;
 
 	/* Initialize expression state for each expression we need */
 	context->exprstates = (ExprState **)
@@ -1795,8 +1864,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++;
 		}
@@ -1816,11 +1897,9 @@ ExecInitPruningContext(PartitionPruneContext *context,
  *
  * Must only be called once per 'prunestate', and only if initial pruning
  * is required.
- *
- * 'nsubplans' must be passed as the total number of unpruned subplans.
  */
-Bitmapset *
-ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubplans)
+static Bitmapset *
+ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate)
 {
 	Bitmapset  *result = NULL;
 	MemoryContext oldcontext;
@@ -1845,14 +1924,20 @@ ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubplans)
 		PartitionedRelPruningData *pprune;
 
 		prunedata = prunestate->partprunedata[i];
+
+		/*
+		 * 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.
+		 */
 		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 */
+		/* Expression eval may have used space in ExprContext too */
 		if (pprune->initial_pruning_steps)
-			ResetExprContext(pprune->initial_context.planstate->ps_ExprContext);
+			ResetExprContext(pprune->initial_context.exprcontext);
 	}
 
 	/* Add in any subplans that partition pruning didn't account for */
@@ -1865,118 +1950,120 @@ ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubplans)
 
 	MemoryContextReset(prunestate->prune_context);
 
+	return result;
+}
+
+/*
+ * ExecPartitionPruneFixSubPlanIndexes
+ *		Fix mapping of partition indexes to subplan indexes contained in
+ *		prunestate by considering the new list of subplans that survived
+ *		initial pruning
+ *
+ * Subplans would be previously indexed 0..(n_total_subplans - 1), though
+ * now should be changed to index range 0..num(initially_valid_subplans).
+ */
+static void
+ExecPartitionPruneFixSubPlanIndexes(PartitionPruneState *prunestate,
+									Bitmapset *initially_valid_subplans,
+									int n_total_subplans)
+{
+	int		   *new_subplan_indexes;
+	Bitmapset  *new_other_subplans;
+	int			i;
+	int			newidx;
+
 	/*
-	 * 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).
+	 * 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.
 	 */
-	if (prunestate->do_exec_prune && bms_num_members(result) < nsubplans)
+	new_subplan_indexes = (int *) palloc0(sizeof(int) * n_total_subplans);
+	newidx = 1;
+	i = -1;
+	while ((i = bms_next_member(initially_valid_subplans, i)) >= 0)
 	{
-		int		   *new_subplan_indexes;
-		Bitmapset  *new_other_subplans;
-		int			i;
-		int			newidx;
+		Assert(i < n_total_subplans);
+		new_subplan_indexes[i] = newidx++;
+	}
 
-		/*
-		 * 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) * nsubplans);
-		newidx = 1;
-		i = -1;
-		while ((i = bms_next_member(result, i)) >= 0)
-		{
-			Assert(i < nsubplans);
-			new_subplan_indexes[i] = newidx++;
-		}
+	/*
+	 * 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;
 
 		/*
-		 * Now we can update each PartitionedRelPruneInfo's subplan_map with
-		 * new subplan indexes.  We must also recompute its present_parts
-		 * bitmap.
+		 * 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 (i = 0; i < prunestate->num_partprunedata; i++)
+		for (j = prunedata->num_partrelprunedata - 1; j >= 0; j--)
 		{
-			PartitionPruningData *prunedata = prunestate->partprunedata[i];
-			int			j;
+			PartitionedRelPruningData *pprune = &prunedata->partrelprunedata[j];
+			int			nparts = pprune->nparts;
+			int			k;
 
-			/*
-			 * 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--)
-			{
-				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;
 
-				/* We just rebuild present_parts from scratch */
-				bms_free(pprune->present_parts);
-				pprune->present_parts = NULL;
+			for (k = 0; k < nparts; k++)
+			{
+				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;
-
-					/*
-					 * 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;
+					Assert(oldidx < n_total_subplans);
+					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);
 }
 
 /*
@@ -2018,11 +2105,16 @@ ExecFindMatchingSubPlans(PartitionPruneState *prunestate)
 		prunedata = prunestate->partprunedata[i];
 		pprune = &prunedata->partrelprunedata[0];
 
+		/*
+		 * 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, false, &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 */
diff --git a/src/backend/executor/nodeAppend.c b/src/backend/executor/nodeAppend.c
index 7937f1c88f..5b6d3eb23b 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.  Initial pruning steps, if any, are
+		 * performed as part of the setup, adding the set of indexes of
+		 * surviving subplans to 'validsubplans'.
+		 */
+		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
diff --git a/src/backend/executor/nodeMergeAppend.c b/src/backend/executor/nodeMergeAppend.c
index 418f89dea8..9a9f29e845 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.  Initial pruning steps, if any, are
+		 * performed as part of the setup, adding the set of indexes of
+		 * surviving subplans to 'validsubplans'.
+		 */
+		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
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..fd5735a946 100644
--- a/src/include/executor/execPartition.h
+++ b/src/include/executor/execPartition.h
@@ -119,10 +119,9 @@ extern ResultRelInfo *ExecFindPartition(ModifyTableState *mtstate,
 										EState *estate);
 extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
 									PartitionTupleRouting *proute);
-extern PartitionPruneState *ExecCreatePartitionPruneState(PlanState *planstate,
-														  PartitionPruneInfo *partitionpruneinfo);
+extern PartitionPruneState *ExecInitPartitionPruning(PlanState *planstate,
+						 int n_total_subplans,
+						 PartitionPruneInfo *pruneinfo,
+						 Bitmapset **initially_valid_subplans);
 extern Bitmapset *ExecFindMatchingSubPlans(PartitionPruneState *prunestate);
-extern Bitmapset *ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate,
-												  int nsubplans);
-
 #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

