fix test failure

Author: huaxingao

native/core/src/execution/datafusion/planner.rs

@@ -1233,10 +1233,38 @@ impl PhysicalPlanner {
     ) -> Result<Arc<dyn AggregateExpr>, ExecutionError> {
         match spark_expr.expr_struct.as_ref().unwrap() {
             AggExprStruct::Count(expr) => {
-                let if_expr = self.convert_count_to_if(&expr.children, schema.clone())?;
+                assert!(!expr.children.is_empty());
+                // Using `count_udaf` from Comet is exceptionally slow for some reason, so
+                // as a workaround we translate it to `SUM(IF(expr IS NOT NULL, 1, 0))`
+                // https://github.com/apache/datafusion-comet/issues/744
+
+                let children = expr
+                    .children
+                    .iter()
+                    .map(|child| self.create_expr(child, schema.clone()))
+                    .collect::<Result<Vec<_>, _>>()?;
+
+                // create `IS NOT NULL expr` and join them with `AND` if there are multiple
+                let not_null_expr: Arc<dyn PhysicalExpr> = children.iter().skip(1).fold(
+                    Arc::new(IsNotNullExpr::new(children[0].clone())) as Arc<dyn PhysicalExpr>,
+                    |acc, child| {
+                        Arc::new(BinaryExpr::new(
+                            acc,
+                            DataFusionOperator::And,
+                            Arc::new(IsNotNullExpr::new(child.clone())),
+                        ))
+                    },
+                );
+
+                let child = Arc::new(IfExpr::new(
+                    not_null_expr,
+                    Arc::new(Literal::new(ScalarValue::Int64(Some(1)))),
+                    Arc::new(Literal::new(ScalarValue::Int64(Some(0)))),
+                ));
+
                 create_aggregate_expr(
                     &sum_udaf(),
-                    &[if_expr],
+                    &[child],
                     &[],
                     &[],
                     &[],
@@ -1491,40 +1519,6 @@ impl PhysicalPlanner {
         }
     }
 
-    fn convert_count_to_if(
-        &self,
-        children: &[Expr],
-        schema: Arc<Schema>,
-    ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
-        assert!(!children.is_empty(), "Children should not be empty");
-
-        // Translate `COUNT` to `SUM(IF(expr IS NOT NULL, 1, 0))`
-        let children_exprs = children
-            .iter()
-            .map(|child| self.create_expr(child, schema.clone()))
-            .collect::<Result<Vec<_>, _>>()?;
-
-        // Create `IS NOT NULL expr` and combine with `AND` for multiple children
-        let not_null_expr = children_exprs.iter().skip(1).fold(
-            Arc::new(IsNotNullExpr::new(children_exprs[0].clone())) as Arc<dyn PhysicalExpr>,
-            |acc, child| {
-                Arc::new(BinaryExpr::new(
-                    acc,
-                    DataFusionOperator::And,
-                    Arc::new(IsNotNullExpr::new(child.clone())),
-                ))
-            },
-        );
-
-        let if_expr = Arc::new(IfExpr::new(
-            not_null_expr,
-            Arc::new(Literal::new(ScalarValue::Int64(Some(1)))),
-            Arc::new(Literal::new(ScalarValue::Int64(Some(0)))),
-        ));
-
-        Ok(if_expr)
-    }
-
     /// Create a DataFusion windows physical expression from Spark physical expression
     fn create_window_expr<'a>(
         &'a self,
@@ -1646,8 +1640,12 @@ impl PhysicalPlanner {
     ) -> Result<(String, Vec<Arc<dyn PhysicalExpr>>), ExecutionError> {
         match &agg_func.expr_struct {
             Some(AggExprStruct::Count(expr)) => {
-                let if_expr = self.convert_count_to_if(&expr.children, schema.clone())?;
-                Ok(("count".to_string(), vec![if_expr]))
+                let children = expr
+                    .children
+                    .iter()
+                    .map(|child| self.create_expr(child, schema.clone()))
+                    .collect::<Result<Vec<_>, _>>()?;
+                Ok(("count".to_string(), children))
             }
             Some(AggExprStruct::Min(expr)) => {
                 let child = self.create_expr(expr.child.as_ref().unwrap(), schema.clone())?;

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

@@ -209,7 +209,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
       expr match {
         case agg: AggregateExpression =>
           agg.aggregateFunction match {
-            case _: Min | _: Max | _: Count | _: Sum | _: Min | _: Max =>
+            case _: Min | _: Max | _: Count | _: Sum =>
               Some(agg)
             case _ =>
               withInfo(windowExpr, "Unsupported aggregate", expr)