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 fa7a134733194..044eb1dfe1270 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; @@ -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_schema_with_large_offsets, max_duplicate_ordinal, + output_batch_rows, restore_group_key_types, }; use super::AggregateTableMetrics; @@ -309,7 +310,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 +327,11 @@ impl AggregateHashTable { } }; - let batch = output.next_batch(batch_size)?; + let rows = output.next_batch_rows(&output_schema, batch_size); + let batch = output + .next_batch(rows)? + .map(|batch| restore_group_key_types(batch, &output_schema)) + .transpose()?; if output.is_exhausted() { self.state = AggregateHashTableState::Done; } else { @@ -362,6 +371,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 +398,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); @@ -643,6 +660,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, @@ -972,6 +998,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] @@ -1017,6 +1044,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/aggregate_hash_table/common_clustered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs index 52ac5f7c761f7..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, + evaluate_group_by, group_schema_with_large_offsets, }; use super::AggregateTableMetrics; @@ -320,6 +320,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 +344,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 +424,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 +451,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/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 1ee6d02d13e8e..43550df7eff06 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_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}; @@ -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, @@ -787,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 { @@ -805,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); @@ -817,6 +827,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); @@ -1046,7 +1062,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..93ca24b90a9bd 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -164,12 +164,19 @@ 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::{ + 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; @@ -178,7 +185,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 +561,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 +716,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 +2863,197 @@ 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_key_types_differ(batch: &RecordBatch, schema: &Schema) -> bool { + batch + .columns() + .iter() + .zip(schema.fields()) + .any(|(column, field)| column.data_type() != field.data_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) +} + +// 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, + 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 { + 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::>>()?; + 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; + } + // 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. + let mut offset = 0; + loop { + let rows = output_batch_rows(&batch, &schema, offset, batch_size); + emitter + .emit(restore_group_key_types(batch.slice(offset, rows), &schema)?) + .await; + offset += rows; + if offset >= batch.num_rows() { + break; + } + } + } + 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 +10190,379 @@ 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 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]); + + 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(()) + } + + #[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(()) + } + + 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 = 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); + // 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.clone())]), + )); + let output = collect(restore_group_key_stream(input, 8192)).await?; + + 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]); + 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]); + } + + 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/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()); 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..009eb306de665 --- /dev/null +++ b/datafusion/sqllogictest/test_files/aggregate_byte_group_keys.slt @@ -0,0 +1,444 @@ +# 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 + +# 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' + +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