Perform runtime initial pruning outside ExecInitNode()

This commit builds on the prior change that moved PartitionPruneInfos
out of individual plan nodes into a list in PlannedStmt, making it
possible to initialize PartitionPruneStates without traversing the
plan tree and perform runtime initial pruning before ExecInitNode()
initializes the plan trees.  These tasks are now handled in a new
routine, ExecDoInitialPruning(), which is called by InitPlan()
before calling ExecInitNode() on various plan trees.

ExecDoInitialPruning() performs the initial pruning and saves the
result -- a Bitmapset of indexes for surviving child subnodes -- in
es_part_prune_results, a list in EState.

PartitionPruneStates created for initial pruning are stored in
es_part_prune_states, another list in EState, for later use during
exec pruning. Both lists are parallel to es_part_prune_infos, which
holds the PartitionPruneInfos from PlannedStmt, enabling shared
indexing.

PartitionPruneStates initialized in ExecDoInitialPruning() now
include only the PartitionPruneContexts for initial pruning steps.
Exec pruning contexts are initialized later in
ExecInitPartitionExecPruning() when the parent plan node is
initialized, as the exec pruning step expressions depend on the parent
node's PlanState.

The existing function PartitionPruneFixSubPlanMap() has been
repurposed for this initialization to avoid duplicating a similar
loop structure for finding PartitionedRelPruningData to initialize
exec pruning contexts for.  It has been renamed to
InitExecPruningContexts() to reflect its new primary responsibility.
The original logic to "fix subplan maps" remains intact but is now
encapsulated within the renamed function.

This commit removes two obsolete Asserts in partkey_datum_from_expr().
The ExprContext used for pruning expression evaluation is now
independent of the parent PlanState, making these Asserts unnecessary.

By centralizing pruning logic and decoupling it from the plan
initialization step (ExecInitNode()), this change sets the stage for
future patches that will use the result of initial pruning to
save the overhead of redundant processing for pruned partitions.

Reviewed-by: Robert Haas <robertmhaas@gmail.com>
Reviewed-by: Tomas Vondra <tomas@vondra.me>
Discussion: https://postgr.es/m/CA+HiwqFGkMSge6TgC9KQzde0ohpAycLQuV7ooitEEpbKB0O_mg@mail.gmail.com

Author: amitlan

src/backend/executor/execMain.c

@@ -46,6 +46,7 @@
 #include "commands/matview.h"
 #include "commands/trigger.h"
 #include "executor/executor.h"
+#include "executor/execPartition.h"
 #include "executor/nodeSubplan.h"
 #include "foreign/fdwapi.h"
 #include "mb/pg_wchar.h"
@@ -855,6 +856,17 @@ InitPlan(QueryDesc *queryDesc, int eflags)
 	estate->es_plannedstmt = plannedstmt;
 	estate->es_part_prune_infos = plannedstmt->partPruneInfos;
 
+	/*
+	 * Perform runtime "initial" pruning to identify which child subplans,
+	 * corresponding to the children of plan nodes that contain
+	 * PartitionPruneInfo such as Append, will not be executed. The results,
+	 * which are bitmapsets of indexes of the child subplans that will be
+	 * executed, are saved in es_part_prune_results.  These results correspond
+	 * to each PartitionPruneInfo entry, and the es_part_prune_results list is
+	 * parallel to es_part_prune_infos.
+	 */
+	ExecDoInitialPruning(estate);
+
 	/*
 	 * Next, build the ExecRowMark array from the PlanRowMark(s), if any.
 	 */

src/backend/executor/execPartition.c

@@ -144,11 +144,11 @@ ExecInitAppend(Append *node, EState *estate, int eflags)
 		 * subplans to initialize (validsubplans) by taking into account the
 		 * result of performing initial pruning if any.
 		 */
-		prunestate = ExecInitPartitionPruning(&appendstate->ps,
-											  list_length(node->appendplans),
-											  node->part_prune_index,
-											  node->apprelids,
-											  &validsubplans);
+		prunestate = ExecInitPartitionExecPruning(&appendstate->ps,
+												  list_length(node->appendplans),
+												  node->part_prune_index,
+												  node->apprelids,
+												  &validsubplans);
 		appendstate->as_prune_state = prunestate;
 		nplans = bms_num_members(validsubplans);
 

src/backend/executor/nodeAppend.c

@@ -92,11 +92,11 @@ ExecInitMergeAppend(MergeAppend *node, EState *estate, int eflags)
 		 * subplans to initialize (validsubplans) by taking into account the
 		 * result of performing initial pruning if any.
 		 */
-		prunestate = ExecInitPartitionPruning(&mergestate->ps,
-											  list_length(node->mergeplans),
-											  node->part_prune_index,
-											  node->apprelids,
-											  &validsubplans);
+		prunestate = ExecInitPartitionExecPruning(&mergestate->ps,
+												  list_length(node->mergeplans),
+												  node->part_prune_index,
+												  node->apprelids,
+												  &validsubplans);
 		mergestate->ms_prune_state = prunestate;
 		nplans = bms_num_members(validsubplans);
 

