From 9fb0e9f9d80591025d524e9fe73b4b6e9e90fa52 Mon Sep 17 00:00:00 2001
From: amit <amitlangote09@gmail.com>
Date: Mon, 4 Mar 2019 15:52:11 +0900
Subject: [PATCH v31 4/8] Adjust inheritance_planner to reuse source child
 tables

Set of source inheritance child tables won't change across repeated
planning of the query for different target child relations.  So,
note down the RangeTblEntrys, AppendRelInfos, and PlanRowMarks of
the source inheritance child tables after the planning for the first
child table is finished.  When planning for the subsequent child
target tables, put the saved objects of source inheritance child
tables into the child PlannerInfos and set contains_inherit_children
to true, so that query_planner doesn't add them again.

This restores the regression test outputs to their original state,
because like before, source inheritance child tables are not added
multiple times.
---
 .../postgres_fdw/expected/postgres_fdw.out    |  46 ++---
 src/backend/nodes/outfuncs.c                  |   1 +
 src/backend/optimizer/plan/planner.c          | 191 +++++++++++-------
 src/backend/optimizer/util/relnode.c          |  48 ++++-
 src/include/nodes/pathnodes.h                 |   7 +
 src/test/regress/expected/partition_prune.out |  66 +++++-
 src/test/regress/expected/rowsecurity.out     |  16 +-
 src/test/regress/sql/partition_prune.sql      |   7 +
 8 files changed, 268 insertions(+), 114 deletions(-)

diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out
index bbc9ead2e9..c912e95fe1 100644
--- a/contrib/postgres_fdw/expected/postgres_fdw.out
+++ b/contrib/postgres_fdw/expected/postgres_fdw.out
@@ -7190,8 +7190,8 @@ select * from bar where f1 in (select f1 from foo) for share;
 -- Check UPDATE with inherited target and an inherited source table
 explain (verbose, costs off)
 update bar set f2 = f2 + 100 where f1 in (select f1 from foo);
-                                            QUERY PLAN                                             
----------------------------------------------------------------------------------------------------
+                                         QUERY PLAN                                          
+---------------------------------------------------------------------------------------------
  Update on public.bar
    Update on public.bar
    Foreign Update on public.bar2
@@ -7214,22 +7214,22 @@ update bar set f2 = f2 + 100 where f1 in (select f1 from foo);
                                  Output: foo2.f1, foo2.ctid, foo2.*, foo2.tableoid
                                  Remote SQL: SELECT f1, f2, f3, ctid FROM public.loct1
    ->  Hash Join
-         Output: bar2.f1, (bar2.f2 + 100), bar2.f3, bar2.ctid, foo_1.ctid, foo_1.*, foo_1.tableoid
+         Output: bar2.f1, (bar2.f2 + 100), bar2.f3, bar2.ctid, foo.ctid, foo.*, foo.tableoid
          Inner Unique: true
-         Hash Cond: (bar2.f1 = foo_1.f1)
+         Hash Cond: (bar2.f1 = foo.f1)
          ->  Foreign Scan on public.bar2
                Output: bar2.f1, bar2.f2, bar2.f3, bar2.ctid
                Remote SQL: SELECT f1, f2, f3, ctid FROM public.loct2 FOR UPDATE
          ->  Hash
-               Output: foo_1.f1, foo_1.ctid, foo_1.*, foo_1.tableoid
+               Output: foo.f1, foo.ctid, foo.*, foo.tableoid
                ->  HashAggregate
-                     Output: foo_1.f1, foo_1.ctid, foo_1.*, foo_1.tableoid
-                     Group Key: foo_1.f1
+                     Output: foo.f1, foo.ctid, foo.*, foo.tableoid
+                     Group Key: foo.f1
                      ->  Append
-                           ->  Seq Scan on public.foo foo_1
-                                 Output: foo_1.f1, foo_1.ctid, foo_1.*, foo_1.tableoid
-                           ->  Foreign Scan on public.foo2 foo2_1
-                                 Output: foo2_1.f1, foo2_1.ctid, foo2_1.*, foo2_1.tableoid
+                           ->  Seq Scan on public.foo
+                                 Output: foo.f1, foo.ctid, foo.*, foo.tableoid
+                           ->  Foreign Scan on public.foo2
+                                 Output: foo2.f1, foo2.ctid, foo2.*, foo2.tableoid
                                  Remote SQL: SELECT f1, f2, f3, ctid FROM public.loct1
 (39 rows)
 
