From 880064faac82d68e9d5047c63a68b00c097c8f19 Mon Sep 17 00:00:00 2001 From: Max Burke Date: Fri, 17 Jul 2026 14:39:52 -0700 Subject: [PATCH 1/5] fix: Allow more than 2 GiB of data in group keys Store Utf8 and Binary group keys with 64-bit offsets. This removes the 2 GiB limit on the combined key data in an aggregate table. Divide the output into batches, then change the keys to their declared types. Each output column must still fit the 32-bit offset limit. Keys inside lists or dictionaries keep their types. Use the same schema helper for each aggregate table. Use the existing stream adapter to divide large batches and restore the output types. --- .../aggregates/aggregate_hash_table/common.rs | 32 ++- .../aggregate_hash_table/common_clustered.rs | 25 +- .../src/aggregates/clustered_final_stream.rs | 3 + .../aggregates/clustered_partial_stream.rs | 3 + .../src/aggregates/clustered_single_stream.rs | 6 + .../src/aggregates/grouped_hash_stream.rs | 39 ++- .../src/aggregates/hash_stream.rs | 7 +- .../physical-plan/src/aggregates/mod.rs | 243 +++++++++++++++++- .../src/aggregates/partial_reduce_stream.rs | 7 +- .../physical-plan/src/aggregates/spill.rs | 4 +- 10 files changed, 344 insertions(+), 25 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 fa7a134733194..54a60c0e46b04 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -29,7 +29,8 @@ use crate::aggregates::group_values::{ use crate::aggregates::order::GroupClustering; use crate::aggregates::{ AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, - group_id_array, max_duplicate_ordinal, + group_id_array, group_keys_with_large_offsets, group_schema_with_large_offsets, + max_duplicate_ordinal, restore_group_key_types, }; use super::AggregateTableMetrics; @@ -200,9 +201,13 @@ impl AggregateHashTable { ) -> Result { let state = self.state.building_mut(); // Outer vec: one per grouping set; inner vec: group-by expressions. - let grouping_set_args = self - .group_by_metrics - .time_group_key_preparation(|| evaluate_group_by(&state.group_by, batch))?; + let grouping_set_args = + self.group_by_metrics.time_group_key_preparation(|| { + evaluate_group_by(&state.group_by, batch)? + .into_iter() + .map(group_keys_with_large_offsets) + .collect::>>() + })?; // The evaluated args for each accumulator. let accumulator_args = self.group_by_metrics.time_aggregate_arguments(|| { @@ -309,7 +314,11 @@ impl AggregateHashTable { Ok::<_, datafusion_common::DataFusionError>(columns) })?; - let batch = RecordBatch::try_new(output_schema, columns)?; + let emit_schema = group_schema_with_large_offsets( + &output_schema, + state.group_by.num_group_exprs(), + ); + let batch = RecordBatch::try_new(emit_schema, columns)?; debug_assert!(batch.num_rows() > 0); MaterializedAggregateOutput::new(batch) } @@ -322,7 +331,10 @@ impl AggregateHashTable { } }; - let batch = output.next_batch(batch_size)?; + let batch = output + .next_batch(batch_size)? + .map(|batch| restore_group_key_types(batch, &output_schema)) + .transpose()?; if output.is_exhausted() { self.state = AggregateHashTableState::Done; } else { @@ -362,6 +374,10 @@ impl AggregateHashTable { /// Unlike normal single aggregation output, this materializes intermediate /// states rather than final values. The states can therefore be merged after /// spilling without finalizing the same group more than once. + /// + /// The keys use 64-bit offsets because this batch can contain more than + /// 2 GiB of key data. + /// Callers divide the batch, then change the keys to the output types. pub(in crate::aggregates) fn take_state_batch( &mut self, ) -> Result> { @@ -385,6 +401,10 @@ impl AggregateHashTable { Ok::<_, datafusion_common::DataFusionError>(output) })?; + let state_schema = group_schema_with_large_offsets( + &state_schema, + state.group_by.num_group_exprs(), + ); let batch = RecordBatch::try_new(state_schema, output)?; debug_assert!(batch.num_rows() > 0); diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs index 52ac5f7c761f7..717b5da3cbd83 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs @@ -22,7 +22,7 @@ use crate::aggregates::group_values::{ use crate::aggregates::order::{GroupClustering, GroupClusteringMode}; use crate::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, - evaluate_group_by, + evaluate_group_by, group_keys_with_large_offsets, group_schema_with_large_offsets, }; use super::AggregateTableMetrics; @@ -234,7 +234,10 @@ impl ClusteredAggregateTable { ) -> Result { let grouping_set_args = self.group_by_metrics.time_group_key_preparation(|| { - evaluate_group_by(&self.buffer.group_by, batch) + evaluate_group_by(&self.buffer.group_by, batch)? + .into_iter() + .map(group_keys_with_large_offsets) + .collect::>>() })?; let accumulator_args = self.group_by_metrics.time_aggregate_arguments(|| { @@ -320,6 +323,10 @@ impl ClusteredAggregateTable { /// active (incomplete) groups. Partial aggregation can pass those states to /// its final stage, while single and final aggregation sort and spill them /// before replay. + /// + /// The keys use 64-bit offsets because this batch can contain more than + /// 2 GiB of key data. + /// Callers divide the batch, then change the keys to the output types. pub(in crate::aggregates) fn take_state_batch( &mut self, ) -> Result> { @@ -340,7 +347,11 @@ impl ClusteredAggregateTable { Ok::<_, datafusion_common::DataFusionError>(output) })?; - let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?; + let state_schema = group_schema_with_large_offsets( + &self.state_schema, + self.buffer.group_by.num_group_exprs(), + ); + let batch = RecordBatch::try_new(state_schema, output)?; debug_assert!(batch.num_rows() > 0); // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the @@ -416,6 +427,8 @@ impl ClusteredAggregateTable { /// Removes the selected groups once and materializes their output columns. /// The caller chooses the completed prefix. + /// + /// The keys keep 64-bit offsets, as in [`Self::take_state_batch`]. pub(super) fn materialize_groups( &mut self, emit_to: EmitTo, @@ -441,7 +454,11 @@ impl ClusteredAggregateTable { Ok::<_, datafusion_common::DataFusionError>(output) })?; - let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; + let output_schema = group_schema_with_large_offsets( + &self.output_schema, + self.buffer.group_by.num_group_exprs(), + ); + let batch = RecordBatch::try_new(output_schema, output)?; debug_assert!(batch.num_rows() > 0); Ok(batch) diff --git a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs index 6eeb37a475ffe..17fc62a0ee6e2 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs @@ -19,6 +19,7 @@ use super::aggregate_hash_table::{ ClusteredAggregateTable, ClusteredAggregateTableMetrics, FinalMarker, }; use super::order::GroupClusteringMode; +use super::restore_group_key_stream; use super::spill::AggregateSpill; use crate::SendableRecordBatchStream; use crate::aggregates::AggregateMode; @@ -272,6 +273,7 @@ impl ClusteredFinalAggregateStream { stage, } = self; let schema = Arc::clone(&context.schema); + let batch_size = context.batch_size; let metrics = context.baseline_metrics.clone(); let stream = async_try_stream(|mut emitter| async move { let mut stage = Some(stage); @@ -296,6 +298,7 @@ impl ClusteredFinalAggregateStream { Ok(()) }); let stream = Box::pin(RecordBatchStreamAdapter::new(schema, stream)); + let stream = restore_group_key_stream(stream, batch_size); Box::pin(ObservedStream::new(stream, metrics, None)) } } diff --git a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs index 29a2e25a34022..eb83a3451ac34 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs @@ -15,6 +15,7 @@ use futures::stream::StreamExt; use super::AggregateExec; use super::aggregate_hash_table::{ClusteredAggregateTable, PartialMarker}; +use super::restore_group_key_stream; use crate::aggregates::AggregateMode; use crate::aggregates::order::{GroupClustering, GroupClusteringMode}; use crate::metrics::{BaselineMetrics, MetricBuilder, SpillMetrics}; @@ -269,6 +270,7 @@ impl ClusteredPartialAggregateStream { stage, } = self; let schema = Arc::clone(&context.schema); + let batch_size = context.batch_size; let metrics = context.baseline_metrics.clone(); let stream = async_try_stream(|mut emitter| async move { let mut stage = Some(stage); @@ -287,6 +289,7 @@ impl ClusteredPartialAggregateStream { Ok(()) }); let stream = Box::pin(RecordBatchStreamAdapter::new(schema, stream)); + let stream = restore_group_key_stream(stream, batch_size); Box::pin(ObservedStream::new(stream, metrics, None)) } } diff --git a/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs index 4163925249b38..5c2f801cdd3ca 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_single_stream.rs @@ -17,6 +17,7 @@ use futures::stream::{Stream, StreamExt}; use super::aggregate_hash_table::{ClusteredAggregateTable, SingleMarker}; use super::order::GroupClusteringMode; +use super::restore_group_key_stream; use super::spill::AggregateSpill; use super::{AggregateExec, create_schema}; use crate::aggregates::AggregateMode; @@ -209,6 +210,11 @@ impl ClusteredSingleAggregateStream { }) } + pub(super) fn into_stream(self) -> SendableRecordBatchStream { + let batch_size = self.batch_size; + restore_group_key_stream(Box::pin(self), batch_size) + } + fn close_input(&mut self) { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 1ee6d02d13e8e..9bea9af8dd6e3 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -33,7 +33,8 @@ use crate::aggregates::order::GroupClusteringFull; use crate::aggregates::{ AggregateInputMode, AggregateMode, AggregateOutputMode, PhysicalGroupBy, aggregate_metric_label, create_schema, evaluate_group_by, evaluate_optional, - group_id_array, max_duplicate_ordinal, + group_id_array, group_keys_with_large_offsets, group_schema_with_large_offsets, + max_duplicate_ordinal, restore_group_key_types, }; use crate::metrics::{BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput}; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; @@ -276,6 +277,9 @@ pub(crate) struct GroupedHashAggregateStream { // the execution. // ======================================================================== schema: SchemaRef, + /// Batches from [`Self::emit`] can contain more key data than the + /// 32-bit offset limit permits. + emit_schema: SchemaRef, input_schema: SchemaRef, input: SendableRecordBatchStream, mode: AggregateMode, @@ -466,12 +470,19 @@ 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, - )?); + // The keys in files on disk keep 64-bit offsets. + let spill_schema = group_schema_with_large_offsets( + &Arc::new(create_schema( + &agg.input().schema(), + &agg_group_by, + &aggregate_exprs, + AggregateMode::Partial, + )?), + group_schema.fields().len(), + ); + + let emit_schema = + group_schema_with_large_offsets(&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 @@ -618,6 +629,7 @@ impl GroupedHashAggregateStream { Ok(GroupedHashAggregateStream { schema: agg_schema, + emit_schema, input_schema: agg.input().schema(), input, mode: agg.mode, @@ -817,6 +829,12 @@ impl Stream for GroupedHashAggregateStream { reduction_factor.add_part(output_batch.num_rows()); } + let output_batch = + match restore_group_key_types(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); @@ -866,6 +884,11 @@ impl GroupedHashAggregateStream { } })?; + let group_by_values = group_by_values + .into_iter() + .map(group_keys_with_large_offsets) + .collect::>>()?; + // Evaluate aggregate arguments and filters together. Per-aggregate // argument timers remain limited to their argument expressions because // filters are evaluated collectively here. @@ -1046,7 +1069,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); diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index ca85ccb59f6bb..21e475a7a74c3 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -25,6 +25,7 @@ use super::aggregate_hash_table::{ PartialSkipMarker, }; use super::order::GroupClusteringMode; +use super::restore_group_key_stream; use super::skip_partial::SkipAggregationProbe; use super::spill::AggregateSpill; use crate::metrics::{ @@ -275,8 +276,12 @@ impl PartialHashAggregateStream { pub(crate) fn into_stream(self) -> SendableRecordBatchStream { let schema = Arc::clone(&self.schema); + let batch_size = self.batch_size; - Box::pin(RecordBatchStreamAdapter::new(schema, self.create_stream())) + restore_group_key_stream( + Box::pin(RecordBatchStreamAdapter::new(schema, self.create_stream())), + batch_size, + ) } /// Entry point for the partial hash aggregate state machine. diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 0943c780e117b..edd4e07ffe360 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -164,10 +164,14 @@ use crate::{ Partitioning, SendableRecordBatchStream, Statistics, }; use datafusion_common::config::ConfigOptions; +use futures::StreamExt; use parking_lot::Mutex; use std::collections::HashMap; -use arrow::array::{ArrayRef, UInt8Array, UInt16Array, UInt32Array, UInt64Array}; +use arrow::array::{ + ArrayRef, MutableArrayData, UInt8Array, UInt16Array, UInt32Array, UInt64Array, + make_array, +}; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; @@ -178,7 +182,7 @@ use datafusion_common::{ ColumnStatistics, Constraint, Constraints, Result, ScalarValue, assert_eq_or_internal_err, internal_err, not_impl_err, }; -use datafusion_execution::TaskContext; +use datafusion_execution::{TaskContext, async_try_stream}; use datafusion_expr::{Accumulator, Aggregate, AggregateMetrics}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use datafusion_physical_expr::equivalence::ProjectionMapping; @@ -554,8 +558,16 @@ impl PhysicalGroupBy { Aggregate::grouping_id_type(self.expr.len(), max_duplicate_ordinal(&self.groups)) } + /// Returns the schema for stored group keys. + /// + /// [`DataType::Utf8`] and [`DataType::Binary`] keys use 64-bit offsets to + /// allow more than 2 GiB of key data. Output batches use the declared types. pub fn group_schema(&self, schema: &Schema) -> Result { - Ok(Arc::new(Schema::new(self.group_fields(schema)?))) + let schema = Arc::new(Schema::new(self.group_fields(schema)?)); + Ok(group_schema_with_large_offsets( + &schema, + schema.fields().len(), + )) } /// Returns the fields that are used as the grouping keys. @@ -701,7 +713,7 @@ impl From for SendableRecordBatchStream { StreamType::SingleHash(stream) => stream.into_stream(), StreamType::ClusteredPartialAggregate(stream) => stream.into_stream(), StreamType::ClusteredFinalAggregate(stream) => stream.into_stream(), - StreamType::ClusteredSingleAggregate(stream) => Box::pin(stream), + StreamType::ClusteredSingleAggregate(stream) => stream.into_stream(), StreamType::GroupedHash(stream) => Box::pin(stream), StreamType::GroupedPriorityQueue(stream) => Box::pin(stream), } @@ -2848,6 +2860,141 @@ impl AggregateExec { } } +fn group_key_type_with_large_offsets(data_type: &DataType) -> Option { + match data_type { + DataType::Utf8 => Some(DataType::LargeUtf8), + DataType::Binary => Some(DataType::LargeBinary), + _ => None, + } +} + +fn group_schema_with_large_offsets( + schema: &SchemaRef, + num_group_columns: usize, +) -> SchemaRef { + if !schema + .fields() + .iter() + .take(num_group_columns) + .any(|field| group_key_type_with_large_offsets(field.data_type()).is_some()) + { + return Arc::clone(schema); + } + + let fields = schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| { + if index < num_group_columns + && let Some(data_type) = + group_key_type_with_large_offsets(field.data_type()) + { + Arc::new(field.as_ref().clone().with_data_type(data_type)) + } else { + Arc::clone(field) + } + }) + .collect::>(); + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) +} + +fn group_keys_with_large_offsets(arrays: Vec) -> Result> { + arrays + .into_iter() + .map( + |array| match group_key_type_with_large_offsets(array.data_type()) { + Some(data_type) => Ok(arrow::compute::cast(&array, &data_type)?), + None => Ok(array), + }, + ) + .collect() +} + +fn group_key_types_differ(batch: &RecordBatch, schema: &Schema) -> bool { + batch + .columns() + .iter() + .zip(schema.fields()) + .any(|(column, field)| column.data_type() != field.data_type()) +} + +// Each output column must fit the byte limit of its declared type. +fn restore_group_key_types( + batch: RecordBatch, + schema: &SchemaRef, +) -> Result { + if !group_key_types_differ(&batch, schema) { + 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 { + // A slice can start beyond the 32-bit offset limit. + // Copy its values so that the first offset is zero. + let data = array.to_data(); + let mut mutable = MutableArrayData::new(vec![&data], false, data.len()); + mutable.try_extend(0, 0, data.len())?; + Ok(arrow::compute::cast( + &make_array(mutable.freeze()), + field.data_type(), + )?) + } + }) + .collect::>>()?; + Ok(RecordBatch::try_new(Arc::clone(schema), columns)?) +} + +fn restore_group_key_stream( + mut input: SendableRecordBatchStream, + batch_size: usize, +) -> SendableRecordBatchStream { + let schema = input.schema(); + if !schema + .fields() + .iter() + .any(|field| group_key_type_with_large_offsets(field.data_type()).is_some()) + { + return input; + } + + let output_schema = Arc::clone(&schema); + let stream = async_try_stream(move |mut emitter| async move { + while let Some(batch) = input.next().await.transpose()? { + if !group_key_types_differ(&batch, &schema) { + emitter.emit(batch).await; + continue; + } + if batch.num_rows() <= batch_size { + emitter.emit(restore_group_key_types(batch, &schema)?).await; + continue; + } + // When memory is not sufficient, the input can return all groups at once. + // Divide the batch before conversion to limit the bytes in each column. + for offset in (0..batch.num_rows()).step_by(batch_size) { + let length = batch_size.min(batch.num_rows() - offset); + emitter + .emit(restore_group_key_types( + batch.slice(offset, length), + &schema, + )?) + .await; + } + } + Ok(()) + }); + Box::pin(crate::stream::RecordBatchStreamAdapter::new( + output_schema, + stream, + )) +} + /// Creates the output schema for an [`AggregateExec`] containing the group by columns followed /// by the aggregate columns. fn create_schema( @@ -9984,4 +10131,92 @@ mod tests { Ok(()) } + + #[test] + fn group_schema_uses_large_offsets() -> Result<()> { + let input_schema = Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("b", DataType::Binary, true), + Field::new("ls", DataType::LargeUtf8, true), + Field::new("v", DataType::Utf8View, true), + Field::new("i", DataType::Int64, true), + ]); + let group_by = PhysicalGroupBy::new_single( + input_schema + .fields() + .iter() + .map(|f| Ok((col(f.name(), &input_schema)?, f.name().clone()))) + .collect::>()?, + ); + + let group_types = group_by + .group_schema(&input_schema)? + .fields() + .iter() + .map(|f| f.data_type().clone()) + .collect::>(); + assert_eq!( + group_types, + vec![ + DataType::LargeUtf8, + DataType::LargeBinary, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Int64, + ] + ); + Ok(()) + } + + #[test] + fn restore_group_key_types_from_slice() -> Result<()> { + let stored_schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::LargeUtf8, true), + Field::new("c", DataType::Int64, false), + ])); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Utf8, true), + Field::new("c", DataType::Int64, false), + ])); + assert_eq!( + group_schema_with_large_offsets(&output_schema, 1), + stored_schema + ); + assert_eq!( + group_schema_with_large_offsets(&output_schema, 0), + output_schema + ); + + let batch = RecordBatch::try_new( + Arc::clone(&stored_schema), + vec![ + Arc::new(arrow::array::LargeStringArray::from(vec![ + Some("aa"), + None, + Some("cccc"), + Some("dd"), + ])), + Arc::new(Int64Array::from(vec![1, 2, 3, 4])), + ], + )?; + + let output = restore_group_key_types(batch.slice(1, 3), &output_schema)?; + assert_eq!(output.schema(), output_schema); + let k = output.column(0).as_string::(); + assert_eq!(k.value_offsets(), &[0, 0, 4, 6]); + assert_eq!(k.value_data(), b"ccccdd"); + assert_eq!( + k.iter().collect::>(), + vec![None, Some("cccc"), Some("dd")] + ); + let c = output.column(1).as_primitive::(); + assert_eq!(c.values(), &[2, 3, 4]); + + let unchanged = restore_group_key_types(batch.clone(), &stored_schema)?; + assert_eq!(unchanged.schema(), stored_schema); + assert!(Arc::ptr_eq(unchanged.column(0), batch.column(0))); + + Ok(()) + } + } diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index 1d003af37121a..0c21191214024 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -16,6 +16,7 @@ use futures::stream::{Stream, StreamExt}; use super::AggregateExec; use super::GroupClusteringMode; use super::aggregate_hash_table::{AggregateHashTable, PartialReduceMarker}; +use super::restore_group_key_stream; use crate::SendableRecordBatchStream; use crate::metrics::{BaselineMetrics, Count, MetricBuilder, RecordOutput, SpillMetrics}; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; @@ -145,8 +146,12 @@ impl PartialReduceHashAggregateStream { pub(crate) fn into_stream(self) -> SendableRecordBatchStream { let schema = Arc::clone(&self.schema); + let batch_size = self.batch_size; - Box::pin(RecordBatchStreamAdapter::new(schema, self.create_stream())) + restore_group_key_stream( + Box::pin(RecordBatchStreamAdapter::new(schema, self.create_stream())), + batch_size, + ) } /// Entry point for the partial reduce hash aggregate. diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index c08a76adc2153..138263f125c83 100644 --- a/datafusion/physical-plan/src/aggregates/spill.rs +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -19,7 +19,7 @@ use futures::StreamExt; use super::aggregate_hash_table::ClusteredAggregateTableMetrics; use super::clustered_final_stream::ClusteredFinalAggregateStream; use super::order::GroupClusteringMode; -use super::{AggregateExec, AggregateMode}; +use super::{AggregateExec, AggregateMode, restore_group_key_types}; use crate::SendableRecordBatchStream; use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::sorts::IncrementalSortIterator; @@ -223,6 +223,7 @@ impl AggregateSpill { return Ok(()); }; + let spill_schema = Arc::clone(self.spill_manager.schema()); let mut max_batch_rows = 0; let sorted_iter = IncrementalSortIterator::new( state_batch, @@ -230,6 +231,7 @@ impl AggregateSpill { self.batch_size, ) .with_max_batch_bytes(SPILL_BATCH_TARGET_BYTES) + .map(|batch| restore_group_key_types(batch?, &spill_schema)) .inspect(|batch| { if let Ok(batch) = batch { max_batch_rows = max_batch_rows.max(batch.num_rows()); From 6e7f2bdad735aaef4cca860ddc13e4ed07d9dfe3 Mon Sep 17 00:00:00 2001 From: Max Burke Date: Tue, 10 Nov 2026 19:52:58 -0700 Subject: [PATCH 2/5] test: Check string and binary group keys in aggregate streams Check that each aggregate stream keeps the declared key types and returns the expected results. Check sorted and unsorted input, one and two stages, grouping sets, and groups that the stream writes to disk. The memory limit allows all four final streams to start. Check output under a memory limit before the input ends. Check streams that can and cannot reserve memory for the output batch. Also check empty batches and batches that need no type conversion. --- .../physical-plan/src/aggregates/mod.rs | 156 +++++++ .../test_files/aggregate_byte_group_keys.slt | 420 ++++++++++++++++++ 2 files changed, 576 insertions(+) create mode 100644 datafusion/sqllogictest/test_files/aggregate_byte_group_keys.slt diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index edd4e07ffe360..85262c53f41cd 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -10219,4 +10219,160 @@ mod tests { Ok(()) } + #[tokio::test] + async fn restore_group_key_stream_splits_large_batches() -> Result<()> { + let output_schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("b", DataType::Binary, true), + Field::new("c", DataType::Int64, false), + ])); + let stored_schema = group_schema_with_large_offsets(&output_schema, 2); + + let keys = (0..10) + .map(|i| (i % 4 != 0).then(|| format!("key-{i}"))) + .collect::>(); + let expected = RecordBatch::try_new( + Arc::clone(&output_schema), + vec![ + Arc::new(StringArray::from(keys.clone())), + Arc::new(arrow::array::BinaryArray::from_iter( + keys.iter().map(|k| k.as_ref().map(|k| k.as_bytes())), + )), + Arc::new(Int64Array::from_iter_values(0..10)), + ], + )?; + let stored = RecordBatch::try_new( + Arc::clone(&stored_schema), + vec![ + arrow::compute::cast(expected.column(0), &DataType::LargeUtf8)?, + arrow::compute::cast(expected.column(1), &DataType::LargeBinary)?, + Arc::clone(expected.column(2)), + ], + )?; + + let input = Box::pin(crate::stream::RecordBatchStreamAdapter::new( + Arc::clone(&output_schema), + futures::stream::iter(vec![ + Ok(RecordBatch::new_empty(stored_schema)), + Ok(stored), + Ok(expected.clone()), + ]), + )); + let output = collect(restore_group_key_stream(input, 4)).await?; + + let rows = output.iter().map(|b| b.num_rows()).collect::>(); + assert_eq!(rows, vec![0, 4, 4, 2, 10]); + assert!(output.iter().all(|batch| batch.schema() == output_schema)); + assert_eq!(concat_batches(&output_schema, &output[1..4])?, expected); + assert_eq!(output[4], expected); + + Ok(()) + } + + #[rstest::rstest] + #[case::reserves_batch_memory(320 * 1024, true)] + #[case::cannot_reserve_batch_memory(64 * 1024, false)] + #[tokio::test] + async fn partial_hash_restores_key_types_with_limited_memory( + #[case] memory_limit: usize, + #[case] has_reserved_memory: bool, + ) -> Result<()> { + let batch_size = 1024; + let num_groups = 4 * batch_size; + + let schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, false), + Field::new("b", DataType::Binary, false), + Field::new("v", DataType::Int64, false), + ])); + let keys = (0..num_groups) + .map(|i| format!("key-{i:04}")) + .collect::>(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from_iter_values(&keys)), + Arc::new(arrow::array::BinaryArray::from_iter_values(&keys)), + Arc::new(Int64Array::from(vec![1; num_groups])), + ], + )?; + + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(memory_limit, 1.0) + .build_arc()?; + // The test must group the input rows to check the key conversion. + let session_config = SessionConfig::new().with_batch_size(batch_size).set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(1.0)), + ); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(Arc::clone(&runtime)), + ); + + // The input waits at the barrier. + // Only the memory limit can cause output before that point. + let input = Arc::new( + crate::test::exec::BarrierExec::new(vec![vec![batch]], Arc::clone(&schema)) + .without_start_barrier() + .with_finish_barrier() + .with_log(false), + ); + let aggregate = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![ + (col("s", &schema)?, "s".to_string()), + (col("b", &schema)?, "b".to_string()), + ]), + vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("v", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("c") + .build()?, + )], + vec![None], + Arc::clone(&input) as Arc, + Arc::clone(&schema), + )?; + let mut stream = + PartialHashAggregateStream::new(&aggregate, &task_ctx, 0)?.into_stream(); + + let mut batches = vec![]; + let mut num_rows = 0; + while num_rows < num_groups { + let batch = + tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("the memory limit must cause output before the input ends") + .expect("the stream ended before all groups were returned")?; + if batches.is_empty() { + assert_eq!( + runtime.memory_pool.reserved() >= batch.get_array_memory_size(), + has_reserved_memory + ); + } + num_rows += batch.num_rows(); + batches.push(batch); + } + input.wait_finish().await; + assert!(stream.next().await.is_none()); + + for batch in &batches { + assert_eq!(batch.schema(), aggregate.schema()); + assert_eq!(batch.num_rows(), batch_size); + } + let output = concat_batches(&aggregate.schema(), &batches)?; + let strings = output.column(0).as_string::(); + assert!(strings.iter().flatten().eq(keys.iter().map(String::as_str))); + let binaries = output.column(1).as_binary::(); + assert!( + binaries + .iter() + .flatten() + .eq(keys.iter().map(String::as_bytes)) + ); + + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/aggregate_byte_group_keys.slt b/datafusion/sqllogictest/test_files/aggregate_byte_group_keys.slt new file mode 100644 index 0000000000000..9f5a921f070b7 --- /dev/null +++ b/datafusion/sqllogictest/test_files/aggregate_byte_group_keys.slt @@ -0,0 +1,420 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +# SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +# Each query checks that Utf8 and Binary group keys keep their declared types. +# The outer query fails if the key types change. +# +# The tests change key types because varchar uses Utf8View. +# The input has 15000 groups. Of these groups, 1500 have a NULL s key. + +statement ok +SET datafusion.execution.batch_size = 128 + +statement ok +SET datafusion.execution.enable_migration_aggregate = true + +statement ok +SET datafusion.optimizer.enable_round_robin_repartition = false + +statement ok +SET datafusion.optimizer.prefer_existing_sort = true + +statement ok +SET datafusion.execution.parquet.schema_force_view_types = false + +statement ok +SET datafusion.explain.physical_plan_only = true + +statement ok +CREATE VIEW byte_keys AS +SELECT + arrow_cast(concat('p', v % 3), 'Utf8') AS p, + arrow_cast(CASE WHEN v % 10 = 0 THEN NULL ELSE concat('s', (v * 7) % 5000) END, 'Utf8') AS s, + arrow_cast(concat('b', (v * 13) % 5000), 'Binary') AS b, + v +FROM generate_series(1, 20000) AS t(v) + +# Sorted input selects the streams that use the group order. +query I +COPY (SELECT * FROM byte_keys ORDER BY p, s, b) +TO 'test_files/scratch/aggregate_byte_group_keys/sorted.parquet' STORED AS PARQUET +---- +20000 + +statement ok +CREATE EXTERNAL TABLE byte_keys_sorted STORED AS PARQUET +LOCATION 'test_files/scratch/aggregate_byte_group_keys/sorted.parquet' +WITH ORDER (p ASC, s ASC, b ASC) + +statement ok +CREATE EXTERNAL TABLE byte_keys_sorted_by_p STORED AS PARQUET +LOCATION 'test_files/scratch/aggregate_byte_group_keys/sorted.parquet' +WITH ORDER (p ASC) + +query TTT +SELECT arrow_typeof(p), arrow_typeof(s), arrow_typeof(b) FROM byte_keys LIMIT 1 +---- +Utf8 Utf8 Binary + +query TTT +SELECT arrow_typeof(p), arrow_typeof(s), arrow_typeof(b) FROM byte_keys_sorted LIMIT 1 +---- +Utf8 Utf8 Binary + +# Reference results with Utf8View and BinaryView keys + +statement ok +SET datafusion.execution.target_partitions = 1 + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM ( + SELECT p, s, b, count(*) AS c, sum(v) AS total + FROM ( + SELECT arrow_cast(p, 'Utf8View') AS p, arrow_cast(s, 'Utf8View') AS s, + arrow_cast(b, 'BinaryView') AS b, v + FROM byte_keys + ) + GROUP BY p, s, b +) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +# One stage with unsorted input + +query TT +EXPLAIN SELECT p, s, b, count(*), sum(v) FROM byte_keys GROUP BY p, s, b +---- +AggregateExec: mode=Single, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys.v)] + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +query TT?II +SELECT p, s, b, count(*), sum(v) FROM byte_keys GROUP BY p, s, b ORDER BY p, s, b LIMIT 5 +---- +p0 s1 6232383539 1 7143 +p0 s1001 6231383539 1 10143 +p0 s1002 6234373138 2 19572 +p0 s1003 6232353737 1 9429 +p0 s1004 62343336 2 18144 + +# A single key column uses a separate group storage method. +query IIITII +SELECT count(*), count(s), sum(length(s)), min(s), sum(c), sum(total) +FROM (SELECT s, count(*) AS c, sum(v) AS total FROM byte_keys GROUP BY s) +---- +4501 4500 21501 s1 20000 200010000 + +query ITTII +SELECT count(*), min(arrow_cast(b, 'Utf8')), max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT b, count(*) AS c, sum(v) AS total FROM byte_keys GROUP BY b) +---- +5000 b0 b999 20000 200010000 + +# The stream writes groups to disk because of the memory limit. +statement ok +SET datafusion.runtime.memory_limit = '300K' + +query TT +EXPLAIN ANALYZE SELECT p, s, b, count(*), sum(v) FROM byte_keys GROUP BY p, s, b +---- +AggregateExec: mode=Single, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys.v)], metrics=[spilled_rows= K, + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +statement ok +RESET datafusion.runtime.memory_limit + +# Two stages with unsorted input + +statement ok +SET datafusion.execution.target_partitions = 4 + +query TT +EXPLAIN SELECT p, s, b, count(*), sum(v) FROM byte_keys GROUP BY p, s, b +---- +AggregateExec: mode=FinalPartitioned, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys.v)] +AggregateExec: mode=Partial, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys.v)] + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +query IIITII +SELECT count(*), count(s), sum(length(s)), min(s), sum(c), sum(total) +FROM (SELECT s, count(*) AS c, sum(v) AS total FROM byte_keys GROUP BY s) +---- +4501 4500 21501 s1 20000 200010000 + +# The final stage writes groups to disk because of the memory limit. +statement ok +SET datafusion.runtime.memory_limit = '1M' + +query TT +EXPLAIN ANALYZE SELECT p, s, b, count(*), sum(v) FROM byte_keys GROUP BY p, s, b +---- +AggregateExec: mode=FinalPartitioned, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys.v)], metrics=[spilled_rows= K, +AggregateExec: mode=Partial, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys.v)], metrics=[spill_count=0, + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +statement ok +RESET datafusion.runtime.memory_limit + +# Grouping sets use NULL for keys outside the selected set. +query IIIIIII +SELECT count(*), count(p), count(s), sum(length(s)), sum(g), sum(c), sum(total) +FROM ( + SELECT p, s, grouping(p, s) AS g, count(*) AS c, sum(v) AS total + FROM byte_keys GROUP BY ROLLUP(p, s) +) +---- +13507 13506 13500 64503 6 60000 600030000 + +query IIIIIII +SELECT count(*), count(p), count(s), sum(length(s)), sum(g), sum(c), sum(total) +FROM ( + SELECT p, s, grouping(p, s) AS g, count(*) AS c, sum(v) AS total + FROM (SELECT arrow_cast(p, 'Utf8View') AS p, arrow_cast(s, 'Utf8View') AS s, v FROM byte_keys) + GROUP BY ROLLUP(p, s) +) +---- +13507 13506 13500 64503 6 60000 600030000 + +# The partial stage returns input rows without groups. +statement ok +SET datafusion.execution.skip_partial_aggregation_probe_rows_threshold = 0 + +statement ok +SET datafusion.execution.skip_partial_aggregation_probe_ratio_threshold = 0.0 + +query TT +EXPLAIN ANALYZE SELECT p, s, b, count(*), sum(v) FROM byte_keys GROUP BY p, s, b +---- +AggregateExec: mode=FinalPartitioned, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys.v)], metrics=[spill_count=0, +AggregateExec: mode=Partial, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys.v)], metrics=[spill_count=0,skipped_aggregation_rows= K, + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +statement ok +RESET datafusion.execution.skip_partial_aggregation_probe_rows_threshold + +statement ok +RESET datafusion.execution.skip_partial_aggregation_probe_ratio_threshold + +# Input sorted by all group keys + +statement ok +SET datafusion.execution.target_partitions = 1 + +query TT +EXPLAIN SELECT p, s, b, count(*), sum(v) FROM byte_keys_sorted GROUP BY p, s, b +---- +AggregateExec: mode=Single, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys_sorted.v)], group_clustering_mode=Full + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys_sorted GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +statement ok +SET datafusion.execution.target_partitions = 2 + +query TT +EXPLAIN SELECT p, s, b, count(*), sum(v) FROM byte_keys_sorted GROUP BY p, s, b +---- +AggregateExec: mode=FinalPartitioned, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys_sorted.v)], group_clustering_mode=Full +AggregateExec: mode=Partial, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys_sorted.v)], group_clustering_mode=Full + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys_sorted GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +# Input sorted by the first group key + +statement ok +SET datafusion.execution.target_partitions = 1 + +query TT +EXPLAIN SELECT p, s, b, count(*), sum(v) FROM byte_keys_sorted_by_p GROUP BY p, s, b +---- +AggregateExec: mode=Single, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys_sorted_by_p.v)], group_clustering_mode=Partial([0]) + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys_sorted_by_p GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +statement ok +SET datafusion.runtime.memory_limit = '300K' + +query TT +EXPLAIN ANALYZE SELECT p, s, b, count(*), sum(v) FROM byte_keys_sorted_by_p GROUP BY p, s, b +---- +AggregateExec: mode=Single, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys_sorted_by_p.v)], group_clustering_mode=Partial([0]), metrics=[spilled_rows= K, + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys_sorted_by_p GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +statement ok +RESET datafusion.runtime.memory_limit + +statement ok +SET datafusion.execution.target_partitions = 2 + +query TT +EXPLAIN SELECT p, s, b, count(*), sum(v) FROM byte_keys_sorted_by_p GROUP BY p, s, b +---- +AggregateExec: mode=FinalPartitioned, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys_sorted_by_p.v)], group_clustering_mode=Partial([0]) +AggregateExec: mode=Partial, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys_sorted_by_p.v)], group_clustering_mode=Partial([0]) + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys_sorted_by_p GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +statement ok +SET datafusion.runtime.memory_limit = '600K' + +query TT +EXPLAIN ANALYZE SELECT p, s, b, count(*), sum(v) FROM byte_keys_sorted_by_p GROUP BY p, s, b +---- +AggregateExec: mode=FinalPartitioned, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys_sorted_by_p.v)], group_clustering_mode=Partial([0]), metrics=[spilled_rows= K, +AggregateExec: mode=Partial, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys_sorted_by_p.v)], group_clustering_mode=Partial([0]), metrics=[spill_count=0, + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys_sorted_by_p GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +statement ok +RESET datafusion.runtime.memory_limit + +# GroupedHashAggregateStream keeps 64-bit offsets in files on disk. + +statement ok +SET datafusion.execution.enable_migration_aggregate = false + +statement ok +SET datafusion.execution.target_partitions = 4 + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys_sorted_by_p GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +statement ok +SET datafusion.execution.target_partitions = 1 + +statement ok +SET datafusion.runtime.memory_limit = '600K' + +query TT +EXPLAIN ANALYZE SELECT p, s, b, count(*), sum(v) FROM byte_keys GROUP BY p, s, b +---- +AggregateExec: mode=Single, gby=[p@0 as p, s@1 as s, b@2 as b], aggr=[count(Int64(1)), sum(byte_keys.v)], metrics=[spilled_rows= K, + + +query IIIITTII +SELECT count(*), count(s), sum(length(p)), sum(length(s)), min(s), + max(arrow_cast(b, 'Utf8')), sum(c), sum(total) +FROM (SELECT p, s, b, count(*) AS c, sum(v) AS total FROM byte_keys GROUP BY p, s, b) +---- +15000 13500 30000 64503 s1 b999 20000 200010000 + +statement ok +RESET datafusion.runtime.memory_limit + +# Set the test options to their initial values. +statement ok +DROP TABLE byte_keys_sorted_by_p + +statement ok +DROP TABLE byte_keys_sorted + +statement ok +DROP VIEW byte_keys + +statement ok +RESET datafusion.execution.batch_size + +statement ok +SET datafusion.execution.target_partitions = 4 + +statement ok +RESET datafusion.execution.enable_migration_aggregate + +statement ok +RESET datafusion.optimizer.enable_round_robin_repartition + +statement ok +RESET datafusion.optimizer.prefer_existing_sort + +statement ok +RESET datafusion.execution.parquet.schema_force_view_types + +statement ok +RESET datafusion.explain.physical_plan_only + +statement ok +RESET datafusion.catalog.create_default_catalog_and_schema From 33de0143377f9d13b9230fcfbe4ac76fed355b7a Mon Sep 17 00:00:00 2001 From: Max Burke Date: Tue, 10 Nov 2026 21:20:08 -0700 Subject: [PATCH 3/5] fix: Keep aggregate output keys within the byte limit A batch can fit the row limit and still contain too much key data for 32-bit offsets. Limit both rows and key bytes before output conversion. Use the same limit check for aggregate tables and streams. Search the offsets only when the requested rows do not fit. Test the 32-bit boundary, array slices, and keys that cannot fit in one output batch. Addresses https://github.com/apache/datafusion/pull/23695#discussion_r3663622881 --- .../aggregates/aggregate_hash_table/common.rs | 41 +++- .../src/aggregates/grouped_hash_stream.rs | 6 +- .../physical-plan/src/aggregates/mod.rs | 180 ++++++++++++++++-- 3 files changed, 207 insertions(+), 20 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 54a60c0e46b04..6531acb7aee4a 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -9,7 +9,7 @@ use std::sync::Arc; use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, new_null_array}; use arrow::compute::prep_null_mask_filter; -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{Schema, SchemaRef}; use arrow::downcast_dictionary_array; use arrow::record_batch::RecordBatch; use datafusion_common::utils::compact_dictionary; @@ -30,7 +30,7 @@ use crate::aggregates::order::GroupClustering; use crate::aggregates::{ AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, group_id_array, group_keys_with_large_offsets, group_schema_with_large_offsets, - max_duplicate_ordinal, restore_group_key_types, + max_duplicate_ordinal, output_batch_rows, restore_group_key_types, }; use super::AggregateTableMetrics; @@ -331,8 +331,9 @@ impl AggregateHashTable { } }; + let rows = output.next_batch_rows(&output_schema, batch_size); let batch = output - .next_batch(batch_size)? + .next_batch(rows)? .map(|batch| restore_group_key_types(batch, &output_schema)) .transpose()?; if output.is_exhausted() { @@ -663,6 +664,15 @@ impl MaterializedAggregateOutput { Self { batch, offset: 0 } } + /// Returns a row count for [`Self::next_batch`] within the row and key byte limits. + pub(super) fn next_batch_rows( + &self, + output_schema: &Schema, + batch_size: usize, + ) -> usize { + output_batch_rows(&self.batch, output_schema, self.offset, batch_size) + } + pub(super) fn next_batch( &mut self, batch_size: usize, @@ -992,6 +1002,7 @@ mod tests { use super::*; use crate::aggregates::group_values::aggregate_sub_metrics; + use crate::aggregates::tests::{GIB, zeroed_group_keys}; use crate::metrics::ExecutionPlanMetricsSet; #[test] @@ -1037,6 +1048,30 @@ mod tests { Ok(()) } + #[test] + fn next_batch_rows_limits_key_bytes() -> Result<()> { + let output_schema = + Schema::new(vec![Field::new("group_col", DataType::Binary, false)]); + let stored_schema = Arc::new(Schema::new(vec![Field::new( + "group_col", + DataType::LargeBinary, + false, + )])); + // The first two keys hold one byte more than `i32::MAX`. + let keys = zeroed_group_keys(&[(DataType::LargeBinary, &[GIB, GIB, 1, 1, 1])]); + let batch = RecordBatch::try_new(stored_schema, keys)?; + let mut output = MaterializedAggregateOutput::new(batch); + + let mut rows = vec![]; + while !output.is_exhausted() { + let next_rows = output.next_batch_rows(&output_schema, 8192); + rows.push(output.next_batch(next_rows)?.unwrap().num_rows()); + } + assert_eq!(rows, vec![1, 4]); + + Ok(()) + } + fn dictionary_batch( keys: PrimitiveArray, num_values: usize, diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 9bea9af8dd6e3..07e6d56e0b103 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -34,7 +34,7 @@ use crate::aggregates::{ AggregateInputMode, AggregateMode, AggregateOutputMode, PhysicalGroupBy, aggregate_metric_label, create_schema, evaluate_group_by, evaluate_optional, group_id_array, group_keys_with_large_offsets, group_schema_with_large_offsets, - max_duplicate_ordinal, restore_group_key_types, + max_duplicate_ordinal, output_batch_rows, restore_group_key_types, }; use crate::metrics::{BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput}; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; @@ -799,7 +799,7 @@ impl Stream for GroupedHashAggregateStream { ExecutionState::ProducingOutput(batch) => { // slice off a part of the batch, if needed let output_batch; - let size = self.batch_size; + let size = output_batch_rows(batch, &self.schema, 0, self.batch_size); (self.exec_state, output_batch) = if batch.num_rows() <= size { ( if self.input_done { @@ -817,8 +817,6 @@ impl Stream for GroupedHashAggregateStream { batch.clone(), ) } else { - // output first batch_size rows - let size = self.batch_size; let num_remaining = batch.num_rows() - size; let remaining = batch.slice(size, num_remaining); let output = batch.slice(0, size); diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 85262c53f41cd..997ffb8f763f8 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -169,8 +169,8 @@ use parking_lot::Mutex; use std::collections::HashMap; use arrow::array::{ - ArrayRef, MutableArrayData, UInt8Array, UInt16Array, UInt32Array, UInt64Array, - make_array, + ArrayRef, AsArray, MutableArrayData, UInt8Array, UInt16Array, UInt32Array, + UInt64Array, make_array, }; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; @@ -2919,7 +2919,40 @@ fn group_key_types_differ(batch: &RecordBatch, schema: &Schema) -> bool { .any(|(column, field)| column.data_type() != field.data_type()) } -// Each output column must fit the byte limit of its declared type. +// A row limit alone does not keep key data within the 32-bit offset limit. +fn output_batch_rows( + batch: &RecordBatch, + schema: &Schema, + offset: usize, + batch_size: usize, +) -> usize { + let max_rows = batch_size.min(batch.num_rows() - offset); + let rows = batch + .columns() + .iter() + .zip(schema.fields()) + .filter(|(column, field)| column.data_type() != field.data_type()) + .filter_map(|(column, _)| match column.data_type() { + DataType::LargeUtf8 => Some(column.as_string::().value_offsets()), + DataType::LargeBinary => Some(column.as_binary::().value_offsets()), + _ => None, + }) + .map(|offsets| { + let offsets = &offsets[offset..=offset + max_rows]; + if offsets[max_rows] - offsets[0] <= i64::from(i32::MAX) { + max_rows + } else { + let limit = offsets[0] + i64::from(i32::MAX); + offsets[1..].partition_point(|end| *end <= limit) + } + }) + .min() + .unwrap_or(max_rows); + // Select at least one row so conversion can report an invalid key. + rows.max(1).min(max_rows) +} + +// `batch` must fit the 32-bit offsets of `schema`: see `output_batch_rows`. fn restore_group_key_types( batch: RecordBatch, schema: &SchemaRef, @@ -2971,20 +3004,18 @@ fn restore_group_key_stream( emitter.emit(batch).await; continue; } - if batch.num_rows() <= batch_size { - emitter.emit(restore_group_key_types(batch, &schema)?).await; - continue; - } // When memory is not sufficient, the input can return all groups at once. // Divide the batch before conversion to limit the bytes in each column. - for offset in (0..batch.num_rows()).step_by(batch_size) { - let length = batch_size.min(batch.num_rows() - offset); + let mut offset = 0; + loop { + let rows = output_batch_rows(&batch, &schema, offset, batch_size); emitter - .emit(restore_group_key_types( - batch.slice(offset, length), - &schema, - )?) + .emit(restore_group_key_types(batch.slice(offset, rows), &schema)?) .await; + offset += rows; + if offset >= batch.num_rows() { + break; + } } } Ok(()) @@ -10269,6 +10300,129 @@ mod tests { Ok(()) } + pub(in crate::aggregates) const GIB: usize = 1 << 30; + + /// Returns columns that share a buffer filled with zeros. + /// Tests can check large offsets without reading the entire buffer. + pub(in crate::aggregates) fn zeroed_group_keys( + columns: &[(DataType, &[usize])], + ) -> Vec { + let num_bytes = columns + .iter() + .map(|(_, lengths)| lengths.iter().sum::()) + .max() + .unwrap_or_default(); + let values = arrow::buffer::Buffer::from(vec![0_u8; num_bytes]); + columns + .iter() + .map(|(data_type, lengths)| { + let offsets = arrow::buffer::OffsetBuffer::::from_lengths( + lengths.iter().copied(), + ); + let values = values.clone(); + match data_type { + DataType::LargeBinary => Arc::new( + arrow::array::LargeBinaryArray::new(offsets, values, None), + ) as ArrayRef, + // SAFETY: zero bytes are valid UTF-8. + DataType::LargeUtf8 => Arc::new(unsafe { + arrow::array::LargeStringArray::new_unchecked( + offsets, values, None, + ) + }), + other => unreachable!("{other} keys do not use 64-bit offsets"), + } + }) + .collect() + } + + #[test] + fn output_batch_rows_fit_32_bit_offsets() -> Result<()> { + let output_schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("b", DataType::Binary, true), + Field::new("c", DataType::Int64, false), + ])); + let stored_schema = group_schema_with_large_offsets(&output_schema, 2); + let stored = |s: &[usize], b: &[usize]| { + let mut columns = zeroed_group_keys(&[ + (DataType::LargeUtf8, s), + (DataType::LargeBinary, b), + ]); + columns.push(Arc::new(Int64Array::from(vec![0; s.len()]))); + RecordBatch::try_new(Arc::clone(&stored_schema), columns) + }; + + let batch = stored(&[7, GIB, GIB - 1, 1, 0, 5], &[1, 1, 1, GIB, GIB, 1])?; + for (offset, batch_size, rows) in [ + // The next key of `s` does not fit after 7 + GIB bytes. + (0, 8192, 2), + // `s` holds exactly `i32::MAX` bytes. + (1, 8192, 2), + // The next key of `b` does not fit after 1 + GIB bytes. + (2, 8192, 2), + // Two keys of `b` hold one byte more than `i32::MAX`. + (3, 8192, 1), + (4, 8192, 2), + (1, 1, 1), + (6, 8192, 0), + ] { + assert_eq!( + output_batch_rows(&batch, &output_schema, offset, batch_size), + rows, + "offset {offset} with batch size {batch_size}" + ); + } + + let slice = batch.slice(1, 5); + assert_eq!(output_batch_rows(&slice, &output_schema, 0, 8192), 2); + assert_eq!(output_batch_rows(&slice, &output_schema, 2, 8192), 1); + + // Output keys with 64-bit offsets do not need the 32-bit limit. + assert_eq!(output_batch_rows(&batch, &stored_schema, 0, 8192), 6); + drop((batch, slice)); + + // An invalid key must reach conversion so that it returns an error. + let batch = stored(&[2 * GIB, 1], &[1, 1])?; + assert_eq!(output_batch_rows(&batch, &output_schema, 0, 8192), 1); + assert_eq!(output_batch_rows(&batch, &output_schema, 1, 8192), 1); + + Ok(()) + } + + #[tokio::test] + async fn restore_group_key_stream_limits_key_bytes() -> Result<()> { + let output_schema = Arc::new(Schema::new(vec![ + Field::new("s", DataType::Utf8, true), + Field::new("c", DataType::Int64, false), + ])); + let stored_schema = group_schema_with_large_offsets(&output_schema, 1); + // The two keys hold one byte more than `i32::MAX`. + let mut columns = zeroed_group_keys(&[(DataType::LargeUtf8, &[GIB, GIB])]); + columns.push(Arc::new(Int64Array::from(vec![1, 2]))); + let stored = RecordBatch::try_new(stored_schema, columns)?; + + let input = Box::pin(crate::stream::RecordBatchStreamAdapter::new( + Arc::clone(&output_schema), + futures::stream::iter(vec![Ok(stored)]), + )); + let mut output = restore_group_key_stream(input, 8192); + + // An output batch owns a copy of its keys, so the test holds one + // batch at a time. + let mut counts = vec![]; + while let Some(batch) = output.next().await.transpose()? { + assert_eq!(batch.schema(), output_schema); + let keys = batch.column(0).as_string::(); + assert_eq!(keys.value_offsets(), &[0, GIB as i32]); + counts + .extend_from_slice(batch.column(1).as_primitive::().values()); + } + assert_eq!(counts, vec![1, 2]); + + Ok(()) + } + #[rstest::rstest] #[case::reserves_batch_memory(320 * 1024, true)] #[case::cannot_reserve_batch_memory(64 * 1024, false)] From 1887985a0dce14ca053ce753c885eb8a8b002bf1 Mon Sep 17 00:00:00 2001 From: Max Burke Date: Tue, 10 Nov 2026 22:55:28 -0700 Subject: [PATCH 4/5] perf: Read string and binary keys without conversion Let group storage with 64-bit offsets read Utf8 and Binary input arrays directly. Reuse the existing comparison and storage methods without a new offset buffer for each input column. Keep input conversion in GroupValuesRows because its row format requires the schema types. Reject extra columns before conversion. Test strings, binary values, nulls, sorted groups, and both offset sizes. Addresses https://github.com/apache/datafusion/pull/23695#discussion_r3663622885 --- .../physical-expr-common/src/binary_map.rs | 92 +++++++++++-- .../aggregates/aggregate_hash_table/common.rs | 14 +- .../aggregate_hash_table/common_clustered.rs | 7 +- .../src/aggregates/group_values/mod.rs | 109 ++++++++++++++- .../group_values/multi_group_by/bytes.rs | 125 ++++++++++++++---- .../src/aggregates/group_values/row.rs | 40 +++++- .../src/aggregates/grouped_hash_stream.rs | 9 +- .../physical-plan/src/aggregates/mod.rs | 12 -- .../test_files/aggregate_byte_group_keys.slt | 24 ++++ 9 files changed, 354 insertions(+), 78 deletions(-) diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 161497569999a..903d08f029ecf 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -9,7 +9,7 @@ use arrow::array::{ Array, ArrayRef, BufferBuilder, GenericBinaryArray, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, cast::AsArray, - types::{ByteArrayType, GenericBinaryType, GenericStringType}, + types::{BinaryType, ByteArrayType, GenericBinaryType, GenericStringType, Utf8Type}, }; use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::DataType; @@ -348,7 +348,9 @@ where /// /// # Arguments: /// - /// `values`: array whose values are inserted + /// `values`: the input array. A map with 64-bit offsets also accepts + /// 32-bit input offsets. The map can store more than 2 GiB of values + /// without input conversion. /// /// `make_payload_fn`: invoked for each value that is not already present /// to create the payload, in order of the values in `values` @@ -374,6 +376,8 @@ where MP: FnMut(Option<&[u8]>) -> V, OP: FnMut(V), { + let reads_32_bit_offsets = O::IS_LARGE + && matches!(values.data_type(), DataType::Utf8 | DataType::Binary); // Sanity array type match self.output_type { OutputType::Binary => { @@ -381,22 +385,38 @@ where values.data_type(), DataType::Binary | DataType::LargeBinary )); - self.insert_if_new_inner::>( - values, - make_payload_fn, - observe_payload_fn, - ) + if reads_32_bit_offsets { + self.insert_if_new_inner::( + values, + make_payload_fn, + observe_payload_fn, + ) + } else { + self.insert_if_new_inner::>( + values, + make_payload_fn, + observe_payload_fn, + ) + } } OutputType::Utf8 => { assert!(matches!( values.data_type(), DataType::Utf8 | DataType::LargeUtf8 )); - self.insert_if_new_inner::>( - values, - make_payload_fn, - observe_payload_fn, - ) + if reads_32_bit_offsets { + self.insert_if_new_inner::( + values, + make_payload_fn, + observe_payload_fn, + ) + } else { + self.insert_if_new_inner::>( + values, + make_payload_fn, + observe_payload_fn, + ) + } } _ => unreachable!("View types should use `ArrowBytesViewMap`"), } @@ -736,7 +756,7 @@ where #[cfg(test)] mod tests { use super::*; - use arrow::array::{BinaryArray, LargeBinaryArray, StringArray}; + use arrow::array::{BinaryArray, LargeBinaryArray, LargeStringArray, StringArray}; use std::collections::HashMap; /// A lower bound on the bytes a hashbrown table holding `entries` entries @@ -1201,6 +1221,52 @@ mod tests { assert_eq!(&set.into_state(), &expected); } + #[test] + fn large_offset_set_inserts_arrays_with_32_bit_offsets() { + let mut set = ArrowBytesSet::::new(OutputType::Utf8); + let values: ArrayRef = Arc::new(StringArray::from(vec![ + Some("a"), + Some("CXCCCCCCCC"), + None, + Some("CXCCCCCCCC"), + ])); + set.insert(&values); + let values: ArrayRef = Arc::new(LargeStringArray::from(vec![ + Some("CXCCCCCCCC"), + Some("b"), + None, + ])); + set.insert(&values); + let expected: ArrayRef = Arc::new(LargeStringArray::from(vec![ + Some("a"), + Some("CXCCCCCCCC"), + None, + Some("b"), + ])); + assert_eq!(&set.into_state(), &expected); + + let mut set = ArrowBytesSet::::new(OutputType::Binary); + let values: ArrayRef = Arc::new(BinaryArray::from_opt_vec(vec![ + Some(b"a"), + Some(b"CXCCCCCCCC"), + None, + Some(b"CXCCCCCCCC"), + ])); + set.insert(&values); + let values: ArrayRef = Arc::new(LargeBinaryArray::from_opt_vec(vec![ + Some(b"CXCCCCCCCC"), + Some(b"b"), + ])); + set.insert(&values); + let expected: ArrayRef = Arc::new(LargeBinaryArray::from_opt_vec(vec![ + Some(b"a"), + Some(b"CXCCCCCCCC"), + None, + Some(b"b"), + ])); + assert_eq!(&set.into_state(), &expected); + } + #[test] #[should_panic( expected = "matches!(values.data_type(), DataType::Utf8 | DataType::LargeUtf8)" 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 6531acb7aee4a..044eb1dfe1270 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -29,8 +29,8 @@ use crate::aggregates::group_values::{ use crate::aggregates::order::GroupClustering; use crate::aggregates::{ AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, - group_id_array, group_keys_with_large_offsets, group_schema_with_large_offsets, - max_duplicate_ordinal, output_batch_rows, restore_group_key_types, + group_id_array, group_schema_with_large_offsets, max_duplicate_ordinal, + output_batch_rows, restore_group_key_types, }; use super::AggregateTableMetrics; @@ -201,13 +201,9 @@ impl AggregateHashTable { ) -> Result { let state = self.state.building_mut(); // Outer vec: one per grouping set; inner vec: group-by expressions. - let grouping_set_args = - self.group_by_metrics.time_group_key_preparation(|| { - evaluate_group_by(&state.group_by, batch)? - .into_iter() - .map(group_keys_with_large_offsets) - .collect::>>() - })?; + let grouping_set_args = self + .group_by_metrics + .time_group_key_preparation(|| evaluate_group_by(&state.group_by, batch))?; // The evaluated args for each accumulator. let accumulator_args = self.group_by_metrics.time_aggregate_arguments(|| { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs index 717b5da3cbd83..61b180da4676d 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs @@ -22,7 +22,7 @@ use crate::aggregates::group_values::{ use crate::aggregates::order::{GroupClustering, GroupClusteringMode}; use crate::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, - evaluate_group_by, group_keys_with_large_offsets, group_schema_with_large_offsets, + evaluate_group_by, group_schema_with_large_offsets, }; use super::AggregateTableMetrics; @@ -234,10 +234,7 @@ impl ClusteredAggregateTable { ) -> Result { let grouping_set_args = self.group_by_metrics.time_group_key_preparation(|| { - evaluate_group_by(&self.buffer.group_by, batch)? - .into_iter() - .map(group_keys_with_large_offsets) - .collect::>>() + evaluate_group_by(&self.buffer.group_by, batch) })?; let accumulator_args = self.group_by_metrics.time_aggregate_arguments(|| { diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index 67f0aa797ff7a..5c0111d9c3e7c 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -148,6 +148,10 @@ pub trait GroupValues: Send { /// /// - Otherwise, the general implementation `GroupValuesRows` will be chosen. /// +/// A schema with [`DataType::LargeUtf8`] or [`DataType::LargeBinary`] keys +/// also accepts [`DataType::Utf8`] or [`DataType::Binary`] input arrays. +/// Only [`GroupValuesRows`] needs to change these input arrays to the schema types. +/// /// `GroupColumn`: crate::aggregates::group_values::multi_group_by::GroupColumn /// `GroupValuesColumn`: crate::aggregates::group_values::multi_group_by::GroupValuesColumn /// `GroupValuesRows`: crate::aggregates::group_values::GroupValuesRows @@ -212,13 +216,15 @@ mod tests { use std::sync::Arc; use arrow::array::{ - ArrayRef, AsArray, BooleanArray, Int32Array, StringArray, StringViewArray, + ArrayRef, AsArray, BinaryArray, BooleanArray, Int32Array, NullArray, StringArray, + StringViewArray, }; + use arrow::compute::cast; use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use datafusion_expr::{EmitTo, GroupSelection}; - use super::new_group_values; - use crate::aggregates::order::GroupClustering; + use super::{multi_group_by, new_group_values}; + use crate::aggregates::order::{GroupClustering, GroupClusteringMode}; #[test] fn preserving_values_keep_group_indices_valid() { @@ -278,6 +284,103 @@ mod tests { assert_eq!(actual[0].as_primitive::(), &expected); } + #[rstest::rstest] + #[case::string(vec![DataType::LargeUtf8], GroupClusteringMode::None)] + #[case::binary(vec![DataType::LargeBinary], GroupClusteringMode::None)] + #[case::columns( + vec![DataType::LargeUtf8, DataType::LargeBinary], + GroupClusteringMode::None + )] + #[case::partially_clustered_columns( + vec![DataType::LargeUtf8, DataType::LargeBinary], + GroupClusteringMode::Partial(vec![0]) + )] + #[case::clustered_columns( + vec![DataType::LargeUtf8, DataType::LargeBinary], + GroupClusteringMode::Full + )] + #[case::rows(vec![DataType::LargeUtf8, DataType::Null], GroupClusteringMode::None)] + fn large_offset_columns_accept_32_bit_input_offsets( + #[case] data_types: Vec, + #[case] group_clustering_mode: GroupClusteringMode, + ) { + let fields = data_types + .iter() + .enumerate() + .map(|(index, data_type)| { + Field::new(format!("group{index}"), data_type.clone(), true) + }) + .collect::>(); + let schema = Arc::new(Schema::new(fields)); + if data_types.contains(&DataType::Null) { + assert!(!multi_group_by::supported_schema(&schema)); + } + + let input = |keys: &[Option<&str>]| -> Vec { + data_types + .iter() + .map(|data_type| match data_type { + DataType::LargeUtf8 => { + Arc::new(StringArray::from(keys.to_vec())) as ArrayRef + } + DataType::LargeBinary => Arc::new(BinaryArray::from_iter( + keys.iter().map(|key| key.map(str::as_bytes)), + )), + _ => Arc::new(NullArray::new(keys.len())), + }) + .collect() + }; + let stored = |input: &[ArrayRef]| -> Vec { + input + .iter() + .zip(schema.fields()) + .map(|(array, field)| cast(array, field.data_type()).unwrap()) + .collect() + }; + + let group_clustering = GroupClustering::try_new(&group_clustering_mode).unwrap(); + let mut group_values = + new_group_values(Arc::clone(&schema), &group_clustering).unwrap(); + let mut expected_group_values = + new_group_values(Arc::clone(&schema), &group_clustering).unwrap(); + let mut groups = vec![]; + + // Equal keys must be next to each other for the implementations that use the group order. + let long_key = "a key longer than sixteen bytes"; + type Step<'a> = (&'a [Option<&'a str>], &'a [usize], Option); + let steps: [Step; 3] = [ + ( + &[Some("a"), Some("a"), None, Some(long_key)], + &[0, 0, 1, 2], + None, + ), + ( + &[Some(long_key), Some("c")], + &[2, 3], + Some(EmitTo::First(2)), + ), + (&[Some("c"), Some("d")], &[1, 2], Some(EmitTo::All)), + ]; + for (keys, expected_groups, emit_to) in steps { + let input = input(keys); + group_values.intern(&input, &mut groups).unwrap(); + assert_eq!(groups, expected_groups); + expected_group_values + .intern(&stored(&input), &mut groups) + .unwrap(); + assert_eq!(groups, expected_groups); + + // After `EmitTo::First`, a single column stores the remaining keys again with 64-bit input offsets. + if let Some(emit_to) = emit_to { + let values = group_values.emit(emit_to).unwrap(); + assert_eq!(values, expected_group_values.emit(emit_to).unwrap()); + for (value, field) in values.iter().zip(schema.fields()) { + assert_eq!(value.data_type(), field.data_type()); + } + } + } + } + #[test] fn preserving_non_nullable_primitive_and_boolean_values() { let schema = Arc::new(Schema::new(vec![ 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 a2247e00a6e2e..01e2d1143fdc3 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 @@ -12,7 +12,9 @@ use arrow::array::{ types::GenericStringType, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; -use arrow::datatypes::{ByteArrayType, DataType, GenericBinaryType}; +use arrow::datatypes::{ + BinaryType, ByteArrayType, DataType, GenericBinaryType, Utf8Type, +}; use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::utils::split_vec_min_alloc; use datafusion_common::{Result, exec_datafusion_err}; @@ -22,6 +24,11 @@ use std::mem::size_of; use std::sync::Arc; use std::vec; +#[inline] +fn uses_smaller_input_offsets(array: &ArrayRef) -> bool { + O::IS_LARGE && matches!(array.data_type(), DataType::Utf8 | DataType::Binary) +} + fn checked_output_offset( current_len: usize, additional_len: usize, @@ -39,6 +46,9 @@ fn checked_output_offset( /// /// 1. Efficient comparison of incoming rows to existing rows /// 2. Efficient construction of the final output array +/// +/// A builder with 64-bit offsets also accepts 32-bit input offsets. +/// This allows more than 2 GiB of stored keys without input conversion. pub struct ByteGroupValueBuilder where O: OffsetSizeTrait, @@ -74,19 +84,54 @@ where } } - fn equal_to_inner(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool + fn equal_to_inner( + &self, + lhs_row: usize, + array: &ArrayRef, + rhs_row: usize, + ) -> bool where - B: ByteArrayType, + Input32: ByteArrayType, + Stored: ByteArrayType, { - let array = array.as_bytes::(); + // The type check also gets the array. A separate check would add a call for each row. + let array = array.as_any(); + if O::IS_LARGE + && let Some(array) = array.downcast_ref::>() + { + return self.do_equal_to_inner(lhs_row, array, rhs_row); + } + let array: &GenericByteArray = array.downcast_ref().expect("byte array"); self.do_equal_to_inner(lhs_row, array, rhs_row) } - fn append_val_inner(&mut self, array: &ArrayRef, row: usize) -> Result<()> + fn append_val_inner( + &mut self, + array: &ArrayRef, + row: usize, + ) -> Result<()> + where + Input32: ByteArrayType, + Stored: ByteArrayType, + { + let array = array.as_any(); + if O::IS_LARGE + && let Some(array) = array.downcast_ref::>() + { + return self.do_append_row_inner(array, row); + } + let array: &GenericByteArray = array.downcast_ref().expect("byte array"); + self.do_append_row_inner(array, row) + } + + fn do_append_row_inner( + &mut self, + arr: &GenericByteArray, + row: usize, + ) -> Result<()> where B: ByteArrayType, { - let arr = array.as_bytes::(); if arr.is_null(row) { self.nulls.append_null(); // nulls need a zero length in the offset buffer @@ -146,7 +191,7 @@ where match all_null_or_non_null { Nulls::Some => { for &row in rows { - self.append_val_inner::(array, row)? + self.do_append_row_inner(arr, row)? } } @@ -230,14 +275,18 @@ where column.data_type(), DataType::Binary | DataType::LargeBinary )); - self.equal_to_inner::>(lhs_row, column, rhs_row) + 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) + self.equal_to_inner::>( + lhs_row, column, rhs_row, + ) } _ => unreachable!("View types should use `ArrowBytesViewMap`"), } @@ -251,14 +300,14 @@ where column.data_type(), DataType::Binary | DataType::LargeBinary )); - self.append_val_inner::>(column, row)? + self.append_val_inner::>(column, row)? } OutputType::Utf8 => { debug_assert!(matches!( column.data_type(), DataType::Utf8 | DataType::LargeUtf8 )); - self.append_val_inner::>(column, row)? + self.append_val_inner::>(column, row)? } _ => unreachable!("View types should use `ArrowBytesViewMap`"), } @@ -280,24 +329,42 @@ where array.data_type(), DataType::Binary | DataType::LargeBinary )); - self.vectorized_equal_to_inner::>( - lhs_rows, - array, - rhs_rows, - equal_to_results, - ); + if uses_smaller_input_offsets::(array) { + self.vectorized_equal_to_inner::( + lhs_rows, + array, + rhs_rows, + equal_to_results, + ); + } else { + 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, - ); + if uses_smaller_input_offsets::(array) { + self.vectorized_equal_to_inner::( + lhs_rows, + array, + rhs_rows, + equal_to_results, + ); + } else { + self.vectorized_equal_to_inner::>( + lhs_rows, + array, + rhs_rows, + equal_to_results, + ); + } } _ => unreachable!("View types should use `ArrowBytesViewMap`"), } @@ -310,14 +377,22 @@ where column.data_type(), DataType::Binary | DataType::LargeBinary )); - self.vectorized_append_inner::>(column, rows)? + if uses_smaller_input_offsets::(column) { + self.vectorized_append_inner::(column, rows)? + } else { + self.vectorized_append_inner::>(column, rows)? + } } OutputType::Utf8 => { debug_assert!(matches!( column.data_type(), DataType::Utf8 | DataType::LargeUtf8 )); - self.vectorized_append_inner::>(column, rows)? + if uses_smaller_input_offsets::(column) { + self.vectorized_append_inner::(column, rows)? + } else { + self.vectorized_append_inner::>(column, rows)? + } } _ => unreachable!("View types should use `ArrowBytesViewMap`"), } diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 29b5eaa66b51e..4fc63189bfe90 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -11,10 +11,10 @@ use arrow::array::{ use arrow::compute::cast; use arrow::datatypes::{DataType, SchemaRef}; use arrow::row::{RowConverter, Rows, SortField}; -use datafusion_common::Result; use datafusion_common::float_semantics; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; +use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::{EmitTo, GroupSelection}; use hashbrown::hash_table::HashTable; @@ -109,13 +109,29 @@ impl GroupValuesRows { impl GroupValues for GroupValuesRows { fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + assert_eq_or_internal_err!( + cols.len(), + self.schema.fields().len(), + "The number of group columns must match the schema" + ); // Normalize -0.0 → +0.0 so RowConverter (IEEE 754 totalOrder) and // primitive hashing both group ±0 together. No-op for non-float // columns. - let normalized_cols: Vec = cols + let normalized_cols = cols .iter() - .map(|col| float_semantics::normalize(Arc::clone(col))) - .collect(); + .zip(self.schema.fields()) + .map(|(col, field)| { + // The row format requires input arrays to match the schema types. + let col = match (col.data_type(), field.data_type()) { + (DataType::Utf8, DataType::LargeUtf8) + | (DataType::Binary, DataType::LargeBinary) => { + cast(col, field.data_type())? + } + _ => Arc::clone(col), + }; + Ok(float_semantics::normalize(col)) + }) + .collect::>>()?; let cols = normalized_cols.as_slice(); // Convert the group keys into the row format @@ -491,6 +507,22 @@ mod tests { )])) } + #[test] + fn rejects_extra_group_columns() -> Result<()> { + let mut group_values = GroupValuesRows::try_new(int32_schema())?; + let input: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let error = group_values + .intern(&[Arc::clone(&input), input], &mut vec![]) + .unwrap_err(); + assert!( + error + .to_string() + .contains("The number of group columns must match the schema") + ); + assert!(group_values.is_empty()); + Ok(()) + } + #[test] fn size_retains_reusable_buffers_after_emit() -> Result<()> { let mut group_values = GroupValuesRows::try_new(int32_schema())?; diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 07e6d56e0b103..43550df7eff06 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -33,8 +33,8 @@ use crate::aggregates::order::GroupClusteringFull; use crate::aggregates::{ AggregateInputMode, AggregateMode, AggregateOutputMode, PhysicalGroupBy, aggregate_metric_label, create_schema, evaluate_group_by, evaluate_optional, - group_id_array, group_keys_with_large_offsets, group_schema_with_large_offsets, - max_duplicate_ordinal, output_batch_rows, restore_group_key_types, + group_id_array, group_schema_with_large_offsets, max_duplicate_ordinal, + output_batch_rows, restore_group_key_types, }; use crate::metrics::{BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput}; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; @@ -882,11 +882,6 @@ impl GroupedHashAggregateStream { } })?; - let group_by_values = group_by_values - .into_iter() - .map(group_keys_with_large_offsets) - .collect::>>()?; - // Evaluate aggregate arguments and filters together. Per-aggregate // argument timers remain limited to their argument expressions because // filters are evaluated collectively here. diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 997ffb8f763f8..b152e76035b7a 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -2899,18 +2899,6 @@ fn group_schema_with_large_offsets( Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) } -fn group_keys_with_large_offsets(arrays: Vec) -> Result> { - arrays - .into_iter() - .map( - |array| match group_key_type_with_large_offsets(array.data_type()) { - Some(data_type) => Ok(arrow::compute::cast(&array, &data_type)?), - None => Ok(array), - }, - ) - .collect() -} - fn group_key_types_differ(batch: &RecordBatch, schema: &Schema) -> bool { batch .columns() diff --git a/datafusion/sqllogictest/test_files/aggregate_byte_group_keys.slt b/datafusion/sqllogictest/test_files/aggregate_byte_group_keys.slt index 9f5a921f070b7..009eb306de665 100644 --- a/datafusion/sqllogictest/test_files/aggregate_byte_group_keys.slt +++ b/datafusion/sqllogictest/test_files/aggregate_byte_group_keys.slt @@ -119,6 +119,30 @@ FROM (SELECT b, count(*) AS c, sum(v) AS total FROM byte_keys GROUP BY b) ---- 5000 b0 b999 20000 200010000 +# A key column of the Null type uses the row format for the group values. +query TT +EXPLAIN SELECT s, n, count(*) FROM (SELECT s, NULL AS n FROM byte_keys) GROUP BY s, n +---- +AggregateExec: mode=Single, gby=[s@0 as s, n@1 as n], aggr=[count(Int64(1))] + + +query TT +SELECT arrow_typeof(s), arrow_typeof(n) +FROM (SELECT s, n, count(*) FROM (SELECT s, NULL AS n FROM byte_keys) GROUP BY s, n) +LIMIT 1 +---- +Utf8 Null + +query IIIII +SELECT count(*), count(s), count(n), sum(length(s)), sum(c) +FROM ( + SELECT s, n, count(*) AS c + FROM (SELECT s, NULL AS n FROM byte_keys) + GROUP BY s, n +) +---- +4501 4500 0 21501 20000 + # The stream writes groups to disk because of the memory limit. statement ok SET datafusion.runtime.memory_limit = '300K' From 8d24422ed9d187c5fb6df666f775abbd859bfaa8 Mon Sep 17 00:00:00 2001 From: Max Burke Date: Wed, 11 Nov 2026 08:01:57 -0700 Subject: [PATCH 5/5] perf: Share stored key data with aggregate output Output batches need 32-bit offsets but can use the stored key bytes. Subtract the first offset and share the required part of the value buffer. This removes the copy of the key data for each output batch. The output keeps the original value buffer in memory. Check that output keys share this buffer, including keys beyond the 32-bit offset limit. --- .../physical-plan/src/aggregates/mod.rs | 102 +++++++++++++----- 1 file changed, 75 insertions(+), 27 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index b152e76035b7a..93ca24b90a9bd 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -169,11 +169,14 @@ use parking_lot::Mutex; use std::collections::HashMap; use arrow::array::{ - ArrayRef, AsArray, MutableArrayData, UInt8Array, UInt16Array, UInt32Array, - UInt64Array, make_array, + Array, ArrayRef, AsArray, GenericByteArray, UInt8Array, UInt16Array, UInt32Array, + UInt64Array, }; +use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::compute::SortOptions; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::datatypes::{ + BinaryType, ByteArrayType, DataType, Field, Schema, SchemaRef, Utf8Type, +}; use arrow::record_batch::RecordBatch; use arrow_schema::FieldRef; use datafusion_common::stats::Precision; @@ -2940,6 +2943,39 @@ fn output_batch_rows( rows.max(1).min(max_rows) } +// The result shares the value buffer, so it keeps that buffer in memory. +fn with_32_bit_offsets(array: &GenericByteArray) -> Result> +where + T: ByteArrayType, + U: ByteArrayType, +{ + let offsets = array.value_offsets(); + let start = offsets[0]; + let num_bytes = offsets[offsets.len() - 1] - start; + if num_bytes > i64::from(i32::MAX) { + return internal_err!( + "Group keys of {num_bytes} bytes do not fit 32-bit offsets" + ); + } + // A slice can start beyond the 32-bit limit. Its new offsets must start at zero. + let offsets = offsets + .iter() + .map(|offset| (offset - start) as i32) + .collect::>(); + let values = array + .values() + .slice_with_length(start as usize, num_bytes as usize); + // SAFETY: the checked length fits i32. Subtraction preserves the offset order + // and value boundaries. The value bytes and null flags do not change. + Ok(unsafe { + GenericByteArray::new_unchecked( + OffsetBuffer::new_unchecked(offsets), + values, + array.nulls().cloned(), + ) + }) +} + // `batch` must fit the 32-bit offsets of `schema`: see `output_batch_rows`. fn restore_group_key_types( batch: RecordBatch, @@ -2957,15 +2993,19 @@ fn restore_group_key_types( if array.data_type() == field.data_type() { Ok(Arc::clone(array)) } else { - // A slice can start beyond the 32-bit offset limit. - // Copy its values so that the first offset is zero. - let data = array.to_data(); - let mut mutable = MutableArrayData::new(vec![&data], false, data.len()); - mutable.try_extend(0, 0, data.len())?; - Ok(arrow::compute::cast( - &make_array(mutable.freeze()), - field.data_type(), - )?) + match (array.data_type(), field.data_type()) { + (DataType::LargeUtf8, DataType::Utf8) => { + let keys = array.as_string::(); + Ok(Arc::new(with_32_bit_offsets::<_, Utf8Type>(keys)?) as _) + } + (DataType::LargeBinary, DataType::Binary) => { + let keys = array.as_binary::(); + Ok(Arc::new(with_32_bit_offsets::<_, BinaryType>(keys)?) as _) + } + (stored_type, output_type) => internal_err!( + "Group keys of type {stored_type} cannot have the output type {output_type}" + ), + } } }) .collect::>>()?; @@ -10228,6 +10268,11 @@ mod tests { k.iter().collect::>(), vec![None, Some("cccc"), Some("dd")] ); + let stored_keys = batch.column(0).as_string::(); + assert!(std::ptr::eq( + k.value_data().as_ptr(), + stored_keys.value_data()[2..].as_ptr() + )); let c = output.column(1).as_primitive::(); assert_eq!(c.values(), &[2, 3, 4]); @@ -10304,9 +10349,7 @@ mod tests { columns .iter() .map(|(data_type, lengths)| { - let offsets = arrow::buffer::OffsetBuffer::::from_lengths( - lengths.iter().copied(), - ); + let offsets = OffsetBuffer::::from_lengths(lengths.iter().copied()); let values = values.clone(); match data_type { DataType::LargeBinary => Arc::new( @@ -10385,28 +10428,33 @@ mod tests { Field::new("c", DataType::Int64, false), ])); let stored_schema = group_schema_with_large_offsets(&output_schema, 1); - // The two keys hold one byte more than `i32::MAX`. - let mut columns = zeroed_group_keys(&[(DataType::LargeUtf8, &[GIB, GIB])]); - columns.push(Arc::new(Int64Array::from(vec![1, 2]))); + // Two keys hold one byte more than `i32::MAX`, and the last key starts + // beyond the 32-bit offset limit. + let mut columns = zeroed_group_keys(&[(DataType::LargeUtf8, &[GIB, GIB, GIB])]); + columns.push(Arc::new(Int64Array::from(vec![1, 2, 3]))); let stored = RecordBatch::try_new(stored_schema, columns)?; + let error = restore_group_key_types(stored.clone(), &output_schema).unwrap_err(); + assert_contains!(error.to_string(), "do not fit 32-bit offsets"); let input = Box::pin(crate::stream::RecordBatchStreamAdapter::new( Arc::clone(&output_schema), - futures::stream::iter(vec![Ok(stored)]), + futures::stream::iter(vec![Ok(stored.clone())]), )); - let mut output = restore_group_key_stream(input, 8192); + let output = collect(restore_group_key_stream(input, 8192)).await?; - // An output batch owns a copy of its keys, so the test holds one - // batch at a time. - let mut counts = vec![]; - while let Some(batch) = output.next().await.transpose()? { + assert_eq!(output.len(), 3); + let stored_keys = stored.column(0).as_string::(); + for (index, batch) in output.iter().enumerate() { assert_eq!(batch.schema(), output_schema); let keys = batch.column(0).as_string::(); assert_eq!(keys.value_offsets(), &[0, GIB as i32]); - counts - .extend_from_slice(batch.column(1).as_primitive::().values()); + assert!(std::ptr::eq( + keys.value_data().as_ptr(), + stored_keys.value_data()[index * GIB..].as_ptr() + )); + let counts = batch.column(1).as_primitive::(); + assert_eq!(counts.values(), &[index as i64 + 1]); } - assert_eq!(counts, vec![1, 2]); Ok(()) }