Improve performance of partition pruning remapping a little.

ExecFindInitialMatchingSubPlans has to update the PartitionPruneState's
subplan mapping data to account for the removal of subplans it prunes.
However, that's only necessary if run-time pruning will also occur,
so we can skip it when that won't happen, which should result in not
needing to do the remapping in many cases.  (We now need it only when
some partitions are potentially startup-time prunable and others are
potentially run-time prunable, which seems like an unusual case.)

Also make some marginal performance improvements in the remapping
itself.  These will mainly win if most partitions got pruned by
the startup-time pruning, which is perhaps a debatable assumption
in this context.

Also fix some bogus comments, and rearrange code to marginally
reduce space consumption in the executor's query-lifespan context.

David Rowley, reviewed by Yoshikazu Imai

Discussion: https://postgr.es/m/CAKJS1f9+m6-di-zyy4B4AGn0y1B9F8UKDRigtBbNviXOkuyOpw@mail.gmail.com

Author: tglsfdc

src/backend/executor/execPartition.c

@@ -1596,10 +1596,13 @@ ExecCreatePartitionPruneState(PlanState *planstate,
 /*
  * ExecFindInitialMatchingSubPlans
  *		Identify the set of subplans that cannot be eliminated by initial
- *		pruning (disregarding any pruning constraints involving PARAM_EXEC
- *		Params).  Also re-map the translation matrix which allows conversion
- *		of partition indexes into subplan indexes to account for the unneeded
- *		subplans having been removed.
+ *		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.
@@ -1613,17 +1616,18 @@ ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubplans)
 	MemoryContext oldcontext;
 	int			i;
 
+	/* Caller error if we get here without do_initial_prune */
 	Assert(prunestate->do_initial_prune);
 
 	/*
 	 * Switch to a temp context to avoid leaking memory in the executor's
-	 * memory context.
+	 * query-lifespan memory context.
 	 */
 	oldcontext = MemoryContextSwitchTo(prunestate->prune_context);
 
 	/*
-	 * For each hierarchy, do the pruning tests, and add deletable subplans'
-	 * indexes to "result".
+	 * For each hierarchy, do the pruning tests, and add nondeletable
+	 * subplans' indexes to "result".
 	 */
 	for (i = 0; i < prunestate->num_partprunedata; i++)
 	{
@@ -1640,22 +1644,27 @@ ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubplans)
 		ResetExprContext(pprune->context.planstate->ps_ExprContext);
 	}
 
+	/* 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);
 
-	/* Add in any subplans that partition pruning didn't account for */
-	result = bms_add_members(result, prunestate->other_subplans);
-
 	MemoryContextReset(prunestate->prune_context);
 
 	/*
-	 * If any subplans were pruned, 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.
+	 * 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).
 	 */
-	if (bms_num_members(result) < nsubplans)
+	if (prunestate->do_exec_prune && bms_num_members(result) < nsubplans)
 	{
 		int		   *new_subplan_indexes;
 		Bitmapset  *new_other_subplans;
@@ -1664,25 +1673,17 @@ ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubplans)
 
 		/*
 		 * First we must build a temporary array which maps old subplan
-		 * indexes to new ones.  While we're at it, also recompute the
-		 * other_subplans set, since indexes in it may change.
+		 * 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 *) palloc(sizeof(int) * nsubplans);
-		new_other_subplans = NULL;
-		newidx = 0;
-		for (i = 0; i < nsubplans; i++)
+		new_subplan_indexes = (int *) palloc0(sizeof(int) * nsubplans);
+		newidx = 1;
+		i = -1;
+		while ((i = bms_next_member(result, i)) >= 0)
 		{
-			if (bms_is_member(i, result))
-				new_subplan_indexes[i] = newidx++;
-			else
-				new_subplan_indexes[i] = -1;	/* Newly pruned */
-
-			if (bms_is_member(i, prunestate->other_subplans))
-				new_other_subplans = bms_add_member(new_other_subplans,
-													new_subplan_indexes[i]);
+			Assert(i < nsubplans);
+			new_subplan_indexes[i] = newidx++;
 		}
-		bms_free(prunestate->other_subplans);
-		prunestate->other_subplans = new_other_subplans;
 
 		/*
 		 * Now we can update each PartitionedRelPruneInfo's subplan_map with
@@ -1699,7 +1700,7 @@ ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubplans)
 			 * 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 that from 'present_parts'.
+			 * can exclude it from the current level's present_parts.
 			 */
 			for (j = prunedata->num_partrelprunedata - 1; j >= 0; j--)
 			{
@@ -1728,9 +1729,9 @@ ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubplans)
 					if (oldidx >= 0)
 					{
 						Assert(oldidx < nsubplans);
-						pprune->subplan_map[k] = new_subplan_indexes[oldidx];
+						pprune->subplan_map[k] = new_subplan_indexes[oldidx] - 1;
 
-						if (new_subplan_indexes[oldidx] >= 0)
+						if (new_subplan_indexes[oldidx] > 0)
 							pprune->present_parts =
 								bms_add_member(pprune->present_parts, k);
 					}
@@ -1748,6 +1749,19 @@ ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubplans)
 			}
 		}
 
