diff --git a/datafusion-examples/examples/udf/advanced_udaf.rs b/datafusion-examples/examples/udf/advanced_udaf.rs index f1651dbf28913..b990740159906 100644 --- a/datafusion-examples/examples/udf/advanced_udaf.rs +++ b/datafusion-examples/examples/udf/advanced_udaf.rs @@ -268,7 +268,6 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&arrow::array::BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 2, "two arguments to merge_batch"); @@ -280,7 +279,7 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator { self.null_state.accumulate( group_indices, partial_counts, - opt_filter, + None, total_num_groups, |group_index, partial_count| { self.counts[group_index] += partial_count; @@ -292,7 +291,7 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator { self.null_state.accumulate( group_indices, partial_prods, - opt_filter, + None, total_num_groups, |group_index, new_value: ::Native| { let prod = &mut self.prods[group_index]; diff --git a/datafusion/core/tests/user_defined/user_defined_aggregates.rs b/datafusion/core/tests/user_defined/user_defined_aggregates.rs index 7d22c5df70dfc..b895cb9c7ce2c 100644 --- a/datafusion/core/tests/user_defined/user_defined_aggregates.rs +++ b/datafusion/core/tests/user_defined/user_defined_aggregates.rs @@ -872,7 +872,6 @@ impl GroupsAccumulator for TestGroupsAccumulator { &mut self, _values: &[ArrayRef], _group_indices: &[usize], - _opt_filter: Option<&arrow::array::BooleanArray>, _total_num_groups: usize, ) -> Result<()> { Ok(()) diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs index da5da384c7b4e..b021674cbec2c 100644 --- a/datafusion/expr-common/src/groups_accumulator.rs +++ b/datafusion/expr-common/src/groups_accumulator.rs @@ -183,12 +183,14 @@ pub trait GroupsAccumulator: Send + std::any::Any { /// /// * `values`: arrays produced from previously calling `state` on other accumulators. /// - /// Other arguments are the same as for [`Self::update_batch`]. + /// Other arguments are the same as for [`Self::update_batch`], except that + /// there is no `opt_filter` — aggregate filters are applied during the + /// partial (update) phase, so by the time intermediate states are merged + /// no per-row filtering is needed. fn merge_batch( &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()>; diff --git a/datafusion/ffi/src/udaf/groups_accumulator.rs b/datafusion/ffi/src/udaf/groups_accumulator.rs index 1600bef39da45..272afdb6abfb1 100644 --- a/datafusion/ffi/src/udaf/groups_accumulator.rs +++ b/datafusion/ffi/src/udaf/groups_accumulator.rs @@ -64,7 +64,6 @@ pub struct FFI_GroupsAccumulator { accumulator: &mut Self, values: SVec, group_indices: SVec, - opt_filter: FFI_Option, total_num_groups: usize, ) -> FFI_Result<()>, @@ -195,21 +194,14 @@ unsafe extern "C" fn merge_batch_fn_wrapper( accumulator: &mut FFI_GroupsAccumulator, values: SVec, group_indices: SVec, - opt_filter: FFI_Option, total_num_groups: usize, ) -> FFI_Result<()> { unsafe { let accumulator = accumulator.inner_mut(); let values = sresult_return!(process_values(values)); let group_indices: Vec = group_indices.into_iter().collect(); - let opt_filter = sresult_return!(process_opt_filter(opt_filter)); - sresult!(accumulator.merge_batch( - &values, - &group_indices, - opt_filter.as_ref(), - total_num_groups - )) + sresult!(accumulator.merge_batch(&values, &group_indices, total_num_groups)) } } @@ -379,7 +371,6 @@ impl GroupsAccumulator for ForeignGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { unsafe { @@ -388,20 +379,11 @@ impl GroupsAccumulator for ForeignGroupsAccumulator { .map(WrappedArray::try_from) .collect::, ArrowError>>()?; let group_indices = group_indices.iter().cloned().collect(); - let opt_filter = opt_filter - .map(|bool_array| to_ffi(&bool_array.to_data())) - .transpose()? - .map(|(array, schema)| WrappedArray { - array, - schema: WrappedSchema(schema), - }) - .into(); df_result!((self.accumulator.merge_batch)( &mut self.accumulator, values.into_iter().collect(), group_indices, - opt_filter, total_num_groups )) } @@ -517,8 +499,7 @@ mod tests { let second_states = vec![make_array(create_array!(Boolean, vec![false]).to_data())]; - let opt_filter = create_array!(Boolean, vec![true]); - foreign_accum.merge_batch(&second_states, &[0], Some(opt_filter.as_ref()), 1)?; + foreign_accum.merge_batch(&second_states, &[0], 1)?; let groups_bool = foreign_accum.evaluate(EmitTo::All)?; assert_eq!(groups_bool.len(), 1); assert_eq!( 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 d370d59c90012..60fe0388c430f 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs @@ -160,7 +160,6 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> datafusion_common::Result<()> { debug_assert_eq!(values.len(), 1); diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index ad2a21bb4733c..b412b4ffe09f2 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -375,13 +375,12 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { self.invoke_per_accumulator( values, group_indices, - opt_filter, + None, total_num_groups, |accumulator, values_to_accumulate| { accumulator.merge_batch(values_to_accumulate)?; diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs index d1d8924a2c3e8..afb1dec24a484 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs @@ -132,11 +132,10 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // update / merge are the same - self.update_batch(values, group_indices, opt_filter, total_num_groups) + self.update_batch(values, group_indices, None, total_num_groups) } fn size(&self) -> usize { diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs index a81b89e1e46f1..474899d8f3c6a 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs @@ -131,11 +131,10 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // update / merge are the same - self.update_batch(values, group_indices, opt_filter, total_num_groups) + self.update_batch(values, group_indices, None, total_num_groups) } /// Converts an input batch directly to a state batch diff --git a/datafusion/functions-aggregate/benches/first_last.rs b/datafusion/functions-aggregate/benches/first_last.rs index 1d18e1c7dcd44..8f28e126a4009 100644 --- a/datafusion/functions-aggregate/benches/first_last.rs +++ b/datafusion/functions-aggregate/benches/first_last.rs @@ -235,7 +235,6 @@ fn merge_bench( Arc::clone(&is_set), ], &group_indices, - opt_filter, num_groups, ) .unwrap(), diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 38b902964f546..0c5a438454092 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -705,15 +705,8 @@ impl GroupsAccumulator for HllGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - // Since aggregate filter should be applied in partial stage, in final stage there should be no filter - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { - assert!( - opt_filter.is_none(), - "aggregate filter should be applied in partial stage, there should be no filter in final stage" - ); - self.ensure_groups(total_num_groups); let states = downcast_value!(values[0], BinaryArray); let mut delta: isize = 0; diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 8ed3fbf8c3d26..1dd111f9182c9 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -732,7 +732,6 @@ impl GroupsAccumulator for ArrayAggGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 1, "one argument to merge_batch"); @@ -2019,7 +2018,7 @@ mod tests { // Merge acc2's state into acc1 let state = acc2.state(EmitTo::All)?; - acc1.merge_batch(&state, &[0, 1], None, 2)?; + acc1.merge_batch(&state, &[0, 1], 2)?; // Another update_batch on acc1 after the merge let values: ArrayRef = Arc::new(Int32Array::from(vec![5, 6])); @@ -2088,7 +2087,7 @@ mod tests { // Feed state into a new accumulator via merge_batch let mut acc2 = ArrayAggGroupsAccumulator::new(DataType::Int32, false); - acc2.merge_batch(&state, &[0, 0, 1], None, 2)?; + acc2.merge_batch(&state, &[0, 0, 1], 2)?; // Group 0 received rows 0 ([1]) and 1 ([NULL]) → [1, NULL] let vals = eval_i32_lists(&mut acc2, EmitTo::All)?; @@ -2118,7 +2117,7 @@ mod tests { // Feed state into a new accumulator via merge_batch let mut acc2 = ArrayAggGroupsAccumulator::new(DataType::Int32, true); - acc2.merge_batch(&state, &[0, 0, 1, 1], None, 2)?; + acc2.merge_batch(&state, &[0, 0, 1, 1], 2)?; // Group 0: received [1] and null (skipped) → [1] let vals = eval_i32_lists(&mut acc2, EmitTo::All)?; diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs index ddeb9b0870a16..06c76946343dc 100644 --- a/datafusion/functions-aggregate/src/average.rs +++ b/datafusion/functions-aggregate/src/average.rs @@ -908,7 +908,6 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 2, "two arguments to merge_batch"); @@ -920,7 +919,7 @@ where self.null_state.accumulate( group_indices, partial_counts, - opt_filter, + None, total_num_groups, |group_index, partial_count| { // SAFETY: group_index is guaranteed to be in bounds @@ -934,7 +933,7 @@ where self.null_state.accumulate( group_indices, partial_sums, - opt_filter, + None, total_num_groups, |group_index, new_value: ::Native| { // SAFETY: group_index is guaranteed to be in bounds diff --git a/datafusion/functions-aggregate/src/correlation.rs b/datafusion/functions-aggregate/src/correlation.rs index 5a95cfe8320fc..7fcf4bb61ffad 100644 --- a/datafusion/functions-aggregate/src/correlation.rs +++ b/datafusion/functions-aggregate/src/correlation.rs @@ -493,7 +493,6 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // Resize vectors to accommodate total number of groups @@ -512,11 +511,6 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator { let partial_sum_xx = values[4].as_primitive::(); let partial_sum_yy = values[5].as_primitive::(); - assert!( - opt_filter.is_none(), - "aggregate filter should be applied in partial stage, there should be no filter in final stage" - ); - accumulate_correlation_states( group_indices, ( diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index eab36d4951a9c..f0ce8c82a1bb2 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -665,8 +665,6 @@ impl GroupsAccumulator for CountGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - // Since aggregate filter should be applied in partial stage, in final stage there should be no filter - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 1, "one argument to merge_batch"); diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index 1935f29c4cfe8..cecb277cb844a 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -671,7 +671,6 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator, total_num_groups: usize, ) -> Result<()> { self.resize_states(total_num_groups); @@ -690,7 +689,7 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator().unwrap(); + // group 0 keeps merged value=1 (ordering=1). + // group 1 keeps merged value=-6 (ordering=-6 < 6, so -6 is "first"). + // group 2 had no merged value (is_set=false), so update_batch value=6 wins. let expect: PrimitiveArray = - Int64Array::from(vec![Some(1), Some(6), Some(6), None]); + Int64Array::from(vec![Some(1), Some(-6), Some(6), None]); assert_eq!(eval_result, &expect); @@ -1680,7 +1677,7 @@ mod tests { group_acc.compute_size_of_orderings() ); - group_acc.merge_batch(&s, &Vec::from_iter(0..s[0].len()), None, 100)?; + group_acc.merge_batch(&s, &Vec::from_iter(0..s[0].len()), 100)?; assert_eq!( group_acc.size_of_orderings, group_acc.compute_size_of_orderings() @@ -1753,12 +1750,7 @@ mod tests { ]; assert_eq!(state, expected_state); - group_acc.merge_batch( - &state, - &[0, 1, 2], - Some(&BooleanArray::from(vec![true, false, false])), - 3, - )?; + group_acc.merge_batch(&state, &[0, 1, 2], 3)?; val_with_orderings.clear(); val_with_orderings.push(Arc::new(Int64Array::from(vec![66, 6]))); @@ -1769,6 +1761,10 @@ mod tests { let binding = group_acc.evaluate(EmitTo::All)?; let eval_result = binding.as_any().downcast_ref::().unwrap(); + // group 0: merged value=1 (ordering=1, is_set=true), update not called. + // group 1: merged value=-6 (ordering=-6, is_set=true); update ordering=66 > -6 + // → LAST_VALUE keeps the higher ordering, so group 1 becomes 66. + // group 2: is_set=false after merge; update_batch sets it to 6. let expect: PrimitiveArray = Int64Array::from(vec![Some(1), Some(66), Some(6), None]); diff --git a/datafusion/functions-aggregate/src/median.rs b/datafusion/functions-aggregate/src/median.rs index e7e7d03937f12..4a0da10f51845 100644 --- a/datafusion/functions-aggregate/src/median.rs +++ b/datafusion/functions-aggregate/src/median.rs @@ -393,8 +393,6 @@ impl GroupsAccumulator for MedianGroupsAccumulator, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 1, "one argument to merge_batch"); diff --git a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs index b56c2106e32b5..7a3c605d82e4d 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs @@ -309,11 +309,10 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // min/max are their own states (no transition needed) - self.update_batch(values, group_indices, opt_filter, total_num_groups) + self.update_batch(values, group_indices, None, total_num_groups) } fn convert_to_state( diff --git a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs index 7c94e7f5738be..10580ac18d3ec 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs @@ -134,11 +134,10 @@ impl GroupsAccumulator for MinMaxStructAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // min/max are their own states (no transition needed) - self.update_batch(values, group_indices, opt_filter, total_num_groups) + self.update_batch(values, group_indices, None, total_num_groups) } fn convert_to_state( diff --git a/datafusion/functions-aggregate/src/percentile_cont.rs b/datafusion/functions-aggregate/src/percentile_cont.rs index 714988bde2acf..e8e6fd127e65d 100644 --- a/datafusion/functions-aggregate/src/percentile_cont.rs +++ b/datafusion/functions-aggregate/src/percentile_cont.rs @@ -537,8 +537,6 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - // Since aggregate filter should be applied in partial stage, in final stage there should be no filter - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 1, "one argument to merge_batch"); diff --git a/datafusion/functions-aggregate/src/stddev.rs b/datafusion/functions-aggregate/src/stddev.rs index 68e38a3b8db07..f0482b23d12a7 100644 --- a/datafusion/functions-aggregate/src/stddev.rs +++ b/datafusion/functions-aggregate/src/stddev.rs @@ -329,11 +329,10 @@ impl GroupsAccumulator for StddevGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&arrow::array::BooleanArray>, total_num_groups: usize, ) -> Result<()> { self.variance - .merge_batch(values, group_indices, opt_filter, total_num_groups) + .merge_batch(values, group_indices, total_num_groups) } fn evaluate(&mut self, emit_to: datafusion_expr::EmitTo) -> Result { diff --git a/datafusion/functions-aggregate/src/string_agg.rs b/datafusion/functions-aggregate/src/string_agg.rs index f0757818afb93..6b0665f479d78 100644 --- a/datafusion/functions-aggregate/src/string_agg.rs +++ b/datafusion/functions-aggregate/src/string_agg.rs @@ -413,11 +413,10 @@ impl GroupsAccumulator for StringAggGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // State is always LargeUtf8, which update_batch already handles. - self.update_batch(values, group_indices, opt_filter, total_num_groups) + self.update_batch(values, group_indices, None, total_num_groups) } fn convert_to_state( @@ -898,7 +897,7 @@ mod tests { // Simulate a second accumulator's state (LargeUtf8 partial strings) let partial_state: ArrayRef = Arc::new(LargeStringArray::from(vec!["c,d", "e"])); - acc.merge_batch(&[partial_state], &[0, 1], None, 2)?; + acc.merge_batch(&[partial_state], &[0, 1], 2)?; let result = evaluate_groups(&mut acc, EmitTo::All); assert_eq!( diff --git a/datafusion/functions-aggregate/src/variance.rs b/datafusion/functions-aggregate/src/variance.rs index d5fddf01f2d52..551fcfe120352 100644 --- a/datafusion/functions-aggregate/src/variance.rs +++ b/datafusion/functions-aggregate/src/variance.rs @@ -528,8 +528,6 @@ impl GroupsAccumulator for VarianceGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - // Since aggregate filter should be applied in partial stage, in final stage there should be no filter - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 3, "two arguments to merge_batch"); @@ -673,8 +671,8 @@ mod tests { Arc::new(Float64Array::from(vec![1.0])), ]; let mut acc = VarianceGroupsAccumulator::new(StatsType::Sample); - acc.merge_batch(&state_1, &[0], None, 1)?; - acc.merge_batch(&state_2, &[0], None, 1)?; + acc.merge_batch(&state_1, &[0], 1)?; + acc.merge_batch(&state_2, &[0], 1)?; let result = acc.evaluate(EmitTo::All)?; let result = result.as_any().downcast_ref::().unwrap(); assert_eq!(result.len(), 1); diff --git a/datafusion/physical-plan/src/aggregates/hash_aggregate.rs b/datafusion/physical-plan/src/aggregates/hash_aggregate.rs index 29d292b215d16..59ee09912f621 100644 --- a/datafusion/physical-plan/src/aggregates/hash_aggregate.rs +++ b/datafusion/physical-plan/src/aggregates/hash_aggregate.rs @@ -993,3 +993,274 @@ impl RecordBatchStream for FinalHashAggregateStream { Arc::clone(&self.schema) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::aggregates::{AggregateMode, PhysicalGroupBy}; + use crate::execution_plan::ExecutionPlan; + use crate::test::TestMemoryExec; + + use arrow::array::{Int32Array, Int64Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::Result; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + use futures::StreamExt; + + #[tokio::test] + async fn test_partial_hash_stream_double_emission_race_condition_bug() -> Result<()> { + // Fix for https://github.com/apache/datafusion/issues/18701 + // This test specifically proves that we have fixed double emission race condition + // where emit_early_if_necessary() and switch_to_skip_aggregation() + // both emit in the same loop iteration, causing data loss + + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + + // Create data that will trigger BOTH conditions in the same iteration: + // 1. More groups than batch_size (triggers early emission when memory pressure hits) + // 2. High cardinality ratio (triggers skip aggregation) + let batch_size = 1024; // We'll set this in session config + let num_groups = batch_size + 100; // Slightly more than batch_size (1124 groups) + + // Create exactly 1 row per group = 100% cardinality ratio + let group_ids: Vec = (0..num_groups as i32).collect(); + let values: Vec = vec![1; num_groups]; + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids)), + Arc::new(Int64Array::from(values)), + ], + )?; + let input_partitions = vec![vec![batch]]; + + // Create constrained memory to trigger early emission but not completely fail + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(1024, 1.0) // small enough to start but will trigger pressure + .build_arc()?; + + let mut task_ctx = TaskContext::default().with_runtime(runtime); + + // Configure to trigger BOTH conditions: + // 1. Low probe threshold (triggers skip probe after few rows) + // 2. Low ratio threshold (triggers skip aggregation immediately) + // 3. Set batch_size to 1024 so our 1124 groups will trigger early emission + // This creates the race condition where both emit paths are triggered + let mut session_config = task_ctx.session_config().clone(); + session_config = session_config.set( + "datafusion.execution.batch_size", + &datafusion_common::ScalarValue::UInt64(Some(1024)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &datafusion_common::ScalarValue::UInt64(Some(50)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &datafusion_common::ScalarValue::Float64(Some(0.8)), + ); + task_ctx = task_ctx.with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // Create aggregate: COUNT(*) GROUP BY group_col + let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())]; + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )]; + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + + // Use Partial mode where the race condition occurs + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(group_expr), + aggr_expr, + vec![None], + exec, + Arc::clone(&schema), + )?; + + // Execute and collect results + let mut stream = + PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + let mut results = Vec::new(); + + while let Some(result) = stream.next().await { + let batch = result?; + results.push(batch); + } + + // Count total groups emitted + let mut total_output_groups = 0; + for batch in &results { + total_output_groups += batch.num_rows(); + } + + assert_eq!( + total_output_groups, num_groups, + "Unexpected number of groups", + ); + + Ok(()) + } + + #[tokio::test] + async fn test_partial_hash_stream_skip_aggregation_probe_not_locked_until_skip() + -> Result<()> { + // Test that the probe is not locked until we actually decide to skip. + // This allows us to continue evaluating the skip condition across multiple batches. + // + // Scenario: + // - Batch 1: Hits rows threshold but NOT ratio threshold (low cardinality) -> don't skip + // - Batch 2: Now hits ratio threshold (high cardinality) -> skip + // + // Without the fix, the probe would be locked after batch 1, preventing the skip + // decision from being made on batch 2. + + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int32, false), + ])); + + // Configure thresholds: + // - probe_rows_threshold: 100 rows + // - probe_ratio_threshold: 0.8 (80%) + let probe_rows_threshold = 100; + let probe_ratio_threshold = 0.8; + + // Batch 1: 100 rows with only 10 unique groups + // Ratio: 10/100 = 0.1 (10%) < 0.8 -> should NOT skip + // This will hit the rows threshold but not the ratio threshold + let batch1_rows = 100; + let batch1_groups = 10; + let mut group_ids_batch1 = Vec::new(); + for i in 0..batch1_rows { + group_ids_batch1.push((i % batch1_groups) as i32); + } + let values_batch1: Vec = vec![1; batch1_rows]; + + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch1)), + Arc::new(Int32Array::from(values_batch1)), + ], + )?; + + // Batch 2: 360 rows with 360 unique NEW groups (starting from group 10) + // After batch 2, total: 460 rows, 370 groups + // Ratio: 370/460 is about 0.804 (80.4%) > 0.8 -> SHOULD decide to skip + let batch2_rows = 360; + let batch2_groups = 360; + let group_ids_batch2: Vec = (batch1_groups..(batch1_groups + batch2_groups)) + .map(|x| x as i32) + .collect(); + let values_batch2: Vec = vec![1; batch2_rows]; + + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch2)), + Arc::new(Int32Array::from(values_batch2)), + ], + )?; + + // Batch 3: This batch should be skipped since we decided to skip after batch 2 + // 100 rows with 100 unique groups (continuing from where batch 2 left off) + let batch3_rows = 100; + let batch3_groups = 100; + let batch3_start_group = batch1_groups + batch2_groups; + let group_ids_batch3: Vec = (batch3_start_group + ..(batch3_start_group + batch3_groups)) + .map(|x| x as i32) + .collect(); + let values_batch3: Vec = vec![1; batch3_rows]; + + let batch3 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch3)), + Arc::new(Int32Array::from(values_batch3)), + ], + )?; + + let input_partitions = vec![vec![batch1, batch2, batch3]]; + + let runtime = RuntimeEnvBuilder::default().build_arc()?; + let mut task_ctx = TaskContext::default().with_runtime(runtime); + + // Configure skip aggregation settings + let mut session_config = task_ctx.session_config().clone(); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &datafusion_common::ScalarValue::UInt64(Some(probe_rows_threshold)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &datafusion_common::ScalarValue::Float64(Some(probe_ratio_threshold)), + ); + task_ctx = task_ctx.with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // Create aggregate: COUNT(*) GROUP BY group_col + let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())]; + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )]; + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + + // Use Partial mode + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(group_expr), + aggr_expr, + vec![None], + exec, + Arc::clone(&schema), + )?; + + // Execute and collect results + let mut stream = + PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + let mut results = Vec::new(); + + while let Some(result) = stream.next().await { + let batch = result?; + results.push(batch); + } + + // Check that skip aggregation actually happened. + // The key metric is skipped_aggregation_rows. + let metrics = aggregate_exec.metrics().unwrap(); + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|m| m.as_usize()) + .unwrap_or(0); + + // We expect batch 3's rows to be skipped (100 rows) + assert_eq!( + skipped_rows, batch3_rows, + "Expected batch 3's rows ({batch3_rows}) to be skipped", + ); + + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/hash_table.rs b/datafusion/physical-plan/src/aggregates/hash_table.rs index 2e5702f750546..e6b2fa22c137f 100644 --- a/datafusion/physical-plan/src/aggregates/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/hash_table.rs @@ -221,12 +221,8 @@ impl HashAggregateAccumulator { total_num_groups: usize, ) -> Result<()> { debug_assert!(values.filter.is_none()); - self.accumulator.merge_batch( - &values.arguments, - group_indices, - None, - total_num_groups, - ) + self.accumulator + .merge_batch(&values.arguments, group_indices, total_num_groups) } fn evaluate_final(&mut self, emit_to: EmitTo) -> Result { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 7d382c231d386..54e44aa86d66c 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3989,6 +3989,8 @@ mod tests { Ok(()) } + // Migrated to PartialHashAggregateStream coverage below; + // kept here for the legacy GroupedHashAggregateStream implementation. #[tokio::test] async fn test_skip_aggregation_after_first_batch() -> Result<()> { let schema = Arc::new(Schema::new(vec![ @@ -4071,6 +4073,8 @@ mod tests { Ok(()) } + // Migrated to PartialHashAggregateStream coverage below; + // kept here for the legacy GroupedHashAggregateStream implementation. #[tokio::test] async fn test_skip_aggregation_after_threshold() -> Result<()> { let schema = Arc::new(Schema::new(vec![ diff --git a/datafusion/physical-plan/src/aggregates/row_hash.rs b/datafusion/physical-plan/src/aggregates/row_hash.rs index c501fe662b76d..d46faf9acc14a 100644 --- a/datafusion/physical-plan/src/aggregates/row_hash.rs +++ b/datafusion/physical-plan/src/aggregates/row_hash.rs @@ -923,7 +923,7 @@ impl GroupedHashAggregateStream { // if aggregation is over intermediate states, // use merge - acc.merge_batch(values, group_indices, None, total_num_groups)?; + acc.merge_batch(values, group_indices, total_num_groups)?; } self.group_by_metrics .aggregation_time @@ -1111,17 +1111,31 @@ impl GroupedHashAggregateStream { // Prime each accumulator for the registered group count with no data. // // We build 1-row null arrays for each aggregate argument and pass them - // with an all-false filter. The filter ensures no row is accumulated - // into any group, which keeps every group in its "zero" initial state - // (NULL for SUM/AVG/MIN/MAX, 0 for COUNT). + // with an all-false filter to update_batch. The filter ensures no row + // is accumulated into any group, which keeps every group in its "zero" + // initial state (NULL for SUM/AVG/MIN/MAX, 0 for COUNT). // // Using a 1-row batch rather than 0 rows is required to avoid a fast // path in `NullState::accumulate` that treats "0 nulls in a 0-row // array" as "all groups have been seen", which would cause SUM to // return 0 instead of NULL. // - // Argument types are inferred directly from the expression metadata so - // we never need to construct a full `RecordBatch`. + // This path always runs in a Raw input mode, so `update_batch` (not + // `merge_batch`) is the right entry point: + // + // - `has_grouping_set()` can only be true for the Partial / Single / + // SinglePartitioned modes, whose `input_mode()` is `Raw`. The final + // modes rebuild their group-by via `PhysicalGroupBy::as_final()`, + // which clears `has_grouping_set`, so this method returns early for + // them and never reaches here. + // + // Since every row is filtered out, the actual data content never + // matters. The assert documents and guards the invariant above. + debug_assert_eq!( + self.mode.input_mode(), + AggregateInputMode::Raw, + "init_empty_grouping_sets must only run in a Raw input mode" + ); let total_groups = self.group_values.len(); let null_args: Vec> = self .aggregate_arguments @@ -1137,11 +1151,7 @@ impl GroupedHashAggregateStream { .collect::>>()?; let false_filter = BooleanArray::from(vec![false]); for (acc, args) in self.accumulators.iter_mut().zip(null_args.iter()) { - if self.mode.input_mode() == AggregateInputMode::Raw { - acc.update_batch(args, &[0], Some(&false_filter), total_groups)?; - } else { - acc.merge_batch(args, &[0], Some(&false_filter), total_groups)?; - } + acc.update_batch(args, &[0], Some(&false_filter), total_groups)?; } } @@ -1394,6 +1404,8 @@ mod tests { use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::col; + // Migrated to PartialHashAggregateStream coverage in hash_aggregate.rs; + // kept here for the legacy GroupedHashAggregateStream implementation. #[tokio::test] async fn test_double_emission_race_condition_bug() -> Result<()> { // Fix for https://github.com/apache/datafusion/issues/18701 @@ -1500,6 +1512,9 @@ mod tests { Ok(()) } + // TODO: migrate to PartialHashAggregateStream when it supports + // InputOrderMode::PartiallySorted; kept here for the legacy + // GroupedHashAggregateStream implementation. #[tokio::test] async fn test_emit_early_with_partially_sorted() -> Result<()> { // Reproducer for #20445: EmitEarly with PartiallySorted panics in diff --git a/datafusion/physical-plan/src/aggregates/skip_partial.rs b/datafusion/physical-plan/src/aggregates/skip_partial.rs index a4306c69b411e..903235f950f58 100644 --- a/datafusion/physical-plan/src/aggregates/skip_partial.rs +++ b/datafusion/physical-plan/src/aggregates/skip_partial.rs @@ -134,6 +134,8 @@ mod tests { use datafusion_physical_expr::expressions::col; use futures::StreamExt; + // Migrated to PartialHashAggregateStream coverage in hash_aggregate.rs; + // kept here for the legacy GroupedHashAggregateStream implementation. #[tokio::test] async fn test_skip_aggregation_probe_not_locked_until_skip() -> Result<()> { // Test that the probe is not locked until we actually decide to skip. diff --git a/datafusion/spark/src/function/aggregate/avg.rs b/datafusion/spark/src/function/aggregate/avg.rs index 5f4d2c253a2dc..6ca3c59309e70 100644 --- a/datafusion/spark/src/function/aggregate/avg.rs +++ b/datafusion/spark/src/function/aggregate/avg.rs @@ -289,7 +289,6 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 2, "two arguments to merge_batch"); @@ -464,7 +463,6 @@ mod tests { acc.merge_batch( &state, &[0, 0, 0], - None, 1, // single group ) .unwrap(); @@ -486,7 +484,7 @@ mod tests { Some(3.0), ]))]; let state = acc.convert_to_state(&input, None).unwrap(); - acc.merge_batch(&state, &[0, 0, 0], None, 1).unwrap(); + acc.merge_batch(&state, &[0, 0, 0], 1).unwrap(); let result = acc.evaluate(EmitTo::All).unwrap(); let result = result.as_primitive::(); diff --git a/datafusion/sqllogictest/test_files/grouping.slt b/datafusion/sqllogictest/test_files/grouping.slt index eac901b2a300f..2c05dd851e61a 100644 --- a/datafusion/sqllogictest/test_files/grouping.slt +++ b/datafusion/sqllogictest/test_files/grouping.slt @@ -232,6 +232,12 @@ SELECT SUM(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING ---- NULL +# grouping_sets_empty_input_avg: AVG returns NULL for the empty group +query R +SELECT AVG(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING SETS(()) +---- +NULL + # grouping_sets_empty_input_count: COUNT returns 0 for the empty group, not a missing row query I SELECT COUNT(*) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING SETS(()) diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 2ebb6952fe04b..abe99846c12d8 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -45,6 +45,58 @@ will now appear as `Decimal128(NULL,10,2)`. Query result values already used human-readable decimal formatting and are unchanged. +### `GroupsAccumulator::merge_batch` no longer takes `opt_filter` + +The `opt_filter` argument has been removed from +`datafusion_expr_common::groups_accumulator::GroupsAccumulator::merge_batch`: + +```diff + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], +- opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()>; +``` + +Aggregate `FILTER` clauses only apply to raw input rows during the partial +(update) phase, so by the time intermediate states are merged there is nothing +left to filter per row. In practice `opt_filter` was always `None` here, so +removing it makes the API self-explanatory and impossible to misuse. + +**Who is affected:** + +- Anyone with a custom `GroupsAccumulator` implementation. +- Anyone calling `merge_batch` directly. + +**Migration guide:** + +Drop the `opt_filter` argument from your `merge_batch` signature and from any +call sites: + +```diff + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], +- opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + // ... + } +``` + +```diff +- acc.merge_batch(values, group_indices, None, total_num_groups)?; ++ acc.merge_batch(values, group_indices, total_num_groups)?; +``` + +If your implementation previously inspected `opt_filter` (for example asserting +it was `None`), that code can simply be deleted. + +See [issue #22775](https://github.com/apache/datafusion/issues/22775) for details. + ### `is_dynamic_physical_expr` is deprecated `datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr` is