feat: let a partial aggregate with a double AVG stop aggregating

native/core/src/execution/jni_api.rs

@@ -61,6 +61,7 @@ use datafusion::{
 };
 use datafusion_comet_common::decode_string_arrays;
 use datafusion_comet_proto::spark_expression::agg_expr::ExprStruct as AggExprStruct;
+use datafusion_comet_proto::spark_expression::data_type::DataTypeId;
 use datafusion_comet_proto::spark_operator::{AggregateMode, Operator, ShuffleScan};
 use datafusion_comet_shuffle::decoded_batch_coalescer::DecodedBatchCoalescer;
 use datafusion_comet_spark_expr::url_funcs::{CometParseUrl, CometTryParseUrl};
@@ -1034,6 +1035,12 @@ fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operato
                         && agg.agg_exprs.iter().all(|expr| {
                             matches!(&expr.expr_struct, Some(AggExprStruct::Count(count))
                                 if count.children.len() == 1)
+                                // A decimal average's accumulator can't convert rows to states.
+                                || matches!(&expr.expr_struct, Some(AggExprStruct::Avg(avg))
+                                    if avg.child.is_some()
+                                        && avg.datatype.as_ref().is_some_and(|datatype| {
+                                            datatype.type_id == DataTypeId::Double as i32
+                                        }))
                         })
                 }
                 // PartialMerge consumes states, not rows.
