Allow count(distinct) in queries with a subquery

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 "

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
+ count 
-DETAIL:  distinct in the outermost query is unsupported
+-------
+     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,

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