Fix an oversight in cbc1279 to handle MERGE correctly

ExecInitModifyTable() forgot to trim MERGE-related lists to exclude
entries for result relations pruned during initial pruning, so fix
that.

While at it, make the function's use of the pruned resultRelations
list, rather than ModifyTable.resultRelations, more consistent.

Reported-by: Alexander Lakhin <exclusion@gmail.com> (via sqlsmith)
Reviewed-by: Junwang Zhao <zhjwpku@gmail.com>
Discussion: https://postgr.es/m/e72c94d9-e5f9-4753-9bc1-69d72bd54b8a@gmail.com

Author: amitlan

src/backend/executor/nodeModifyTable.c

@@ -3643,14 +3643,15 @@ ExecMergeNotMatched(ModifyTableContext *context, ResultRelInfo *resultRelInfo,
 void
 ExecInitMerge(ModifyTableState *mtstate, EState *estate)
 {
-	ModifyTable *node = (ModifyTable *) mtstate->ps.plan;
+	List	   *mergeActionLists = mtstate->mt_mergeActionLists;
+	List	   *mergeJoinConditions = mtstate->mt_mergeJoinConditions;
 	ResultRelInfo *rootRelInfo = mtstate->rootResultRelInfo;
 	ResultRelInfo *resultRelInfo;
 	ExprContext *econtext;
 	ListCell   *lc;
 	int			i;
 
-	if (node->mergeActionLists == NIL)
+	if (mergeActionLists == NIL)
 		return;
 
 	mtstate->mt_merge_subcommands = 0;
@@ -3667,14 +3668,14 @@ ExecInitMerge(ModifyTableState *mtstate, EState *estate)
 	 * anything here, do so there too.
 	 */
 	i = 0;
-	foreach(lc, node->mergeActionLists)
+	foreach(lc, mergeActionLists)
 	{
 		List	   *mergeActionList = lfirst(lc);
 		Node	   *joinCondition;
 		TupleDesc	relationDesc;
 		ListCell   *l;
 
-		joinCondition = (Node *) list_nth(node->mergeJoinConditions, i);
+		joinCondition = (Node *) list_nth(mergeJoinConditions, i);
 		resultRelInfo = mtstate->resultRelInfo + i;
 		i++;
 		relationDesc = RelationGetDescr(resultRelInfo->ri_RelationDesc);
@@ -4475,6 +4476,8 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	List	   *withCheckOptionLists = NIL;
 	List	   *returningLists = NIL;
 	List	   *updateColnosLists = NIL;
+	List	   *mergeActionLists = NIL;
+	List	   *mergeJoinConditions = NIL;
 	ResultRelInfo *resultRelInfo;
 	List	   *arowmarks;
 	ListCell   *l;
@@ -4518,6 +4521,18 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 
 				updateColnosLists = lappend(updateColnosLists, updateColnosList);
 			}
+			if (node->mergeActionLists)
+			{
+				List	   *mergeActionList = list_nth(node->mergeActionLists, i);
+
+				mergeActionLists = lappend(mergeActionLists, mergeActionList);
+			}
+			if (node->mergeJoinConditions)
+			{
+				List	   *mergeJoinCondition = list_nth(node->mergeJoinConditions, i);
+
+				mergeJoinConditions = lappend(mergeJoinConditions, mergeJoinCondition);
+			}
 		}
 		i++;
 	}
@@ -4544,6 +4559,8 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	mtstate->mt_merge_updated = 0;
 	mtstate->mt_merge_deleted = 0;
 	mtstate->mt_updateColnosLists = updateColnosLists;
+	mtstate->mt_mergeActionLists = mergeActionLists;
+	mtstate->mt_mergeJoinConditions = mergeJoinConditions;
 
 	/*----------
 	 * Resolve the target relation. This is the same as:
@@ -4556,7 +4573,8 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	 * If it's a partitioned or inherited table, the root partition or
 	 * appendrel RTE doesn't appear elsewhere in the plan and its RT index is
 	 * given explicitly in node->rootRelation.  Otherwise, the target relation
-	 * is the sole relation in the node->resultRelations list.
+	 * is the sole relation in the node->resultRelations list and, since it can
+	 * never be pruned, also in the resultRelations list constructed above.
 	 *----------
 	 */
 	if (node->rootRelation > 0)
@@ -4569,9 +4587,10 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	else
 	{
 		Assert(list_length(node->resultRelations) == 1);
+		Assert(list_length(resultRelations) == 1);
 		mtstate->rootResultRelInfo = mtstate->resultRelInfo;
 		ExecInitResultRelation(estate, mtstate->resultRelInfo,
-							   linitial_int(node->resultRelations));
+							   linitial_int(resultRelations));
 	}
 
 	/* set up epqstate with dummy subplan data for the moment */
@@ -4599,8 +4618,8 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 		Index		resultRelation = lfirst_int(l);
 		List	   *mergeActions = NIL;
 
