Fix the avg accumulator

Author: masonh22

datafusion/functions-aggregate/src/average.rs

@@ -653,8 +653,15 @@ impl Accumulator for AvgAccumulator {
     }
 
     fn state(&mut self) -> Result<Vec<ScalarValue>> {
+        // With no non-NULL values both fields are NULL, matching
+        // `AvgGroupsAccumulator::state`
+        let (count, sum) = if self.count == 0 {
+            (None, None)
+        } else {
+            (Some(self.count), self.sum)
+        };
         Ok(vec![
-            ScalarValue::from(self.count),
+            ScalarValue::UInt64(count),
             ScalarValue::Float64(self.sum),
         ])
     }
@@ -833,9 +840,16 @@ where
     }
 
     fn state(&mut self) -> Result<Vec<ScalarValue>> {
+        // With no non-NULL values both fields are NULL, matching
+        // `AvgGroupsAccumulator::state`
+        let (count, sum) = if self.count == 0 {
+            (None, None)
+        } else {
+            (Some(self.count), self.sum)
+        };
         Ok(vec![
-            ScalarValue::from(self.count),
-            ScalarValue::new_primitive::<S>(self.sum, &self.sum_data_type)?,
+            ScalarValue::UInt64(count),
+            ScalarValue::new_primitive::<S>(sum, &self.sum_data_type)?,
         ])
     }
 
@@ -916,14 +930,21 @@ impl Accumulator for DurationAvgAccumulator {
     }
 
     fn state(&mut self) -> Result<Vec<ScalarValue>> {
+        // With no non-NULL values both fields are NULL, matching
+        // `AvgGroupsAccumulator::state`
+        let (count, sum) = if self.count == 0 {
+            (None, None)
+        } else {
+            (Some(self.count), self.sum)
+        };
         let duration_value = match self.time_unit {
-            TimeUnit::Second => ScalarValue::DurationSecond(self.sum),
-            TimeUnit::Millisecond => ScalarValue::DurationMillisecond(self.sum),
-            TimeUnit::Microsecond => ScalarValue::DurationMicrosecond(self.sum),
-            TimeUnit::Nanosecond => ScalarValue::DurationNanosecond(self.sum),
+            TimeUnit::Second => ScalarValue::DurationSecond(sum),
+            TimeUnit::Millisecond => ScalarValue::DurationMillisecond(sum),
+            TimeUnit::Microsecond => ScalarValue::DurationMicrosecond(sum),
+            TimeUnit::Nanosecond => ScalarValue::DurationNanosecond(sum),
         };
 
-        Ok(vec![ScalarValue::from(self.count), duration_value])
+        Ok(vec![ScalarValue::UInt64(count), duration_value])
     }
 
     fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {