Improve run-time partition pruning to handle any stable expression.

The initial coding of the run-time-pruning feature only coped with cases
where the partition key(s) are compared to Params.  That is a bit silly;
we can allow it to work with any non-Var-containing stable expression, as
long as we take special care with expressions containing PARAM_EXEC Params.
The code is hardly any longer this way, and it's considerably clearer
(IMO at least).  Per gripe from Pavel Stehule.

David Rowley, whacked around a bit by me

Discussion: https://postgr.es/m/CAFj8pRBjrufA3ocDm8o4LPGNye9Y+pm1b9kCwode4X04CULG3g@mail.gmail.com

src/backend/executor/nodeAppend.c

@@ -133,29 +133,27 @@ ExecInitAppend(Append *node, EState *estate, int eflags)
 	{
 		PartitionPruneState *prunestate;
 
+		/* We may need an expression context to evaluate partition exprs */
 		ExecAssignExprContext(estate, &appendstate->ps);
 
 		prunestate = ExecSetupPartitionPruneState(&appendstate->ps,
 												  node->part_prune_infos);
 
-		/*
-		 * When there are external params matching the partition key we may be
-		 * able to prune away Append subplans now.
-		 */
-		if (!bms_is_empty(prunestate->extparams))
+		/* Perform an initial partition prune, if required. */
+		if (prunestate->do_initial_prune)
 		{
-			/* Determine which subplans match the external params */
+			/* Determine which subplans survive initial pruning */
 			validsubplans = ExecFindInitialMatchingSubPlans(prunestate,
 															list_length(node->appendplans));
 
 			/*
-			 * If no subplans match the given parameters then we must handle
-			 * this case in a special way.  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 as_whichplan to NO_MATCHING_SUBPLANS
-			 * to indicate that we don't need to scan any subnodes.
+			 * 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 as_whichplan to NO_MATCHING_SUBPLANS to indicate that
+			 * we don't really need to scan any subnodes.
 			 */
 			if (bms_is_empty(validsubplans))
 			{
@@ -175,14 +173,13 @@ ExecInitAppend(Append *node, EState *estate, int eflags)
 		}
 
 		/*
-		 * If there's no exec params then no further pruning can be done, we
-		 * can just set the valid subplans to all remaining subplans.
+		 * If no runtime pruning is required, we can fill as_valid_subplans
+		 * immediately, preventing later calls to ExecFindMatchingSubPlans.
 		 */
-		if (bms_is_empty(prunestate->execparams))
+		if (!prunestate->do_exec_prune)
 			appendstate->as_valid_subplans = bms_add_range(NULL, 0, nplans - 1);
 
 		appendstate->as_prune_state = prunestate;
-
 	}
 	else
 	{
@@ -190,7 +187,7 @@ ExecInitAppend(Append *node, EState *estate, int eflags)
 
 		/*
 		 * When run-time partition pruning is not enabled we can just mark all
-		 * subplans as valid, they must also all be initialized.
+		 * subplans as valid; they must also all be initialized.
 		 */
 		appendstate->as_valid_subplans = validsubplans =
 			bms_add_range(NULL, 0, nplans - 1);
@@ -341,13 +338,13 @@ ExecReScanAppend(AppendState *node)
 	int			i;
 
 	/*
-	 * If any of the parameters being used for partition pruning have changed,
-	 * then we'd better unset the valid subplans so that they are reselected
-	 * for the new parameter values.
+	 * 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->as_prune_state &&
 		bms_overlap(node->ps.chgParam,
-					node->as_prune_state->execparams))
+					node->as_prune_state->execparamids))
 	{
 		bms_free(node->as_valid_subplans);
 		node->as_valid_subplans = NULL;
@@ -531,9 +528,9 @@ choose_next_subplan_for_leader(AppendState *node)
 		node->as_whichplan = node->as_nplans - 1;
 
 		/*
-		 * If we've yet to determine the valid subplans for these parameters
-		 * then do so now.  If run-time pruning is disabled then the valid
-		 * subplans will always be set to all subplans.
+		 * 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.
 		 */
 		if (node->as_valid_subplans == NULL)
 		{
@@ -606,9 +603,9 @@ choose_next_subplan_for_worker(AppendState *node)
 		node->as_pstate->pa_finished[node->as_whichplan] = true;
 
 	/*
-	 * If we've yet to determine the valid subplans for these parameters then
-	 * do so now.  If run-time pruning is disabled then the valid subplans
-	 * will always be set to all subplans.
+	 * 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.
 	 */
 	else if (node->as_valid_subplans == NULL)
 	{

src/backend/nodes/copyfuncs.c

@@ -2175,10 +2175,13 @@ _copyPartitionPruneInfo(const PartitionPruneInfo *from)
 	COPY_NODE_FIELD(pruning_steps);
 	COPY_BITMAPSET_FIELD(present_parts);
 	COPY_SCALAR_FIELD(nparts);
+	COPY_SCALAR_FIELD(nexprs);
 	COPY_POINTER_FIELD(subnode_map, from->nparts * sizeof(int));
 	COPY_POINTER_FIELD(subpart_map, from->nparts * sizeof(int));
-	COPY_BITMAPSET_FIELD(extparams);
-	COPY_BITMAPSET_FIELD(execparams);
+	COPY_POINTER_FIELD(hasexecparam, from->nexprs * sizeof(bool));
+	COPY_SCALAR_FIELD(do_initial_prune);
+	COPY_SCALAR_FIELD(do_exec_prune);
+	COPY_BITMAPSET_FIELD(execparamids);
 
 	return newnode;
 }

src/backend/nodes/outfuncs.c

@@ -1742,6 +1742,7 @@ _outPartitionPruneInfo(StringInfo str, const PartitionPruneInfo *node)
 	WRITE_NODE_FIELD(pruning_steps);
 	WRITE_BITMAPSET_FIELD(present_parts);
 	WRITE_INT_FIELD(nparts);
+	WRITE_INT_FIELD(nexprs);
 
 	appendStringInfoString(str, " :subnode_map");
 	for (i = 0; i < node->nparts; i++)
@@ -1751,8 +1752,13 @@ _outPartitionPruneInfo(StringInfo str, const PartitionPruneInfo *node)
 	for (i = 0; i < node->nparts; i++)
 		appendStringInfo(str, " %d", node->subpart_map[i]);
 
-	WRITE_BITMAPSET_FIELD(extparams);
-	WRITE_BITMAPSET_FIELD(execparams);
+	appendStringInfoString(str, " :hasexecparam");
+	for (i = 0; i < node->nexprs; i++)
+		appendStringInfo(str, " %s", booltostr(node->hasexecparam[i]));
+
+	WRITE_BOOL_FIELD(do_initial_prune);
+	WRITE_BOOL_FIELD(do_exec_prune);
+	WRITE_BITMAPSET_FIELD(execparamids);
 }
 
 /*****************************************************************************

src/backend/nodes/readfuncs.c

@@ -1363,10 +1363,13 @@ _readPartitionPruneInfo(void)
 	READ_NODE_FIELD(pruning_steps);
 	READ_BITMAPSET_FIELD(present_parts);
 	READ_INT_FIELD(nparts);
+	READ_INT_FIELD(nexprs);
 	READ_INT_ARRAY(subnode_map, local_node->nparts);
 	READ_INT_ARRAY(subpart_map, local_node->nparts);
-	READ_BITMAPSET_FIELD(extparams);
-	READ_BITMAPSET_FIELD(execparams);
+	READ_BOOL_ARRAY(hasexecparam, local_node->nexprs);
+	READ_BOOL_FIELD(do_initial_prune);
+	READ_BOOL_FIELD(do_exec_prune);
+	READ_BITMAPSET_FIELD(execparamids);
 
 	READ_DONE();
 }

src/include/executor/execPartition.h

@@ -127,15 +127,16 @@ typedef struct PartitionTupleRouting
  * subpart_map					An array containing the offset into the
  *								'partprunedata' array in PartitionPruning, or
  *								-1 if there is no such element in that array.
- * present_parts				A Bitmapset of the partition index that we have
- *								subnodes mapped for.
+ * present_parts				A Bitmapset of the partition indexes that we
+ *								have subnodes mapped for.
  * context						Contains the context details required to call
  *								the partition pruning code.
- * pruning_steps				Contains a list of PartitionPruneStep used to
+ * pruning_steps				List of PartitionPruneSteps used to
  *								perform the actual pruning.
- * extparams					Contains paramids of external params found
- *								matching partition keys in 'pruning_steps'.
- * allparams					As 'extparams' but also including exec params.
+ * do_initial_prune				true if pruning should be performed during
+ *								executor startup.
+ * do_exec_prune				true if pruning should be performed during
+ *								executor run.
  *-----------------------
  */
 typedef struct PartitionPruningData
@@ -145,15 +146,14 @@ typedef struct PartitionPruningData
 	Bitmapset  *present_parts;
 	PartitionPruneContext context;
 	List	   *pruning_steps;
-	Bitmapset  *extparams;
-	Bitmapset  *allparams;
+	bool		do_initial_prune;
+	bool		do_exec_prune;
 } PartitionPruningData;
 
 /*-----------------------
  * PartitionPruneState - State object required for executor nodes to perform
  * partition pruning elimination of their subnodes.  This encapsulates a
- * flattened hierarchy of PartitionPruningData structs and also stores all
- * paramids which were found to match the partition keys of each partition.
+ * flattened hierarchy of PartitionPruningData structs.
  * This struct can be attached to node types which support arbitrary Lists of
  * subnodes containing partitions to allow subnodes to be eliminated due to
  * the clauses being unable to match to any tuple that the subnode could
@@ -163,24 +163,24 @@ typedef struct PartitionPruningData
  *						partitioned relation. First element contains the
  *						details for the target partitioned table.
  * num_partprunedata	Number of items in 'partprunedata' array.
+ * do_initial_prune		true if pruning should be performed during executor
+ *						startup (at any hierarchy level).
+ * do_exec_prune		true if pruning should be performed during
+ *						executor run (at any hierarchy level).
  * prune_context		A memory context which can be used to call the query
  *						planner's partition prune functions.
- * extparams			All PARAM_EXTERN paramids which were found to match a
- *						partition key in each of the contained
- *						PartitionPruningData structs.
- * execparams			As above but for PARAM_EXEC.
- * allparams			Union of 'extparams' and 'execparams', saved to avoid
- *						recalculation.
+ * execparamids			Contains paramids of PARAM_EXEC Params found within
+ *						any of the partprunedata structs.
  *-----------------------
  */
 typedef struct PartitionPruneState
 {
 	PartitionPruningData *partprunedata;
 	int			num_partprunedata;
+	bool		do_initial_prune;
+	bool		do_exec_prune;
 	MemoryContext prune_context;
-	Bitmapset  *extparams;
-	Bitmapset  *execparams;
-	Bitmapset  *allparams;
+	Bitmapset  *execparamids;
 } PartitionPruneState;
 
 extern PartitionTupleRouting *ExecSetupPartitionTupleRouting(ModifyTableState *mtstate,

src/include/nodes/primnodes.h

@@ -1597,11 +1597,17 @@ typedef struct PartitionPruneInfo
 	List	   *pruning_steps;	/* List of PartitionPruneStep */
 	Bitmapset  *present_parts;	/* Indexes of all partitions which subnodes
 								 * are present for. */
-	int			nparts;			/* The length of the following two arrays */
+	int			nparts;			/* Length of subnode_map[] and subpart_map[] */
+	int			nexprs;			/* Length of hasexecparam[] */
 	int		   *subnode_map;	/* subnode index by partition id, or -1 */
 	int		   *subpart_map;	/* subpart index by partition id, or -1 */
-	Bitmapset  *extparams;		/* All external paramids seen in prunesteps */
-	Bitmapset  *execparams;		/* All exec paramids seen in prunesteps */
+	bool	   *hasexecparam;	/* true if corresponding pruning_step contains
+								 * any PARAM_EXEC Params. */
+	bool		do_initial_prune;	/* true if pruning should be performed
+									 * during executor startup. */
+	bool		do_exec_prune;	/* true if pruning should be performed during
+								 * executor run. */
+	Bitmapset  *execparamids;	/* All PARAM_EXEC Param IDs in pruning_steps */
 } PartitionPruneInfo;
 
 #endif							/* PRIMNODES_H */

src/include/partitioning/partprune.h

@@ -40,23 +40,27 @@ typedef struct PartitionPruneContext
 	PartitionBoundInfo boundinfo;
 
 	/*
-	 * Can be set when the context is used from the executor to allow params
-	 * found matching the partition key to be evaluated.
+	 * This will be set when the context is used from the executor, to allow
+	 * Params to be evaluated.
 	 */
 	PlanState  *planstate;
 
-	/*
-	 * Parameters that are safe to be used for partition pruning. execparams
-	 * are not safe to use until the executor is running.
-	 */
-	Bitmapset  *safeparams;
-
 	/*
 	 * Array of ExprStates, indexed as per PruneCtxStateIdx; one for each
 	 * partkey in each pruning step.  Allocated if planstate is non-NULL,
 	 * otherwise NULL.
 	 */
 	ExprState **exprstates;
+
+	/*
+	 * Similar array of flags, each true if corresponding 'exprstate'
+	 * expression contains any PARAM_EXEC Params.  (Can be NULL if planstate
+	 * is NULL.)
+	 */
+	bool	   *exprhasexecparam;
+
+	/* true if it's safe to evaluate PARAM_EXEC Params */
+	bool		evalexecparams;
 } PartitionPruneContext;
 
 #define PruneCxtStateIdx(partnatts, step_id, keyno) \