+		/*
+		 * 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;
+
 		pfree(new_subplan_indexes);
 	}
 
@@ -1768,15 +1782,22 @@ ExecFindMatchingSubPlans(PartitionPruneState *prunestate)
 	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.
+	 */
+	Assert(prunestate->do_exec_prune);
+
 	/*
 	 * Switch to a temp context to avoid leaking memory in the executor's
-	 * memory context.
+	 * query-lifespan memory context.
 	 */
 	oldcontext = MemoryContextSwitchTo(prunestate->prune_context);
 
 	/*
-	 * For each hierarchy, do the pruning tests, and add deletable subplans'
-	 * indexes to "result".
+	 * For each hierarchy, do the pruning tests, and add nondeletable
+	 * subplans' indexes to "result".
 	 */
 	for (i = 0; i < prunestate->num_partprunedata; i++)
 	{
@@ -1792,14 +1813,14 @@ ExecFindMatchingSubPlans(PartitionPruneState *prunestate)
 		ResetExprContext(pprune->context.planstate->ps_ExprContext);
 	}
 
+	/* 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);
 
-	/* Add in any subplans that partition pruning didn't account for */
-	result = bms_add_members(result, prunestate->other_subplans);
-
 	MemoryContextReset(prunestate->prune_context);
 
 	return result;

src/test/regress/expected/partition_prune.out

@@ -2472,11 +2472,61 @@ select * from (select * from ab where a = 1 union all (values(10,5)) union all s
          Filter: (b = $0)
 (39 rows)
 
+-- Another UNION ALL test, but containing a mix of exec init and exec run-time pruning.
+create table xy_1 (x int, y int);
+insert into xy_1 values(100,-10);
+set enable_bitmapscan = 0;
+set enable_indexscan = 0;
+set plan_cache_mode = 'force_generic_plan';
+prepare ab_q6 as
+select * from (
+	select tableoid::regclass,a,b from ab
+union all
+	select tableoid::regclass,x,y from xy_1
+union all
+	select tableoid::regclass,a,b from ab
+) ab where a = $1 and b = (select -10);
+-- Ensure the xy_1 subplan is not pruned.
+explain (analyze, costs off, summary off, timing off) execute ab_q6(1);
+                       QUERY PLAN                       
+--------------------------------------------------------
+ Append (actual rows=0 loops=1)
+   InitPlan 1 (returns $0)
+     ->  Result (actual rows=1 loops=1)
+   Subplans Removed: 12
+   ->  Seq Scan on ab_a1_b1 (never executed)
+         Filter: ((a = $1) AND (b = $0))
+   ->  Seq Scan on ab_a1_b2 (never executed)
+         Filter: ((a = $1) AND (b = $0))
+   ->  Seq Scan on ab_a1_b3 (never executed)
+         Filter: ((a = $1) AND (b = $0))
+   ->  Seq Scan on xy_1 (actual rows=0 loops=1)
+         Filter: ((x = $1) AND (y = $0))
+         Rows Removed by Filter: 1
+   ->  Seq Scan on ab_a1_b1 ab_a1_b1_1 (never executed)
+         Filter: ((a = $1) AND (b = $0))
+   ->  Seq Scan on ab_a1_b2 ab_a1_b2_1 (never executed)
+         Filter: ((a = $1) AND (b = $0))
+   ->  Seq Scan on ab_a1_b3 ab_a1_b3_1 (never executed)
+         Filter: ((a = $1) AND (b = $0))
+(19 rows)
+
+-- Ensure we see just the xy_1 row.
+execute ab_q6(100);
+ tableoid |  a  |  b  
+----------+-----+-----
+ xy_1     | 100 | -10
+(1 row)
+
+reset enable_bitmapscan;
+reset enable_indexscan;
+reset plan_cache_mode;
 deallocate ab_q1;
 deallocate ab_q2;
 deallocate ab_q3;
 deallocate ab_q4;
 deallocate ab_q5;
+deallocate ab_q6;
 -- UPDATE on a partition subtree has been seen to have problems.
 insert into ab values (1,2);
 explain (analyze, costs off, summary off, timing off)

src/test/regress/sql/partition_prune.sql

@@ -548,11 +548,39 @@ select * from (select * from ab where a = 1 union all select * from ab) ab where
 explain (analyze, costs off, summary off, timing off)
 select * from (select * from ab where a = 1 union all (values(10,5)) union all select * from ab) ab where b = (select 1);
 
+-- Another UNION ALL test, but containing a mix of exec init and exec run-time pruning.
+create table xy_1 (x int, y int);
+insert into xy_1 values(100,-10);
+
+set enable_bitmapscan = 0;
+set enable_indexscan = 0;
+set plan_cache_mode = 'force_generic_plan';
+
+prepare ab_q6 as
+select * from (
+	select tableoid::regclass,a,b from ab
+union all
+	select tableoid::regclass,x,y from xy_1
+union all
+	select tableoid::regclass,a,b from ab
+) ab where a = $1 and b = (select -10);
+
+-- Ensure the xy_1 subplan is not pruned.
+explain (analyze, costs off, summary off, timing off) execute ab_q6(1);
+
+-- Ensure we see just the xy_1 row.
+execute ab_q6(100);
+
+reset enable_bitmapscan;
+reset enable_indexscan;
+reset plan_cache_mode;
+
 deallocate ab_q1;
 deallocate ab_q2;
 deallocate ab_q3;
 deallocate ab_q4;
 deallocate ab_q5;
+deallocate ab_q6;
 
 -- UPDATE on a partition subtree has been seen to have problems.
 insert into ab values (1,2);