-		if (node->mergeActionLists)
-			mergeActions = list_nth(node->mergeActionLists, i);
+		if (mergeActionLists)
+			mergeActions = list_nth(mergeActionLists, i);
 
 		if (resultRelInfo != mtstate->rootResultRelInfo)
 		{

src/include/nodes/execnodes.h

@@ -1448,10 +1448,13 @@ typedef struct ModifyTableState
 	double		mt_merge_deleted;
 
 	/*
-	 * List of valid updateColnosLists.  Contains only those belonging to
-	 * unpruned relations from ModifyTable.updateColnosLists.
+	 * Lists of valid updateColnosLists, mergeActionLists, and
+	 * mergeJoinConditions.  These contain only entries for unpruned
+	 * relations, filtered from the corresponding lists in ModifyTable.
 	 */
 	List	   *mt_updateColnosLists;
+	List	   *mt_mergeActionLists;
+	List	   *mt_mergeJoinConditions;
 } ModifyTableState;
 
 /* ----------------

src/test/regress/expected/partition_prune.out

@@ -4512,6 +4512,83 @@ explain (costs off) execute update_part_abc_view (2, 'a');
 execute update_part_abc_view (2, 'a');
 ERROR:  new row violates check option for view "part_abc_view"
 DETAIL:  Failing row contains (2, a, t).
+-- All pruned.
+explain (costs off) execute update_part_abc_view (3, 'a');
+         QUERY PLAN          
+-----------------------------
+ Update on part_abc
+   ->  Append
+         Subplans Removed: 2
+(3 rows)
+
+execute update_part_abc_view (3, 'a');
+ a | b | c 
+---+---+---
+(0 rows)
+
 deallocate update_part_abc_view;
+-- Runtime pruning on MERGE using a stable function
+create function stable_one() returns int as $$ begin return 1; end; $$ language plpgsql stable;
+explain (costs off)
+merge into part_abc_view pt
+using (select stable_one() as pid) as q join part_abc_1 pt1 on (q.pid = pt1.a) on pt.a = pt1.a
+when matched then delete returning pt.a;
+                              QUERY PLAN                               
+-----------------------------------------------------------------------
+ Merge on part_abc
+   Merge on part_abc_1
+   ->  Nested Loop
+         ->  Append
+               Subplans Removed: 1
+               ->  Seq Scan on part_abc_1
+                     Filter: ((b <> 'a'::text) AND (a = stable_one()))
+         ->  Materialize
+               ->  Seq Scan on part_abc_1 pt1
+                     Filter: (a = stable_one())
+(10 rows)
+
+merge into part_abc_view pt
+using (select stable_one() as pid) as q join part_abc_1 pt1 on (q.pid = pt1.a) on pt.a = pt1.a
+when matched then delete returning pt.a;
+ a 
+---
+ 1
+(1 row)
+
+table part_abc_view;
+ a | b | c 
+---+---+---
+ 2 | c | t
+(1 row)
+
+-- All pruned.
+explain (costs off)
+merge into part_abc_view pt
+using (select stable_one() + 2 as pid) as q join part_abc_1 pt1 on (q.pid = pt1.a) on pt.a = pt1.a
+when matched then delete returning pt.a;
+                      QUERY PLAN                      
+------------------------------------------------------
+ Merge on part_abc
+   ->  Nested Loop
+         ->  Append
+               Subplans Removed: 2
+         ->  Materialize
+               ->  Seq Scan on part_abc_1 pt1
+                     Filter: (a = (stable_one() + 2))
+(7 rows)
+
+merge into part_abc_view pt
+using (select stable_one() + 2 as pid) as q join part_abc_1 pt1 on (q.pid = pt1.a) on pt.a = pt1.a
+when matched then delete returning pt.a;
+ a 
+---
+(0 rows)
+
+table part_abc_view;
+ a | b | c 
+---+---+---
+ 2 | c | t
+(1 row)
+
 drop view part_abc_view;
 drop table part_abc;

src/test/regress/sql/partition_prune.sql

@@ -1371,6 +1371,31 @@ explain (costs off) execute update_part_abc_view (1, 'd');
 execute update_part_abc_view (1, 'd');
 explain (costs off) execute update_part_abc_view (2, 'a');
 execute update_part_abc_view (2, 'a');
+-- All pruned.
+explain (costs off) execute update_part_abc_view (3, 'a');
+execute update_part_abc_view (3, 'a');
 deallocate update_part_abc_view;
+
+-- Runtime pruning on MERGE using a stable function
+create function stable_one() returns int as $$ begin return 1; end; $$ language plpgsql stable;
+explain (costs off)
+merge into part_abc_view pt
+using (select stable_one() as pid) as q join part_abc_1 pt1 on (q.pid = pt1.a) on pt.a = pt1.a
+when matched then delete returning pt.a;
+merge into part_abc_view pt
+using (select stable_one() as pid) as q join part_abc_1 pt1 on (q.pid = pt1.a) on pt.a = pt1.a
+when matched then delete returning pt.a;
+table part_abc_view;
+
+-- All pruned.
+explain (costs off)
+merge into part_abc_view pt
+using (select stable_one() + 2 as pid) as q join part_abc_1 pt1 on (q.pid = pt1.a) on pt.a = pt1.a
+when matched then delete returning pt.a;
+merge into part_abc_view pt
+using (select stable_one() + 2 as pid) as q join part_abc_1 pt1 on (q.pid = pt1.a) on pt.a = pt1.a
+when matched then delete returning pt.a;
+table part_abc_view;
+
 drop view part_abc_view;
 drop table part_abc;