Set partitioned_rels appropriately when UNION ALL is used.

In most cases, this omission won't matter, because the appropriate
locks will have been acquired during parse/plan or by AcquireExecutorLocks.
But it's a bug all the same.

Report by Ashutosh Bapat.  Patch by me, reviewed by Amit Langote.

Discussion: http://postgr.es/m/CAFjFpRdHb_ZnoDTuBXqrudWXh3H1ibLkr6nHsCFT96fSK4DXtA@mail.gmail.com

2 files changed, 40 insertions, 8 deletions

diff --git a/src/backend/optimizer/path/allpaths.c b/src/backend/optimizer/path/allpaths.c
index f087ddb61db..4a82636734e 100644
--- a/src/backend/optimizer/path/allpaths.c
+++ b/src/backend/optimizer/path/allpaths.c
@@ -1287,13 +1287,34 @@ add_paths_to_append_rel(PlannerInfo *root, RelOptInfo *rel,
 	ListCell   *l;
 	List	   *partitioned_rels = NIL;
 	RangeTblEntry *rte;
+	bool		build_partitioned_rels = false;
 
+	/*
+	 * A plain relation will already have a PartitionedChildRelInfo if it is
+	 * partitioned.  For a subquery RTE, no PartitionedChildRelInfo exists; we
+	 * collect all partitioned_rels associated with any child.  (This assumes
+	 * that we don't need to look through multiple levels of subquery RTEs; if
+	 * we ever do, we could create a PartitionedChildRelInfo with the
+	 * accumulated list of partitioned_rels which would then be found when
+	 * populated our parent rel with paths.  For the present, that appears to
+	 * be unnecessary.)
+	 */
 	rte = planner_rt_fetch(rel->relid, root);
-	if (rte->relkind == RELKIND_PARTITIONED_TABLE)
+	switch (rte->rtekind)
 	{
-		partitioned_rels = get_partitioned_child_rels(root, rel->relid);
-		/* The root partitioned table is included as a child rel */
-		Assert(list_length(partitioned_rels) >= 1);
+		case RTE_RELATION:
+			if (rte->relkind == RELKIND_PARTITIONED_TABLE)
+			{
+				partitioned_rels =
+					get_partitioned_child_rels(root, rel->relid);
+				Assert(list_length(partitioned_rels) >= 1);
+			}
+			break;
+		case RTE_SUBQUERY:
+			build_partitioned_rels = true;
+			break;
+		default:
+			elog(ERROR, "unexpcted rtekind: %d", (int) rte->rtekind);
 	}
 
 	/*
@@ -1307,6 +1328,19 @@ add_paths_to_append_rel(PlannerInfo *root, RelOptInfo *rel,
 		ListCell   *lcp;
 
 		/*
+		 * If we need to build partitioned_rels, accumulate the partitioned
+		 * rels for this child.
+		 */
+		if (build_partitioned_rels)
+		{
+			List	   *cprels;
+
+			cprels = get_partitioned_child_rels(root, childrel->relid);
+			partitioned_rels = list_concat(partitioned_rels,
+										   list_copy(cprels));
+		}
+
+		/*
 		 * If child has an unparameterized cheapest-total path, add that to
 		 * the unparameterized Append path we are constructing for the parent.
 		 * If not, there's no workable unparameterized path.
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 3e93eee6d8e..793242759d1 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -6068,7 +6068,8 @@ plan_cluster_use_sort(Oid tableOid, Oid indexOid)
  *		Returns a list of the RT indexes of the partitioned child relations
  *		with rti as the root parent RT index.
  *
- * Note: Only call this function on RTEs known to be partitioned tables.
+ * Note: This function might get called even for range table entries that
+ * are not partitioned tables; in such a case, it will simply return NIL.
  */
 List *
 get_partitioned_child_rels(PlannerInfo *root, Index rti)
@@ -6087,8 +6088,5 @@ get_partitioned_child_rels(PlannerInfo *root, Index rti)
 		}
 	}
 
-	/* The root partitioned table is included as a child rel */
-	Assert(list_length(result) >= 1);
-
 	return result;
 }