fixes

native/core/src/execution/planner.rs

@@ -108,6 +108,7 @@ use jni::objects::GlobalRef;
 use num::{BigInt, ToPrimitive};
 use std::cmp::max;
 use std::{collections::HashMap, sync::Arc};
+use datafusion_expr::test::function_stub::count_udaf;
 
 // For clippy error on type_complexity.
 type PhyAggResult = Result<Vec<AggregateFunctionExpr>, ExecutionError>;

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -433,6 +433,11 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
           None
         }
       case Count(children) =>
+        if (children.length > 1 && aggExpr.isDistinct) {
+          withInfo(aggExpr, "no support for count distinct with multiple expressions")
+          return None
+        }
+
         val exprChildren = children.map(exprToProto(_, inputs, binding))
 
         if (exprChildren.forall(_.isDefined)) {

spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala

@@ -108,7 +108,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
         val cometShuffles = collect(df2.queryExecution.executedPlan) {
           case _: CometShuffleExchangeExec => true
         }
-        if (shuffleMode == "jvm") {
+        if (shuffleMode == "jvm" || shuffleMode == "auto") {
           assert(cometShuffles.length == 1)
         } else {
           // we fall back to Spark for shuffle because we do not support
@@ -608,7 +608,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
               withView("v") {
                 sql("CREATE TEMP VIEW v AS SELECT _1, _2 FROM tbl ORDER BY _1")
                 checkSparkAnswer(
-                  "SELECT _2, SUM(_1), SUM(DISTINCT _1), MIN(_1), MAX(_1), COUNT(_1)," +
+                  "SELECT _2, SUM(_1), SUM(DISTINCT _2), MIN(_1), MAX(_1), COUNT(_1)," +
                     " COUNT(DISTINCT _2), AVG(_1), FIRST(_1), LAST(_1) FROM v GROUP BY _2 ORDER BY _2")
               }
             }