Fix EPQ run-time pruning using the parent's PARAM_EXEC values

8741e48e5 made EvalPlanQualStart() borrow the parent EState's
es_part_prune_states.  However, the ExprContext in each
PartitionPruneState belongs to the parent EState, so exec pruning during
an EPQ recheck used the parent's PARAM_EXEC values, i.e. those from the
old row version.  This could prune away the partition the recheck
needed, causing the row to be silently skipped, e.g. with a correlated
subquery on a partitioned table in the WHERE clause of an UPDATE.

Here we fix this by giving the EPQ EState its own PartitionPruneStates,
while still borrowing the parent's es_part_prune_results so that the
same subplans are initialized.  As the states are no longer shared, the
"initialized" field added by 1b5dd3a24 is no longer needed, so remove
it.  This also undoes the ABI break that field caused in 18.

Backpatch to 18, where 8741e48e5 and 1b5dd3a24 were backpatched.

Author: Zsolt Parragi <zsolt.parragi@percona.com>
Reviewed-by: David Rowley <dgrowleyml@gmail.com>
Discussion: https://postgr.es/m/CAN4CZFPZ6qKnv6FPmq1RQZ49cHALCwLBu_NJdqUrkZNg0z81JQ%40mail.gmail.com
Backpatch-through: 18

src/backend/executor/execMain.c

@@ -3120,12 +3120,17 @@ EvalPlanQualStart(EPQState *epqstate, Plan *planTree)
 	 * and MergeAppend subplans as the parent did.
 	 */
 	rcestate->es_part_prune_infos = parentestate->es_part_prune_infos;
-	rcestate->es_part_prune_states = parentestate->es_part_prune_states;
 	rcestate->es_part_prune_results = parentestate->es_part_prune_results;
 
 	/* We'll also borrow the es_partition_directory from the parent state */
 	rcestate->es_partition_directory = parentestate->es_partition_directory;
 
+	/*
+	 * The PartitionPruneStates can't be borrowed, though, since exec pruning
+	 * must see our own es_param_exec_vals.
+	 */
+	ExecCreatePartitionPruneStates(rcestate);
+
 	/*
 	 * Initialize private state information for each SubPlan.  We must do this
 	 * before running ExecInitNode on the main query tree, since

src/backend/executor/execPartition.c

@@ -1856,6 +1856,36 @@ ExecDoInitialPruning(EState *estate)
 	}
 }
 
+/*
+ * ExecCreatePartitionPruneStates
+ *		Create a PartitionPruneState for each PartitionPruneInfo entry in
+ *		estate->es_part_prune_infos, without performing initial pruning.
+ *
+ * This is for an EState that must reuse the initial pruning results of
+ * another EState, such as an EPQ recheck's EState, for which the caller must
+ * have set es_part_prune_results.  The PartitionPruneStates can't be shared
+ * with the other EState, since exec pruning must evaluate PARAM_EXEC Params
+ * using this EState's es_param_exec_vals.
+ */
+void
+ExecCreatePartitionPruneStates(EState *estate)
+{
+	ListCell   *lc;
+
+	Assert(estate->es_part_prune_states == NIL);
+
+	foreach(lc, estate->es_part_prune_infos)
+	{
+		PartitionPruneInfo *pruneinfo = lfirst_node(PartitionPruneInfo, lc);
+		Bitmapset  *all_leafpart_rtis = NULL;
+
+		estate->es_part_prune_states =
+			lappend(estate->es_part_prune_states,
+					CreatePartitionPruneState(estate, pruneinfo,
+											  &all_leafpart_rtis));
+	}
+}
+
 /*
  * ExecInitPartitionExecPruning
  *		Initialize the data structures needed for runtime "exec" partition
@@ -2007,8 +2037,6 @@ CreatePartitionPruneState(EState *estate, PartitionPruneInfo *pruneinfo,
 	prunestate->other_subplans = bms_copy(pruneinfo->other_subplans);
 	prunestate->do_initial_prune = false;	/* may be set below */
 	prunestate->do_exec_prune = false;	/* may be set below */
