Expand run-time partition pruning to work with MergeAppend

This expands the support for the run-time partition pruning which was added
for Append in 499be013de to also allow unneeded subnodes of a MergeAppend
to be removed.

Author: David Rowley
Discussion: https://www.postgresql.org/message-id/CAKJS1f_F_V8D7Wu-HVdnH7zCUxhoGK8XhLLtd%3DCu85qDZzXrgg%40mail.gmail.com

6 files changed, 160 insertions, 29 deletions

diff --git a/src/backend/executor/nodeAppend.c b/src/backend/executor/nodeAppend.c
index 5ce4fb43e1a..e7188b2d310 100644
--- a/src/backend/executor/nodeAppend.c
+++ b/src/backend/executor/nodeAppend.c
@@ -151,7 +151,7 @@ ExecInitAppend(Append *node, EState *estate, int eflags)
 			/*
 			 * The case where no subplans survive pruning must be handled
 			 * specially.  The problem here is that code in explain.c requires
-			 * an Append to have at least one subplan in order for it to
+			 * a MergeAppend to have at least one subplan in order for it to
 			 * properly determine the Vars in that subplan's targetlist.  We
 			 * sidestep this issue by just initializing the first subplan and
 			 * setting as_whichplan to NO_MATCHING_SUBPLANS to indicate that
diff --git a/src/backend/executor/nodeMergeAppend.c b/src/backend/executor/nodeMergeAppend.c
index 118f4ef07df..ec8a49c3a84 100644
--- a/src/backend/executor/nodeMergeAppend.c
+++ b/src/backend/executor/nodeMergeAppend.c
@@ -39,6 +39,7 @@
 #include "postgres.h"
 
 #include "executor/execdebug.h"
+#include "executor/execPartition.h"
 #include "executor/nodeMergeAppend.h"
 #include "lib/binaryheap.h"
 #include "miscadmin.h"
@@ -65,8 +66,10 @@ ExecInitMergeAppend(MergeAppend *node, EState *estate, int eflags)
 {
 	MergeAppendState *mergestate = makeNode(MergeAppendState);
 	PlanState **mergeplanstates;
+	Bitmapset  *validsubplans;
 	int			nplans;
-	int			i;
+	int			i,
+				j;
 	ListCell   *lc;
 
 	/* check for unsupported flags */
@@ -79,18 +82,79 @@ ExecInitMergeAppend(MergeAppend *node, EState *estate, int eflags)
 	ExecLockNonLeafAppendTables(node->partitioned_rels, estate);
 
 	/*
-	 * Set up empty vector of subplan states
-	 */
-	nplans = list_length(node->mergeplans);
-
-	mergeplanstates = (PlanState **) palloc0(nplans * sizeof(PlanState *));
-
-	/*
 	 * create new MergeAppendState for our node
 	 */
 	mergestate->ps.plan = (Plan *) node;
 	mergestate->ps.state = estate;
 	mergestate->ps.ExecProcNode = ExecMergeAppend;
+	mergestate->ms_noopscan = false;
+
+	/* If run-time partition pruning is enabled, then set that up now */
+	if (node->part_prune_infos != NIL)
+	{
+		PartitionPruneState *prunestate;
+
+		/* We may need an expression context to evaluate partition exprs */
+		ExecAssignExprContext(estate, &mergestate->ps);
+
+		prunestate = ExecCreatePartitionPruneState(&mergestate->ps,
+												   node->part_prune_infos);
+		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));
+
+			/*
+			 * The case where no subplans survive pruning must be handled
+			 * specially.  The problem here is that code in explain.c requires
+			 * an Append to have at least one subplan in order for it to
+			 * properly determine the Vars in that subplan's targetlist.  We
+			 * sidestep this issue by just initializing the first subplan and
+			 * setting ms_noopscan to true to indicate that we don't really
+			 * need to scan any subnodes.
+			 */
+			if (bms_is_empty(validsubplans))
+			{
+				mergestate->ms_noopscan = true;
+
+				/* Mark the first as valid so that it's initialized below */
+				validsubplans = bms_make_singleton(0);
+			}
+
+			nplans = bms_num_members(validsubplans);
+		}
+		else
+		{
+			/* We'll need to initialize all subplans */
+			nplans = list_length(node->mergeplans);
+			validsubplans = bms_add_range(NULL, 0, nplans - 1);
+		}
+
+		/*
+		 * If no runtime pruning is required, we can fill ms_valid_subplans
+		 * immediately, preventing later calls to ExecFindMatchingSubPlans.
+		 */
+		if (!prunestate->do_exec_prune)
+			mergestate->ms_valid_subplans = bms_add_range(NULL, 0, nplans - 1);
+	}
+	else
+	{
+		nplans = list_length(node->mergeplans);
+
+		/*
+		 * When run-time partition pruning is not enabled we can just mark all
+		 * subplans as valid; they must also all be initialized.
+		 */
+		mergestate->ms_valid_subplans = validsubplans =
+			bms_add_range(NULL, 0, nplans - 1);
+		mergestate->ms_prune_state = NULL;
+	}
+
+	mergeplanstates = (PlanState **) palloc(nplans * sizeof(PlanState *));
 	mergestate->mergeplans = mergeplanstates;
 	mergestate->ms_nplans = nplans;
 
