3 files changed, 45 insertions, 14 deletions

diff --git a/src/backend/optimizer/path/allpaths.c b/src/backend/optimizer/path/allpaths.c
index 9ebe69ebad..8873a298d2 100644
--- a/src/backend/optimizer/path/allpaths.c
+++ b/src/backend/optimizer/path/allpaths.c
@@ -1964,7 +1964,7 @@ set_subquery_pathlist(PlannerInfo *root, RelOptInfo *rel,
 											 subpath->pathkeys,
 											 make_tlist_from_pathtarget(subpath->pathtarget));
 
-		if (subroot->distribution && subroot->distribution->distributionExpr)
+		if (subpath->distribution && subpath->distribution->distributionExpr)
 		{
 			ListCell *lc;
 
@@ -1976,14 +1976,14 @@ set_subquery_pathlist(PlannerInfo *root, RelOptInfo *rel,
 			 * be from the rel, need conversion.
 			 */
 			distribution = makeNode(Distribution);
-			distribution->distributionType = subroot->distribution->distributionType;
-			distribution->nodes = bms_copy(subroot->distribution->nodes);
-			distribution->restrictNodes = bms_copy(subroot->distribution->restrictNodes);
+			distribution->distributionType = subpath->distribution->distributionType;
+			distribution->nodes = bms_copy(subpath->distribution->nodes);
+			distribution->restrictNodes = bms_copy(subpath->distribution->restrictNodes);
 
 			foreach(lc, targetlist)
 			{
 				TargetEntry *tle = (TargetEntry *) lfirst(lc);
-				if (equal(tle->expr, subroot->distribution->distributionExpr))
+				if (equal(tle->expr, subpath->distribution->distributionExpr))
 				{
 					distribution->distributionExpr = (Node *)
 							makeVarFromTargetEntry(rel->relid, tle);
@@ -1992,7 +1992,7 @@ set_subquery_pathlist(PlannerInfo *root, RelOptInfo *rel,
 			}
 		}
 		else
-			distribution = subroot->distribution;
+			distribution = subpath->distribution;
 
 		/* Generate outer path using this subpath */
 		add_path(rel, (Path *)
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index c216680889..bce977ec81 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -368,6 +368,9 @@ standard_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 	final_rel = fetch_upper_rel(root, UPPERREL_FINAL, NULL);
 	best_path = get_cheapest_fractional_path(final_rel, tuple_fraction);
 
+	if (!root->distribution)
+		root->distribution = best_path->distribution;
+
 	top_plan = create_plan(root, best_path);
 #ifdef XCP
 	if (root->distribution)
@@ -6915,23 +6918,45 @@ equal_distributions(PlannerInfo *root, Distribution *dst1,
 	return false;
 }
 
+/*
+ * adjust_path_distribution
+ *		Adjust distribution of the path to match what's expected by ModifyTable.
+ *
+ * We use root->distribution to communicate distribution expected by a ModifyTable.
+ * Currently it's set either in preprocess_targetlist() for simple target relations,
+ * or in inheritance_planner() for targets that are inheritance trees.
+ *
+ * If root->distribution is NULL, we don't need to do anything and we can leave the
+ * path distribution as it is. This happens when there is no ModifyTable node, for
+ * example.
+ *
+ * If the root->distribution is set, we need to inspect it and redistribute the data
+ * if needed (when it root->distribution does not match path->distribution).
+ *
+ * We also detect DML (e.g. correlated UPDATE/DELETE or updates of distribution key)
+ * that we can't handle at this point.
+ *
+ * XXX We must not update root->distribution here, because we need to do this on all
+ * paths considered by grouping_planner(), and there's no obvious guarantee all the
+ * paths will share the same distribution. Postgres-XL 9.5 was allowed to do that,
+ * because prior to the pathification (in PostgreSQL 9.6) grouping_planner() picked
+ * before the distributions were adjusted.
+ */
 static Path *
 adjust_path_distribution(PlannerInfo *root, Query *parse, Path *path)
 {
-	/* if the root distribution is NULL, set it to path distribution */
+	/* if there is no root distribution, no redistribution is needed */
 	if (!root->distribution)
-	{
-		root->distribution = path->distribution;
-		return path;
-	}
-
-	/* don't touch paths without distribution attached (catalogs etc.) */
-	if ((path->distribution == NULL) && (root->distribution == NULL))
 		return path;
 
+	/* and also skip dummy paths */
 	if (IS_DUMMY_PATH(path))
 		return path;
 
+	/*
+	 * Both the path and root have distribution. Let's see if they differ,
+	 * and do a redistribution if not.
+	 */
 	if (equal_distributions(root, root->distribution, path->distribution))
 	{
 		if (IsLocatorReplicated(path->distribution->distributionType) &&
diff --git a/src/backend/optimizer/plan/subselect.c b/src/backend/optimizer/plan/subselect.c
index beeebd3a57..e7f82bf9b5 100644
--- a/src/backend/optimizer/plan/subselect.c
+++ b/src/backend/optimizer/plan/subselect.c
@@ -548,6 +548,9 @@ make_subplan(PlannerInfo *root, Query *orig_subquery,
 	final_rel = fetch_upper_rel(subroot, UPPERREL_FINAL, NULL);
 	best_path = get_cheapest_fractional_path(final_rel, tuple_fraction);
 
+	if (!subroot->distribution)
+		subroot->distribution = best_path->distribution;
+
 	plan = create_plan(subroot, best_path);
 
 #ifdef XCP
@@ -1223,6 +1226,9 @@ SS_process_ctes(PlannerInfo *root)
 		final_rel = fetch_upper_rel(subroot, UPPERREL_FINAL, NULL);
 		best_path = final_rel->cheapest_total_path;
 
+		if (!subroot->distribution)
+			subroot->distribution = best_path->distribution;
+
 		plan = create_plan(subroot, best_path);
 
 #ifdef XCP