-	prunestate->initialized = false;	/* set in
-										 * InitExecPartitionPruneContexts */
 	prunestate->num_partprunedata = n_part_hierarchies;
 
 	/*
@@ -2356,16 +2384,6 @@ InitExecPartitionPruneContexts(PartitionPruneState *prunestate,
 	Assert(parent_plan != NULL);
 	estate = parent_plan->state;
 
-	/*
-	 * PartitionPruneStates are sometimes shared and this one may have been
-	 * initialized already.  Sharing of states occurs for EPQ, for example.
-	 * See EvalPlanQualStart().
-	 */
-	if (prunestate->initialized)
-		return;
-
-	prunestate->initialized = true;
-
 	/*
 	 * No need to fix subplans maps if initial pruning didn't eliminate any
 	 * subplans.

src/include/executor/execPartition.h

@@ -113,8 +113,6 @@ typedef struct PartitionPruningData
  *						startup (at any hierarchy level).
  * do_exec_prune		true if pruning should be performed during
  *						executor run (at any hierarchy level).
- * initialized			true if InitExecPartitionPruneContexts has been called
- *						on this PartitionPruneState
  * num_partprunedata	Number of items in "partprunedata" array.
  * partprunedata		Array of PartitionPruningData pointers for the plan's
  *						partitioned relation(s), one for each partitioning
@@ -128,12 +126,12 @@ typedef struct PartitionPruneState
 	MemoryContext prune_context;
 	bool		do_initial_prune;
 	bool		do_exec_prune;
-	bool		initialized;
 	int			num_partprunedata;
 	PartitionPruningData *partprunedata[FLEXIBLE_ARRAY_MEMBER];
 } PartitionPruneState;
 
 extern void ExecDoInitialPruning(EState *estate);
+extern void ExecCreatePartitionPruneStates(EState *estate);
 extern PartitionPruneState *ExecInitPartitionExecPruning(PlanState *planstate,
 														 int n_total_subplans,
 														 int part_prune_index,

src/test/isolation/expected/eval-plan-qual-partition-prune.out

@@ -28,3 +28,19 @@ a|b
 2|3
 (2 rows)
 
+
+starting permutation: s1_begin s1_update_lt s2_update_lt s1_commit
+step s1_begin: BEGIN;
+step s1_update_lt: UPDATE lt SET k = 2, v = 2 WHERE id = 1;
+step s2_update_lt: 
+  UPDATE lt SET done = true
+  WHERE lt.v = (SELECT lp.a FROM lp WHERE lp.a = lt.k)
+  RETURNING *;
+ <waiting ...>
+step s1_commit: COMMIT;
+step s2_update_lt: <... completed>
+id|k|v|done
+--+-+-+----
+ 1|2|2|t   
+(1 row)
+

src/test/isolation/specs/eval-plan-qual-partition-prune.spec

@@ -8,16 +8,19 @@ setup
   INSERT INTO lp VALUES (1, 1), (2, 1);
   CREATE TABLE lk (id int PRIMARY KEY, n int);
   INSERT INTO lk VALUES (1, 0);
+  CREATE TABLE lt (id int PRIMARY KEY, k int, v int, done bool DEFAULT false);
+  INSERT INTO lt VALUES (1, 1, 1);
 }
 
 teardown
 {
-  DROP TABLE lp, lk;
+  DROP TABLE lp, lk, lt;
 }
 
 session s1
 step s1_begin	{ BEGIN; }
 step s1_update	{ UPDATE lk SET n = n + 1 WHERE id = 1; }
+step s1_update_lt	{ UPDATE lt SET k = 2, v = 2 WHERE id = 1; }
 step s1_commit	{ COMMIT; }
 
 session s2
@@ -38,6 +41,14 @@ step s2_update
   )
   SELECT * FROM upd ORDER BY a;
 }
+# The correlated subquery's exec pruning must use the EPQ's param values
+step s2_update_lt
+{
+  UPDATE lt SET done = true
+  WHERE lt.v = (SELECT lp.a FROM lp WHERE lp.a = lt.k)
+  RETURNING *;
+}
 step s2_select	{ SELECT * FROM lp ORDER BY a; }
 
 permutation s1_begin s1_update s2_update s1_commit s2_select
+permutation s1_begin s1_update_lt s2_update_lt s1_commit