From ea6b98fda42fe0a6d5070844177449a35fcb87c4 Mon Sep 17 00:00:00 2001
From: Marco Slot <marco.slot@gmail.com>
Date: Sat, 9 Dec 2017 16:23:14 +0100
Subject: [PATCH] Allow count(distinct) in queries with a subquery

---
 .../planner/multi_logical_optimizer.c         | 12 ++-----
 src/test/regress/expected/multi_subquery.out  | 34 +++++++++++++++++--
 src/test/regress/sql/multi_subquery.sql       | 26 +++++++++++++-
 3 files changed, 58 insertions(+), 14 deletions(-)

diff --git a/src/backend/distributed/planner/multi_logical_optimizer.c b/src/backend/distributed/planner/multi_logical_optimizer.c
index 290d83a20..958c0b262 100644
--- a/src/backend/distributed/planner/multi_logical_optimizer.c
+++ b/src/backend/distributed/planner/multi_logical_optimizer.c
@@ -2754,15 +2754,6 @@ ErrorIfUnsupportedAggregateDistinct(Aggref *aggregateExpression,
 
 	AggregateType aggregateType = GetAggregateType(aggregateExpression->aggfnoid);
 
-	/* check if logical plan includes a subquery */
-	List *subqueryMultiTableList = SubqueryMultiTableList(logicalPlanNode);
-	if (subqueryMultiTableList != NIL)
-	{
-		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-						errmsg("cannot push down this subquery"),
-						errdetail("distinct in the outermost query is unsupported")));
-	}
-
 	/*
 	 * We partially support count(distinct) in subqueries, other distinct aggregates in
 	 * subqueries are not supported yet.
@@ -2790,7 +2781,8 @@ ErrorIfUnsupportedAggregateDistinct(Aggref *aggregateExpression,
 		foreach(multiTableNodeCell, multiTableNodeList)
 		{
 			MultiTable *multiTable = (MultiTable *) lfirst(multiTableNodeCell);
-			if (multiTable->relationId == SUBQUERY_RELATION_ID)
+			if (multiTable->relationId == SUBQUERY_RELATION_ID ||
+				multiTable->relationId == SUBQUERY_PUSHDOWN_RELATION_ID)
 			{
 				ereport(ERROR, (errmsg("cannot compute aggregate (distinct)"),
 								errdetail("Only count(distinct) aggregate is "
diff --git a/src/test/regress/expected/multi_subquery.out b/src/test/regress/expected/multi_subquery.out
index b7c7da2e9..8bd197c4d 100644
--- a/src/test/regress/expected/multi_subquery.out
+++ b/src/test/regress/expected/multi_subquery.out
@@ -152,7 +152,7 @@ ERROR:  cannot push down this subquery
 DETAIL:  Limit in subquery without limit in the outermost query is unsupported
 -- reset the flag for next query
 SET citus.subquery_pushdown to OFF;
--- Check that we error out if the outermost query is a distinct clause.
+-- Check that we support count distinct with a subquery
 SELECT
 	count(DISTINCT a)
 FROM (
@@ -163,8 +163,36 @@ FROM (
 	GROUP BY
 	   l_orderkey
 ) z;
-ERROR:  cannot push down this subquery
-DETAIL:  distinct in the outermost query is unsupported
+ count 
+-------
+     7
+(1 row)
+
+-- We do not support distinct aggregates other than count distinct with a subquery
+SELECT
+	sum(DISTINCT a)
+FROM (
+	SELECT
+		count(*) a
+	FROM
+		lineitem_subquery
+	GROUP BY
+	   l_orderkey
+) z;
+ERROR:  cannot compute aggregate (distinct)
+DETAIL:  Only count(distinct) aggregate is supported in subqueries
+SELECT
+	avg(DISTINCT a)
+FROM (
+	SELECT
+		count(*) a
+	FROM
+		lineitem_subquery
+	GROUP BY
+	   l_orderkey
+) z;
+ERROR:  cannot compute aggregate (distinct)
+DETAIL:  Only count(distinct) aggregate is supported in subqueries
 -- Check supported subquery types.
 SELECT
 	o_custkey,
diff --git a/src/test/regress/sql/multi_subquery.sql b/src/test/regress/sql/multi_subquery.sql
index 99887f318..e9468f10a 100644
--- a/src/test/regress/sql/multi_subquery.sql
+++ b/src/test/regress/sql/multi_subquery.sql
@@ -145,7 +145,7 @@ FROM
 -- reset the flag for next query
 SET citus.subquery_pushdown to OFF;
 
--- Check that we error out if the outermost query is a distinct clause.
+-- Check that we support count distinct with a subquery
 
 SELECT
 	count(DISTINCT a)
@@ -158,6 +158,30 @@ FROM (
 	   l_orderkey
 ) z;
 
+-- We do not support distinct aggregates other than count distinct with a subquery
+
+SELECT
+	sum(DISTINCT a)
+FROM (
+	SELECT
+		count(*) a
+	FROM
+		lineitem_subquery
+	GROUP BY
+	   l_orderkey
+) z;
+
+SELECT
+	avg(DISTINCT a)
+FROM (
+	SELECT
+		count(*) a
+	FROM
+		lineitem_subquery
+	GROUP BY
+	   l_orderkey
+) z;
+
 -- Check supported subquery types.
 
 SELECT