From b00b1f9165801724d0330286235b4dca258f4fd1 Mon Sep 17 00:00:00 2001
From: jcoleman <jtc331@gmail.com>
Date: Thu, 15 Apr 2021 01:48:36 +0000
Subject: [PATCH v2] Fix find_em_expr_usable_for_sorting_rel

---
 src/backend/optimizer/path/equivclass.c       | 44 ++++++++++++++++++-
 .../regress/expected/incremental_sort.out     | 26 +++++++++++
 src/test/regress/sql/incremental_sort.sql     |  7 +++
 3 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/src/backend/optimizer/path/equivclass.c b/src/backend/optimizer/path/equivclass.c
index 0188c1e9a1..27867e16e8 100644
--- a/src/backend/optimizer/path/equivclass.c
+++ b/src/backend/optimizer/path/equivclass.c
@@ -68,6 +68,7 @@ static Bitmapset *get_eclass_indexes_for_relids(PlannerInfo *root,
 												Relids relids);
 static Bitmapset *get_common_eclass_indexes(PlannerInfo *root, Relids relids1,
 											Relids relids2);
+static Expr* expr_list_member_ignore_relabel(Expr *node, List *exprlist);
 
 
 /*
@@ -798,6 +799,27 @@ find_em_expr_for_rel(EquivalenceClass *ec, RelOptInfo *rel)
 	return NULL;
 }
 
+static Expr*
+expr_list_member_ignore_relabel(Expr *node, List *exprlist)
+{
+	ListCell   *temp;
+
+	while (node && IsA(node, RelabelType))
+		node = ((RelabelType *) node)->arg;
+
+	foreach(temp, exprlist)
+	{
+		Expr	   *tlexpr = (Expr *) lfirst(temp);
+
+		while (tlexpr && IsA(tlexpr, RelabelType))
+			tlexpr = ((RelabelType *) tlexpr)->arg;
+
+		if (equal(node, tlexpr))
+			return node;
+	}
+	return NULL;
+}
+
 /*
  * Find an equivalence class member expression that can be safely used to build
  * a sort node using the provided relation. The rules are a subset of those
@@ -819,6 +841,9 @@ find_em_expr_usable_for_sorting_rel(PlannerInfo *root, EquivalenceClass *ec,
 	{
 		EquivalenceMember *em = lfirst(lc_em);
 		Expr	   *em_expr = em->em_expr;
+		List	   *exprvars;
+		ListCell   *k;
+		PathTarget *target = rel->reltarget;
 
 		/*
 		 * We shouldn't be trying to sort by an equivalence class that
@@ -858,8 +883,23 @@ find_em_expr_usable_for_sorting_rel(PlannerInfo *root, EquivalenceClass *ec,
 		 * valuable here to expend the work trying to find a match in the
 		 * target's exprs since such a sort will have to be postponed anyway.
 		 */
-		if (!ec->ec_has_volatile)
-			return em->em_expr;
+		if (ec->ec_has_volatile)
+			return NULL;
+
+
+		exprvars = pull_var_clause((Node *) em_expr,
+								   PVC_INCLUDE_AGGREGATES |
+								   PVC_INCLUDE_WINDOWFUNCS |
+								   PVC_INCLUDE_PLACEHOLDERS);
+		foreach(k, exprvars)
+		{
+			if (!expr_list_member_ignore_relabel(lfirst(k), target->exprs))
+				break;
+		}
+		list_free(exprvars);
+
+		if (!k) /* If we made it through exprvars finding a tlist member for each */
+			return em->em_expr; /* found usable expression */
 	}
 
 	/* We didn't find any suitable equivalence class expression */
diff --git a/src/test/regress/expected/incremental_sort.out b/src/test/regress/expected/incremental_sort.out
index a417b566d9..fd18af4c0c 100644
--- a/src/test/regress/expected/incremental_sort.out
+++ b/src/test/regress/expected/incremental_sort.out
@@ -1579,6 +1579,32 @@ order by 1, 2;
    ->  Function Scan on generate_series
 (7 rows)
 
+-- Parallel sort with an aggregate that be safely generated partially in parallel,
+-- but we can't sort by the partial aggregates.
+explain (costs off) select count(*)
+from tenk1 t1
+join tenk1 t2 on t1.unique1 = t2.unique2
+join tenk1 t3 on t2.unique1 = t3.unique1
+order by count(*);
+                                          QUERY PLAN                                           
+-----------------------------------------------------------------------------------------------
+ Sort
+   Sort Key: (count(*))
+   ->  Finalize Aggregate
+         ->  Gather
+               Workers Planned: 2
+               ->  Partial Aggregate
+                     ->  Parallel Hash Join
+                           Hash Cond: (t2.unique1 = t3.unique1)
+                           ->  Parallel Hash Join
+                                 Hash Cond: (t1.unique1 = t2.unique2)
+                                 ->  Parallel Index Only Scan using tenk1_unique1 on tenk1 t1
+                                 ->  Parallel Hash
+                                       ->  Parallel Index Scan using tenk1_unique2 on tenk1 t2
+                           ->  Parallel Hash
+                                 ->  Parallel Index Only Scan using tenk1_unique1 on tenk1 t3
+(15 rows)
+
 -- Parallel sort but with expression (correlated subquery) that
 -- is prohibited in parallel plans.
 explain (costs off) select distinct
diff --git a/src/test/regress/sql/incremental_sort.sql b/src/test/regress/sql/incremental_sort.sql
index 81429164d4..67748bb64b 100644
--- a/src/test/regress/sql/incremental_sort.sql
+++ b/src/test/regress/sql/incremental_sort.sql
@@ -257,6 +257,13 @@ from tenk1, lateral (select tenk1.unique1 from generate_series(1, 1000)) as sub;
 explain (costs off) select sub.unique1, md5(stringu1)
 from tenk1, lateral (select tenk1.unique1 from generate_series(1, 1000)) as sub
 order by 1, 2;
+-- Parallel sort with an aggregate that be safely generated partially in parallel,
+-- but we can't sort by the partial aggregates.
+explain (costs off) select count(*)
+from tenk1 t1
+join tenk1 t2 on t1.unique1 = t2.unique2
+join tenk1 t3 on t2.unique1 = t3.unique1
+order by count(*);
 -- Parallel sort but with expression (correlated subquery) that
 -- is prohibited in parallel plans.
 explain (costs off) select distinct
-- 
2.20.1