@@ -2712,7 +2719,7 @@ mod tests {
     use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
     use datafusion::physical_plan::ExecutionPlan;
     use datafusion_comet_proto::spark_expression;
-    use datafusion_comet_proto::spark_expression::{AggExpr, Count, Expr, Sum};
+    use datafusion_comet_proto::spark_expression::{AggExpr, Avg, Count, Expr, Sum};
     use datafusion_comet_proto::spark_operator::{HashAggregate, ShuffleWriter};
     use std::cell::Cell;
     use std::future::Future;
@@ -2729,6 +2736,17 @@ mod tests {
             expr_struct: Some(AggExprStruct::Sum(Sum::default())),
             ..Default::default()
         };
+        let avg = |child: Option<Expr>, type_id: DataTypeId| AggExpr {
+            expr_struct: Some(AggExprStruct::Avg(Avg {
+                child,
+                datatype: Some(spark_expression::DataType {
+                    type_id: type_id as i32,
+                    type_info: None,
+                }),
+                ..Default::default()
+            })),
+            ..Default::default()
+        };
         let partial = HashAggregate {
             grouping_exprs: vec![Expr::default()],
             agg_exprs: vec![count.clone()],
@@ -2773,6 +2791,13 @@ mod tests {
                 agg_exprs: vec![],
                 ..partial.clone()
             },
+            HashAggregate {
+                agg_exprs: vec![
+                    count.clone(),
+                    avg(Some(Expr::default()), DataTypeId::Double),
+                ],
+                ..partial.clone()
+            },
             HashAggregate {
                 agg_exprs: vec![count.clone(), count],
                 ..partial.clone()
@@ -2797,6 +2822,14 @@ mod tests {
                 agg_exprs: vec![AggExpr::default()],
                 ..partial.clone()
             },
+            HashAggregate {
+                agg_exprs: vec![avg(None, DataTypeId::Double)],
+                ..partial.clone()
+            },
+            HashAggregate {
+                agg_exprs: vec![avg(Some(Expr::default()), DataTypeId::Decimal)],
+                ..partial.clone()
+            },
             HashAggregate {
                 agg_exprs: vec![AggExpr {
                     expr_struct: Some(AggExprStruct::Count(Count {

native/spark-expr/src/agg_funcs/avg.rs

@@ -334,12 +334,33 @@ where
         ])
     }
 
+    /// Each row's own partial state, as `state` emits it: the value and a count of one, or a
+    /// zero sum and count for a null or filtered-out row.
     fn convert_to_state(
         &self,
-        _values: &[ArrayRef],
-        _opt_filter: Option<&arrow::array::BooleanArray>,
+        values: &[ArrayRef],
+        opt_filter: Option<&arrow::array::BooleanArray>,
     ) -> Result<Vec<ArrayRef>> {
-        not_impl_err!("Input batch conversion to state not implemented")
+        let values = values[0].as_primitive::<T>();
+        let mut sums = Vec::with_capacity(values.len());
+        let mut counts = Vec::with_capacity(values.len());
+        for idx in 0..values.len() {
+            let included = values.is_valid(idx)
+                && opt_filter.is_none_or(|filter| filter.is_valid(idx) && filter.value(idx));
+            if included {
+                sums.push(values.value(idx));
+                counts.push(1);
+            } else {
+                sums.push(T::default_value());
+                counts.push(0);
+            }
+        }
+        let sums = PrimitiveArray::<T>::new(sums.into(), None)
+            .with_data_type(self.return_data_type.clone());
+        Ok(vec![
+            Arc::new(sums) as ArrayRef,
+            Arc::new(Int64Array::new(counts.into(), None)) as ArrayRef,
+        ])
     }
 
     fn size(&self) -> usize {
@@ -365,4 +386,32 @@ mod tests {
         assert_eq!(partial.evaluate()?, ScalarValue::Float64(None));
         Ok(())
     }
+
+    /// Converting rows straight to partial states, as a partial aggregate that stops aggregating
+    /// does, gives each row the state that aggregating it alone into a group would.
+    #[test]
+    fn convert_to_state_matches_one_row_groups() -> Result<()> {
+        use arrow::array::BooleanArray;
+        let values: ArrayRef = Arc::new(Float64Array::from(vec![
+            Some(1.5),
+            None,
+            Some(-2.0),
+            Some(4.0),
+        ]));
+        let filter = BooleanArray::from(vec![Some(true), Some(true), Some(true), Some(false)]);
+        let new_accumulator = || {
+            AvgGroupsAccumulator::<arrow::datatypes::Float64Type, _>::new(
+                &Float64,
+                |sum: f64, count: i64| Ok(sum / count as f64),
+            )
+        };
+
+        let converted =
+            new_accumulator().convert_to_state(&[Arc::clone(&values)], Some(&filter))?;
+
+        let mut grouped = new_accumulator();
+        grouped.update_batch(&[values], &[0, 1, 2, 3], Some(&filter), 4)?;
+        assert_eq!(converted, grouped.state(EmitTo::All)?);
+        Ok(())
+    }
 }

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -710,10 +710,11 @@ object CometConf extends ShimCometConf {
         "Experimental opt-in: let a native partial aggregate stop aggregating once its input " +
           "looks mostly distinct, and send the rest of the task's rows to the shuffle " +
           "unaggregated. Only applies to partial aggregates that feed Comet native shuffle " +
-          "and whose aggregate functions, if any, are all single-argument COUNT. The check " +
-          "starts after the first 100,000 input rows of a task, and once aggregation stops " +
-          "it does not resume, so a task whose keys repeat after a mostly distinct start can " +
-          s"shuffle many times more rows than it would with this disabled. $TUNING_GUIDE.")
+          "and whose aggregate functions, if any, are all single-argument COUNT or AVG of a " +
+          "non-decimal value. The check starts after the first 100,000 input rows of a task, " +
+          "and once aggregation stops it does not resume, so a task whose keys repeat after a " +
+          "mostly distinct start can shuffle many times more rows than it would with this " +
+          s"disabled. $TUNING_GUIDE.")
       .booleanConf
       .createWithDefault(false)
 

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

@@ -3254,7 +3254,17 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
           }
           withSQLConf(
             "spark.comet.datafusion.execution.skip_partial_aggregation_probe_ratio_threshold" -> "0.8") {
-            for (expression <- Seq("sum(v)", "count(v, w)", "count(*) + sum(v)")) {
+            // A non-decimal AVG converts rows to partial states, including null rows.
+            assert(
+              aggregates("SELECT k, avg(v), avg(w) FROM skip_partial_eligibility GROUP BY k")
+                .map(skipped)
+                .sum > 0L)
+            for (expression <- Seq(
+                "sum(v)",
+                "count(v, w)",
+                "count(*) + sum(v)",
+                // A decimal AVG's accumulator can't convert rows to partial states.
+                "avg(CAST(v AS DECIMAL(30, 2)))")) {
               val result = aggregates(
                 "SELECT sum(n) FROM " +
                   s"(SELECT k, $expression n FROM skip_partial_eligibility GROUP BY k)")