feat: Implement state conversion for remaining group accumulators

datafusion-examples/examples/udf/advanced_udaf.rs

@@ -23,8 +23,10 @@ use datafusion::{arrow::datatypes::DataType, logical_expr::Volatility};
 use std::sync::Arc;
 
 use arrow::array::{
-    ArrayRef, AsArray, Float32Array, PrimitiveArray, PrimitiveBuilder, UInt32Array,
+    Array, ArrayRef, AsArray, BooleanArray, Float32Array, PrimitiveArray,
+    PrimitiveBuilder, UInt32Array,
 };
+use arrow::buffer::NullBuffer;
 use arrow::datatypes::{ArrowNativeTypeOp, ArrowPrimitiveType, Float64Type, UInt32Type};
 use arrow::record_batch::RecordBatch;
 use arrow_schema::FieldRef;
@@ -237,7 +239,7 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        opt_filter: Option<&arrow::array::BooleanArray>,
+        opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
         assert_eq!(values.len(), 1, "single argument to update_batch");
@@ -359,6 +361,43 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator {
         ])
     }
 
+    fn convert_to_state(
+        &self,
+        values: &[ArrayRef],
+        opt_filter: Option<&BooleanArray>,
+    ) -> Result<Vec<ArrayRef>> {
+        assert_eq!(values.len(), 1, "single argument to convert_to_state");
+
+        let prods = values[0]
+            .as_primitive::<Float64Type>()
+            .clone()
+            .with_data_type(self.prod_data_type.clone());
+        let counts = UInt32Array::from_value(1, prods.len());
+
+        let filter_nulls = opt_filter.map(|filter| {
+            let validity = match filter.nulls() {
+                Some(nulls) => filter.values() & nulls.inner(),
+                None => filter.values().clone(),
+            };
+            NullBuffer::new(validity)
+        });
+        let nulls = NullBuffer::union(filter_nulls.as_ref(), prods.nulls());
+
+        let prods =
+            PrimitiveArray::<Float64Type>::new(prods.values().clone(), nulls.clone())
+                .with_data_type(self.prod_data_type.clone());
+        let counts = UInt32Array::new(counts.values().clone(), nulls);
+
+        Ok(vec![
+            Arc::new(prods) as ArrayRef,
+            Arc::new(counts) as ArrayRef,
+        ])
+    }
+
+    fn supports_convert_to_state(&self) -> bool {
+        true
+    }
+
     fn size(&self) -> usize {
         self.counts.capacity() * size_of::<u32>()
             + self.prods.capacity() * size_of::<Float64Type>()

datafusion/core/tests/user_defined/user_defined_aggregates.rs

@@ -877,6 +877,22 @@ impl GroupsAccumulator for TestGroupsAccumulator {
         Ok(())
     }
 
+    fn convert_to_state(
+        &self,
+        values: &[ArrayRef],
+        _opt_filter: Option<&arrow::array::BooleanArray>,
+    ) -> Result<Vec<ArrayRef>> {
+        let len = values.first().map_or(0, |value| value.len());
+        Ok(vec![
+            Arc::new(PrimitiveArray::<UInt64Type>::from_value(self.result, len))
+                as ArrayRef,
+        ])
+    }
+
+    fn supports_convert_to_state(&self) -> bool {
+        true
+    }
+
     fn size(&self) -> usize {
         size_of::<u64>()
     }

datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs

@@ -16,7 +16,8 @@
 // under the License.
 
 use arrow::array::{
-    ArrayRef, AsArray, BooleanArray, Int64Array, ListArray, PrimitiveArray,
+    Array, ArrayRef, AsArray, BooleanArray, Int64Array, ListArray, ListBuilder,
+    PrimitiveArray, PrimitiveBuilder,
 };
 use arrow::buffer::{OffsetBuffer, ScalarBuffer};
 use arrow::datatypes::{ArrowPrimitiveType, Field};
@@ -182,9 +183,127 @@ where
         Ok(())
     }
 
