diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 30d15291e3..c4244e6d29 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -1802,7 +1802,9 @@ inheritance_planner(PlannerInfo *root)
 	 * When performing UPDATE/DELETE on a partitioned table, if the query has
 	 * a WHERE clause which supports it, we may be able to perform run-time
 	 * partition pruning.  The following code sets things up to allow this to
-	 * be possible.
+	 * be possible using the information from partition_root that was used
+	 * during planning of the SELECT version of this query which we performed
+	 * above.
 	 */
 	if (partition_root && !dummy_update)
 	{
@@ -1810,36 +1812,29 @@ inheritance_planner(PlannerInfo *root)
 		int			i;
 
 		/*
-		 * Fetch the target partitioned table from the SELECT version of
-		 * the query which we performed above.  This may have the base quals
+		 * Fetch the target partitioned table which may have the base quals
 		 * which could allow the run-time pruning to work.
 		 */
 		parent_rel = partition_root->simple_rel_array[top_parentRTindex];
-
 		final_rel->baserestrictinfo = parent_rel->baserestrictinfo;
 
-		/* build a list of partitioned rels */
+		/* Collect all non-leaf tables in the partition tree being updated. */
 		i = -1;
 		while ((i = bms_next_member(parent_relids, i)) > 0)
 			partitioned_rels = lappend_int(partitioned_rels, i);
 
-
 		/*
-		 * In order to build the run-time pruning data we'll need append rels
-		 * any sub-partitioned tables.  If there are some of those and the
-		 * append_rel_array is not already allocated, then do that now.
+		 * There can only be a single partition tree, the one whose root is
+		 * the query's main target table.
 		 */
-		if (list_length(partitioned_rels) > 1 &&
-			root->append_rel_array == NULL)
-			root->append_rel_array = palloc0(sizeof(AppendRelInfo *) *
-											 root->simple_rel_array_size);
+		partitioned_rels = list_make1(partitioned_rels);
 
 		/*
-		 * There can only be a single partition hierarchy, so it's fine to
-		 * just make a single element list of the partitioned_rels.
+		 * Update simple_rel_array and append_rel_array so that runtime
+		 * pruning setup logic can find the relavant partitioned relations.
+		 * Just use the one that the planning of SELECT version of the query
+		 * would have created.
 		 */
-		partitioned_rels = list_make1(partitioned_rels);
-
 		i = -1;
 		while ((i = bms_next_member(parent_relids, i)) >= 0)
 		{
@@ -1847,11 +1842,7 @@ inheritance_planner(PlannerInfo *root)
 
 			root->simple_rel_array[i] = partition_root->simple_rel_array[i];
 
-			/*
-			 * The root partition won't have an append rel entry, so we can
-			 * skip that.  We'll need to take the partition_root's version for
-			 * any sub-partitioned table's
-			 */
+			/* Root partitioned table doesn't have an AppendRelInfo. */
 			if (i != top_parentRTindex)
 			{
 				Assert(root->append_rel_array[i] == NULL);
diff --git a/src/backend/partitioning/partprune.c b/src/backend/partitioning/partprune.c
index ca53f684c2..020bc60fed 100644
--- a/src/backend/partitioning/partprune.c
+++ b/src/backend/partitioning/partprune.c
@@ -213,9 +213,6 @@ static void partkey_datum_from_expr(PartitionPruneContext *context,
  *
  * 'parentrel' is the RelOptInfo for an appendrel, and 'subpaths' is the list
  * of scan paths for its child rels.
- *
- * If 'resultRelations' is non-NIL, then this List of relids is used to build
- * the mapping structures.  Otherwise the 'subpaths' List is used.
  *
   * 'partitioned_rels' is a List containing Lists of relids of partitioned
  * tables (a/k/a non-leaf partitions) that are parents of some of the child
@@ -227,6 +224,12 @@ static void partkey_datum_from_expr(PartitionPruneContext *context,
  * that set into the PartitionPruneInfo's 'other_subplans' field.  Callers
  * will likely never want to prune subplans which are mentioned in this field.
  *
+ * 'subpaths' (of an Append/MergeAppend for SELECT) and 'resultRelations'
+ * (of a ModifyTable for UPDATE/DELETE) are provided to map a given partition's
+ * index to their corresponding subpath's and result relation's index, resp.
+ * Having these maps allows the executor to easily skip subplans or result
+ * relations based on the indexes of partitions that are pruned.
+ *
  * 'prunequal' is a list of potential pruning quals.
  */
 PartitionPruneInfo *
@@ -249,10 +252,6 @@ make_partition_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 	 */
 	relid_subplan_map = palloc0(sizeof(int) * root->simple_rel_array_size);
 
-	/*
-	 * If 'resultRelations' are present then map these, otherwise we map the
-	 * 'subpaths' List.
-	 */
 	if (resultRelations != NIL)
 	{
 		i = 1;