@@ -101,26 +165,24 @@ ExecInitMergeAppend(MergeAppend *node, EState *estate, int eflags)
 	/*
 	 * Miscellaneous initialization
 	 *
-	 * MergeAppend plans don't have expression contexts because they never
-	 * call ExecQual or ExecProject.
-	 */
-
-	/*
 	 * MergeAppend nodes do have Result slots, which hold pointers to tuples,
 	 * so we have to initialize them.
 	 */
 	ExecInitResultTupleSlotTL(estate, &mergestate->ps);
 
 	/*
-	 * call ExecInitNode on each of the plans to be executed and save the
-	 * results into the array "mergeplans".
+	 * call ExecInitNode on each of the valid plans to be executed and save
+	 * the results into the mergeplanstates array.
 	 */
-	i = 0;
+	j = i = 0;
 	foreach(lc, node->mergeplans)
 	{
-		Plan	   *initNode = (Plan *) lfirst(lc);
+		if (bms_is_member(i, validsubplans))
+		{
+			Plan	   *initNode = (Plan *) lfirst(lc);
 
-		mergeplanstates[i] = ExecInitNode(initNode, estate, eflags);
+			mergeplanstates[j++] = ExecInitNode(initNode, estate, eflags);
+		}
 		i++;
 	}
 
@@ -178,11 +240,25 @@ ExecMergeAppend(PlanState *pstate)
 
 	if (!node->ms_initialized)
 	{
+		/* Nothing to do if all subplans were pruned */
+		if (node->ms_noopscan)
+			return ExecClearTuple(node->ps.ps_ResultTupleSlot);
+
 		/*
-		 * First time through: pull the first tuple from each subplan, and set
-		 * up the heap.
+		 * If we've yet to determine the valid subplans then do so now.  If
+		 * run-time pruning is disabled then the valid subplans will always be
+		 * set to all subplans.
 		 */
-		for (i = 0; i < node->ms_nplans; i++)
+		if (node->ms_valid_subplans == NULL)
+			node->ms_valid_subplans =
+				ExecFindMatchingSubPlans(node->ms_prune_state);
+
+		/*
+		 * First time through: pull the first tuple from each valid subplan,
+		 * and set up the heap.
+		 */
+		i = -1;
+		while ((i = bms_next_member(node->ms_valid_subplans, i)) >= 0)
 		{
 			node->ms_slots[i] = ExecProcNode(node->mergeplans[i]);
 			if (!TupIsNull(node->ms_slots[i]))
@@ -288,6 +364,12 @@ ExecEndMergeAppend(MergeAppendState *node)
 	 */
 	for (i = 0; i < nplans; i++)
 		ExecEndNode(mergeplans[i]);
+
+	/*
+	 * release any resources associated with run-time pruning
+	 */
+	if (node->ms_prune_state)
+		ExecDestroyPartitionPruneState(node->ms_prune_state);
 }
 
 void
@@ -295,6 +377,19 @@ ExecReScanMergeAppend(MergeAppendState *node)
 {
 	int			i;
 
+	/*
+	 * If any PARAM_EXEC Params used in pruning expressions have changed, then
+	 * we'd better unset the valid subplans so that they are reselected for
+	 * the new parameter values.
+	 */
+	if (node->ms_prune_state &&
+		bms_overlap(node->ps.chgParam,
+					node->ms_prune_state->execparamids))
+	{
+		bms_free(node->ms_valid_subplans);
+		node->ms_valid_subplans = NULL;
+	}
+
 	for (i = 0; i < node->ms_nplans; i++)
 	{
 		PlanState  *subnode = node->mergeplans[i];
diff --git a/src/backend/nodes/copyfuncs.c b/src/backend/nodes/copyfuncs.c
index 1c12075b017..96836ef19cd 100644
--- a/src/backend/nodes/copyfuncs.c
+++ b/src/backend/nodes/copyfuncs.c
@@ -273,6 +273,7 @@ _copyMergeAppend(const MergeAppend *from)
 	COPY_POINTER_FIELD(sortOperators, from->numCols * sizeof(Oid));
 	COPY_POINTER_FIELD(collations, from->numCols * sizeof(Oid));
 	COPY_POINTER_FIELD(nullsFirst, from->numCols * sizeof(bool));
+	COPY_NODE_FIELD(part_prune_infos);
 
 	return newnode;
 }
diff --git a/src/backend/nodes/outfuncs.c b/src/backend/nodes/outfuncs.c
index a88c0aecd09..a6454ce28b3 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -434,6 +434,8 @@ _outMergeAppend(StringInfo str, const MergeAppend *node)
 	appendStringInfoString(str, " :nullsFirst");
 	for (i = 0; i < node->numCols; i++)
 		appendStringInfo(str, " %s", booltostr(node->nullsFirst[i]));
+
+	WRITE_NODE_FIELD(part_prune_infos);
 }
 
 static void
diff --git a/src/backend/nodes/readfuncs.c b/src/backend/nodes/readfuncs.c
index 42aff7f57a3..9a01eb6b639 100644
--- a/src/backend/nodes/readfuncs.c
+++ b/src/backend/nodes/readfuncs.c
@@ -1634,6 +1634,7 @@ _readMergeAppend(void)
 	READ_OID_ARRAY(sortOperators, local_node->numCols);
 	READ_OID_ARRAY(collations, local_node->numCols);
 	READ_BOOL_ARRAY(nullsFirst, local_node->numCols);
+	READ_NODE_FIELD(part_prune_infos);
 
 	READ_DONE();
 }
diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index 165a9e9b8e7..0a0bec3bfcc 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -1068,6 +1068,11 @@ create_append_plan(PlannerInfo *root, AppendPath *best_path)
 		subplans = lappend(subplans, subplan);
 	}
 
+	/*
+	 * If any quals exist, they may be useful to perform further partition
+	 * pruning during execution.  Gather information needed by the executor
+	 * to do partition pruning.
+	 */
 	if (enable_partition_pruning &&
 		rel->reloptkind == RELOPT_BASEREL &&
 		best_path->partitioned_rels != NIL)
@@ -1078,7 +1083,6 @@ create_append_plan(PlannerInfo *root, AppendPath *best_path)
 
 		if (best_path->path.param_info)
 		{
-
 			List	   *prmquals = best_path->path.param_info->ppi_clauses;
 
 			prmquals = extract_actual_clauses(prmquals, false);
@@ -1088,12 +1092,6 @@ create_append_plan(PlannerInfo *root, AppendPath *best_path)
 			prunequal = list_concat(prunequal, prmquals);
 		}
 
-		/*
-		 * If any quals exist, they may be useful to perform further partition
-		 * pruning during execution.  Generate a PartitionPruneInfo for each
-		 * partitioned rel to store these quals and allow translation of
-		 * partition indexes into subpath indexes.
-		 */
 		if (prunequal != NIL)
 			partpruneinfos =
 				make_partition_pruneinfo(root,
@@ -1133,6 +1131,8 @@ create_merge_append_plan(PlannerInfo *root, MergeAppendPath *best_path)
 	List	   *pathkeys = best_path->path.pathkeys;
 	List	   *subplans = NIL;
 	ListCell   *subpaths;
+	RelOptInfo *rel = best_path->path.parent;
+	List	   *partpruneinfos = NIL;
 
 	/*
 	 * We don't have the actual creation of the MergeAppend node split out
@@ -1218,8 +1218,40 @@ create_merge_append_plan(PlannerInfo *root, MergeAppendPath *best_path)
 		subplans = lappend(subplans, subplan);
 	}
 
+	/*
+	 * If any quals exist, they may be useful to perform further partition
+	 * pruning during execution.  Gather information needed by the executor
+	 * to do partition pruning.
+	 */
+	if (enable_partition_pruning &&
+		rel->reloptkind == RELOPT_BASEREL &&
+		best_path->partitioned_rels != NIL)
+	{
+		List	   *prunequal;
+
+		prunequal = extract_actual_clauses(rel->baserestrictinfo, false);
+
+		if (best_path->path.param_info)
+		{
+
+			List	   *prmquals = best_path->path.param_info->ppi_clauses;
+
+			prmquals = extract_actual_clauses(prmquals, false);
+			prmquals = (List *) replace_nestloop_params(root,
+														(Node *) prmquals);
+
+			prunequal = list_concat(prunequal, prmquals);
+		}
+
+		if (prunequal != NIL)
+			partpruneinfos = make_partition_pruneinfo(root,
+													  best_path->partitioned_rels,
+													  best_path->subpaths, prunequal);
+	}
+
 	node->partitioned_rels = best_path->partitioned_rels;
 	node->mergeplans = subplans;
+	node->part_prune_infos = partpruneinfos;
 
 	return (Plan *) node;
 }