From d44c2e892f1d56543343bd46ca81ab7540497848 Mon Sep 17 00:00:00 2001 From: linfeng Date: Tue, 30 Jun 2026 21:34:20 +0800 Subject: [PATCH 1/3] Implement state conversion for group accumulators --- .../user_defined/user_defined_aggregates.rs | 16 ++ .../src/aggregate/count_distinct/groups.rs | 121 ++++++++++++++- .../src/approx_distinct.rs | 129 ++++++++++++++++ .../functions-aggregate/src/correlation.rs | 141 ++++++++++++++++++ datafusion/functions-aggregate/src/stddev.rs | 16 +- .../functions-aggregate/src/variance.rs | 123 +++++++++++++++ 6 files changed, 543 insertions(+), 3 deletions(-) diff --git a/datafusion/core/tests/user_defined/user_defined_aggregates.rs b/datafusion/core/tests/user_defined/user_defined_aggregates.rs index b895cb9c7ce2c..1d4b22230147f 100644 --- a/datafusion/core/tests/user_defined/user_defined_aggregates.rs +++ b/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> { + let len = values.first().map_or(0, |value| value.len()); + Ok(vec![ + Arc::new(PrimitiveArray::::from_value(self.result, len)) + as ArrayRef, + ]) + } + + fn supports_convert_to_state(&self) -> bool { + true + } + fn size(&self) -> usize { size_of::() } diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs index 60fe0388c430f..986d4ec0d71ae 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs +++ b/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> { + debug_assert_eq!(values.len(), 1); + let arr = values[0].as_primitive::(); + + let values_builder = PrimitiveBuilder::::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.seen.capacity() * (size_of::<(usize, T::Native)>() + size_of::()) + self.counts.capacity() * size_of::() } } + +#[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::::new(); + direct.update_batch( + std::slice::from_ref(&values), + &group_indices, + Some(&filter), + 2, + )?; + let direct = direct.evaluate(EmitTo::All)?; + + let converter = PrimitiveDistinctCountGroupsAccumulator::::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::::new(); + merged.merge_batch(&state, &group_indices, 2)?; + let merged = merged.evaluate(EmitTo::All)?; + + assert_eq!( + direct.as_any().downcast_ref::().unwrap(), + merged.as_any().downcast_ref::().unwrap() + ); + Ok(()) + } + + #[test] + fn convert_to_state_preserves_empty_and_filtered_rows() -> Result<()> { + let converter = PrimitiveDistinctCountGroupsAccumulator::::new(); + let empty_values = + Arc::new(Int32Array::from(Vec::>::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::(); + for row in 0..list_state.len() { + assert_eq!(list_state.value_length(row), 0); + } + + let mut merged = PrimitiveDistinctCountGroupsAccumulator::::new(); + merged.merge_batch(&state, &group_indices, 2)?; + let result = merged.evaluate(EmitTo::All)?; + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &Int64Array::from(vec![0, 0]) + ); + Ok(()) + } +} diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 1062a478b7bea..77c85705227b7 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/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> { + 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::() + 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::() + .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::() + .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::>::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::().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::() + .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). diff --git a/datafusion/functions-aggregate/src/correlation.rs b/datafusion/functions-aggregate/src/correlation.rs index 7fcf4bb61ffad..2e90cac6d9298 100644 --- a/datafusion/functions-aggregate/src/correlation.rs +++ b/datafusion/functions-aggregate/src/correlation.rs @@ -489,6 +489,61 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator { ]) } + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + assert_eq!(values.len(), 2, "two arguments to convert_to_state"); + let array_x = downcast_array::(&values[0]); + let array_y = downcast_array::(&values[1]); + + let len = array_x.len(); + let mut counts = Vec::with_capacity(len); + let mut sum_x = Vec::with_capacity(len); + let mut sum_y = Vec::with_capacity(len); + let mut sum_xy = Vec::with_capacity(len); + let mut sum_xx = Vec::with_capacity(len); + let mut sum_yy = Vec::with_capacity(len); + + for row in 0..len { + let included = array_x.is_valid(row) + && array_y.is_valid(row) + && opt_filter + .is_none_or(|filter| filter.is_valid(row) && filter.value(row)); + if included { + let x = array_x.value(row); + let y = array_y.value(row); + counts.push(1); + sum_x.push(x); + sum_y.push(y); + sum_xy.push(x * y); + sum_xx.push(x * x); + sum_yy.push(y * y); + } else { + counts.push(0); + sum_x.push(0.0); + sum_y.push(0.0); + sum_xy.push(0.0); + sum_xx.push(0.0); + sum_yy.push(0.0); + } + } + + Ok(vec![ + Arc::new(UInt64Array::from(counts)), + Arc::new(Float64Array::from(sum_x)), + Arc::new(Float64Array::from(sum_y)), + Arc::new(Float64Array::from(sum_xy)), + Arc::new(Float64Array::from(sum_xx)), + Arc::new(Float64Array::from(sum_yy)), + ]) + } + + fn supports_convert_to_state(&self) -> bool { + true + } + fn merge_batch( &mut self, values: &[ArrayRef], @@ -593,4 +648,90 @@ mod tests { }); assert!(result.is_err()); } + + #[test] + fn convert_to_state_roundtrips_through_merge() -> Result<()> { + let x = Arc::new(Float64Array::from(vec![ + Some(1.0), + Some(2.0), + None, + Some(4.0), + Some(8.0), + Some(16.0), + Some(32.0), + ])) as ArrayRef; + let y = Arc::new(Float64Array::from(vec![ + Some(2.0), + Some(4.0), + Some(6.0), + None, + Some(16.0), + Some(32.0), + Some(64.0), + ])) as ArrayRef; + let filter = BooleanArray::from(vec![ + Some(true), + Some(false), + Some(true), + Some(true), + None, + Some(true), + Some(true), + ]); + let values = vec![x, y]; + let group_indices = vec![0, 1, 0, 1, 0, 0, 0]; + + let mut direct = CorrelationGroupsAccumulator::new(); + direct.update_batch(&values, &group_indices, Some(&filter), 2)?; + let direct = direct.evaluate(EmitTo::All)?; + + let converter = CorrelationGroupsAccumulator::new(); + let state = converter.convert_to_state(&values, Some(&filter))?; + let mut merged = CorrelationGroupsAccumulator::new(); + merged.merge_batch(&state, &group_indices, 2)?; + let merged = merged.evaluate(EmitTo::All)?; + + assert_eq!( + direct.as_any().downcast_ref::().unwrap(), + merged.as_any().downcast_ref::().unwrap() + ); + Ok(()) + } + + #[test] + fn convert_to_state_preserves_empty_and_filtered_rows() -> Result<()> { + let converter = CorrelationGroupsAccumulator::new(); + let empty_values = vec![ + Arc::new(Float64Array::from(Vec::>::new())) as ArrayRef, + Arc::new(Float64Array::from(Vec::>::new())) as ArrayRef, + ]; + let state = converter.convert_to_state(&empty_values, None)?; + for state_array in &state { + assert_eq!(state_array.len(), 0); + assert_eq!(state_array.null_count(), 0); + } + + let values = vec![ + Arc::new(Float64Array::from(vec![Some(1.0), Some(2.0), None])) as ArrayRef, + Arc::new(Float64Array::from(vec![Some(2.0), None, Some(4.0)])) as ArrayRef, + ]; + let filter = BooleanArray::from(vec![Some(false), None, Some(false)]); + let group_indices = vec![0, 1, 0]; + let state = converter.convert_to_state(&values, Some(&filter))?; + for state_array in &state { + assert_eq!(state_array.len(), values[0].len()); + assert_eq!(state_array.null_count(), 0); + } + + let counts = state[0].as_any().downcast_ref::().unwrap(); + assert_eq!(counts, &UInt64Array::from(vec![0, 0, 0])); + + let mut merged = CorrelationGroupsAccumulator::new(); + merged.merge_batch(&state, &group_indices, 2)?; + let result = merged.evaluate(EmitTo::All)?; + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result.len(), 2); + assert_eq!(result.null_count(), 2); + Ok(()) + } } diff --git a/datafusion/functions-aggregate/src/stddev.rs b/datafusion/functions-aggregate/src/stddev.rs index f0482b23d12a7..a31517b93e003 100644 --- a/datafusion/functions-aggregate/src/stddev.rs +++ b/datafusion/functions-aggregate/src/stddev.rs @@ -22,7 +22,7 @@ use std::hash::Hash; use std::mem::align_of_val; use std::sync::Arc; -use arrow::array::Float64Array; +use arrow::array::{BooleanArray, Float64Array}; use arrow::datatypes::FieldRef; use arrow::{array::ArrayRef, datatypes::DataType, datatypes::Field}; use datafusion_common::ScalarValue; @@ -318,7 +318,7 @@ impl GroupsAccumulator for StddevGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&arrow::array::BooleanArray>, + opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { self.variance @@ -345,6 +345,18 @@ impl GroupsAccumulator for StddevGroupsAccumulator { self.variance.state(emit_to) } + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + self.variance.convert_to_state(values, opt_filter) + } + + fn supports_convert_to_state(&self) -> bool { + true + } + fn size(&self) -> usize { self.variance.size() } diff --git a/datafusion/functions-aggregate/src/variance.rs b/datafusion/functions-aggregate/src/variance.rs index 551fcfe120352..0278ce2c233e4 100644 --- a/datafusion/functions-aggregate/src/variance.rs +++ b/datafusion/functions-aggregate/src/variance.rs @@ -580,6 +580,44 @@ impl GroupsAccumulator for VarianceGroupsAccumulator { ]) } + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + assert_eq!(values.len(), 1, "single argument to convert_to_state"); + let values = as_float64_array(&values[0])?; + + let len = values.len(); + let mut counts = Vec::with_capacity(len); + let mut means = Vec::with_capacity(len); + let mut m2s = Vec::with_capacity(len); + + for row in 0..len { + if values.is_valid(row) + && opt_filter + .is_none_or(|filter| filter.is_valid(row) && filter.value(row)) + { + counts.push(1); + means.push(values.value(row)); + } else { + counts.push(0); + means.push(0.0); + } + m2s.push(0.0); + } + + Ok(vec![ + Arc::new(UInt64Array::new(counts.into(), None)), + Arc::new(Float64Array::new(means.into(), None)), + Arc::new(Float64Array::new(m2s.into(), None)), + ]) + } + + fn supports_convert_to_state(&self) -> bool { + true + } + fn size(&self) -> usize { self.m2s.capacity() * size_of::() + self.means.capacity() * size_of::() @@ -679,4 +717,89 @@ mod tests { assert_eq!(result.value(0), 1.0); Ok(()) } + + #[test] + fn convert_to_state_roundtrips_through_merge() -> Result<()> { + let values = Arc::new(Float64Array::from(vec![ + Some(1.0), + Some(2.0), + None, + Some(4.0), + Some(8.0), + Some(16.0), + Some(32.0), + ])) as ArrayRef; + let filter = BooleanArray::from(vec![ + Some(true), + Some(false), + Some(true), + None, + Some(true), + Some(true), + Some(true), + ]); + let group_indices = vec![0, 1, 0, 1, 0, 0, 0]; + + let mut direct = VarianceGroupsAccumulator::new(StatsType::Sample); + direct.update_batch( + std::slice::from_ref(&values), + &group_indices, + Some(&filter), + 2, + )?; + let direct = direct.evaluate(EmitTo::All)?; + + let converter = VarianceGroupsAccumulator::new(StatsType::Sample); + let state = + converter.convert_to_state(std::slice::from_ref(&values), Some(&filter))?; + let mut merged = VarianceGroupsAccumulator::new(StatsType::Sample); + merged.merge_batch(&state, &group_indices, 2)?; + let merged = merged.evaluate(EmitTo::All)?; + + let direct = direct.as_any().downcast_ref::().unwrap(); + let merged = merged.as_any().downcast_ref::().unwrap(); + assert_eq!(direct.len(), merged.len()); + for row in 0..direct.len() { + assert_eq!(direct.is_null(row), merged.is_null(row)); + if direct.is_valid(row) { + assert!((direct.value(row) - merged.value(row)).abs() < 1e-12); + } + } + Ok(()) + } + + #[test] + fn convert_to_state_preserves_empty_and_filtered_rows() -> Result<()> { + let converter = VarianceGroupsAccumulator::new(StatsType::Sample); + let empty_values = + Arc::new(Float64Array::from(Vec::>::new())) as ArrayRef; + let state = + converter.convert_to_state(std::slice::from_ref(&empty_values), None)?; + for state_array in &state { + assert_eq!(state_array.len(), 0); + assert_eq!(state_array.null_count(), 0); + } + + let values = + Arc::new(Float64Array::from(vec![Some(1.0), Some(2.0), None])) as ArrayRef; + let filter = BooleanArray::from(vec![Some(false), None, Some(false)]); + let group_indices = vec![0, 1, 0]; + let state = + converter.convert_to_state(std::slice::from_ref(&values), Some(&filter))?; + for state_array in &state { + assert_eq!(state_array.len(), values.len()); + assert_eq!(state_array.null_count(), 0); + } + + let counts = state[0].as_any().downcast_ref::().unwrap(); + assert_eq!(counts, &UInt64Array::from(vec![0, 0, 0])); + + let mut merged = VarianceGroupsAccumulator::new(StatsType::Sample); + merged.merge_batch(&state, &group_indices, 2)?; + let result = merged.evaluate(EmitTo::All)?; + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result.len(), 2); + assert_eq!(result.null_count(), 2); + Ok(()) + } } From 35585f8ea0c1af92d21a86948ebb3a31dd14e30b Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Wed, 1 Jul 2026 14:21:45 +0800 Subject: [PATCH 2/3] add convert_to_state to advanced UDAF example --- .../examples/udf/advanced_udaf.rs | 43 ++++++++++++++++++- 1 file changed, 41 insertions(+), 2 deletions(-) diff --git a/datafusion-examples/examples/udf/advanced_udaf.rs b/datafusion-examples/examples/udf/advanced_udaf.rs index b990740159906..096753d2b5d7b 100644 --- a/datafusion-examples/examples/udf/advanced_udaf.rs +++ b/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> { + assert_eq!(values.len(), 1, "single argument to convert_to_state"); + + let prods = values[0] + .as_primitive::() + .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::::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::() + self.prods.capacity() * size_of::() From 2c28d6c497fd2500cc553c6ea0dc58ade20f995a Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Fri, 10 Jul 2026 10:01:28 +0800 Subject: [PATCH 3/3] test: add skip partial SLT coverage for state conversion --- .../test_files/aggregate_skip_partial.slt | 40 ++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/test_files/aggregate_skip_partial.slt b/datafusion/sqllogictest/test_files/aggregate_skip_partial.slt index a10417f232409..20c23a44924d7 100644 --- a/datafusion/sqllogictest/test_files/aggregate_skip_partial.slt +++ b/datafusion/sqllogictest/test_files/aggregate_skip_partial.slt @@ -220,7 +220,7 @@ e true false NULL statement ok reset datafusion.execution.batch_size; -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; @@ -636,6 +636,44 @@ FROM aggregate_test_100_null GROUP BY c2 ORDER BY c2; 4 3 5 6 +# Test variance and stddev with nullable fields and filters +query IRRR +SELECT c2, + var_samp(c11) FILTER (WHERE c3 > 0), + stddev_samp(c11) FILTER (WHERE c3 > 0), + stddev_pop(c11) FILTER (WHERE c3 > 0) +FROM aggregate_test_100_null GROUP BY c2 ORDER BY c2; +---- +1 0.085786074994 0.29289259976 0.276141791266 +2 0.070769227104 0.266024861816 0.24884347232 +3 0.087515365779 0.295829960922 0.270054571305 +4 0.038697229167 0.196716113134 0.182123731489 +5 0.081817232141 0.286037116719 0.23354832782 + +# Test corr with nullable fields and filters +query IR +SELECT c2, + corr(c3, c11) FILTER (WHERE c5 > 0) +FROM aggregate_test_100_null GROUP BY c2 ORDER BY c2; +---- +1 -0.38515658251 +2 0.414489249329 +3 -0.796447131429 +4 -0.938568248748 +5 -0.051058146743 + +# Test count distinct with nullable fields and filters +query II +SELECT c2, + count(distinct c3) FILTER (WHERE c11 > 0.5) +FROM aggregate_test_100_null GROUP BY c2 ORDER BY c2; +---- +1 10 +2 6 +3 3 +4 3 +5 6 + # Test median with nullable fields and filter query IRR SELECT c2,