2 files changed, 77 insertions, 6 deletions

diff --git a/src/backend/optimizer/path/allpaths.c b/src/backend/optimizer/path/allpaths.c
index 0eeff804bcf..09a0ee24d60 100644
--- a/src/backend/optimizer/path/allpaths.c
+++ b/src/backend/optimizer/path/allpaths.c
@@ -2756,7 +2756,8 @@ get_useful_pathkeys_for_relation(PlannerInfo *root, RelOptInfo *rel)
 	/*
 	 * Considering query_pathkeys is always worth it, because it might allow
 	 * us to avoid a total sort when we have a partially presorted path
-	 * available.
+	 * available or to push the total sort into the parallel portion of the
+	 * query.
 	 */
 	if (root->query_pathkeys)
 	{
@@ -2769,17 +2770,17 @@ get_useful_pathkeys_for_relation(PlannerInfo *root, RelOptInfo *rel)
 			EquivalenceClass *pathkey_ec = pathkey->pk_eclass;
 
 			/*
-			 * We can only build an Incremental Sort for pathkeys which
-			 * contain an EC member in the current relation, so ignore any
-			 * suffix of the list as soon as we find a pathkey without an EC
-			 * member the relation.
+			 * We can only build a sort for pathkeys which contain an EC
+			 * member in the current relation's target, so ignore any suffix
+			 * of the list as soon as we find a pathkey without an EC member
+			 * in the relation.
 			 *
 			 * By still returning the prefix of the pathkeys list that does
 			 * meet criteria of EC membership in the current relation, we
 			 * enable not just an incremental sort on the entirety of
 			 * query_pathkeys but also incremental sort below a JOIN.
 			 */
-			if (!find_em_expr_for_rel(pathkey_ec, rel))
+			if (!find_em_expr_usable_for_sorting_rel(pathkey_ec, rel))
 				break;
 
 			npathkeys++;
diff --git a/src/backend/optimizer/path/equivclass.c b/src/backend/optimizer/path/equivclass.c
index 823422edad0..130f482eb62 100644
--- a/src/backend/optimizer/path/equivclass.c
+++ b/src/backend/optimizer/path/equivclass.c
@@ -795,6 +795,76 @@ find_em_expr_for_rel(EquivalenceClass *ec, RelOptInfo *rel)
 }
 
 /*
+ * Find an equivalence class member expression that can be safely used by a
+ * sort node on top of the provided relation. The rules here must match those
+ * applied in prepare_sort_from_pathkeys.
+ */
+Expr *
+find_em_expr_usable_for_sorting_rel(EquivalenceClass *ec, RelOptInfo *rel)
+{
+	ListCell   *lc_em;
+
+	/*
+	 * If there is more than one equivalence member matching these
+	 * requirements we'll be content to choose any one of them.
+	 */
+	foreach(lc_em, ec->ec_members)
+	{
+		EquivalenceMember *em = lfirst(lc_em);
+		Expr	   *em_expr = em->em_expr;
+		PathTarget *target = rel->reltarget;
+		ListCell   *lc_target_expr;
+
+		/*
+		 * We shouldn't be trying to sort by an equivalence class that
+		 * contains a constant, so no need to consider such cases any further.
+		 */
+		if (em->em_is_const)
+			continue;
+
+		/*
+		 * Any Vars in the equivalence class member need to come from this
+		 * relation. This is a superset of prepare_sort_from_pathkeys ignoring
+		 * child members unless they belong to the rel being sorted.
+		 */
+		if (!bms_is_subset(em->em_relids, rel->relids))
+			continue;
+
+		/*
+		 * As long as the expression isn't volatile then
+		 * prepare_sort_from_pathkeys is able to generate a new target entry,
+		 * so there's no need to verify that one already exists.
+		 */
+		if (!ec->ec_has_volatile)
+			return em->em_expr;
+
+		/*
+		 * If, however, it's volatile, we have to verify that the
+		 * equivalence member's expr is already generated in the
+		 * relation's target (we do strip relabels first from both
+		 * expressions, which is cheap and might allow us to match
+		 * more expressions).
+		 */
+		while (em_expr && IsA(em_expr, RelabelType))
+			em_expr = ((RelabelType *) em_expr)->arg;
+
+		foreach(lc_target_expr, target->exprs)
+		{
+			Expr	   *target_expr = lfirst(lc_target_expr);
+
+			while (target_expr && IsA(target_expr, RelabelType))
+				target_expr = ((RelabelType *) target_expr)->arg;
+
+			if (equal(target_expr, em_expr))
+				return em->em_expr;
+		}
+	}
+
+	/* We didn't find any suitable equivalence class expression */
+	return NULL;
+}
+
+/*
  * generate_base_implied_equalities
  *	  Generate any restriction clauses that we can deduce from equivalence
  *	  classes.