src/test/regress/expected/partition_prune.out

@@ -1726,8 +1726,8 @@ explain (analyze, costs off, summary off, timing off) execute ab_q1 (2, 4);
          Filter: ((a >= $1) AND (a <= $2) AND (b < 3))
 (10 rows)
 
--- Ensure a mix of external and exec params work together at different
--- levels of partitioning.
+-- Ensure a mix of PARAM_EXTERN and PARAM_EXEC Params work together at
+-- different levels of partitioning.
 prepare ab_q2 (int, int) as
 select a from ab where a between $1 and $2 and b < (select 3);
 execute ab_q2 (1, 8);
@@ -1770,7 +1770,7 @@ explain (analyze, costs off, summary off, timing off) execute ab_q2 (2, 2);
          Filter: ((a >= $1) AND (a <= $2) AND (b < $0))
 (10 rows)
 
--- As above, but with swap the exec param to the first partition level
+-- As above, but swap the PARAM_EXEC Param to the first partition level
 prepare ab_q3 (int, int) as
 select a from ab where b between $1 and $2 and a < (select 3);
 execute ab_q3 (1, 8);
@@ -1835,6 +1835,54 @@ fetch backward all from cur;
 (2 rows)
 
 commit;
+begin;
+-- Test run-time pruning using stable functions
+create function list_part_fn(int) returns int as $$ begin return $1; end;$$ language plpgsql stable;
+-- Ensure pruning works using a stable function containing no Vars
+explain (analyze, costs off, summary off, timing off) select * from list_part where a = list_part_fn(1);
+                      QUERY PLAN                      
+------------------------------------------------------
+ Append (actual rows=1 loops=1)
+   Subplans Removed: 3
+   ->  Seq Scan on list_part1 (actual rows=1 loops=1)
+         Filter: (a = list_part_fn(1))
+(4 rows)
+
+-- Ensure pruning does not take place when the function has a Var parameter
+explain (analyze, costs off, summary off, timing off) select * from list_part where a = list_part_fn(a);
+                      QUERY PLAN                      
+------------------------------------------------------
+ Append (actual rows=4 loops=1)
+   ->  Seq Scan on list_part1 (actual rows=1 loops=1)
+         Filter: (a = list_part_fn(a))
+   ->  Seq Scan on list_part2 (actual rows=1 loops=1)
+         Filter: (a = list_part_fn(a))
+   ->  Seq Scan on list_part3 (actual rows=1 loops=1)
+         Filter: (a = list_part_fn(a))
+   ->  Seq Scan on list_part4 (actual rows=1 loops=1)
+         Filter: (a = list_part_fn(a))
+(9 rows)
+
+-- Ensure pruning does not take place when the expression contains a Var.
+explain (analyze, costs off, summary off, timing off) select * from list_part where a = list_part_fn(1) + a;
+                      QUERY PLAN                      
+------------------------------------------------------
+ Append (actual rows=0 loops=1)
+   ->  Seq Scan on list_part1 (actual rows=0 loops=1)
+         Filter: (a = (list_part_fn(1) + a))
+         Rows Removed by Filter: 1
+   ->  Seq Scan on list_part2 (actual rows=0 loops=1)
+         Filter: (a = (list_part_fn(1) + a))
+         Rows Removed by Filter: 1
+   ->  Seq Scan on list_part3 (actual rows=0 loops=1)
+         Filter: (a = (list_part_fn(1) + a))
+         Rows Removed by Filter: 1
+   ->  Seq Scan on list_part4 (actual rows=0 loops=1)
+         Filter: (a = (list_part_fn(1) + a))
+         Rows Removed by Filter: 1
+(13 rows)
+
+rollback;
 drop table list_part;
 -- Parallel append
 -- Suppress the number of loops each parallel node runs for.  This is because
