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"
@@ -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)

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
+ 1 |   1
-   5 |
+ 5 |
-   3 |   2
+ 3 |   2
-   7 |   4
+ 7 |   4
-   6 |
+ 6 |
-   2 | 1.5
+ 2 | 1.5
-   9 |   0
+ 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;