src/backend/executor/nodeMergeAppend.c

@@ -3783,13 +3783,8 @@ partkey_datum_from_expr(PartitionPruneContext *context,
 		/*
 		 * We should never see a non-Const in a step unless the caller has
 		 * passed a valid ExprContext.
-		 *
-		 * When context->planstate is valid, context->exprcontext is same as
-		 * context->planstate->ps_ExprContext.
 		 */
-		Assert(context->planstate != NULL || context->exprcontext != NULL);
-		Assert(context->planstate == NULL ||
-			   (context->exprcontext == context->planstate->ps_ExprContext));
+		Assert(context->exprcontext != NULL);
 
 		exprstate = context->exprstates[stateidx];
 		ectx = context->exprcontext;

src/backend/partitioning/partprune.c

@@ -42,6 +42,9 @@ extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
  * PartitionedRelPruneInfo (see plannodes.h); though note that here,
  * subpart_map contains indexes into PartitionPruningData.partrelprunedata[].
  *
+ * partrel						Partitioned table Relation; obtained by
+ * 								ExecGetRangeTableRelation(estate, rti), where
+ *								rti is PartitionedRelPruneInfo.rtindex.
  * nparts						Length of subplan_map[] and subpart_map[].
  * subplan_map					Subplan index by partition index, or -1.
  * subpart_map					Subpart index by partition index, or -1.
@@ -58,6 +61,7 @@ extern void ExecCleanupTupleRouting(ModifyTableState *mtstate,
  */
 typedef struct PartitionedRelPruningData
 {
+	Relation	partrel;
 	int			nparts;
 	int		   *subplan_map;
 	int		   *subpart_map;
@@ -90,6 +94,8 @@ typedef struct PartitionPruningData
  * the clauses being unable to match to any tuple that the subplan could
  * possibly produce.
  *
+ * econtext				Standalone ExprContext to evaluate expressions in
+ *						the pruning steps
  * execparamids			Contains paramids of PARAM_EXEC Params found within
  *						any of the partprunedata structs.  Pruning must be
  *						done again each time the value of one of these
@@ -112,6 +118,7 @@ typedef struct PartitionPruningData
  */
 typedef struct PartitionPruneState
 {
+	ExprContext *econtext;
 	Bitmapset  *execparamids;
 	Bitmapset  *other_subplans;
 	MemoryContext prune_context;
@@ -121,11 +128,12 @@ typedef struct PartitionPruneState
 	PartitionPruningData *partprunedata[FLEXIBLE_ARRAY_MEMBER];
 } PartitionPruneState;
 
-extern PartitionPruneState *ExecInitPartitionPruning(PlanState *planstate,
-													 int n_total_subplans,
-													 int part_prune_index,
-													 Bitmapset *relids,
-													 Bitmapset **initially_valid_subplans);
+extern void ExecDoInitialPruning(EState *estate);
+extern PartitionPruneState *ExecInitPartitionExecPruning(PlanState *planstate,
+														 int n_total_subplans,
+														 int part_prune_index,
+														 Bitmapset *relids,
+														 Bitmapset **initially_valid_subplans);
 extern Bitmapset *ExecFindMatchingSubPlans(PartitionPruneState *prunestate,
 										   bool initial_prune);
 

src/include/executor/execPartition.h

@@ -656,6 +656,8 @@ typedef struct EState
 	List	   *es_rteperminfos;	/* List of RTEPermissionInfo */
 	PlannedStmt *es_plannedstmt;	/* link to top of plan tree */
 	List	   *es_part_prune_infos;	/* List of PartitionPruneInfo */
+	List	   *es_part_prune_states;	/* List of PartitionPruneState */
+	List	   *es_part_prune_results;	/* List of Bitmapset */
 	const char *es_sourceText;	/* Source text from QueryDesc */
 
 	JunkFilter *es_junkFilter;	/* top-level junk filter, if any */

src/include/nodes/execnodes.h

@@ -281,7 +281,11 @@ typedef struct Append
 	 */
 	int			first_partial_plan;
 
-	/* Index to PlannerInfo.partPruneInfos or -1 if no run-time pruning */
+	/*
+	 * Index into PlannedStmt.partPruneInfos and parallel lists in EState:
+	 * es_part_prune_states and es_part_prune_results. Set to -1 if no
+	 * run-time pruning is used.
+	 */
 	int			part_prune_index;
 } Append;
 
@@ -316,7 +320,11 @@ typedef struct MergeAppend
 	/* NULLS FIRST/LAST directions */
 	bool	   *nullsFirst pg_node_attr(array_size(numCols));
 
-	/* Index to PlannerInfo.partPruneInfos or -1 if no run-time pruning */
+	/*
+	 * Index into PlannedStmt.partPruneInfos and parallel lists in EState:
+	 * es_part_prune_states and es_part_prune_results. Set to -1 if no
+	 * run-time pruning is used.
+	 */
 	int			part_prune_index;
 } MergeAppend;