Replace the relid in some missing fields during SJE

Reported-by: Alexander Lakhin
Discussion: https://postgr.es/m/a89f480f-8143-0965-f22d-0a892777f501%40gmail.com
Author: Andrei Lepikhov

src/backend/optimizer/plan/analyzejoins.c

@@ -1866,6 +1866,8 @@ remove_self_join_rel(PlannerInfo *root, PlanRowMark *kmark, PlanRowMark *rmark,
 	/* Replace varno in all the query structures */
 	query_tree_walker(root->parse, replace_varno_walker, &ctx,
 					  QTW_EXAMINE_SORTGROUP);
+	if (root->parse->resultRelation == toRemove->relid)
+		root->parse->resultRelation = toKeep->relid;
 
 	/* Replace links in the planner info */
 	remove_rel_from_query(root, toRemove, toKeep->relid, NULL, NULL);
@@ -1875,6 +1877,9 @@ remove_self_join_rel(PlannerInfo *root, PlanRowMark *kmark, PlanRowMark *rmark,
 				  toRemove->relid, toKeep->relid);
 	replace_varno((Node *) root->processed_groupClause,
 				  toRemove->relid, toKeep->relid);
+	replace_relid(root->all_result_relids, toRemove->relid, toKeep->relid);
+	replace_relid(root->leaf_result_relids, toRemove->relid, toKeep->relid);
+
 
 	/*
 	 * There may be references to the rel in root->fkey_list, but if so,

src/test/regress/expected/join.out

@@ -6868,6 +6868,18 @@ select * from emp1 t1
          ->  Seq Scan on emp1 t3
 (6 rows)
 
+-- Check that SJE replaces target relation correctly
+explain (costs off)
+WITH t1 AS (SELECT * FROM emp1)
+UPDATE emp1 SET code = t1.code + 1 FROM t1
+WHERE t1.id = emp1.id RETURNING emp1.id, emp1.code;
+            QUERY PLAN            
+----------------------------------
+ Update on emp1
+   ->  Seq Scan on emp1
+         Filter: (id IS NOT NULL)
+(3 rows)
+
 -- We can remove the join even if we find the join can't duplicate rows and
 -- the base quals of each side are different.  In the following case we end up
 -- moving quals over to s1 to make it so it can't match any rows.
@@ -7030,7 +7042,7 @@ EXPLAIN (COSTS OFF) -- A RowMark exists for the table being kept
 UPDATE sj sq SET b = 1 FROM sj as sz WHERE sq.a = sz.a;
            QUERY PLAN            
 ---------------------------------
- Update on sj sq
+ Update on sj sz
    ->  Seq Scan on sj sz
          Filter: (a IS NOT NULL)
 (3 rows)

src/test/regress/sql/join.sql

@@ -2616,6 +2616,12 @@ select * from emp1 t1
    inner join emp1 t2 on t1.id = t2.id
     left join emp1 t3 on t1.id > 1 and t1.id < 2;
 
+-- Check that SJE replaces target relation correctly
+explain (costs off)
+WITH t1 AS (SELECT * FROM emp1)
+UPDATE emp1 SET code = t1.code + 1 FROM t1
+WHERE t1.id = emp1.id RETURNING emp1.id, emp1.code;
+
 -- We can remove the join even if we find the join can't duplicate rows and
 -- the base quals of each side are different.  In the following case we end up
 -- moving quals over to s1 to make it so it can't match any rows.