Document and enforce compatibility between an aggregate's accumulators

datafusion/expr-common/src/accumulator.rs

@@ -97,7 +97,7 @@ impl Drop for AggregateMetricTimer<'_> {
 /// `Accumulator`s are stateful objects that implement a single group. They
 /// aggregate values from multiple rows together into a final output aggregate.
 ///
-/// [`GroupsAccumulator]` is an additional more performant (but also complex) API
+/// [`GroupsAccumulator`] is an additional more performant (but also complex) API
 /// that manages state for multiple groups at once.
 ///
 /// An accumulator knows how to:
@@ -118,6 +118,7 @@ impl Drop for AggregateMetricTimer<'_> {
 /// [`state`]: Self::state
 /// [`evaluate`]: Self::evaluate
 /// [`merge_batch`]: Self::merge_batch
+/// [`GroupsAccumulator`]: crate::groups_accumulator::GroupsAccumulator
 /// [window function]: https://en.wikipedia.org/wiki/Window_function_(SQL)
 pub trait Accumulator: Send + Sync + Debug + std::any::Any {
     /// Supplies optional metrics owned by this aggregate expression.

datafusion/functions-aggregate/Cargo.toml

@@ -112,3 +112,4 @@ harness = false
 
 [features]
 force_hash_collisions = ["datafusion-common/force_hash_collisions"]
+testing = []

datafusion/functions-aggregate/src/approx_distinct.rs

@@ -126,6 +126,65 @@ impl<A: Accumulator> Accumulator for ApproxDistinctBitmapWrapper<A> {
     }
 }
 
+/// A validated, zero-copy view of a serialized `approx_distinct` partial state,
+/// as produced by [`GroupHll::serialize`] or by the per-group [`Accumulator`]s.
+enum SerializedHll<'a> {
+    /// The raw [`NUM_REGISTERS`] registers of a dense sketch.
+    Dense(&'a [u8; NUM_REGISTERS]),
+    /// Little-endian hashes of at most [`SPARSE_LIMIT`] distinct values. An
+    /// empty state decodes as an empty sparse state.
+    Sparse(&'a [[u8; size_of::<u64>()]]),
+}
+
+impl<'a> SerializedHll<'a> {
+    fn decode(bytes: &'a [u8]) -> Result<Self> {
+        if let Ok(registers) = <&[u8; NUM_REGISTERS]>::try_from(bytes) {
+            return Ok(Self::Dense(registers));
+        }
+        let (chunks, rest) = bytes.as_chunks::<{ size_of::<u64>() }>();
+        if !rest.is_empty() {
+            return internal_err!(
+                "approx_distinct: malformed sparse state: length {} is not a multiple of {}",
+                bytes.len(),
+                size_of::<u64>()
+            );
+        }
+        if chunks.len() > SPARSE_LIMIT {
+            return internal_err!(
+                "approx_distinct: malformed sparse state: length {} exceeds sparse limit {}",
+                bytes.len(),
+                SPARSE_LIMIT * size_of::<u64>()
+            );
+        }
+        Ok(Self::Sparse(chunks))
+    }
+}
+
+/// Merge the serialized partial states in `states` into `hll`.
+fn merge_states<T: Hash + ?Sized>(
+    hll: &mut HyperLogLog<T>,
+    states: &[ArrayRef],
+) -> Result<()> {
+    assert_eq!(1, states.len(), "expect only 1 element in the states");
+    let binary_array = downcast_value!(states[0], BinaryArray);
+    for v in binary_array.iter() {
+        let v = v.ok_or_else(|| {
+            internal_datafusion_err!("Impossibly got empty binary array from states")
+        })?;
+        match SerializedHll::decode(v)? {
+            SerializedHll::Dense(registers) => {
+                hll.merge(&HyperLogLog::new_with_registers(*registers));
+            }
+            SerializedHll::Sparse(chunks) => {
+                for chunk in chunks {
+                    hll.add_hashed(u64::from_le_bytes(*chunk));
+                }
+            }
+        }
+    }
+    Ok(())
+}
+
 #[derive(Debug)]
 struct HLLAccumulator {
     hll: HyperLogLog<u8>,
@@ -166,16 +225,7 @@ impl Accumulator for HLLAccumulator {
     }
 
     fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
-        assert_eq!(1, states.len(), "expect only 1 element in the states");
-        let binary_array = downcast_value!(states[0], BinaryArray);
-        for v in binary_array.iter() {
-            let v = v.ok_or_else(|| {
-                internal_datafusion_err!("Impossibly got empty binary array from states")
-            })?;
-            let other = v.try_into()?;
-            self.hll.merge(&other);
-        }
-        Ok(())
+        merge_states(&mut self.hll, states)
     }
 
     fn state(&mut self) -> Result<Vec<ScalarValue>> {
@@ -226,16 +276,7 @@ where
     }
 
     fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
-        assert_eq!(1, states.len(), "expect only 1 element in the states");
-        let binary_array = downcast_value!(states[0], BinaryArray);
-        for v in binary_array.iter() {
-            let v = v.ok_or_else(|| {
-                internal_datafusion_err!("Impossibly got empty binary array from states")
-            })?;
-            let other = v.try_into()?;
-            self.hll.merge(&other);
-        }
-        Ok(())
+        merge_states(&mut self.hll, states)
     }
 
     fn state(&mut self) -> Result<Vec<ScalarValue>> {
@@ -335,36 +376,17 @@ impl GroupHll {
         }
     }
 
