From 7f573fbb2d061c6fc2d1c943d609d222710545df Mon Sep 17 00:00:00 2001
From: "dgrowley@gmail.com" <dgrowley@gmail.com>
Date: Fri, 7 Sep 2018 18:20:52 +1200
Subject: [PATCH v1] Improve performance of run-time partition pruning

Previously when performing the initial run-time prune during initplan,
when resequencing the subplan indexes we performed a loop over each subplan
checking if it had been newly pruned or not.  If it was, we set the
new_subplan_indexes element to -1.  It's a fairly common case that all but
one partition will be pruned, so it seems better to make this array 1-based
and have 0 represent newly pruned.  This allows the array to be initialized
with palloc0 and also allows us to just loop over the set of unpruned
results instead of looping once for each of the original subplans.

The recalculation of the other_subplans set can be moved to a separate
loop.  Generally this is an empty set, so its more efficient to loop over
what members are there rather than doing the work in the first loop and
checking for the presence of an other_subplan with bms_is_member().
---
 src/backend/executor/execPartition.c          | 34 ++++++++++--------
 src/test/regress/expected/partition_prune.out | 50 +++++++++++++++++++++++++++
 src/test/regress/sql/partition_prune.sql      | 28 +++++++++++++++
 3 files changed, 98 insertions(+), 14 deletions(-)

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index 1a9943c3aa..1a3dc6d1fc 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -1689,23 +1689,29 @@ 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
+		 * a 1-based array and leave newly pruned items as 0.
 		 */
-		new_subplan_indexes = (int *) palloc(sizeof(int) * nsubplans);
+		new_subplan_indexes = (int *) palloc0(sizeof(int) * nsubplans);
 		new_other_subplans = NULL;
-		newidx = 0;
-		for (i = 0; i < nsubplans; i++)
+		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 */
+			Assert(i < nsubplans);
 
-			if (bms_is_member(i, prunestate->other_subplans))
-				new_other_subplans = bms_add_member(new_other_subplans,
-													new_subplan_indexes[i]);
+			new_subplan_indexes[i] = newidx++;
 		}
+
+		/*
+		 * We must also recompute the other_subplans set, since indexes in
+		 * it may change.
+		 */
+		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;
 
@@ -1753,9 +1759,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);
 					}
diff --git a/src/test/regress/expected/partition_prune.out b/src/test/regress/expected/partition_prune.out
index 24313e8c78..120b651bf5 100644
--- a/src/test/regress/expected/partition_prune.out
+++ b/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)
diff --git a/src/test/regress/sql/partition_prune.sql b/src/test/regress/sql/partition_prune.sql
index eca1a7c5ac..dc327caffd 100644
--- a/src/test/regress/sql/partition_prune.sql
+++ b/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);
-- 
2.16.2.windows.1

