From ef9b92c238e457230d2a544fd13334d081a3766b Mon Sep 17 00:00:00 2001 From: Max Burke Date: Fri, 17 Jul 2026 14:39:52 -0700 Subject: [PATCH] Aggregation of byte types uses i64 internally then narrows output slices back to i32-offset-types as necessary. --- .../aggregates/aggregate_hash_table/common.rs | 33 ++- .../aggregate_hash_table/common_ordered.rs | 23 +- .../aggregate_hash_table/partial_table.rs | 1 + .../group_values/multi_group_by/bytes.rs | 259 +++++++----------- .../group_values/multi_group_by/mod.rs | 34 +-- .../src/aggregates/grouped_hash_stream.rs | 102 ++++++- .../physical-plan/src/aggregates/mod.rs | 130 ++++++++- 7 files changed, 382 insertions(+), 200 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index eaf39929ced62..9430fed7cbac3 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -33,6 +33,9 @@ use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, }; +use crate::aggregates::{ + narrow_group_key_columns, widen_group_key_arrays, widen_group_key_schema, +}; /// Marker for raw rows -> partial state aggregation. pub(in crate::aggregates) struct PartialMarker; @@ -82,6 +85,13 @@ pub(in crate::aggregates) struct AggregateHashTable { /// Output schema: group columns followed by aggregate state or final values. pub(super) output_schema: SchemaRef, + /// `output_schema` with string/binary group key columns widened to the + /// internal 64-bit offset representation the group values are emitted + /// with (see `PhysicalGroupBy::group_schema`). The materialized output is + /// built against this schema and narrowed back to `output_schema` after + /// being sliced to `batch_size` rows. + pub(super) emit_schema: SchemaRef, + /// Maximum rows per emitted output batch, from config `batch_size`. pub(super) batch_size: usize, @@ -127,12 +137,15 @@ impl AggregateHashTable { .collect::>()?; let group_schema = agg.group_by.group_schema(&input_schema)?; + let emit_schema = + widen_group_key_schema(&output_schema, group_schema.fields().len()); let group_values = new_group_values(group_schema, &GroupOrdering::None)?; Ok(Self { group_by_metrics: GroupByMetrics::new(&agg.metrics, partition), input_schema, output_schema, + emit_schema, batch_size, state: AggregateHashTableState::Building(AggregateHashTableBuffer { group_by: Arc::clone(&agg.group_by), @@ -184,7 +197,14 @@ impl AggregateHashTable { batch: &RecordBatch, aggregate_fn: AggregateBatchFn, ) -> Result<()> { - let evaluated_batch = self.evaluate_batch(batch)?; + let mut evaluated_batch = self.evaluate_batch(batch)?; + // Widen string/binary group keys to the internal representation that + // `group_values` interns (see `PhysicalGroupBy::group_schema`). + evaluated_batch.grouping_set_args = evaluated_batch + .grouping_set_args + .into_iter() + .map(widen_group_key_arrays) + .collect::>()?; let state = self.state.building_mut(); let _timer = self.group_by_metrics.aggregation_time.timer(); @@ -221,6 +241,7 @@ impl AggregateHashTable { materialize_accumulator_fn: MaterializeAccumulatorFn, ) -> Result> { let output_schema = Arc::clone(&self.output_schema); + let emit_schema = Arc::clone(&self.emit_schema); let batch_size = self.batch_size; let mut output = @@ -240,7 +261,7 @@ impl AggregateHashTable { } drop(timer); - let batch = RecordBatch::try_new(output_schema, columns)?; + let batch = RecordBatch::try_new(emit_schema, columns)?; debug_assert!(batch.num_rows() > 0); MaterializedAggregateOutput::new(batch) } @@ -253,7 +274,13 @@ impl AggregateHashTable { } }; - let batch = output.next_batch(batch_size); + // Narrow internally-widened group key columns back to the output + // schema types. Each slice holds at most `batch_size` rows, so unlike + // the full materialized batch, the narrow representation fits. + let batch = output + .next_batch(batch_size) + .map(|batch| narrow_group_key_columns(batch, &output_schema)) + .transpose()?; if output.is_exhausted() { self.state = AggregateHashTableState::Done; } else { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index c83303c51d6e8..97f70e48a12dc 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -37,6 +37,9 @@ use crate::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, }; +use crate::aggregates::{ + narrow_group_key_columns, widen_group_key_arrays, widen_group_key_schema, +}; use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; @@ -81,6 +84,13 @@ pub(in crate::aggregates) struct OrderedAggregateTable { /// Output schema: group columns followed by aggregate state or final values. pub(super) output_schema: SchemaRef, + /// `output_schema` with string/binary group key columns widened to the + /// internal 64-bit offset representation the group values are emitted + /// with (see `PhysicalGroupBy::group_schema`). Emitted batches are built + /// against this schema and narrowed back to `output_schema`; they are + /// clamped to `batch_size` rows, where the narrow representation fits. + pub(super) emit_schema: SchemaRef, + /// Maximum rows per emitted output batch, from config `batch_size`. pub(super) batch_size: usize, @@ -144,6 +154,8 @@ impl OrderedAggregateTable { let group_ordering = GroupOrdering::try_new(input_order_mode)?; let group_schema = agg.group_by.group_schema(input_schema)?; + let emit_schema = + widen_group_key_schema(&output_schema, group_schema.fields().len()); let group_values = new_group_values(group_schema, &group_ordering)?; let aggregate_arguments = aggregate_expressions( &agg.aggr_expr, @@ -168,6 +180,7 @@ impl OrderedAggregateTable { Ok(Self { output_schema, + emit_schema, batch_size, group_by_metrics: GroupByMetrics::new(&agg.metrics, partition), buffer: OrderedAggregateTableBuffer { @@ -265,6 +278,10 @@ impl OrderedAggregateTable { is_final: bool, ) -> Result<()> { for group_values in &evaluated_batch.grouping_set_args { + // Widen string/binary group keys to the internal representation + // that `group_values` interns (see `PhysicalGroupBy::group_schema`) + let group_values = widen_group_key_arrays(group_values.clone())?; + let group_values = &group_values; let starting_num_groups = self.buffer.group_values.len(); self.buffer .group_values @@ -349,7 +366,11 @@ impl OrderedAggregateTable { } drop(timer); - let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; + // Emitted group keys use the widened internal representation; narrow + // them back to the output schema. The emission is clamped to + // `batch_size` rows above, so the narrow representation always fits. + let batch = RecordBatch::try_new(Arc::clone(&self.emit_schema), output)?; + let batch = narrow_group_key_columns(batch, &self.output_schema)?; debug_assert!(batch.num_rows() > 0); Ok(Some(batch)) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index ffac42feaa3b3..0347cde520fd9 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -95,6 +95,7 @@ impl AggregateHashTable { group_by_metrics: self.group_by_metrics.clone(), input_schema: Arc::clone(&self.input_schema), output_schema: Arc::clone(&self.output_schema), + emit_schema: Arc::clone(&self.emit_schema), batch_size: self.batch_size, state: AggregateHashTableState::Building(AggregateHashTableBuffer { group_by: Arc::clone(&state.group_by), diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs index c83b1da4049bc..d477e1c81f356 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs @@ -20,14 +20,14 @@ use crate::aggregates::group_values::multi_group_by::{ }; use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; use arrow::array::{ - Array, ArrayRef, AsArray, BooleanBufferBuilder, BufferBuilder, GenericBinaryArray, - GenericByteArray, GenericStringArray, OffsetSizeTrait, types::GenericStringType, + Array, ArrayRef, AsArray, BooleanBufferBuilder, BufferBuilder, GenericByteArray, + LargeBinaryArray, LargeStringArray, types::GenericStringType, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{ByteArrayType, DataType, GenericBinaryType}; +use datafusion_common::Result; use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::utils::split_vec_min_alloc; -use datafusion_common::{Result, exec_datafusion_err}; use datafusion_physical_expr_common::binary_map::{INITIAL_BUFFER_CAPACITY, OutputType}; use std::mem::size_of; use std::sync::Arc; @@ -40,38 +40,55 @@ use std::vec; /// /// 1. Efficient comparison of incoming rows to existing rows /// 2. Efficient construction of the final output array -pub struct ByteGroupValueBuilder -where - O: OffsetSizeTrait, -{ +/// +/// Group values are always accumulated with 64-bit offsets and emitted as +/// `LargeUtf8` / `LargeBinary` arrays, regardless of the offset width of the +/// input arrays, so the concatenation of all distinct group values is not +/// limited to `i32::MAX` bytes. Input arrays may use either offset width; +/// each call dispatches on the array's actual data type. +pub struct ByteGroupValueBuilder { output_type: OutputType, buffer: BufferBuilder, /// Offsets into `buffer` for each distinct value. These offsets as used /// directly to create the final `GenericBinaryArray`. The `i`th string is /// stored in the range `offsets[i]..offsets[i+1]` in `buffer`. Null values /// are stored as a zero length string. - offsets: Vec, + offsets: Vec, /// Nulls nulls: MaybeNullBufferBuilder, - /// The maximum size of the buffer for `0` - max_buffer_size: usize, } -impl ByteGroupValueBuilder -where - O: OffsetSizeTrait, -{ +/// Dispatches to `$inner::($($args),*)` with `B` resolved from the +/// builder's output type and the input array's actual offset width. +macro_rules! dispatch_input_type { + ($self:ident, $column:ident, $inner:ident($($args:expr),*)) => { + match ($self.output_type, $column.data_type()) { + (OutputType::Utf8, DataType::Utf8) => { + $self.$inner::>($($args),*) + } + (OutputType::Utf8, DataType::LargeUtf8) => { + $self.$inner::>($($args),*) + } + (OutputType::Binary, DataType::Binary) => { + $self.$inner::>($($args),*) + } + (OutputType::Binary, DataType::LargeBinary) => { + $self.$inner::>($($args),*) + } + (output_type, dt) => unreachable!( + "ByteGroupValueBuilder: unexpected input type {dt} for output type {output_type:?}" + ), + } + }; +} + +impl ByteGroupValueBuilder { pub fn new(output_type: OutputType) -> Self { Self { output_type, buffer: BufferBuilder::new(INITIAL_BUFFER_CAPACITY), - offsets: vec![O::default()], + offsets: vec![0], nulls: MaybeNullBufferBuilder::new(), - max_buffer_size: if O::IS_LARGE { - i64::MAX as usize - } else { - i32::MAX as usize - }, } } @@ -92,10 +109,10 @@ where self.nulls.append(true); // nulls need a zero length in the offset buffer let offset = self.buffer.len(); - self.offsets.push(O::usize_as(offset)); + self.offsets.push(offset as i64); } else { self.nulls.append(false); - self.do_append_val_inner(arr, row)?; + self.do_append_val_inner(arr, row); } Ok(()) @@ -154,7 +171,7 @@ where Nulls::None => { self.nulls.append_n(rows.len(), false); for &row in rows { - self.do_append_val_inner(arr, row)?; + self.do_append_val_inner(arr, row); } } @@ -163,7 +180,7 @@ where let new_len = self.offsets.len() + rows.len(); let offset = self.buffer.len(); - self.offsets.resize(new_len, O::usize_as(offset)); + self.offsets.resize(new_len, offset as i64); } } @@ -188,83 +205,31 @@ where self.value(lhs_row) == (array.value(rhs_row).as_ref() as &[u8]) } - fn do_append_val_inner( - &mut self, - array: &GenericByteArray, - row: usize, - ) -> Result<()> + fn do_append_val_inner(&mut self, array: &GenericByteArray, row: usize) where B: ByteArrayType, { let value: &[u8] = array.value(row).as_ref(); self.buffer.append_slice(value); - - if self.buffer.len() > self.max_buffer_size { - return Err(exec_datafusion_err!( - "offset overflow, buffer size > {}", - self.max_buffer_size - )); - } - - self.offsets.push(O::usize_as(self.buffer.len())); - Ok(()) + self.offsets.push(self.buffer.len() as i64); } /// return the current value of the specified row irrespective of null pub fn value(&self, row: usize) -> &[u8] { - let l = self.offsets[row].as_usize(); - let r = self.offsets[row + 1].as_usize(); + let l = self.offsets[row] as usize; + let r = self.offsets[row + 1] as usize; // Safety: the offsets are constructed correctly and never decrease unsafe { self.buffer.as_slice().get_unchecked(l..r) } } } -impl GroupColumn for ByteGroupValueBuilder -where - O: OffsetSizeTrait, -{ +impl GroupColumn for ByteGroupValueBuilder { fn equal_to(&self, lhs_row: usize, column: &ArrayRef, rhs_row: usize) -> bool { - // Sanity array type - match self.output_type { - OutputType::Binary => { - debug_assert!(matches!( - column.data_type(), - DataType::Binary | DataType::LargeBinary - )); - self.equal_to_inner::>(lhs_row, column, rhs_row) - } - OutputType::Utf8 => { - debug_assert!(matches!( - column.data_type(), - DataType::Utf8 | DataType::LargeUtf8 - )); - self.equal_to_inner::>(lhs_row, column, rhs_row) - } - _ => unreachable!("View types should use `ArrowBytesViewMap`"), - } + dispatch_input_type!(self, column, equal_to_inner(lhs_row, column, rhs_row)) } fn append_val(&mut self, column: &ArrayRef, row: usize) -> Result<()> { - // Sanity array type - match self.output_type { - OutputType::Binary => { - debug_assert!(matches!( - column.data_type(), - DataType::Binary | DataType::LargeBinary - )); - self.append_val_inner::>(column, row)? - } - OutputType::Utf8 => { - debug_assert!(matches!( - column.data_type(), - DataType::Utf8 | DataType::LargeUtf8 - )); - self.append_val_inner::>(column, row)? - } - _ => unreachable!("View types should use `ArrowBytesViewMap`"), - }; - - Ok(()) + dispatch_input_type!(self, column, append_val_inner(column, row)) } fn vectorized_equal_to( @@ -274,56 +239,15 @@ where rhs_rows: &[usize], equal_to_results: &mut BooleanBufferBuilder, ) { - // Sanity array type - match self.output_type { - OutputType::Binary => { - debug_assert!(matches!( - array.data_type(), - DataType::Binary | DataType::LargeBinary - )); - self.vectorized_equal_to_inner::>( - lhs_rows, - array, - rhs_rows, - equal_to_results, - ); - } - OutputType::Utf8 => { - debug_assert!(matches!( - array.data_type(), - DataType::Utf8 | DataType::LargeUtf8 - )); - self.vectorized_equal_to_inner::>( - lhs_rows, - array, - rhs_rows, - equal_to_results, - ); - } - _ => unreachable!("View types should use `ArrowBytesViewMap`"), - } + dispatch_input_type!( + self, + array, + vectorized_equal_to_inner(lhs_rows, array, rhs_rows, equal_to_results) + ) } fn vectorized_append(&mut self, column: &ArrayRef, rows: &[usize]) -> Result<()> { - match self.output_type { - OutputType::Binary => { - debug_assert!(matches!( - column.data_type(), - DataType::Binary | DataType::LargeBinary - )); - self.vectorized_append_inner::>(column, rows)? - } - OutputType::Utf8 => { - debug_assert!(matches!( - column.data_type(), - DataType::Utf8 | DataType::LargeUtf8 - )); - self.vectorized_append_inner::>(column, rows)? - } - _ => unreachable!("View types should use `ArrowBytesViewMap`"), - }; - - Ok(()) + dispatch_input_type!(self, column, vectorized_append_inner(column, rows)) } fn len(&self) -> usize { @@ -342,20 +266,19 @@ where mut buffer, offsets, nulls, - .. } = *self; let null_buffer = nulls.build(); // SAFETY: the offsets were constructed correctly in `insert_if_new` -- - // monotonically increasing, overflows were checked. + // monotonically increasing, i64 offsets cannot overflow. let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(offsets)) }; let values = buffer.finish(); match output_type { OutputType::Binary => { // SAFETY: the offsets were constructed correctly Arc::new(unsafe { - GenericBinaryArray::new_unchecked(offsets, values, null_buffer) + LargeBinaryArray::new_unchecked(offsets, values, null_buffer) }) } OutputType::Utf8 => { @@ -366,7 +289,7 @@ where // all the values that went in were valid (e.g. utf8) so are all // the values that come out Arc::new(unsafe { - GenericStringArray::new_unchecked(offsets, values, null_buffer) + LargeStringArray::new_unchecked(offsets, values, null_buffer) }) } _ => unreachable!("View types should use `ArrowBytesViewMap`"), @@ -376,7 +299,7 @@ where fn take_n(&mut self, n: usize) -> ArrayRef { debug_assert!(self.len() >= n); let null_buffer = self.nulls.take_n(n); - let first_remaining_offset = O::as_usize(self.offsets[n]); + let first_remaining_offset = self.offsets[n] as usize; // Given offsets like [0, 2, 4, 5] and n = 1, we expect to get // offsets [0, 2, 3]. We first create two offsets for first_n as [0, 2] and the remaining as [2, 4, 5]. @@ -384,11 +307,11 @@ where let offset_n = self.offsets[n]; let mut first_n_offsets = split_vec_min_alloc(&mut self.offsets, n); // After the split, self.offsets[0] == offset_n in both branches; normalize in-place. - self.offsets.iter_mut().for_each(|o| *o = o.sub(offset_n)); + self.offsets.iter_mut().for_each(|o| *o -= offset_n); first_n_offsets.push(offset_n); // SAFETY: the offsets were constructed correctly in `insert_if_new` -- - // monotonically increasing, overflows were checked. + // monotonically increasing, i64 offsets cannot overflow. let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(first_n_offsets)) }; @@ -405,7 +328,7 @@ where OutputType::Binary => { // SAFETY: the offsets were constructed correctly Arc::new(unsafe { - GenericBinaryArray::new_unchecked(offsets, values, null_buffer) + LargeBinaryArray::new_unchecked(offsets, values, null_buffer) }) } OutputType::Utf8 => { @@ -416,7 +339,7 @@ where // thus since all the values that went in were valid (e.g. utf8) // so are all the values that come out Arc::new(unsafe { - GenericStringArray::new_unchecked(offsets, values, null_buffer) + LargeStringArray::new_unchecked(offsets, values, null_buffer) }) } _ => unreachable!("View types should use `ArrowBytesViewMap`"), @@ -429,8 +352,11 @@ mod tests { use std::sync::Arc; use crate::aggregates::group_values::multi_group_by::bytes::ByteGroupValueBuilder; - use arrow::array::{ArrayRef, BooleanBufferBuilder, NullBufferBuilder, StringArray}; - use datafusion_common::DataFusionError; + use arrow::array::{ + ArrayRef, AsArray, BooleanBufferBuilder, LargeStringArray, NullBufferBuilder, + StringArray, + }; + use arrow::datatypes::DataType; use datafusion_physical_expr::binary_map::OutputType; use super::GroupColumn; @@ -446,30 +372,32 @@ mod tests { } #[test] - fn test_byte_group_value_builder_overflow() { - let mut builder = ByteGroupValueBuilder::::new(OutputType::Utf8); + fn test_byte_group_value_builder_exceeds_i32_offsets() { + let mut builder = ByteGroupValueBuilder::new(OutputType::Utf8); let large_string = "a".repeat(1024 * 1024); let array = Arc::new(StringArray::from(vec![Some(large_string.as_str())])) as ArrayRef; - // Append items until our buffer length is i32::MAX as usize - for _ in 0..2047 { + // Append items until the buffer length exceeds i32::MAX; this used to + // fail with an "offset overflow" error when the builder used i32 + // offsets for Utf8 input. + for _ in 0..2049 { builder.append_val(&array, 0).unwrap(); } - assert!(matches!( - builder.append_val(&array, 0), - Err(DataFusionError::Execution(e)) if e.contains("offset overflow") - )); + assert_eq!(builder.value(2048), large_string.as_bytes()); - assert_eq!(builder.value(2046), large_string.as_bytes()); + let output = Box::new(builder).build(); + assert_eq!(output.data_type(), &DataType::LargeUtf8); + assert_eq!(output.len(), 2049); + assert_eq!(output.as_string::().value(2048), large_string); } #[test] fn test_byte_take_n() { - let mut builder = ByteGroupValueBuilder::::new(OutputType::Utf8); + let mut builder = ByteGroupValueBuilder::new(OutputType::Utf8); let array = Arc::new(StringArray::from(vec![Some("a"), None])) as ArrayRef; // a, null, null builder.append_val(&array, 0).unwrap(); @@ -478,7 +406,9 @@ mod tests { // (a, null) remaining: null let output = builder.take_n(2); - assert_eq!(&output, &array); + let expected = + Arc::new(LargeStringArray::from(vec![Some("a"), None])) as ArrayRef; + assert_eq!(&output, &expected); // null, a, null, a builder.append_val(&array, 0).unwrap(); @@ -487,8 +417,9 @@ mod tests { // (null, a) remaining: (null, a) let output = builder.take_n(2); - let array = Arc::new(StringArray::from(vec![None, Some("a")])) as ArrayRef; - assert_eq!(&output, &array); + let expected = + Arc::new(LargeStringArray::from(vec![None, Some("a")])) as ArrayRef; + assert_eq!(&output, &expected); let array = Arc::new(StringArray::from(vec![ Some("a"), @@ -503,18 +434,18 @@ mod tests { // (null, a, longstringfortest, null) remaining: (null) let output = builder.take_n(4); - let array = Arc::new(StringArray::from(vec![ + let expected = Arc::new(LargeStringArray::from(vec![ None, Some("a"), Some("longstringfortest"), None, ])) as ArrayRef; - assert_eq!(&output, &array); + assert_eq!(&output, &expected); } #[test] fn test_byte_equal_to() { - let append = |builder: &mut ByteGroupValueBuilder, + let append = |builder: &mut ByteGroupValueBuilder, builder_array: &ArrayRef, append_rows: &[usize]| { for &index in append_rows { @@ -523,7 +454,7 @@ mod tests { }; let equal_to = - |builder: &ByteGroupValueBuilder, + |builder: &ByteGroupValueBuilder, lhs_rows: &[usize], input_array: &ArrayRef, rhs_rows: &[usize], @@ -540,7 +471,7 @@ mod tests { #[test] fn test_byte_vectorized_equal_to() { - let append = |builder: &mut ByteGroupValueBuilder, + let append = |builder: &mut ByteGroupValueBuilder, builder_array: &ArrayRef, append_rows: &[usize]| { builder @@ -549,7 +480,7 @@ mod tests { }; let equal_to = - |builder: &ByteGroupValueBuilder, + |builder: &ByteGroupValueBuilder, lhs_rows: &[usize], input_array: &ArrayRef, rhs_rows: &[usize], @@ -570,7 +501,7 @@ mod tests { // Test the special `all nulls` or `not nulls` input array case // for vectorized append and equal to - let mut builder = ByteGroupValueBuilder::::new(OutputType::Utf8); + let mut builder = ByteGroupValueBuilder::new(OutputType::Utf8); // All nulls input array let all_nulls_input_array = Arc::new(StringArray::from(vec![ @@ -629,9 +560,9 @@ mod tests { fn test_byte_equal_to_internal(mut append: A, mut equal_to: E) where - A: FnMut(&mut ByteGroupValueBuilder, &ArrayRef, &[usize]), + A: FnMut(&mut ByteGroupValueBuilder, &ArrayRef, &[usize]), E: FnMut( - &ByteGroupValueBuilder, + &ByteGroupValueBuilder, &[usize], &ArrayRef, &[usize], @@ -647,7 +578,7 @@ mod tests { // - exist not null, input not null; values equal // Define ByteGroupValueBuilder - let mut builder = ByteGroupValueBuilder::::new(OutputType::Utf8); + let mut builder = ByteGroupValueBuilder::new(OutputType::Utf8); let builder_array = Arc::new(StringArray::from(vec![ None, None, diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index f275d777c3279..e3339720f45b1 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -1034,25 +1034,11 @@ fn make_group_column(field: &Field) -> Result> { DataType::Decimal128(_, _) => { instantiate_primitive!(v, nullable, Decimal128Type, data_type) } - DataType::Utf8 => { - v.push(Box::new(ByteGroupValueBuilder::::new( - OutputType::Utf8, - ))); + DataType::Utf8 | DataType::LargeUtf8 => { + v.push(Box::new(ByteGroupValueBuilder::new(OutputType::Utf8))); } - DataType::LargeUtf8 => { - v.push(Box::new(ByteGroupValueBuilder::::new( - OutputType::Utf8, - ))); - } - DataType::Binary => { - v.push(Box::new(ByteGroupValueBuilder::::new( - OutputType::Binary, - ))); - } - DataType::LargeBinary => { - v.push(Box::new(ByteGroupValueBuilder::::new( - OutputType::Binary, - ))); + DataType::Binary | DataType::LargeBinary => { + v.push(Box::new(ByteGroupValueBuilder::new(OutputType::Binary))); } DataType::Utf8View => { v.push(Box::new(ByteViewGroupValueBuilder::::new())); @@ -1259,7 +1245,9 @@ enum Nulls { mod tests { use std::{collections::HashMap, sync::Arc}; - use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray, StringViewArray}; + use arrow::array::{ + ArrayRef, Int64Array, LargeStringArray, RecordBatch, StringArray, StringViewArray, + }; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::{compute::concat_batches, util::pretty::pretty_format_batches}; use datafusion_common::utils::proxy::HashTableAllocExt; @@ -1791,10 +1779,12 @@ mod tests { Arc::new(col3) as _, ]; - // Expected batch + // Expected batch. Note the `Utf8` input column is emitted as + // `LargeUtf8`: string group keys are always accumulated with + // 64-bit offsets. let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::Int64, true), - Field::new("b", DataType::Utf8, true), + Field::new("b", DataType::LargeUtf8, true), Field::new("c", DataType::Utf8View, true), ])); @@ -1820,7 +1810,7 @@ mod tests { Some(34212), ]); - let col2 = StringArray::from(vec![ + let col2 = LargeStringArray::from(vec![ // Repeated rows in batch Some("string1"), None, diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 0d00e5c4d0d86..173ca26892ad0 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -29,7 +29,8 @@ use crate::aggregates::order::GroupOrderingFull; use crate::aggregates::{ AggregateInputMode, AggregateMode, AggregateOutputMode, PhysicalGroupBy, create_schema, evaluate_group_by, evaluate_many, evaluate_optional, group_id_array, - max_duplicate_ordinal, + max_duplicate_ordinal, narrow_group_key_columns, widen_group_key_arrays, + widen_group_key_schema, }; use crate::metrics::{BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput}; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; @@ -279,6 +280,13 @@ pub(crate) struct GroupedHashAggregateStream { // the execution. // ======================================================================== schema: SchemaRef, + /// `schema` with string/binary group key columns widened to their + /// internal 64-bit offset representation (see + /// [`internal_group_key_type`](crate::aggregates::internal_group_key_type)). + /// + /// Batches built by [`Self::emit`] use this schema; they are narrowed + /// back to `schema` after being sliced to `batch_size` rows for output. + emit_schema: SchemaRef, input_schema: SchemaRef, input: SendableRecordBatchStream, mode: AggregateMode, @@ -440,12 +448,23 @@ impl GroupedHashAggregateStream { // Therefore, when we spill these intermediate states or pass them to another // aggregation operator, we must use a schema that includes both the group // columns **and** the partial-state columns. - let spill_schema = Arc::new(create_schema( - &agg.input().schema(), - &agg_group_by, - &aggregate_exprs, - AggregateMode::Partial, - )?); + // + // Spilled batches hold the group keys as emitted by `group_values`, + // so the group key columns use the widened internal representation. + let spill_schema = widen_group_key_schema( + &Arc::new(create_schema( + &agg.input().schema(), + &agg_group_by, + &aggregate_exprs, + AggregateMode::Partial, + )?), + group_schema.fields().len(), + ); + + // Batches built by `emit()` hold the group keys as emitted by + // `group_values` and are narrowed to `agg_schema` on output. + let emit_schema = + widen_group_key_schema(&agg_schema, group_schema.fields().len()); // Need to update the GROUP BY expressions to point to the correct column after schema change let merging_group_by_expr = agg_group_by @@ -590,6 +609,7 @@ impl GroupedHashAggregateStream { Ok(GroupedHashAggregateStream { schema: agg_schema, + emit_schema, input_schema: agg.input().schema(), input, mode: agg.mode, @@ -806,6 +826,16 @@ impl Stream for GroupedHashAggregateStream { reduction_factor.add_part(output_batch.num_rows()); } + // Narrow internally-widened group key columns back to the + // output schema types. Each slice holds at most + // `batch_size` rows, so unlike the full emitted batch, the + // narrow representation fits. + let output_batch = + match narrow_group_key_columns(output_batch, &self.schema) { + Ok(batch) => batch, + Err(e) => return Poll::Ready(Some(Err(e))), + }; + // Empty record batches should not be emitted. // They need to be treated as [`Option`]es and handled separately debug_assert!(output_batch.num_rows() > 0); @@ -850,6 +880,13 @@ impl GroupedHashAggregateStream { evaluate_group_by(&self.group_by, batch)? }; + // Widen string/binary group keys to the internal representation that + // `group_values` interns (see `PhysicalGroupBy::group_schema`). + let group_by_values = group_by_values + .into_iter() + .map(widen_group_key_arrays) + .collect::>>()?; + // Only create the timer if there are actual aggregate arguments to evaluate let timer = match ( self.spill_state.is_stream_merging, @@ -1026,7 +1063,7 @@ impl GroupedHashAggregateStream { let schema = if spilling { Arc::clone(&self.spill_state.spill_schema) } else { - self.schema() + Arc::clone(&self.emit_schema) }; if self.group_values.is_empty() { return Ok(None); @@ -1415,6 +1452,55 @@ mod tests { // Migrated to PartialHashAggregateStream coverage in hash_stream.rs; // kept here for the legacy GroupedHashAggregateStream implementation. + + #[test] + fn narrow_group_key_columns_rebases_sliced_offsets() -> Result<()> { + let wide_schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::LargeUtf8, true), + Field::new("c", DataType::Int64, false), + ])); + let narrow_schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Utf8, true), + Field::new("c", DataType::Int64, false), + ])); + assert_eq!(widen_group_key_schema(&narrow_schema, 1), wide_schema); + // Group keys past the group column prefix are left alone + assert_eq!(widen_group_key_schema(&narrow_schema, 0), narrow_schema); + + let batch = RecordBatch::try_new( + Arc::clone(&wide_schema), + vec![ + Arc::new(LargeStringArray::from(vec![ + Some("aa"), + None, + Some("cccc"), + Some("dd"), + ])), + Arc::new(Int64Array::from(vec![1, 2, 3, 4])), + ], + )?; + + // Slice so the string column's offsets do not start at zero, as + // happens when an emitted batch is sliced into `batch_size` chunks + let narrowed = + narrow_group_key_columns(batch.slice(1, 3), &Arc::clone(&narrow_schema))?; + assert_eq!(narrowed.schema(), narrow_schema); + let k = narrowed.column(0).as_string::(); + assert!(k.is_null(0)); + assert_eq!(k.value(1), "cccc"); + assert_eq!(k.value(2), "dd"); + let c = narrowed + .column(1) + .as_primitive::(); + assert_eq!(c.values(), &[2, 3, 4]); + + // A batch already matching the target schema passes through untouched + let untouched = narrow_group_key_columns(batch.clone(), &wide_schema)?; + assert_eq!(untouched.schema(), wide_schema); + + Ok(()) + } + #[tokio::test] async fn test_double_emission_race_condition_bug() -> Result<()> { // Fix for https://github.com/apache/datafusion/issues/18701 diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 732da32ab0391..5de258db499be 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -172,7 +172,10 @@ use datafusion_physical_expr::utils::collect_columns; use parking_lot::Mutex; use std::collections::{HashMap, HashSet}; -use arrow::array::{ArrayRef, UInt8Array, UInt16Array, UInt32Array, UInt64Array}; +use arrow::array::{ + ArrayRef, MutableArrayData, UInt8Array, UInt16Array, UInt32Array, UInt64Array, + make_array, +}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use arrow_schema::FieldRef; @@ -552,8 +555,27 @@ impl PhysicalGroupBy { Aggregate::grouping_id_type(self.expr.len(), max_duplicate_ordinal(&self.groups)) } + /// Returns the schema of the group values as they are interned and + /// emitted by the group values machinery of the aggregation streams. + /// + /// This differs from the output schema in that string and binary group + /// keys are widened to their internal 64-bit offset representation (see + /// `internal_group_key_type`). Output batches are narrowed back to the + /// declared output types after being sliced to `batch_size` rows. pub fn group_schema(&self, schema: &Schema) -> Result { - Ok(Arc::new(Schema::new(self.group_fields(schema)?))) + let fields = self + .group_fields(schema)? + .into_iter() + .map(|field| { + let wide = internal_group_key_type(field.data_type()); + if wide == *field.data_type() { + field + } else { + Arc::new(field.as_ref().clone().with_data_type(wide)) + } + }) + .collect::>(); + Ok(Arc::new(Schema::new(fields))) } /// Returns the fields that are used as the grouping keys. @@ -2187,6 +2209,110 @@ impl ExecutionPlan for AggregateExec { /// Creates the output schema for an [`AggregateExec`] containing the group by columns followed /// by the aggregate columns. +/// The type used internally to intern and emit a group-by key of the given +/// input type. +/// +/// String and binary group keys are accumulated with 64-bit offsets +/// (`LargeUtf8` / `LargeBinary`) so that the single array holding all +/// distinct group keys is not limited to `i32::MAX` total bytes. This is an +/// implementation detail of the aggregation streams: output batches are +/// narrowed back to the declared output types after being sliced to +/// `batch_size` rows, where the narrow representation always fits. +pub(crate) fn internal_group_key_type(input_type: &DataType) -> DataType { + match input_type { + DataType::Utf8 => DataType::LargeUtf8, + DataType::Binary => DataType::LargeBinary, + other => other.clone(), + } +} + +/// Returns `schema` with any string/binary fields among the leading +/// `num_group_columns` group key fields widened to the internal group key +/// representation (see [`internal_group_key_type`]). +/// +/// Returns the original schema when nothing needs widening. +pub(crate) fn widen_group_key_schema( + schema: &SchemaRef, + num_group_columns: usize, +) -> SchemaRef { + let mut changed = false; + let fields = schema + .fields() + .iter() + .enumerate() + .map(|(idx, field)| { + if idx < num_group_columns { + let wide = internal_group_key_type(field.data_type()); + if wide != *field.data_type() { + changed = true; + return Arc::new(field.as_ref().clone().with_data_type(wide)); + } + } + Arc::clone(field) + }) + .collect::>(); + + if changed { + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) + } else { + Arc::clone(schema) + } +} + +/// Widens string/binary group key arrays to the internal representation +/// interned by `group_values` (see `PhysicalGroupBy::group_schema`). +pub(crate) fn widen_group_key_arrays(arrays: Vec) -> Result> { + arrays + .into_iter() + .map(|array| { + let wide = internal_group_key_type(array.data_type()); + if wide == *array.data_type() { + Ok(array) + } else { + Ok(arrow::compute::cast(&array, &wide)?) + } + }) + .collect() +} + +/// Narrows internally-widened group key columns of `batch` back to the types +/// declared in `schema`. +/// +/// `batch` is expected to hold at most `batch_size` rows (it is narrowed +/// after the emitted batch is sliced for output); if a single column slice +/// still holds more than `i32::MAX` bytes an error is returned. +pub(crate) fn narrow_group_key_columns( + batch: RecordBatch, + schema: &SchemaRef, +) -> Result { + if batch.schema().fields() == schema.fields() { + return Ok(batch); + } + + let columns = batch + .columns() + .iter() + .zip(schema.fields()) + .map(|(array, field)| { + if array.data_type() == field.data_type() { + Ok(Arc::clone(array)) + } else { + // The array is typically a slice referencing the buffers of + // the full emitted batch; compact it so its offsets are + // rebased into a minimal buffer before converting them to the + // narrow offset width. + let data = array.to_data(); + let mut mutable = MutableArrayData::new(vec![&data], false, data.len()); + mutable.try_extend(0, 0, data.len())?; + let compacted = make_array(mutable.freeze()); + Ok(arrow::compute::cast(&compacted, field.data_type())?) + } + }) + .collect::>>()?; + + Ok(RecordBatch::try_new(Arc::clone(schema), columns)?) +} + fn create_schema( input_schema: &Schema, group_by: &PhysicalGroupBy,