From 35b4063af8d0be4dfa1bd66b58c72b83e74c4a54 Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Thu, 5 Nov 2026 16:22:03 -0800 Subject: [PATCH 1/2] feat: native GroupsAccumulator for first_value without ORDER BY --- .../functions-aggregate/src/first_last.rs | 463 ++++++++++++++++-- 1 file changed, 426 insertions(+), 37 deletions(-) diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index e7f2d8b7eac47..26f7a7fc09ac5 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -72,6 +72,12 @@ fn create_groups_accumulator_helper( state: S, ) -> Result> { let Some(ordering) = LexOrdering::new(args.order_bys.to_vec()) else { + if is_first { + return Ok(Box::new(UnorderedFirstValueGroupsAccumulator::new( + state, + args.ignore_nulls, + ))); + } return internal_err!("Groups accumulator must have an ordering."); }; @@ -191,44 +197,47 @@ fn create_groups_accumulator( } } +/// Returns true when a value state exists for the aggregate's return type. +/// +/// `LAST_VALUE` also needs an ordering; `FIRST_VALUE` has a grouped +/// accumulator with or without one. fn groups_accumulator_supported(args: &AccumulatorArgs) -> bool { use DataType::*; - !args.order_bys.is_empty() - && matches!( - args.return_field.data_type(), - Int8 | Int16 - | Int32 - | Int64 - | UInt8 - | UInt16 - | UInt32 - | UInt64 - | Float16 - | Float32 - | Float64 - | Decimal32(_, _) - | Decimal64(_, _) - | Decimal128(_, _) - | Decimal256(_, _) - | Date32 - | Date64 - | Time32(_) - | Time64(_) - | Timestamp(_, _) - | Utf8 - | LargeUtf8 - | Utf8View - | Binary - | LargeBinary - | BinaryView - | List(_) - | LargeList(_) - | ListView(_) - | LargeListView(_) - | FixedSizeList(_, _) - | Struct(_) - | Map(_, _) - ) + matches!( + args.return_field.data_type(), + Int8 | Int16 + | Int32 + | Int64 + | UInt8 + | UInt16 + | UInt32 + | UInt64 + | Float16 + | Float32 + | Float64 + | Decimal32(_, _) + | Decimal64(_, _) + | Decimal128(_, _) + | Decimal256(_, _) + | Date32 + | Date64 + | Time32(_) + | Time64(_) + | Timestamp(_, _) + | Utf8 + | LargeUtf8 + | Utf8View + | Binary + | LargeBinary + | BinaryView + | List(_) + | LargeList(_) + | ListView(_) + | LargeListView(_) + | FixedSizeList(_, _) + | Struct(_) + | Map(_, _) + ) } #[user_doc( @@ -908,6 +917,135 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator { + state: S, + /// `is_sets[group_idx]` is true once the group holds its first value. + is_sets: BooleanBufferBuilder, + ignore_nulls: bool, +} + +impl UnorderedFirstValueGroupsAccumulator { + fn new(state: S, ignore_nulls: bool) -> Self { + Self { + state, + is_sets: BooleanBufferBuilder::new(0), + ignore_nulls, + } + } + + fn resize(&mut self, total_num_groups: usize) { + if self.is_sets.len() < total_num_groups { + self.state.resize(total_num_groups); + self.is_sets.resize(total_num_groups); + } + } + + /// Stores row `idx` of `values` as the group's value unless the group + /// already has one. + fn set_if_unset( + &mut self, + group_idx: usize, + values: &ArrayRef, + idx: usize, + ) -> Result<()> { + if !self.is_sets.get_bit(group_idx) { + self.state.update(group_idx, values, idx)?; + self.is_sets.set_bit(group_idx, true); + } + Ok(()) + } +} + +impl GroupsAccumulator + for UnorderedFirstValueGroupsAccumulator +{ + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + self.resize(total_num_groups); + let values = &values[0]; + for (idx, &group_idx) in group_indices.iter().enumerate() { + // A row passes the FILTER clause only when the predicate is + // `true`; rows whose predicate evaluates to `null` are excluded. + if let Some(filter) = opt_filter + && !(filter.is_valid(idx) && filter.value(idx)) + { + continue; + } + if self.ignore_nulls && values.is_null(idx) { + continue; + } + self.set_if_unset(group_idx, values, idx)?; + } + Ok(()) + } + + fn evaluate(&mut self, emit_to: EmitTo) -> Result { + let values = self.state.take(emit_to)?; + state::take_need(&mut self.is_sets, emit_to); + Ok(values) + } + + fn state(&mut self, emit_to: EmitTo) -> Result> { + let values = self.state.take(emit_to)?; + let is_sets = state::take_need(&mut self.is_sets, emit_to); + Ok(vec![values, Arc::new(BooleanArray::new(is_sets, None))]) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + let [values, is_set] = values else { + return internal_err!( + "first_value: expected 2 state columns, got {}", + values.len() + ); + }; + let is_set = as_boolean_array(is_set)?; + validate_is_set_flags(is_set, "first_value")?; + self.resize(total_num_groups); + for (idx, &group_idx) in group_indices.iter().enumerate() { + if is_set.value(idx) { + self.set_if_unset(group_idx, values, idx)?; + } + } + Ok(()) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + let values = &values[0]; + let is_set = BooleanArray::from_iter((0..values.len()).map(|idx| { + let passed_filter = + opt_filter.is_none_or(|f| f.is_valid(idx) && f.value(idx)); + Some(passed_filter && !(self.ignore_nulls && values.is_null(idx))) + })); + Ok(vec![Arc::clone(values), Arc::new(is_set)]) + } + + fn size(&self) -> usize { + self.state.size() + self.is_sets.capacity() / 8 + } +} + /// This accumulator is used when there is no ordering specified for the /// `FIRST_VALUE` aggregation. It simply returns the first value it sees /// according to the pre-existing ordering of the input data, and provides @@ -1289,7 +1427,7 @@ impl AggregateUDFImpl for LastValue { } fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - groups_accumulator_supported(&args) + !args.order_bys.is_empty() && groups_accumulator_supported(&args) } fn create_groups_accumulator( @@ -2994,4 +3132,255 @@ mod tests { assert_eq!(merged.size(), updated.size()); Ok(()) } + + #[test] + fn test_unordered_first_group_acc() -> Result<()> { + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + false, + ); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(1), Some(2)])); + group_acc.update_batch(&[values], &[0, 0, 1], None, 3)?; + + let values: ArrayRef = Arc::new(Int64Array::from(vec![5, 6, 7])); + group_acc.update_batch(&[values], &[0, 1, 2], None, 4)?; + + // RESPECT NULLS keeps the leading null as group 0's first value. + // Groups 0 and 1 already have a first value, so only group 2 changes. + let state = group_acc.state(EmitTo::All)?; + let expected_state: Vec = vec![ + Arc::new(Int64Array::from(vec![None, Some(2), Some(7), None])), + Arc::new(BooleanArray::from(vec![true, true, true, false])), + ]; + assert_eq!(state, expected_state); + + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_ignore_nulls() -> Result<()> { + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + true, + ); + + let values: ArrayRef = + Arc::new(Int64Array::from(vec![None, Some(1), None, Some(3)])); + group_acc.update_batch(&[values], &[0, 0, 1, 0], None, 2)?; + + // Group 1 only saw a null, so it stays unset. + let state = group_acc.state(EmitTo::All)?; + let expected_state: Vec = vec![ + Arc::new(Int64Array::from(vec![Some(1), None])), + Arc::new(BooleanArray::from(vec![true, false])), + ]; + assert_eq!(state, expected_state); + + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_filter_null_predicate() -> Result<()> { + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + false, + ); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![10, 20, 30])); + + // Row 0: predicate is null with its value bit set + // Row 1: predicate is false + // Row 2: predicate is true + let filter = BooleanArray::new( + BooleanBuffer::from(vec![false, true, false, true]), + Some(NullBuffer::from(BooleanBuffer::from(vec![ + true, false, true, true, + ]))), + ) + .slice(1, 3); + assert_eq!(filter.offset(), 1); + + group_acc.update_batch(&[values], &[0, 0, 1], Some(&filter), 2)?; + + let binding = group_acc.evaluate(EmitTo::All)?; + let eval_result = binding.as_primitive::(); + + // Group 0 has no row with a `true` predicate, so it must stay unset. + let expect = Int64Array::from(vec![None, Some(30)]); + assert_eq!(eval_result, &expect); + + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_merge() -> Result<()> { + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + false, + ); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![5])); + group_acc.update_batch(&[values], &[0], None, 1)?; + + let state: Vec = vec![ + Arc::new(Int64Array::from(vec![100, 200, 300, 400])), + Arc::new(BooleanArray::from(vec![true, true, false, true])), + ]; + group_acc.merge_batch(&state, &[2, 0, 1, 1], 3)?; + + let binding = group_acc.evaluate(EmitTo::All)?; + let eval_result = binding.as_primitive::(); + + // group 0 keeps its value from update_batch. + // group 1 skips 300 (is_set=false) and takes 400. + // group 2 takes 100. + let expect = Int64Array::from(vec![5, 400, 100]); + assert_eq!(eval_result, &expect); + + Ok(()) + } + + /// `convert_to_state` folds the FILTER clause and IGNORE NULLS into + /// `is_set`, so `merge_batch` keeps only rows `update_batch` would keep. + #[test] + fn test_unordered_first_group_acc_convert_to_state() -> Result<()> { + let group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + true, + ); + + let values: ArrayRef = + Arc::new(Int64Array::from(vec![None, Some(10), Some(20), Some(30)])); + + // Row 0: null value, skipped by IGNORE NULLS + // Row 1: predicate is null with its value bit set + // Row 2: predicate is false + // Row 3: predicate is true + let filter = BooleanArray::new( + BooleanBuffer::from(vec![true, true, false, true]), + Some(NullBuffer::from(BooleanBuffer::from(vec![ + true, false, true, true, + ]))), + ); + + let state = group_acc.convert_to_state(&[values], Some(&filter))?; + assert_eq!(state.len(), 2); + assert_eq!( + state[1].as_boolean(), + &BooleanArray::from(vec![false, false, false, true]) + ); + + let mut merging_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + true, + ); + merging_acc.merge_batch(&state, &[0, 0, 0, 0], 1)?; + + let binding = merging_acc.evaluate(EmitTo::All)?; + let eval_result = binding.as_primitive::(); + let expect = Int64Array::from(vec![Some(30)]); + assert_eq!(eval_result, &expect); + + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_emit_first() -> Result<()> { + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + false, + ); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3])); + group_acc.update_batch(&[values], &[0, 1, 2], None, 3)?; + + let binding = group_acc.evaluate(EmitTo::First(2))?; + assert_eq!( + binding.as_primitive::(), + &Int64Array::from(vec![1, 2]) + ); + + // Old group 2 is now group 0 and keeps its value; group 1 is new. + let values: ArrayRef = Arc::new(Int64Array::from(vec![9, 8])); + group_acc.update_batch(&[values], &[0, 1], None, 2)?; + + let binding = group_acc.evaluate(EmitTo::All)?; + assert_eq!( + binding.as_primitive::(), + &Int64Array::from(vec![3, 8]) + ); + + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_utf8() -> Result<()> { + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + BytesValueState::try_new(DataType::Utf8)?, + false, + ); + + let values: ArrayRef = Arc::new(StringArray::from(vec![ + "first", + "late", + "group-one", + "group-two", + ])); + group_acc.update_batch(&[values], &[0, 0, 1, 2], None, 4)?; + + let state = group_acc.state(EmitTo::All)?; + assert_eq!(state.len(), 2); + assert_eq!( + state[0].as_string::(), + &StringArray::from(vec![ + Some("first"), + Some("group-one"), + Some("group-two"), + None, + ]) + ); + assert_eq!( + state[1].as_boolean(), + &BooleanArray::from(vec![true, true, true, false]) + ); + + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_list() -> Result<()> { + let value_type = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + GenericValueState::new(value_type), + false, + ); + + let values = ListArray::from_iter_primitive::([ + Some(vec![Some(1), Some(2)]), + Some(vec![Some(9)]), + Some(vec![Some(3)]), + Some(vec![Some(4), Some(5)]), + ]); + group_acc.update_batch(&[Arc::new(values)], &[0, 0, 1, 2], None, 4)?; + + let state = group_acc.state(EmitTo::All)?; + assert_eq!(state.len(), 2); + assert_eq!( + state[0].as_list::(), + &ListArray::from_iter_primitive::([ + Some(vec![Some(1), Some(2)]), + Some(vec![Some(3)]), + Some(vec![Some(4), Some(5)]), + None, + ]) + ); + assert_eq!( + state[1].as_boolean(), + &BooleanArray::from(vec![true, true, true, false]) + ); + + Ok(()) + } } From c36da6edb8d959259b926f7851188cd5852fd723 Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Fri, 6 Nov 2026 12:51:49 -0800 Subject: [PATCH 2/2] test: check FirstValue creates the unordered groups accumulator --- .../functions-aggregate/src/first_last.rs | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index 26f7a7fc09ac5..40f961157efe2 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -3383,4 +3383,38 @@ mod tests { Ok(()) } + + #[test] + fn test_first_value_unordered_groups_accumulator() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]); + let exprs = [col("a", &schema)?]; + let expr_fields = [Arc::new(Field::new("a", DataType::Int64, true))]; + let args = || AccumulatorArgs { + return_field: Arc::new(Field::new("f", DataType::Int64, true)), + schema: &schema, + expr_fields: &expr_fields, + ignore_nulls: false, + order_bys: &[], + is_reversed: false, + name: "first_value(a)", + is_distinct: false, + exprs: &exprs, + }; + + // Without ORDER BY, FIRST_VALUE uses its own grouped accumulator + // rather than falling back to `GroupsAccumulatorAdapter`. + let first_value = FirstValue::new(); + assert!(first_value.groups_accumulator_supported(args())); + let group_acc = first_value.create_groups_accumulator(args())?; + let group_acc: &dyn std::any::Any = group_acc.as_ref(); + assert!( + group_acc + .is::>>() + ); + + // LAST_VALUE still requires an ordering. + assert!(!LastValue::new().groups_accumulator_supported(args())); + + Ok(()) + } }