Merge pull request #3371 from citusdata/fix-task-tracker-row-gather-subquery

Fix row-gather for subqueries being handled by task-tracker

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"
@@ -4220,9 +4221,20 @@ MapTaskList(MapMergeJob *mapMergeJob, List *filterTaskList)
 		partitionColumnName = groupByTargetEntry->resname;
 	}
 	else
+	{
+		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)
 	{

src/test/regress/expected/aggregate_support.out

@@ -383,10 +383,10 @@ 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 |
@@ -397,6 +397,16 @@ select * FROM (
  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)

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;