From 281aacce9b6a05191a5137e96e67a26ed4284a57 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philip=20Dub=C3=A9?= <gunkmute@gmail.com>
Date: Thu, 9 Jan 2020 23:34:39 +0000
Subject: [PATCH] Fix row-gather for subqueries being handled by task-tracker

task-tracker has specific logic for MultiPartition when GROUP BY is missing

We were ending up in this code path because row-gather removes GROUP BY
---
 .../planner/multi_physical_planner.c          | 14 +++++++++-
 .../regress/expected/aggregate_support.out    | 28 +++++++++++++------
 src/test/regress/sql/aggregate_support.sql    |  8 +++++-
 3 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/src/backend/distributed/planner/multi_physical_planner.c b/src/backend/distributed/planner/multi_physical_planner.c
index 67f786cfd..a7e5401e8 100644
--- a/src/backend/distributed/planner/multi_physical_planner.c
+++ b/src/backend/distributed/planner/multi_physical_planner.c
@@ -65,6 +65,7 @@
 #include "optimizer/var.h"
 #endif
 #include "optimizer/restrictinfo.h"
+#include "optimizer/tlist.h"
 #include "parser/parse_relation.h"
 #include "parser/parsetree.h"
 #include "utils/builtins.h"
@@ -4221,7 +4222,18 @@ MapTaskList(MapMergeJob *mapMergeJob, List *filterTaskList)
 	}
 	else
 	{
-		partitionColumnName = ColumnName(partitionColumn, rangeTableList);
+		TargetEntry *targetEntry = tlist_member((Expr *) partitionColumn,
+												filterQuery->targetList);
+		if (targetEntry != NULL)
+		{
+			/* targetEntry->resname may be NULL */
+			partitionColumnName = targetEntry->resname;
+		}
+
+		if (partitionColumnName == NULL)
+		{
+			partitionColumnName = ColumnName(partitionColumn, rangeTableList);
+		}
 	}
 
 	foreach(filterTaskCell, filterTaskList)
diff --git a/src/test/regress/expected/aggregate_support.out b/src/test/regress/expected/aggregate_support.out
index 1bfb788db..470721d2e 100644
--- a/src/test/regress/expected/aggregate_support.out
+++ b/src/test/regress/expected/aggregate_support.out
@@ -383,20 +383,30 @@ select * FROM (
 (5 rows)
 
 select * FROM (
-    SELECT key, avg(distinct floor(agg1.val/2)) m from aggdata agg1
+    SELECT key k, avg(distinct floor(agg1.val/2)) m from aggdata agg1
     group by key
 ) subq;
- key |  m
+ k |  m
 ---------------------------------------------------------------------
-   1 |   1
-   5 |
-   3 |   2
-   7 |   4
-   6 |
-   2 | 1.5
-   9 |   0
+ 1 |   1
+ 5 |
+ 3 |   2
+ 7 |   4
+ 6 |
+ 2 | 1.5
+ 9 |   0
 (7 rows)
 
+-- Test TransformsSubqueryNode with group by not in FROM (failed in past)
+select count(*) FROM (
+    SELECT avg(distinct floor(agg1.val/2)) m from aggdata agg1
+    group by key
+) subq;
+ count
+---------------------------------------------------------------------
+     7
+(1 row)
+
 RESET citus.task_executor_type;
 -- This fails due to table types not being managed properly
 select key, count(distinct aggdata)
diff --git a/src/test/regress/sql/aggregate_support.sql b/src/test/regress/sql/aggregate_support.sql
index ba7dcf5b0..4241e3c77 100644
--- a/src/test/regress/sql/aggregate_support.sql
+++ b/src/test/regress/sql/aggregate_support.sql
@@ -184,7 +184,13 @@ select * FROM (
 ) subq ORDER BY 2, 1 LIMIT 5;
 
 select * FROM (
-    SELECT key, avg(distinct floor(agg1.val/2)) m from aggdata agg1
+    SELECT key k, avg(distinct floor(agg1.val/2)) m from aggdata agg1
+    group by key
+) subq;
+
+-- Test TransformsSubqueryNode with group by not in FROM (failed in past)
+select count(*) FROM (
+    SELECT avg(distinct floor(agg1.val/2)) m from aggdata agg1
     group by key
 ) subq;