@@ -7276,8 +7276,8 @@ where bar.f1 = ss.f1;
                ->  Seq Scan on public.bar
                      Output: bar.f1, bar.f2, bar.ctid
    ->  Merge Join
-         Output: bar2.f1, (bar2.f2 + 100), bar2.f3, bar2.ctid, (ROW(foo_2.f1))
-         Merge Cond: (bar2.f1 = foo_2.f1)
+         Output: bar2.f1, (bar2.f2 + 100), bar2.f3, bar2.ctid, (ROW(foo.f1))
+         Merge Cond: (bar2.f1 = foo.f1)
          ->  Sort
                Output: bar2.f1, bar2.f2, bar2.f3, bar2.ctid
                Sort Key: bar2.f1
@@ -7285,18 +7285,18 @@ where bar.f1 = ss.f1;
                      Output: bar2.f1, bar2.f2, bar2.f3, bar2.ctid
                      Remote SQL: SELECT f1, f2, f3, ctid FROM public.loct2 FOR UPDATE
          ->  Sort
-               Output: foo_2.f1, (ROW(foo_2.f1))
-               Sort Key: foo_2.f1
+               Output: foo.f1, (ROW(foo.f1))
+               Sort Key: foo.f1
                ->  Append
-                     ->  Seq Scan on public.foo foo_2
-                           Output: foo_2.f1, ROW(foo_2.f1)
-                     ->  Foreign Scan on public.foo2 foo2_2
-                           Output: foo2_2.f1, ROW(foo2_2.f1)
+                     ->  Seq Scan on public.foo
+                           Output: foo.f1, ROW(foo.f1)
+                     ->  Foreign Scan on public.foo2
+                           Output: foo2.f1, ROW(foo2.f1)
                            Remote SQL: SELECT f1 FROM public.loct1
-                     ->  Seq Scan on public.foo foo_3
-                           Output: (foo_3.f1 + 3), ROW((foo_3.f1 + 3))
-                     ->  Foreign Scan on public.foo2 foo2_3
-                           Output: (foo2_3.f1 + 3), ROW((foo2_3.f1 + 3))
+                     ->  Seq Scan on public.foo foo_1
+                           Output: (foo_1.f1 + 3), ROW((foo_1.f1 + 3))
+                     ->  Foreign Scan on public.foo2 foo2_1
+                           Output: (foo2_1.f1 + 3), ROW((foo2_1.f1 + 3))
                            Remote SQL: SELECT f1 FROM public.loct1
 (45 rows)
 
diff --git a/src/backend/nodes/outfuncs.c b/src/backend/nodes/outfuncs.c
index 69179a07c3..cd4741d36b 100644
--- a/src/backend/nodes/outfuncs.c
+++ b/src/backend/nodes/outfuncs.c
@@ -2220,6 +2220,7 @@ _outPlannerInfo(StringInfo str, const PlannerInfo *node)
 	WRITE_BITMAPSET_FIELD(curOuterRels);
 	WRITE_NODE_FIELD(curOuterParams);
 	WRITE_BOOL_FIELD(partColsUpdated);
