From e28952653ac4bdf69a6ac64cddd72e7a7fa83b0d Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Thu, 5 Nov 2026 16:37:35 -0800 Subject: [PATCH 1/2] feat: native `GroupsAccumulator` for unordered `first_value` first_value only had a GroupsAccumulator when it had an ORDER BY. Without one, grouped aggregation fell back to GroupsAccumulatorAdapter, which keeps a boxed TrivialFirstValueAccumulator per group. DISTINCT ON without ORDER BY plans as one unordered first_value per selected column, so queries with many distinct keys paid that per-group cost in both time and memory. Add UnorderedFirstValueGroupsAccumulator. It stores values in the existing ValueState implementations plus one is_set bit per group. A group takes the first row that passes the FILTER clause (and, with IGNORE NULLS, is not null) and is never written again, which is the same row TrivialFirstValueAccumulator keeps. The state layout stays [first_value, is_set], so partial and final aggregation are unchanged. Also implements evaluate_preserving and state_preserving. LAST_VALUE and ordered FIRST_VALUE are unchanged. --- .../functions-aggregate/src/first_last.rs | 514 ++++++++++++++++-- 1 file changed, 477 insertions(+), 37 deletions(-) diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index a4236333fd98a..0980a86c34aeb 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -82,6 +82,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."); }; @@ -196,44 +202,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( @@ -812,6 +821,169 @@ 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(()) + } + + fn validate_selection(&self, selection: GroupSelection<'_>) -> Result<()> { + selection.validate_num_groups(self.state.num_groups())?; + selection.validate_num_groups(self.is_sets.len()) + } +} + +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 evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + self.validate_selection(selection)?; + self.state.build_preserving(selection) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + + 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 state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + self.validate_selection(selection)?; + let values = self.state.build_preserving(selection)?; + let mut is_sets = BooleanBufferBuilder::new(selection.len()); + for group_idx in selection.iter() { + is_sets.append(self.is_sets.get_bit(group_idx)); + } + Ok(vec![ + values, + Arc::new(BooleanArray::new(is_sets.finish(), None)), + ]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + + 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 @@ -1191,7 +1363,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( @@ -2636,4 +2808,272 @@ mod tests { } 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)?; + + // RESPECT NULLS keeps the leading null as group 0's first value. + let state = group_acc.state_preserving(GroupSelection::all(3))?; + let expected_state: Vec = vec![ + Arc::new(Int64Array::from(vec![None, Some(2), None])), + Arc::new(BooleanArray::from(vec![true, true, false])), + ]; + assert_eq!(state, expected_state); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![5, 6, 7])); + group_acc.update_batch(&[values], &[0, 1, 2], None, 4)?; + + // Groups 0 and 1 already have a first value, so only group 2 changes. + let binding = group_acc.evaluate(EmitTo::All)?; + let eval_result = binding.as_primitive::(); + let expect = Int64Array::from(vec![None, Some(2), Some(7), None]); + assert_eq!(eval_result, &expect); + + 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_preserving_reads_utf8() -> Result<()> { + let mut acc = UnorderedFirstValueGroupsAccumulator::new( + BytesValueState::try_new(DataType::Utf8)?, + false, + ); + + let values: ArrayRef = Arc::new(StringArray::from(vec![ + "first", + "late", + "group-one", + "group-two", + ])); + acc.update_batch(&[values], &[0, 0, 1, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 2, 3], 4)?; + let expected_values = StringArray::from(vec![ + Some("group-two"), + Some("first"), + Some("group-two"), + None, + ]); + for _ in 0..2 { + assert_eq!( + acc.evaluate_preserving(selection)?.as_string::(), + &expected_values + ); + } + + let state = acc.state_preserving(selection)?; + assert_eq!(state.len(), 2); + assert_eq!(state[0].as_string::(), &expected_values); + assert_eq!( + state[1].as_boolean(), + &BooleanArray::from(vec![true, true, true, false]) + ); + assert!(acc.supports_evaluate_preserving()); + assert!(acc.supports_state_preserving()); + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_preserving_reads_list() -> Result<()> { + let value_type = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + let mut 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)]), + ]); + acc.update_batch(&[Arc::new(values)], &[0, 0, 1, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 2, 3], 4)?; + let expected_values = ListArray::from_iter_primitive::([ + Some(vec![Some(4), Some(5)]), + Some(vec![Some(1), Some(2)]), + Some(vec![Some(4), Some(5)]), + None, + ]); + for _ in 0..2 { + assert_eq!( + acc.evaluate_preserving(selection)?.as_list::(), + &expected_values + ); + } + + let state = acc.state_preserving(selection)?; + assert_eq!(state.len(), 2); + assert_eq!(state[0].as_list::(), &expected_values); + assert_eq!( + state[1].as_boolean(), + &BooleanArray::from(vec![true, true, true, false]) + ); + Ok(()) + } } From bcbb675511eb2a53bb86fe8b4498b6da4a620b86 Mon Sep 17 00:00:00 2001 From: Rohan Krishnaswamy <47869999+rkrishn7@users.noreply.github.com> Date: Fri, 6 Nov 2026 12:51:52 -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 0980a86c34aeb..e954c3fe6e848 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -3076,4 +3076,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(()) + } }