Refactor and cleanup runtime partition prune code a little

* 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 CreatePartitionPruneState()
as a local sub-routine of ExecInitPartitionPruning().

* Likewise, ExecFindInitialMatchingSubPlans() that was in 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
initial_prune passed to the common subroutine
find_matching_subplans_recurse(), it seems better to remove it and add
an initial_prune argument to ExecFindMatchingSubPlans().

* 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.

Author: Amit Langote <amitlangote09@gmail.com>
Discussion: https://postgr.es/m/CA+HiwqEYCpEqh2LMDOp9mT+4-QoVe8HgFMKBjntEMCTZLpcCCA@mail.gmail.com

1 files changed, 239 insertions, 220 deletions

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index aca42ca5b8c..615bd809735 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 void ExecInitPruningContext(PartitionPruneContext *context,
-								   List *pruning_steps,
-								   PartitionDesc partdesc,
-								   PartitionKey partkey,
-								   PlanState *planstate);
+static PartitionPruneState *CreatePartitionPruneState(PlanState *planstate,
+													  PartitionPruneInfo *pruneinfo);
+static void InitPartitionPruneContext(PartitionPruneContext *context,
+									  List *pruning_steps,
+									  PartitionDesc partdesc,
+									  PartitionKey partkey,
+									  PlanState *planstate,
+									  ExprContext *econtext);
+static void PartitionPruneFixSubPlanMap(PartitionPruneState *prunestate,
+										Bitmapset *initially_valid_subplans,
+										int n_total_subplans);
 static void find_matching_subplans_recurse(PartitionPruningData *prunedata,
 										   PartitionedRelPruningData *pprune,
 										   bool initial_prune,
@@ -1590,64 +1596,121 @@ 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 subplans to initialize considering the result of performing initial
+ *		pruning steps if any.  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 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.  No need to do this if
+	 * no steps were removed.
+	 */
+	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);
+	}
+
+	return prunestate;
+}
+
+/*
+ * CreatePartitionPruneState
+ *		Build the data structure required for calling ExecFindMatchingSubPlans
  *
  * 'planstate' is the parent plan node's execution state.
  *
- * 'partitionpruneinfo' is a PartitionPruneInfo as generated by
+ * 'pruneinfo' is a PartitionPruneInfo as generated by
  * make_partition_pruneinfo.  Here we build a PartitionPruneState containing a
  * PartitionPruningData for each partitioning hierarchy (i.e., each sublist of
- * partitionpruneinfo->prune_infos), each of which contains a
- * PartitionedRelPruningData for each PartitionedRelPruneInfo appearing in
- * that sublist.  This two-level system is needed to keep from confusing the
- * different hierarchies when a UNION ALL contains multiple partitioned tables
- * as children.  The data stored in each PartitionedRelPruningData can be
- * re-used each time we re-evaluate which partitions match the pruning steps
- * provided in each PartitionedRelPruneInfo.
+ * pruneinfo->prune_infos), each of which contains a PartitionedRelPruningData
+ * for each PartitionedRelPruneInfo appearing in that sublist.  This two-level
+ * system is needed to keep from confusing the different hierarchies when a
+ * UNION ALL contains multiple partitioned tables as children.  The data
+ * stored in each PartitionedRelPruningData can be re-used each time we
+ * re-evaluate which partitions match the pruning steps provided in each
+ * PartitionedRelPruneInfo.
  */
-PartitionPruneState *
-ExecCreatePartitionPruneState(PlanState *planstate,
-							  PartitionPruneInfo *partitionpruneinfo)
+static PartitionPruneState *
+CreatePartitionPruneState(PlanState *planstate, PartitionPruneInfo *pruneinfo)
 {
 	EState	   *estate = planstate->state;
 	PartitionPruneState *prunestate;
 	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)
 		estate->es_partition_directory =
 			CreatePartitionDirectory(estate->es_query_cxt, false);
 
-	n_part_hierarchies = list_length(partitionpruneinfo->prune_infos);
+	n_part_hierarchies = list_length(pruneinfo->prune_infos);
 	Assert(n_part_hierarchies > 0);
 
 	/*
@@ -1659,7 +1722,7 @@ ExecCreatePartitionPruneState(PlanState *planstate,
 
 	prunestate->execparamids = NULL;
 	/* other_subplans can change at runtime, so we need our own copy */
-	prunestate->other_subplans = bms_copy(partitionpruneinfo->other_subplans);
+	prunestate->other_subplans = bms_copy(pruneinfo->other_subplans);
 	prunestate->do_initial_prune = false;	/* may be set below */
 	prunestate->do_exec_prune = false;	/* may be set below */
 	prunestate->num_partprunedata = n_part_hierarchies;
@@ -1676,7 +1739,7 @@ ExecCreatePartitionPruneState(PlanState *planstate,
 							  ALLOCSET_DEFAULT_SIZES);
 
 	i = 0;
