v2-do-not-prune-all-ModifyTable-result-relations.patch

diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
new file mode 100644
index 19ffcc2..967a5f8
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -4698,7 +4698,9 @@ show_modifytable_info(ModifyTableState *
 	/* Should we explicitly label target relations? */
 	labeltargets = (mtstate->mt_nrels > 1 ||
 					(mtstate->mt_nrels == 1 &&
-					 mtstate->resultRelInfo[0].ri_RangeTableIndex != node->nominalRelation));
+					 mtstate->resultRelInfo[0].ri_RangeTableIndex != node->nominalRelation &&
+					 bms_is_member(mtstate->resultRelInfo[0].ri_RangeTableIndex,
+								   mtstate->ps.state->es_unpruned_relids)));
 
 	if (labeltargets)
 		ExplainOpenGroup("Target Tables", "Target Tables", false, es);
diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
new file mode 100644
index 0493b7d..e9bd98c
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -1006,7 +1006,7 @@ InitPlan(QueryDesc *queryDesc, int eflag
 				case ROW_MARK_SHARE:
 				case ROW_MARK_KEYSHARE:
 				case ROW_MARK_REFERENCE:
-					relation = ExecGetRangeTableRelation(estate, rc->rti);
+					relation = ExecGetRangeTableRelation(estate, rc->rti, false);
 					break;
 				case ROW_MARK_COPY:
 					/* no physical table access is required */
diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
new file mode 100644
index 5cd5e2e..6ac165a
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -2076,7 +2076,7 @@ CreatePartitionPruneState(EState *estate
 			 * because that entry will be held open and locked for the
 			 * duration of this executor run.
 			 */
-			partrel = ExecGetRangeTableRelation(estate, pinfo->rtindex);
+			partrel = ExecGetRangeTableRelation(estate, pinfo->rtindex, false);
 
 			/* Remember for InitExecPartitionPruneContext(). */
 			pprune->partrel = partrel;
diff --git a/src/backend/executor/execUtils.c b/src/backend/executor/execUtils.c
new file mode 100644
index 39d6f4d..e1521e2
--- a/src/backend/executor/execUtils.c
+++ b/src/backend/executor/execUtils.c
@@ -746,7 +746,7 @@ ExecOpenScanRelation(EState *estate, Ind
 	Relation	rel;
 
 	/* Open the relation. */
-	rel = ExecGetRangeTableRelation(estate, scanrelid);
+	rel = ExecGetRangeTableRelation(estate, scanrelid, false);
 
 	/*
 	 * Complain if we're attempting a scan of an unscannable relation, except
@@ -815,19 +815,24 @@ ExecInitRangeTable(EState *estate, List
  *
  * The Relations will be closed in ExecEndPlan().
  *
- * Note: The caller must ensure that 'rti' refers to an unpruned relation
- * (i.e., it is a member of estate->es_unpruned_relids) before calling this
- * function. Attempting to open a pruned relation will result in an error.
+ * If isResultRel is true, the relation is being used as a result relation.
+ * Such a relation might have been pruned, which is OK for result relations,
+ * but not for scan relations.  If isResultRel is false, the caller must
+ * ensure that 'rti' refers to an unpruned relation (i.e., it is a member of
+ * estate->es_unpruned_relids) before calling this function. Attempting to
+ * open a pruned relation for scanning will result in an error.
  */
 Relation
-ExecGetRangeTableRelation(EState *estate, Index rti)
+ExecGetRangeTableRelation(EState *estate, Index rti, bool isResultRel)
 {
+	bool		pruned;
 	Relation	rel;
 
 	Assert(rti > 0 && rti <= estate->es_range_table_size);
 
-	if (!bms_is_member(rti, estate->es_unpruned_relids))
-		elog(ERROR, "trying to open a pruned relation");
+	pruned = !bms_is_member(rti, estate->es_unpruned_relids);
+	if (!isResultRel && pruned)
+		elog(ERROR, "trying to open a pruned relation for scanning");
 
 	rel = estate->es_relations[rti - 1];
 	if (rel == NULL)
@@ -837,7 +842,15 @@ ExecGetRangeTableRelation(EState *estate
 
 		Assert(rte->rtekind == RTE_RELATION);
 
-		if (!IsParallelWorker())
+		if (isResultRel && pruned)
+		{
+			/*
+			 * A pruned result relation might not have been locked yet, so we
+			 * must lock it now.
+			 */
+			rel = table_open(rte->relid, rte->rellockmode);
+		}
+		else if (!IsParallelWorker())
 		{
 			/*
 			 * In a normal query, we should already have the appropriate lock,
@@ -880,7 +893,7 @@ ExecInitResultRelation(EState *estate, R
 {
 	Relation	resultRelationDesc;
 
-	resultRelationDesc = ExecGetRangeTableRelation(estate, rti);
+	resultRelationDesc = ExecGetRangeTableRelation(estate, rti, true);
 	InitResultRelInfo(resultRelInfo,
 					  resultRelationDesc,
 					  rti,
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
new file mode 100644
index b0fe500..f2cf999
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -4471,6 +4471,7 @@ ExecInitModifyTable(ModifyTable *node, E
 	ModifyTableState *mtstate;
 	Plan	   *subplan = outerPlan(node);
 	CmdType		operation = node->operation;
+	int			total_nrels = list_length(node->resultRelations);
 	int			nrels;
 	List	   *resultRelations = NIL;
 	List	   *withCheckOptionLists = NIL;
@@ -4490,13 +4491,25 @@ ExecInitModifyTable(ModifyTable *node, E
 	/*
 	 * Only consider unpruned relations for initializing their ResultRelInfo
 	 * struct and other fields such as withCheckOptions, etc.
+	 *
+	 * Note, however, that we must avoid pruning every result relation.  This
+	 * is important for MERGE, since even if every result relation is pruned
+	 * from the subplan, there might still be NOT MATCHED rows, for which
+	 * there may be INSERT actions to perform.  To allow these actions to be
+	 * found, at least one result relation must be kept.  In addition, when
+	 * inserting into a partitioned table, ExecInitPartitionInfo() needs a
+	 * ResultRelInfo struct as a reference for building the ResultRelInfo of
+	 * the target partition.  In either case, it doesn't matter which result
+	 * relation is kept, so we just keep the last one, if all others have been
+	 * pruned.
 	 */
 	i = 0;
 	foreach(l, node->resultRelations)
 	{
 		Index		rti = lfirst_int(l);
 
-		if (bms_is_member(rti, estate->es_unpruned_relids))
+		if (bms_is_member(rti, estate->es_unpruned_relids) ||
+			(i == total_nrels - 1 && resultRelations == NIL))
 		{
 			resultRelations = lappend_int(resultRelations, rti);
 			if (node->withCheckOptionLists)
@@ -4537,6 +4550,7 @@ ExecInitModifyTable(ModifyTable *node, E
 		i++;
 	}
 	nrels = list_length(resultRelations);
+	Assert(nrels > 0);
 
 	/*
 	 * create state structure
@@ -4735,7 +4749,7 @@ ExecInitModifyTable(ModifyTable *node, E
 	 */
 	mtstate->mt_resultOidAttno =
 		ExecFindJunkAttributeInTlist(subplan->targetlist, "tableoid");
-	Assert(AttributeNumberIsValid(mtstate->mt_resultOidAttno) || nrels == 1);
+	Assert(AttributeNumberIsValid(mtstate->mt_resultOidAttno) || total_nrels == 1);
 	mtstate->mt_lastResultOid = InvalidOid; /* force lookup at first tuple */
 	mtstate->mt_lastResultIndex = 0;	/* must be zero if no such attr */
 
@@ -4832,7 +4846,7 @@ ExecInitModifyTable(ModifyTable *node, E
 	if (node->onConflictAction != ONCONFLICT_NONE)
 	{
 		/* insert may only have one relation, inheritance is not expanded */
-		Assert(nrels == 1);
+		Assert(total_nrels == 1);
 		resultRelInfo->ri_onConflictArbiterIndexes = node->arbiterIndexes;
 	}
 
@@ -4979,7 +4993,7 @@ ExecInitModifyTable(ModifyTable *node, E
 	if (operation == CMD_INSERT)
 	{
 		/* insert may only have one relation, inheritance is not expanded */
-		Assert(nrels == 1);
+		Assert(total_nrels == 1);
 		resultRelInfo = mtstate->resultRelInfo;
 		if (!resultRelInfo->ri_usesFdwDirectModify &&
 			resultRelInfo->ri_FdwRoutine != NULL &&
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
new file mode 100644
index 0d2ffab..0db5d18
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -680,7 +680,8 @@ exec_rt_fetch(Index rti, EState *estate)
 	return (RangeTblEntry *) list_nth(estate->es_range_table, rti - 1);
 }
 
-extern Relation ExecGetRangeTableRelation(EState *estate, Index rti);
+extern Relation ExecGetRangeTableRelation(EState *estate, Index rti,
+										  bool isResultRel);
 extern void ExecInitResultRelation(EState *estate, ResultRelInfo *resultRelInfo,
 								   Index rti);
 
diff --git a/src/test/regress/expected/partition_prune.out b/src/test/regress/expected/partition_prune.out
new file mode 100644
index 8097f4e..0bf3526
--- a/src/test/regress/expected/partition_prune.out
+++ b/src/test/regress/expected/partition_prune.out
@@ -4662,6 +4662,88 @@ table part_abc_view;
  2 | c | t
 (1 row)
 
+-- MERGE ... INSERT when all pruned from MERGE source.
+begin;
+explain (costs off)
+merge into part_abc_view pt
+using (select stable_one() + 1 as pid) as q join part_abc_2 pt2 on (q.pid = pt2.a) on pt.a = stable_one() + 2
+when not matched then insert values (1, 'd', false) returning pt.a;
+                   QUERY PLAN                   
+------------------------------------------------
+ Merge on part_abc
+   ->  Nested Loop Left Join
+         ->  Seq Scan on part_abc_2 pt2
+               Filter: ((stable_one() + 1) = a)
+         ->  Materialize
+               ->  Append
+                     Subplans Removed: 2
+(7 rows)
+
+merge into part_abc_view pt
+using (select stable_one() + 1 as pid) as q join part_abc_2 pt2 on (q.pid = pt2.a) on pt.a = stable_one() + 2
+when not matched then insert values (1, 'd', false) returning pt.a;
+ a 
+---
+ 1
+(1 row)
+
+table part_abc_view;
+ a | b | c 
+---+---+---
+ 1 | d | f
+ 2 | c | t
+(2 rows)
+
+rollback;
+-- A case with multiple ModifyTable nodes.
+begin;
+create table part_abc_log (action text, a int, b text, c bool);
+explain (costs off)
+with t as (
+  merge into part_abc_view pt
+  using (select stable_one() + 1 as pid) as q join part_abc_2 pt2 on (q.pid = pt2.a) on pt.a = stable_one() + 2
+  when not matched then insert values (1, 'd', false) returning merge_action(), pt.*
+)
+insert into part_abc_log select * from t returning *;
+                       QUERY PLAN                       
+--------------------------------------------------------
+ Insert on part_abc_log
+   CTE t
+     ->  Merge on part_abc
+           ->  Nested Loop Left Join
+                 ->  Seq Scan on part_abc_2 pt2
+                       Filter: ((stable_one() + 1) = a)
+                 ->  Materialize
+                       ->  Append
+                             Subplans Removed: 2
+   ->  CTE Scan on t
+(10 rows)
+
+with t as (
+  merge into part_abc_view pt
+  using (select stable_one() + 1 as pid) as q join part_abc_2 pt2 on (q.pid = pt2.a) on pt.a = stable_one() + 2
+  when not matched then insert values (1, 'd', false) returning merge_action(), pt.*
+)
+insert into part_abc_log select * from t returning *;
+ action | a | b | c 
+--------+---+---+---
+ INSERT | 1 | d | f
+(1 row)
+
+table part_abc_view;
+ a | b | c 
+---+---+---
+ 1 | d | f
+ 2 | c | t
+(2 rows)
+
+table part_abc_log;
+ action | a | b | c 
+--------+---+---+---
+ INSERT | 1 | d | f
+(1 row)
+
+rollback;
 -- A case with nested MergeAppend with its own PartitionPruneInfo.
 create index on part_abc (a);
 alter table part_abc add d int;
diff --git a/src/test/regress/sql/partition_prune.sql b/src/test/regress/sql/partition_prune.sql
new file mode 100644
index 4a2c74b..f6db947
--- a/src/test/regress/sql/partition_prune.sql
+++ b/src/test/regress/sql/partition_prune.sql
@@ -1401,6 +1401,38 @@ using (select stable_one() + 2 as pid) a
 when matched then delete returning pt.a;
 table part_abc_view;
 
+-- MERGE ... INSERT when all pruned from MERGE source.
+begin;
+explain (costs off)
+merge into part_abc_view pt
+using (select stable_one() + 1 as pid) as q join part_abc_2 pt2 on (q.pid = pt2.a) on pt.a = stable_one() + 2
+when not matched then insert values (1, 'd', false) returning pt.a;
+merge into part_abc_view pt
+using (select stable_one() + 1 as pid) as q join part_abc_2 pt2 on (q.pid = pt2.a) on pt.a = stable_one() + 2
+when not matched then insert values (1, 'd', false) returning pt.a;
+table part_abc_view;
+rollback;
+
+-- A case with multiple ModifyTable nodes.
+begin;
+create table part_abc_log (action text, a int, b text, c bool);
+explain (costs off)
+with t as (
+  merge into part_abc_view pt
+  using (select stable_one() + 1 as pid) as q join part_abc_2 pt2 on (q.pid = pt2.a) on pt.a = stable_one() + 2
+  when not matched then insert values (1, 'd', false) returning merge_action(), pt.*
+)
+insert into part_abc_log select * from t returning *;
+with t as (
+  merge into part_abc_view pt
+  using (select stable_one() + 1 as pid) as q join part_abc_2 pt2 on (q.pid = pt2.a) on pt.a = stable_one() + 2
+  when not matched then insert values (1, 'd', false) returning merge_action(), pt.*
+)
+insert into part_abc_log select * from t returning *;
+table part_abc_view;
+table part_abc_log;
+rollback;
+
 -- A case with nested MergeAppend with its own PartitionPruneInfo.
 create index on part_abc (a);
 alter table part_abc add d int;