+    fn convert_to_state(
+        &self,
+        values: &[ArrayRef],
+        opt_filter: Option<&BooleanArray>,
+    ) -> datafusion_common::Result<Vec<ArrayRef>> {
+        debug_assert_eq!(values.len(), 1);
+        let arr = values[0].as_primitive::<T>();
+
+        let values_builder = PrimitiveBuilder::<T>::with_capacity(arr.len());
+        let mut builder = ListBuilder::new(values_builder)
+            .with_field(Arc::new(Field::new_list_field(T::DATA_TYPE, true)));
+
+        for row in 0..arr.len() {
+            let included = arr.is_valid(row)
+                && opt_filter
+                    .is_none_or(|filter| filter.is_valid(row) && filter.value(row));
+            if included {
+                builder.values().append_value(arr.value(row));
+            }
+            builder.append(true);
+        }
+
+        Ok(vec![Arc::new(builder.finish())])
+    }
+
+    fn supports_convert_to_state(&self) -> bool {
+        true
+    }
+
     fn size(&self) -> usize {
         size_of::<Self>()
             + self.seen.capacity() * (size_of::<(usize, T::Native)>() + size_of::<u64>())
             + self.counts.capacity() * size_of::<i64>()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::Int32Array;
+    use arrow::datatypes::Int32Type;
+    use datafusion_common::Result;
+
+    #[test]
+    fn convert_to_state_roundtrips_through_merge() -> Result<()> {
+        let values = Arc::new(Int32Array::from(vec![
+            Some(1),
+            Some(2),
+            Some(2),
+            None,
+            Some(3),
+            Some(4),
+            Some(5),
+            Some(5),
+        ])) as ArrayRef;
+        let filter = BooleanArray::from(vec![
+            Some(true),
+            Some(true),
+            Some(true),
+            Some(true),
+            None,
+            Some(true),
+            Some(true),
+            Some(true),
+        ]);
+        let group_indices = vec![0usize, 1, 0, 1, 0, 0, 0, 0];
+
+        let mut direct = PrimitiveDistinctCountGroupsAccumulator::<Int32Type>::new();
+        direct.update_batch(
+            std::slice::from_ref(&values),
+            &group_indices,
+            Some(&filter),
+            2,
+        )?;
+        let direct = direct.evaluate(EmitTo::All)?;
+
+        let converter = PrimitiveDistinctCountGroupsAccumulator::<Int32Type>::new();
+        let state =
+            converter.convert_to_state(std::slice::from_ref(&values), Some(&filter))?;
+        assert_eq!(state[0].null_count(), 0);
+        let mut merged = PrimitiveDistinctCountGroupsAccumulator::<Int32Type>::new();
+        merged.merge_batch(&state, &group_indices, 2)?;
+        let merged = merged.evaluate(EmitTo::All)?;
+
+        assert_eq!(
+            direct.as_any().downcast_ref::<Int64Array>().unwrap(),
+            merged.as_any().downcast_ref::<Int64Array>().unwrap()
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn convert_to_state_preserves_empty_and_filtered_rows() -> Result<()> {
+        let converter = PrimitiveDistinctCountGroupsAccumulator::<Int32Type>::new();
+        let empty_values =
+            Arc::new(Int32Array::from(Vec::<Option<i32>>::new())) as ArrayRef;
+        let state =
+            converter.convert_to_state(std::slice::from_ref(&empty_values), None)?;
+        assert_eq!(state[0].len(), 0);
+        assert_eq!(state[0].null_count(), 0);
+
+        let values = Arc::new(Int32Array::from(vec![Some(1), Some(2), None])) as ArrayRef;
+        let filter = BooleanArray::from(vec![Some(false), None, Some(false)]);
+        let group_indices = vec![0usize, 1, 0];
+
+        let state =
+            converter.convert_to_state(std::slice::from_ref(&values), Some(&filter))?;
+        assert_eq!(state[0].len(), values.len());
+        assert_eq!(state[0].null_count(), 0);
+        let list_state = state[0].as_list::<i32>();
+        for row in 0..list_state.len() {
+            assert_eq!(list_state.value_length(row), 0);
+        }
+
+        let mut merged = PrimitiveDistinctCountGroupsAccumulator::<Int32Type>::new();
+        merged.merge_batch(&state, &group_indices, 2)?;
+        let result = merged.evaluate(EmitTo::All)?;
+        assert_eq!(
+            result.as_any().downcast_ref::<Int64Array>().unwrap(),
+            &Int64Array::from(vec![0, 0])
+        );
+        Ok(())
+    }
+}

datafusion/functions-aggregate/src/approx_distinct.rs

@@ -581,6 +581,43 @@ impl GroupsAccumulator for HllGroupsAccumulator {
         Ok(vec![Arc::new(builder.finish())])
     }
 
+    fn convert_to_state(
+        &self,
+        values: &[ArrayRef],
+        opt_filter: Option<&BooleanArray>,
+    ) -> Result<Vec<ArrayRef>> {
+        assert_eq!(values.len(), 1, "single argument to convert_to_state");
+        let array = values[0].as_ref();
+        let mut hashes = vec![0; array.len()];
+        create_hashes([array], &HLL_HASH_STATE, &mut hashes)?;
+
+        let filter_nulls = opt_filter.map(filter_to_nulls);
+        let value_nulls = array.logical_nulls();
+        let combined_nulls =
+            NullBuffer::union(filter_nulls.as_ref(), value_nulls.as_ref());
+
+        let mut builder = BinaryBuilder::new();
+        let mut scratch = Vec::new();
+        for (row, hash) in hashes.into_iter().enumerate() {
+            if combined_nulls
+                .as_ref()
+                .is_none_or(|nulls| nulls.is_valid(row))
+            {
+                scratch.clear();
+                scratch.extend_from_slice(&hash.to_le_bytes());
+                builder.append_value(&scratch);
+            } else {
+                builder.append_value([]);
+            }
+        }
+
+        Ok(vec![Arc::new(builder.finish())])
+    }
+
+    fn supports_convert_to_state(&self) -> bool {
+        true
+    }
+
     fn size(&self) -> usize {
         self.groups.capacity() * size_of::<GroupHll>()
             + self.allocated_bytes
@@ -1019,6 +1056,98 @@ mod tests {
             assert_eq!(counts.value(0), expected);
         }
 
+        #[test]
+        fn groups_convert_to_state_roundtrips_through_merge() {
+            let values: ArrayRef = Arc::new(Int64Array::from(vec![
+                Some(1),
+                Some(2),
+                Some(2),
+                None,
+                Some(3),
+            ]));
+            let filter = BooleanArray::from(vec![
+                Some(true),
+                Some(true),
+                Some(true),
+                Some(true),
+                None,
+            ]);
+            let group_indices = vec![0usize, 1, 0, 1, 0];
+
+            let mut direct = HllGroupsAccumulator::new();
+            direct
+                .update_batch(
+                    std::slice::from_ref(&values),
+                    &group_indices,
+                    Some(&filter),
+                    2,
+                )
+                .unwrap();
+            let direct = direct
+                .evaluate(EmitTo::All)
+                .unwrap()
+                .as_any()
+                .downcast_ref::<UInt64Array>()
+                .unwrap()
+                .clone();
+
+            let converter = HllGroupsAccumulator::new();
+            let state = converter
+                .convert_to_state(std::slice::from_ref(&values), Some(&filter))
+                .unwrap();
+            assert_eq!(state[0].null_count(), 0);
+            let mut merged = HllGroupsAccumulator::new();
+            merged.merge_batch(&state, &group_indices, 2).unwrap();
+            let merged = merged
+                .evaluate(EmitTo::All)
+                .unwrap()
+                .as_any()
+                .downcast_ref::<UInt64Array>()
+                .unwrap()
+                .clone();
+
+            assert_eq!(direct, merged);
+        }
+
+        #[test]
+        fn groups_convert_to_state_preserves_empty_and_filtered_rows() {
+            let converter = HllGroupsAccumulator::new();
+            let empty_values: ArrayRef =
+                Arc::new(Int64Array::from(Vec::<Option<i64>>::new()));
+            let state = converter
+                .convert_to_state(std::slice::from_ref(&empty_values), None)
+                .unwrap();
+            assert_eq!(state[0].len(), 0);
+            assert_eq!(state[0].null_count(), 0);
+
+            let values: ArrayRef =
+                Arc::new(Int64Array::from(vec![Some(1), Some(2), None]));
+            let filter = BooleanArray::from(vec![Some(false), None, Some(false)]);
+            let group_indices = vec![0usize, 1, 0];
+            let state = converter
+                .convert_to_state(std::slice::from_ref(&values), Some(&filter))
+                .unwrap();
+            assert_eq!(state[0].len(), values.len());
+            assert_eq!(state[0].null_count(), 0);
+            let state = state[0].as_any().downcast_ref::<BinaryArray>().unwrap();
+            for row in 0..state.len() {
+                assert_eq!(state.value(row), b"");
+            }
+
+            let mut merged = HllGroupsAccumulator::new();
+            merged
+                .merge_batch(&[Arc::new(state.clone())], &group_indices, 2)
+                .unwrap();
+            let result = merged
+                .evaluate(EmitTo::All)
+                .unwrap()
+                .as_any()
+                .downcast_ref::<UInt64Array>()
+                .unwrap()
+                .clone();
+            assert_eq!(result, UInt64Array::from(vec![0, 0]));
+        }
+
         /// Regression: a short (≤ 12-byte) Utf8View string must hash identically
         /// in an all-inline batch and in a mixed batch that also contains a long
         /// string (which forces a data buffer).