fix: Average expression in Comet Final should handle all null inputs …

…from partial Spark aggregation (#261)

core/src/execution/datafusion/expressions/avg.rs

@@ -176,9 +176,15 @@ impl Accumulator for AvgAccumulator {
     }
 
     fn evaluate(&mut self) -> Result<ScalarValue> {
-        Ok(ScalarValue::Float64(
-            self.sum.map(|f| f / self.count as f64),
-        ))
+        if self.count == 0 {
+            // If all input are nulls, count will be 0 and we will get null after the division.
+            // This is consistent with Spark Average implementation.
+            Ok(ScalarValue::Float64(None))
+        } else {
+            Ok(ScalarValue::Float64(
+                self.sum.map(|f| f / self.count as f64),
+            ))
+        }
     }
 
     fn size(&self) -> usize {

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

@@ -40,6 +40,21 @@ import org.apache.comet.CometSparkSessionExtensions.isSpark34Plus
 class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
   import testImplicits._
 
+  test(
+    "Average expression in Comet Final should handle " +
+      "all null inputs from partial Spark aggregation") {
+    withTempView("allNulls") {
+      allNulls.createOrReplaceTempView("allNulls")
+      withSQLConf(
+        CometConf.COMET_ENABLED.key -> "true",
+        CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
+        CometConf.COMET_COLUMNAR_SHUFFLE_ENABLED.key -> "true") {
+        val df = sql("select sum(a), avg(a) from allNulls")
+        checkSparkAnswer(df)
+      }
+    }
+  }
+
   test("Aggregation without aggregate expressions should use correct result expressions") {
     withSQLConf(
       CometConf.COMET_ENABLED.key -> "true",