fix: incorrect result with aggregate expression with filter

apache · Apr 18, 2024 · a11a622 · a11a622
1 parent 9321be6
commit a11a622
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -1843,6 +1843,11 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde {
           return None
         }
 
+        // Aggregate expressions with filter are not supported yet.
+        if (aggregateExpressions.exists(_.filter.isDefined)) {
+          return None
+        }
+
         val groupingExprs = groupingExpressions.map(exprToProto(_, child.output))
 
         // In some of the cases, the aggregateExpressions could be empty.

diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala
@@ -40,6 +40,19 @@ import org.apache.comet.CometSparkSessionExtensions.isSpark34Plus
 class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
   import testImplicits._
 
+  test("count with aggregation filter") {
+    withSQLConf(
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
+      CometConf.COMET_COLUMNAR_SHUFFLE_ENABLED.key -> "true") {
+      val df1 = sql("SELECT count(DISTINCT 2), count(DISTINCT 2,3)")
+      checkSparkAnswer(df1)
+
+      val df2 = sql("SELECT count(DISTINCT 2), count(DISTINCT 3,2)")
+      checkSparkAnswer(df2)
+    }
+  }
+
   test("lead/lag should return the default value if the offset row does not exist") {
     withSQLConf(
       CometConf.COMET_ENABLED.key -> "true",