+	WRITE_BOOL_FIELD(contains_inherit_children);
 }
 
 static void
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 2b77fa6d90..c564dd8928 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -646,6 +646,7 @@ subquery_planner(PlannerGlobal *glob, Query *parse,
 		root->wt_param_id = -1;
 	root->non_recursive_path = NULL;
 	root->partColsUpdated = false;
+	root->contains_inherit_children = false;
 
 	/*
 	 * If there is a WITH list, process each WITH query and either convert it
@@ -1197,7 +1198,6 @@ inheritance_planner(PlannerInfo *root)
 	Query	   *parse = root->parse;
 	int			top_parentRTindex = parse->resultRelation;
 	Bitmapset  *subqueryRTindexes;
-	Bitmapset  *modifiableARIindexes;
 	int			nominalRelation = -1;
 	Index		rootRelation = 0;
 	List	   *final_rtable = NIL;
@@ -1219,6 +1219,11 @@ inheritance_planner(PlannerInfo *root)
 	Bitmapset  *parent_relids = bms_make_singleton(top_parentRTindex);
 	PlannerInfo **parent_roots = NULL;
 	List	   *orig_append_rel_list = list_copy(root->append_rel_list);
+	List	   *orig_rtable = list_copy(root->parse->rtable);
+	List	   *source_appinfos = NIL;
+	List	   *source_child_rowmarks = NIL;
+	List	   *source_child_rtes = NIL;
+	bool		is_first_child;
 
 	Assert(parse->commandType != CMD_INSERT);
 
@@ -1262,10 +1267,12 @@ inheritance_planner(PlannerInfo *root)
 	 * management of the rowMarks list.
 	 *
 	 * To begin with, generate a bitmapset of the relids of the subquery RTEs.
+	 * We use orig_rtable, not parse->rtable, because we wouldn't need to
+	 * consider any newly added RTEs as they all must be RTE_RELATION entries.
 	 */
 	subqueryRTindexes = NULL;
 	rti = 1;
-	foreach(lc, parse->rtable)
+	foreach(lc, orig_rtable)
 	{
 		RangeTblEntry *rte = lfirst_node(RangeTblEntry, lc);
 
@@ -1274,32 +1281,6 @@ inheritance_planner(PlannerInfo *root)
 		rti++;
 	}
 
-	/*
-	 * Next, we want to identify which AppendRelInfo items contain references
-	 * to any of the aforesaid subquery RTEs.  These items will need to be
-	 * copied and modified to adjust their subquery references; whereas the
-	 * other ones need not be touched.  It's worth being tense over this
-	 * because we can usually avoid processing most of the AppendRelInfo
-	 * items, thereby saving O(N^2) space and time when the target is a large
-	 * inheritance tree.  We can identify AppendRelInfo items by their
-	 * child_relid, since that should be unique within the list.
-	 */
-	modifiableARIindexes = NULL;
-	if (subqueryRTindexes != NULL)
-	{
-		foreach(lc, root->append_rel_list)
-		{
-			AppendRelInfo *appinfo = lfirst_node(AppendRelInfo, lc);
-
-			if (bms_is_member(appinfo->parent_relid, subqueryRTindexes) ||
-				bms_is_member(appinfo->child_relid, subqueryRTindexes) ||
-				bms_overlap(pull_varnos((Node *) appinfo->translated_vars),
-							subqueryRTindexes))
-				modifiableARIindexes = bms_add_member(modifiableARIindexes,
-													  appinfo->child_relid);
-		}
-	}
-
 	/*
 	 * If the parent RTE is a partitioned table, we should use that as the
 	 * nominal target relation, because the RTEs added for partitioned tables
@@ -1329,6 +1310,7 @@ inheritance_planner(PlannerInfo *root)
 	/*
 	 * And now we can get on with generating a plan for each child table.
 	 */
+	is_first_child = true;
 	foreach(lc, root->append_rel_list)
 	{
 		AppendRelInfo *appinfo = lfirst_node(AppendRelInfo, lc);
@@ -1357,9 +1339,6 @@ inheritance_planner(PlannerInfo *root)
 		subroot = makeNode(PlannerInfo);
 		memcpy(subroot, parent_root, sizeof(PlannerInfo));
 
-		/* grouping_planner doesn't need to see the target children. */
-		subroot->append_rel_list = list_copy(orig_append_rel_list);
-
 		/*
 		 * Generate modified query with this rel as target.  We first apply
 		 * adjust_appendrel_attrs, which copies the Query and changes
@@ -1442,31 +1421,39 @@ inheritance_planner(PlannerInfo *root)
 		subroot->rowMarks = copyObject(parent_root->rowMarks);
 
 		/*
-		 * The append_rel_list likewise might contain references to subquery
-		 * RTEs (if any subqueries were flattenable UNION ALLs).  So prepare
-		 * to apply ChangeVarNodes to that, too.  As explained above, we only
-		 * want to copy items that actually contain such references; the rest
-		 * can just get linked into the subroot's append_rel_list.
-		 *
-		 * If we know there are no such references, we can just use the outer
-		 * append_rel_list unmodified.
+		 * Prepare to apply ChangeVarNodes to orig_append_rel_list by copying
+		 * the AppendRelInfos contained in it.  Note that orig_append_rel_list
+		 * contains only the AppendRelInfos pertaining to flattened UNION ALL
+		 * subqueries, so we'll be applying ChangeVarNodes to all of them.
+		 * We don't need to look at other members of
+		 * parent_root->append_rel_list, which are those corresponding to
+		 * child target relations, because they won't contain any subquery
+		 * references.
 		 */
-		if (modifiableARIindexes != NULL)
-		{
-			ListCell   *lc2;
+		subroot->append_rel_list = copyObject(orig_append_rel_list);
 
-			subroot->append_rel_list = NIL;
-			foreach(lc2, parent_root->append_rel_list)
-			{
-				AppendRelInfo *appinfo2 = lfirst_node(AppendRelInfo, lc2);
+		/*
+		 * If this isn't the first child query, then we can use the child
+		 * objects for source child relations created during the planning of
+		 * 1st child query.  IOW, this planning run doesn't need to create the
+		 * child objects again, indicated by setting contains_inherit_children
+		 * of the child PlannerInfo.  We use list_copy() on various
+		 * source_* lists, because we may add more entries to it below.
+		 *
+		 * We won't need to add source_appinfos into subroot yet, because we
+		 * won't need to modify any of its entries.  Only orig_append_rel_list
+		 * will contain contain references to the subquery RTEs that we will
+		 * make copies of below.
+		 */
+		if (!is_first_child)
+			subroot->rowMarks =
+				list_concat(subroot->rowMarks,
+							list_copy(source_child_rowmarks));
 
-				if (bms_is_member(appinfo2->child_relid, modifiableARIindexes))
-					appinfo2 = copyObject(appinfo2);
-
-				subroot->append_rel_list = lappend(subroot->append_rel_list,
-												   appinfo2);
-			}
-		}
+		if (!is_first_child)
+			subroot->parse->rtable =
+				list_concat(subroot->parse->rtable,
+							list_copy(source_child_rtes));
 
 		/*
 		 * Add placeholders to the child Query's rangetable list to fill the
@@ -1481,17 +1468,16 @@ inheritance_planner(PlannerInfo *root)
 		/*
 		 * If this isn't the first child Query, generate duplicates of all
 		 * subquery RTEs, and adjust Var numbering to reference the
-		 * duplicates. To simplify the loop logic, we scan the original rtable
-		 * not the copy just made by adjust_appendrel_attrs; that should be OK
-		 * since subquery RTEs couldn't contain any references to the target
-		 * rel.
+		 * duplicates.  Note that we scan the original rtable before any
+		 * child target relations were added, which is OK, because no other
+		 * RTEs would contain references to subquery rels being modified.
 		 */
 		if (final_rtable != NIL && subqueryRTindexes != NULL)
 		{
 			ListCell   *lr;
 
 			rti = 1;
-			foreach(lr, parent_parse->rtable)
+			foreach(lr, orig_rtable)
 			{
 				RangeTblEntry *rte = lfirst_node(RangeTblEntry, lr);
 
@@ -1508,20 +1494,8 @@ inheritance_planner(PlannerInfo *root)
 					newrti = list_length(subroot->parse->rtable) + 1;
 					ChangeVarNodes((Node *) subroot->parse, rti, newrti, 0);
 					ChangeVarNodes((Node *) subroot->rowMarks, rti, newrti, 0);
-					/* Skip processing unchanging parts of append_rel_list */
-					if (modifiableARIindexes != NULL)
-					{
-						ListCell   *lc2;
-
-						foreach(lc2, subroot->append_rel_list)
-						{
-							AppendRelInfo *appinfo2 = lfirst_node(AppendRelInfo, lc2);
-
-							if (bms_is_member(appinfo2->child_relid,
-											  modifiableARIindexes))
-								ChangeVarNodes((Node *) appinfo2, rti, newrti, 0);
-						}
-					}
+					ChangeVarNodes((Node *) subroot->append_rel_list, rti,
+								   newrti, 0);
 					rte = copyObject(rte);
 					ChangeVarNodes((Node *) rte->securityQuals, rti, newrti, 0);
 					subroot->parse->rtable = lappend(subroot->parse->rtable,
@@ -1531,6 +1505,23 @@ inheritance_planner(PlannerInfo *root)
 			}
 		}
 
+		/*
+		 * Concatenate source AppendRelInfos, so the source child objects that
+		 * we hooked into subroot above are discoverable.
+		 */
+		if (!is_first_child)
+		{
+			subroot->append_rel_list = list_concat(subroot->append_rel_list,
+												   source_appinfos);
+
+			/*
+			 * We have added child RTEs and row marks and now also the
+			 * AppendRelInfos needed to find them children, so the next
+			 * query_planner doesn't need to add them again.
+			 */
+			subroot->contains_inherit_children = true;
+		}
+
 		/* There shouldn't be any OJ info to translate, as yet */
 		Assert(subroot->join_info_list == NIL);
 		/* and we haven't created PlaceHolderInfos, either */
@@ -1548,6 +1539,62 @@ inheritance_planner(PlannerInfo *root)
 		set_cheapest(sub_final_rel);
 		subpath = sub_final_rel->cheapest_total_path;
 
+		/*
+		 * If we finished planning our first child query, copy the source
+		 * child objects that were added during its planning.
+		 */
+		if (is_first_child && subroot->append_rel_list)
+		{
+			int		num_skip_appinfos = list_length(orig_append_rel_list);
+			int		num_skip_rowmarks = list_length(parent_root->rowMarks);
+			int		num_skip_rtes = list_length(parent_parse->rtable);
+			ListCell *lc2;
+
+			Assert(source_appinfos == NIL);
+			source_appinfos = list_copy_tail(subroot->append_rel_list,
+											 num_skip_appinfos);
+			Assert(source_child_rowmarks == NIL);
+			source_child_rowmarks = list_copy_tail(subroot->rowMarks,
+												   num_skip_rowmarks);
+			Assert(source_child_rtes == NIL);
+			source_child_rtes = list_copy_tail(subroot->parse->rtable,
+											   num_skip_rtes);
+
+			/*
+			 * Original parent PlanRowMark is modified when adding the
+			 * child PlanRowMarks.  Copy those changes so that the planning
+			 * of subsequent child queries works correctly.  That is
+			 * necessary, because we won't be adding child objects again,
+			 * so there won't be an opportunity to modify the parent
+			 * PlanRowMark as desired.
+			 *
+			 * All the original parent row marks are the beginning of
+			 * subroot->rowMarks, skip the rest.
+			 */
+			foreach(lc2, parent_root->rowMarks)
+			{
+				PlanRowMark *oldrc = lfirst_node(PlanRowMark, lc2);
+				ListCell *lc3;
+
+				foreach(lc3, subroot->rowMarks)
+				{
+					PlanRowMark *newrc = lfirst_node(PlanRowMark, lc3);
+
+					if (oldrc->rti == newrc->rti)
+					{
+						/* modified in expand_single_inheritance_child() */
+						oldrc->allMarkTypes = newrc->allMarkTypes;
+						/* modified in expand_inherited_rtentry() */
+						oldrc->isParent = newrc->isParent;
+					}
+					else
+						break;
+				}
+			}
+
+			is_first_child = false;
+		}
+
 		/*
 		 * If this child rel was excluded by constraint exclusion, exclude it
 		 * from the result plan.
diff --git a/src/backend/optimizer/util/relnode.c b/src/backend/optimizer/util/relnode.c
index 9dcc3167e0..2d475e3297 100644
--- a/src/backend/optimizer/util/relnode.c
+++ b/src/backend/optimizer/util/relnode.c
@@ -339,21 +339,39 @@ add_appendrel_other_rels(PlannerInfo *root, RangeTblEntry *rte, Index rti)
 {
 	ListCell	   *l;
 	RelOptInfo	   *rel;
+	int				i;
 
 	/*
-	 * Add inheritance children to the query. For child tables that are
-	 * themselves partitioned, their children will be added recursively.
+	 * Add inheritance children to the query if not already done. For child
+	 * tables that are themselves partitioned, their children will be added
+	 * recursively.
 	 */
-	if (rte->rtekind == RTE_RELATION)
+	if (rte->rtekind == RTE_RELATION && !root->contains_inherit_children)
 	{
 		expand_inherited_rtentry(root, rte, rti);
 		return;
 	}
 
-	/* Add other rels of flattened UNION ALL children. */
-	Assert(rte->rtekind == RTE_SUBQUERY);
+	/*
+	 * Add other rels.  We get here for non-subquery RTEs only if root
+	 * already contains inheritance childern but we still need to create
+	 * RelOptInfos for them.
+	 */
+	Assert(rte->rtekind == RTE_SUBQUERY || rte->rtekind == RTE_RELATION);
 	rel = find_base_rel(root, rti);
 
+	/*
+	 * For partitioned tables, we need to store the child RelOptInfos in the
+	 * rel->part_rels array too.
+	 */
+	if (rel->part_scheme)
+	{
+		Assert(rel->nparts > 0);
+		rel->part_rels = (RelOptInfo **)
+				palloc0(sizeof(RelOptInfo *) * rel->nparts);
+	}
+
+	i = 0;
 	foreach(l, root->append_rel_list)
 	{
 		AppendRelInfo  *appinfo = lfirst(l);
@@ -369,10 +387,30 @@ add_appendrel_other_rels(PlannerInfo *root, RangeTblEntry *rte, Index rti)
 		Assert(childrte != NULL);
 		childrel = build_simple_rel(root, childRTindex, rel);
 
+		/*
+		 * For partitioned parents, we also need to add childrel to its
+		 * part_rels array.  The order in which child tables appear in
+		 * append_rel_list is same as the order in which they appear in the
+		 * parent's PartitionDesc, so assigning partitions like this works.
+		 */
+		if (rel->part_scheme != NULL)
+		{
+			Assert(i < rel->nparts);
+			rel->part_rels[i] = childrel;
+		}
+
+		i++;
+
 		/* Child may itself be an inherited relation. */
 		if (childrte->inh)
 			add_appendrel_other_rels(root, childrte, childRTindex);
 	}
+
+	/*
+	 * For a partitioned table with non-zero number of partitions, we must
+	 * have assigned all elements of its part_rels array.
+	 */
+	Assert(rel->nparts == 0 || i == rel->nparts);
 }
 
 /*
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index 253e0b7e48..31811f9adb 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -350,6 +350,13 @@ struct PlannerInfo
 
 	/* Does this query modify any partition key columns? */
 	bool		partColsUpdated;
+
+	/*
+	 * Does this PlannerInfo and its Query object contain *all* inheritance
+	 * children?  If true, the RTEs, the AppendRelInfos, and the PlanRowMarks
+	 * of all the children are assumed to be present.
+	 */
+	bool		contains_inherit_children;
 };
 
 
diff --git a/src/test/regress/expected/partition_prune.out b/src/test/regress/expected/partition_prune.out
index 71942394ba..f1115c561e 100644
--- a/src/test/regress/expected/partition_prune.out
+++ b/src/test/regress/expected/partition_prune.out
@@ -2559,16 +2559,16 @@ update ab_a1 set b = 3 from ab where ab.a = 1 and ab.a = ab_a1.a;
                            Index Cond: (a = 1)
    ->  Nested Loop (actual rows=1 loops=1)
          ->  Append (actual rows=1 loops=1)
-               ->  Bitmap Heap Scan on ab_a1_b1 ab_a1_b1_2 (actual rows=0 loops=1)
+               ->  Bitmap Heap Scan on ab_a1_b1 ab_a1_b1_1 (actual rows=0 loops=1)
                      Recheck Cond: (a = 1)
                      ->  Bitmap Index Scan on ab_a1_b1_a_idx (actual rows=0 loops=1)
                            Index Cond: (a = 1)
-               ->  Bitmap Heap Scan on ab_a1_b2 ab_a1_b2_2 (actual rows=1 loops=1)
+               ->  Bitmap Heap Scan on ab_a1_b2 ab_a1_b2_1 (actual rows=1 loops=1)
                      Recheck Cond: (a = 1)
                      Heap Blocks: exact=1
                      ->  Bitmap Index Scan on ab_a1_b2_a_idx (actual rows=1 loops=1)
                            Index Cond: (a = 1)
-               ->  Bitmap Heap Scan on ab_a1_b3 ab_a1_b3_2 (actual rows=0 loops=1)
+               ->  Bitmap Heap Scan on ab_a1_b3 ab_a1_b3_1 (actual rows=0 loops=1)
                      Recheck Cond: (a = 1)
                      ->  Bitmap Index Scan on ab_a1_b3_a_idx (actual rows=1 loops=1)
                            Index Cond: (a = 1)
@@ -2580,16 +2580,16 @@ update ab_a1 set b = 3 from ab where ab.a = 1 and ab.a = ab_a1.a;
                            Index Cond: (a = 1)
    ->  Nested Loop (actual rows=0 loops=1)
          ->  Append (actual rows=1 loops=1)
-               ->  Bitmap Heap Scan on ab_a1_b1 ab_a1_b1_3 (actual rows=0 loops=1)
+               ->  Bitmap Heap Scan on ab_a1_b1 ab_a1_b1_1 (actual rows=0 loops=1)
                      Recheck Cond: (a = 1)
                      ->  Bitmap Index Scan on ab_a1_b1_a_idx (actual rows=0 loops=1)
                            Index Cond: (a = 1)
-               ->  Bitmap Heap Scan on ab_a1_b2 ab_a1_b2_3 (actual rows=1 loops=1)
+               ->  Bitmap Heap Scan on ab_a1_b2 ab_a1_b2_1 (actual rows=1 loops=1)
                      Recheck Cond: (a = 1)
                      Heap Blocks: exact=1
                      ->  Bitmap Index Scan on ab_a1_b2_a_idx (actual rows=1 loops=1)
                            Index Cond: (a = 1)
-               ->  Bitmap Heap Scan on ab_a1_b3 ab_a1_b3_3 (actual rows=0 loops=1)
+               ->  Bitmap Heap Scan on ab_a1_b3 ab_a1_b3_1 (actual rows=0 loops=1)
                      Recheck Cond: (a = 1)
                      ->  Bitmap Index Scan on ab_a1_b3_a_idx (actual rows=1 loops=1)
                            Index Cond: (a = 1)
@@ -2606,6 +2606,60 @@ table ab;
  1 | 3
 (1 row)
 
+-- Test UPDATE where source relation has run-time pruning enabled
+truncate ab;
+insert into ab values (1, 1), (1, 2), (1, 3), (2, 1);
+explain (analyze, costs off, summary off, timing off)
+update ab_a1 set b = 3 from ab_a2 where ab_a2.b = (select 1);
+                              QUERY PLAN                              
+----------------------------------------------------------------------
+ Update on ab_a1 (actual rows=0 loops=1)
+   Update on ab_a1_b1
+   Update on ab_a1_b2
+   Update on ab_a1_b3
+   InitPlan 1 (returns $0)
+     ->  Result (actual rows=1 loops=1)
+   ->  Nested Loop (actual rows=1 loops=1)
+         ->  Seq Scan on ab_a1_b1 (actual rows=1 loops=1)
+         ->  Materialize (actual rows=1 loops=1)
+               ->  Append (actual rows=1 loops=1)
+                     ->  Seq Scan on ab_a2_b1 (actual rows=1 loops=1)
+                           Filter: (b = $0)
+                     ->  Seq Scan on ab_a2_b2 (never executed)
+                           Filter: (b = $0)
+                     ->  Seq Scan on ab_a2_b3 (never executed)
+                           Filter: (b = $0)
+   ->  Nested Loop (actual rows=1 loops=1)
+         ->  Seq Scan on ab_a1_b2 (actual rows=1 loops=1)
+         ->  Materialize (actual rows=1 loops=1)
+               ->  Append (actual rows=1 loops=1)
+                     ->  Seq Scan on ab_a2_b1 (actual rows=1 loops=1)
+                           Filter: (b = $0)
+                     ->  Seq Scan on ab_a2_b2 (never executed)
+                           Filter: (b = $0)
+                     ->  Seq Scan on ab_a2_b3 (never executed)
+                           Filter: (b = $0)
+   ->  Nested Loop (actual rows=1 loops=1)
+         ->  Seq Scan on ab_a1_b3 (actual rows=1 loops=1)
+         ->  Materialize (actual rows=1 loops=1)
+               ->  Append (actual rows=1 loops=1)
+                     ->  Seq Scan on ab_a2_b1 (actual rows=1 loops=1)
+                           Filter: (b = $0)
+                     ->  Seq Scan on ab_a2_b2 (never executed)
+                           Filter: (b = $0)
+                     ->  Seq Scan on ab_a2_b3 (never executed)
+                           Filter: (b = $0)
+(36 rows)
+
+select tableoid::regclass, * from ab;
+ tableoid | a | b 
+----------+---+---
+ ab_a1_b3 | 1 | 3
+ ab_a1_b3 | 1 | 3
+ ab_a1_b3 | 1 | 3
+ ab_a2_b1 | 2 | 1
+(4 rows)
+
 drop table ab, lprt_a;
 -- Join
 create table tbl1(col1 int);
diff --git a/src/test/regress/expected/rowsecurity.out b/src/test/regress/expected/rowsecurity.out
index a6a499ed4a..2e170497c9 100644
--- a/src/test/regress/expected/rowsecurity.out
+++ b/src/test/regress/expected/rowsecurity.out
@@ -1815,26 +1815,26 @@ AND f_leak(t1_1.b) AND f_leak(t1_2.b) RETURNING *, t1_1, t1_2;
                ->  Seq Scan on t3 t1_2_2
                      Filter: ((a = 4) AND ((a % 2) = 0) AND f_leak(b))
    ->  Nested Loop
-         Join Filter: (t1_1_1.b = t1_2_3.b)
+         Join Filter: (t1_1_1.b = t1_2.b)
          ->  Seq Scan on t2 t1_1_1
                Filter: ((a = 4) AND ((a % 2) = 0) AND f_leak(b))
          ->  Append
-               ->  Seq Scan on t1 t1_2_3
+               ->  Seq Scan on t1 t1_2
                      Filter: ((a = 4) AND ((a % 2) = 0) AND f_leak(b))
-               ->  Seq Scan on t2 t1_2_4
+               ->  Seq Scan on t2 t1_2_1
                      Filter: ((a = 4) AND ((a % 2) = 0) AND f_leak(b))
-               ->  Seq Scan on t3 t1_2_5
+               ->  Seq Scan on t3 t1_2_2
                      Filter: ((a = 4) AND ((a % 2) = 0) AND f_leak(b))
    ->  Nested Loop
-         Join Filter: (t1_1_2.b = t1_2_6.b)
+         Join Filter: (t1_1_2.b = t1_2.b)
          ->  Seq Scan on t3 t1_1_2
                Filter: ((a = 4) AND ((a % 2) = 0) AND f_leak(b))
          ->  Append
-               ->  Seq Scan on t1 t1_2_6
+               ->  Seq Scan on t1 t1_2
                      Filter: ((a = 4) AND ((a % 2) = 0) AND f_leak(b))
-               ->  Seq Scan on t2 t1_2_7
+               ->  Seq Scan on t2 t1_2_1
                      Filter: ((a = 4) AND ((a % 2) = 0) AND f_leak(b))
-               ->  Seq Scan on t3 t1_2_8
+               ->  Seq Scan on t3 t1_2_2
                      Filter: ((a = 4) AND ((a % 2) = 0) AND f_leak(b))
 (37 rows)
 
diff --git a/src/test/regress/sql/partition_prune.sql b/src/test/regress/sql/partition_prune.sql
index dc327caffd..e15afc2dfd 100644
--- a/src/test/regress/sql/partition_prune.sql
+++ b/src/test/regress/sql/partition_prune.sql
@@ -588,6 +588,13 @@ explain (analyze, costs off, summary off, timing off)
 update ab_a1 set b = 3 from ab where ab.a = 1 and ab.a = ab_a1.a;
 table ab;
 
+-- Test UPDATE where source relation has run-time pruning enabled
+truncate ab;
+insert into ab values (1, 1), (1, 2), (1, 3), (2, 1);
+explain (analyze, costs off, summary off, timing off)
+update ab_a1 set b = 3 from ab_a2 where ab_a2.b = (select 1);
+select tableoid::regclass, * from ab;
+
 drop table ab, lprt_a;
 
 -- Join
-- 
2.17.2 (Apple Git-113)