-    /// Merge a serialized state (produced by [`Self::serialize`] or by the
-    /// per-group [`Accumulator`]) into this sketch.
+    /// Merge a serialized state (see [`SerializedHll`]) into this sketch.
     fn merge_serialized(&mut self, bytes: &[u8]) -> Result<isize> {
-        if bytes.is_empty() {
-            return Ok(0);
-        }
-        if bytes.len() == NUM_REGISTERS {
-            let other: HyperLogLog<u8> = bytes.try_into()?;
-            Ok(self.merge_dense(&other))
-        } else {
-            if !bytes.len().is_multiple_of(size_of::<u64>()) {
-                return internal_err!(
-                    "approx_distinct: malformed sparse state: length {} is not a multiple of {}",
-                    bytes.len(),
-                    size_of::<u64>()
-                );
+        Ok(match SerializedHll::decode(bytes)? {
+            SerializedHll::Dense(registers) => {
+                self.merge_dense(&HyperLogLog::new_with_registers(*registers))
             }
-            if bytes.len() > SPARSE_LIMIT * size_of::<u64>() {
-                return internal_err!(
-                    "approx_distinct: malformed sparse state: length {} exceeds sparse limit {}",
-                    bytes.len(),
-                    SPARSE_LIMIT * size_of::<u64>()
-                );
-            }
-            let mut delta = 0;
-            for chunk in bytes.as_chunks::<{ size_of::<u64>() }>().0 {
-                delta += self.add_hash(u64::from_le_bytes(*chunk));
-            }
-            Ok(delta)
-        }
+            SerializedHll::Sparse(chunks) => chunks
+                .iter()
+                .map(|chunk| self.add_hash(u64::from_le_bytes(*chunk)))
+                .sum(),
+        })
     }
 
     /// Merge a dense sketch into this one, promoting to dense if necessary.

datafusion/functions-aggregate/src/average.rs

@@ -653,10 +653,14 @@ impl Accumulator for AvgAccumulator {
     }
 
     fn state(&mut self) -> Result<Vec<ScalarValue>> {
-        Ok(vec![
-            ScalarValue::from(self.count),
-            ScalarValue::Float64(self.sum),
-        ])
+        // 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::UInt64(count), ScalarValue::Float64(sum)])
     }
 
     fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
@@ -833,9 +837,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 +927,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<()> {

datafusion/functions-aggregate/src/lib.rs

@@ -91,6 +91,8 @@ pub mod sum;
 pub mod variance;
 
 pub mod planner;
+#[cfg(any(test, feature = "testing"))]
+pub mod testing;
 mod utils;
 
 use crate::approx_percentile_cont::approx_percentile_cont_udaf;

datafusion/functions-aggregate/src/testing.rs

@@ -0,0 +1,31 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Utilities for testing aggregate function implementations, including user
+//! defined aggregate functions.
+//!
+//! This module requires the `testing` feature, which is intended to be enabled
+//! only in `[dev-dependencies]`:
+//!
+//! ```toml
+//! [dev-dependencies]
+//! datafusion-functions-aggregate = { version = "...", features = ["testing"] }
+//! ```
+
+mod state_compat;
+
+pub use state_compat::check_state_compatibility;