-	foreach(lc, partitionpruneinfo->prune_infos)
+	foreach(lc, pruneinfo->prune_infos)
 	{
 		List	   *partrelpruneinfos = lfirst_node(List, lc);
 		int			npartrelpruneinfos = list_length(partrelpruneinfos);
@@ -1812,18 +1875,20 @@ ExecCreatePartitionPruneState(PlanState *planstate,
 			pprune->initial_pruning_steps = pinfo->initial_pruning_steps;
 			if (pinfo->initial_pruning_steps)
 			{
-				ExecInitPruningContext(&pprune->initial_context,
-									   pinfo->initial_pruning_steps,
-									   partdesc, partkey, planstate);
+				InitPartitionPruneContext(&pprune->initial_context,
+										  pinfo->initial_pruning_steps,
+										  partdesc, partkey, planstate,
+										  econtext);
 				/* Record whether initial pruning is needed at any level */
 				prunestate->do_initial_prune = true;
 			}
 			pprune->exec_pruning_steps = pinfo->exec_pruning_steps;
 			if (pinfo->exec_pruning_steps)
 			{
-				ExecInitPruningContext(&pprune->exec_context,
-									   pinfo->exec_pruning_steps,
-									   partdesc, partkey, planstate);
+				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;
 			}
@@ -1847,11 +1912,12 @@ ExecCreatePartitionPruneState(PlanState *planstate,
  * Initialize a PartitionPruneContext for the given list of pruning steps.
  */
 static void
-ExecInitPruningContext(PartitionPruneContext *context,
-					   List *pruning_steps,
-					   PartitionDesc partdesc,
-					   PartitionKey partkey,
-					   PlanState *planstate)
+InitPartitionPruneContext(PartitionPruneContext *context,
+						  List *pruning_steps,
+						  PartitionDesc partdesc,
+						  PartitionKey partkey,
+						  PlanState *planstate,
+						  ExprContext *econtext)
 {
 	int			n_steps;
 	int			partnatts;
@@ -1872,6 +1938,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 +1967,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 +1988,119 @@ ExecInitPruningContext(PartitionPruneContext *context,
 }
 
 /*
- * ExecFindInitialMatchingSubPlans
- *		Identify the set of subplans that cannot be eliminated by initial
- *		pruning, disregarding any pruning constraints involving PARAM_EXEC
- *		Params.
+ * PartitionPruneFixSubPlanMap
+ *		Fix mapping of partition indexes to subplan indexes contained in
+ *		prunestate by considering the new list of subplans that survived
+ *		initial pruning
  *
- * 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.
- *
- * '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
+PartitionPruneFixSubPlanMap(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 +2108,24 @@ 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.
+ * Pass initial_prune if PARAM_EXEC Params cannot yet be evaluated.  This
+ * differentiates the initial executor-time pruning step from later
+ * runtime pruning.
  */
 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.
+	 * Either we're here on the initial prune done during pruning
+	 * initialization, or we're at a point where PARAM_EXEC Params can be
+	 * evaluated *and* there are steps in which to do so.
 	 */
-	Assert(prunestate->do_exec_prune);
+	Assert(initial_prune || prunestate->do_exec_prune);
 
 	/*
 	 * Switch to a temp context to avoid leaking memory in the executor's
@@ -2117,17 +2139,21 @@ ExecFindMatchingSubPlans(PartitionPruneState *prunestate)
 	 */
 	for (i = 0; i < prunestate->num_partprunedata; i++)
 	{
-		PartitionPruningData *prunedata;
+		PartitionPruningData *prunedata = prunestate->partprunedata[i];
 		PartitionedRelPruningData *pprune;
 
-		prunedata = prunestate->partprunedata[i];
+		/*
+		 * We pass the zeroth 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];
+		find_matching_subplans_recurse(prunedata, pprune, initial_prune,
+									   &result);
 
-		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 */
@@ -2145,8 +2171,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
  */
@@ -2162,25 +2187,19 @@ find_matching_subplans_recurse(PartitionPruningData *prunedata,
 	/* Guard against stack overflow due to overly deep partition hierarchy. */
 	check_stack_depth();
 
-	/* Only prune if pruning would be useful at this level. */
+	/*
+	 * Prune as appropriate, if we have pruning steps matching the current
+	 * execution context.  Otherwise just include all partitions at this
+	 * level.
+	 */
 	if (initial_prune && pprune->initial_pruning_steps)
-	{
 		partset = get_matching_partitions(&pprune->initial_context,
 										  pprune->initial_pruning_steps);
-	}
 	else if (!initial_prune && pprune->exec_pruning_steps)
-	{
 		partset = get_matching_partitions(&pprune->exec_context,
 										  pprune->exec_pruning_steps);
-	}
 	else
-	{
-		/*
-		 * If no pruning is to be done, just include all partitions at this
-		 * level.
-		 */
 		partset = pprune->present_parts;
-	}
 
 	/* Translate partset into subplan indexes */
 	i = -1;