@@ -2007,7 +2055,7 @@ select explain_parallel_append('execute ab_q5 (33, 44, 55)');
                            Filter: ((b < 4) AND (a = ANY (ARRAY[$1, $2, $3])))
 (9 rows)
 
--- Test Parallel Append with exec params
+-- Test Parallel Append with PARAM_EXEC Params
 select explain_parallel_append('select count(*) from ab where (a = (select 1) or a = (select 3)) and b = 2');
                          explain_parallel_append                         
 -------------------------------------------------------------------------
@@ -2079,6 +2127,40 @@ select explain_parallel_append('select avg(ab.a) from ab inner join lprt_a a on
                                  Index Cond: (a = a.a)
 (27 rows)
 
+-- Ensure the same partitions are pruned when we make the nested loop
+-- parameter an Expr rather than a plain Param.
+select explain_parallel_append('select avg(ab.a) from ab inner join lprt_a a on ab.a = a.a + 0 where a.a in(0, 0, 1)');
+                                      explain_parallel_append                                      
+---------------------------------------------------------------------------------------------------
+ Finalize Aggregate (actual rows=1 loops=1)
+   ->  Gather (actual rows=2 loops=1)
+         Workers Planned: 1
+         Workers Launched: 1
+         ->  Partial Aggregate (actual rows=1 loops=2)
+               ->  Nested Loop (actual rows=0 loops=2)
+                     ->  Parallel Seq Scan on lprt_a a (actual rows=51 loops=N)
+                           Filter: (a = ANY ('{0,0,1}'::integer[]))
+                     ->  Append (actual rows=0 loops=102)
+                           ->  Index Scan using ab_a1_b1_a_idx on ab_a1_b1 (actual rows=0 loops=2)
+                                 Index Cond: (a = (a.a + 0))
+                           ->  Index Scan using ab_a1_b2_a_idx on ab_a1_b2 (actual rows=0 loops=2)
+                                 Index Cond: (a = (a.a + 0))
+                           ->  Index Scan using ab_a1_b3_a_idx on ab_a1_b3 (actual rows=0 loops=2)
+                                 Index Cond: (a = (a.a + 0))
+                           ->  Index Scan using ab_a2_b1_a_idx on ab_a2_b1 (never executed)
+                                 Index Cond: (a = (a.a + 0))
+                           ->  Index Scan using ab_a2_b2_a_idx on ab_a2_b2 (never executed)
+                                 Index Cond: (a = (a.a + 0))
+                           ->  Index Scan using ab_a2_b3_a_idx on ab_a2_b3 (never executed)
+                                 Index Cond: (a = (a.a + 0))
+                           ->  Index Scan using ab_a3_b1_a_idx on ab_a3_b1 (never executed)
+                                 Index Cond: (a = (a.a + 0))
+                           ->  Index Scan using ab_a3_b2_a_idx on ab_a3_b2 (never executed)
+                                 Index Cond: (a = (a.a + 0))
+                           ->  Index Scan using ab_a3_b3_a_idx on ab_a3_b3 (never executed)
+                                 Index Cond: (a = (a.a + 0))
+(27 rows)
+
 insert into lprt_a values(3),(3);
 select explain_parallel_append('select avg(ab.a) from ab inner join lprt_a a on ab.a = a.a where a.a in(1, 0, 3)');
                                       explain_parallel_append                                      

src/test/regress/sql/partition_prune.sql

@@ -348,8 +348,8 @@ execute ab_q1 (1, 8);
 explain (analyze, costs off, summary off, timing off) execute ab_q1 (2, 2);
 explain (analyze, costs off, summary off, timing off) execute ab_q1 (2, 4);
 
--- Ensure a mix of external and exec params work together at different
--- levels of partitioning.
+-- Ensure a mix of PARAM_EXTERN and PARAM_EXEC Params work together at
+-- different levels of partitioning.
 prepare ab_q2 (int, int) as
 select a from ab where a between $1 and $2 and b < (select 3);
 
@@ -361,7 +361,7 @@ execute ab_q2 (1, 8);
 
 explain (analyze, costs off, summary off, timing off) execute ab_q2 (2, 2);
 
--- As above, but with swap the exec param to the first partition level
+-- As above, but swap the PARAM_EXEC Param to the first partition level
 prepare ab_q3 (int, int) as
 select a from ab where b between $1 and $2 and a < (select 3);
 
@@ -396,6 +396,22 @@ fetch backward all from cur;
 
 commit;
 
+begin;
+
+-- Test run-time pruning using stable functions
+create function list_part_fn(int) returns int as $$ begin return $1; end;$$ language plpgsql stable;
+
+-- Ensure pruning works using a stable function containing no Vars
+explain (analyze, costs off, summary off, timing off) select * from list_part where a = list_part_fn(1);
+
+-- Ensure pruning does not take place when the function has a Var parameter
+explain (analyze, costs off, summary off, timing off) select * from list_part where a = list_part_fn(a);
+
+-- Ensure pruning does not take place when the expression contains a Var.
+explain (analyze, costs off, summary off, timing off) select * from list_part where a = list_part_fn(1) + a;
+
+rollback;
+
 drop table list_part;
 
 -- Parallel append
@@ -458,7 +474,7 @@ select explain_parallel_append('execute ab_q5 (2, 3, 3)');
 -- We'll still get a single subplan in this case, but it should not be scanned.
 select explain_parallel_append('execute ab_q5 (33, 44, 55)');
 
--- Test Parallel Append with exec params
+-- Test Parallel Append with PARAM_EXEC Params
 select explain_parallel_append('select count(*) from ab where (a = (select 1) or a = (select 3)) and b = 2');
 
 -- Test pruning during parallel nested loop query
@@ -486,6 +502,10 @@ set enable_mergejoin = 0;
 
 select explain_parallel_append('select avg(ab.a) from ab inner join lprt_a a on ab.a = a.a where a.a in(0, 0, 1)');
 
+-- Ensure the same partitions are pruned when we make the nested loop
+-- parameter an Expr rather than a plain Param.
+select explain_parallel_append('select avg(ab.a) from ab inner join lprt_a a on ab.a = a.a + 0 where a.a in(0, 0, 1)');
+
 insert into lprt_a values(3),(3);
 
 select explain_parallel_append('select avg(ab.a) from ab inner join lprt_a a on ab.a = a.a where a.a in(1, 0, 3)');