feat: Support sum in window function

Author: huaxingao

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

@@ -1531,12 +1531,17 @@ impl PhysicalPlanner {
         partition_by: &[Arc<dyn PhysicalExpr>],
         sort_exprs: &[PhysicalSortExpr],
     ) -> Result<Arc<dyn WindowExpr>, ExecutionError> {
-        let (mut window_func_name, mut window_func_args) = (String::new(), Vec::new());
+        let window_func_name: String;
+        let window_args: Vec<Arc<dyn PhysicalExpr>>;
         if let Some(func) = &spark_expr.built_in_window_function {
             match &func.expr_struct {
                 Some(ExprStruct::ScalarFunc(f)) => {
-                    window_func_name.clone_from(&f.func);
-                    window_func_args.clone_from(&f.args);
+                    window_func_name = f.func.clone();
+                    window_args = f
+                        .args
+                        .iter()
+                        .map(|expr| self.create_expr(expr, input_schema.clone()))
+                        .collect::<Result<Vec<_>, ExecutionError>>()?;
                 }
                 other => {
                     return Err(ExecutionError::GeneralError(format!(
@@ -1545,9 +1550,9 @@ impl PhysicalPlanner {
                 }
             };
         } else if let Some(agg_func) = &spark_expr.agg_func {
-            let result = Self::process_agg_func(agg_func)?;
+            let result = self.process_agg_func(agg_func, input_schema.clone())?;
             window_func_name = result.0;
-            window_func_args = result.1;
+            window_args = result.1;
         } else {
             return Err(ExecutionError::GeneralError(
                 "Both func and agg_func are not set".to_string(),
@@ -1563,11 +1568,6 @@ impl PhysicalPlanner {
             }
         };
 
-        let window_args = window_func_args
-            .iter()
-            .map(|expr| self.create_expr(expr, input_schema.clone()))
-            .collect::<Result<Vec<_>, ExecutionError>>()?;
-
         let spark_window_frame = match spark_expr
             .spec
             .as_ref()
@@ -1637,32 +1637,33 @@ impl PhysicalPlanner {
         .map_err(|e| ExecutionError::DataFusionError(e.to_string()))
     }
 
-    fn process_agg_func(agg_func: &AggExpr) -> Result<(String, Vec<Expr>), ExecutionError> {
-        fn optional_expr_to_vec(expr_option: &Option<Expr>) -> Vec<Expr> {
-            expr_option
-                .as_ref()
-                .cloned()
-                .map_or_else(Vec::new, |e| vec![e])
-        }
-
-        fn int_to_stats_type(value: i32) -> Option<StatsType> {
-            match value {
-                0 => Some(StatsType::Sample),
-                1 => Some(StatsType::Population),
-                _ => None,
-            }
-        }
-
+    fn process_agg_func(
+        &self,
+        agg_func: &AggExpr,
+        schema: SchemaRef,
+    ) -> Result<(String, Vec<Arc<dyn PhysicalExpr>>), ExecutionError> {
         match &agg_func.expr_struct {
             Some(AggExprStruct::Count(expr)) => {
-                let args = &expr.children;
-                Ok(("count".to_string(), args.to_vec()))
+                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)) => {
-                Ok(("min".to_string(), optional_expr_to_vec(&expr.child)))
+                let child = self.create_expr(expr.child.as_ref().unwrap(), schema.clone())?;
+                Ok(("min".to_string(), vec![child]))
             }
             Some(AggExprStruct::Max(expr)) => {
-                Ok(("max".to_string(), optional_expr_to_vec(&expr.child)))
+                let child = self.create_expr(expr.child.as_ref().unwrap(), schema.clone())?;
+                Ok(("max".to_string(), vec![child]))
+            }
+            Some(AggExprStruct::Sum(expr)) => {
+                let child = self.create_expr(expr.child.as_ref().unwrap(), schema.clone())?;
+                let datatype = to_arrow_datatype(expr.datatype.as_ref().unwrap());
+                let child = Arc::new(CastExpr::new(child, datatype.clone(), None));
+                Ok(("sum".to_string(), vec![child]))
             }
             other => Err(ExecutionError::GeneralError(format!(
                 "{other:?} not supported for window function"

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 =>
+            case _: Min | _: Max | _: Count | _: Sum =>
               Some(agg)
             case _ =>
               withInfo(windowExpr, "Unsupported aggregate", expr)

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

@@ -1510,7 +1510,13 @@ class CometExecSuite extends CometTestBase {
         SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled) {
         withParquetTable((0 until 10).map(i => (i, 10 - i)), "t1") { // TODO: test nulls
           val aggregateFunctions =
-            List("COUNT(_1)", "COUNT(*)", "MAX(_1)", "MIN(_1)") // TODO: Test all the aggregates
+            List(
+              "COUNT(_1)",
+              "COUNT(*)",
+              "MAX(_1)",
+              "MIN(_1)",
+              "SUM(_1)"
+            ) // TODO: Test all the aggregates
 
           aggregateFunctions.foreach { function =>
             val queries = Seq(