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 e6e690c4d1e08..45db47f978923 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -27,7 +27,10 @@ use datafusion_expr::{EmitTo, GroupsAccumulator}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use crate::PhysicalExpr; -use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::group_values::{ + GroupByMetrics, GroupValues, cast_group_values_to_schema, + group_value_emit_batch_size, new_group_values, schema_with_group_values, +}; use crate::aggregates::grouped_hash_stream::create_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ @@ -205,6 +208,72 @@ impl AggregateHashTable { Ok(()) } + /// Materializes all partial state once, then returns schema-stable slices. + /// + /// A materialized dictionary group array may need a wider key type for all + /// groups. Each returned partial batch is therefore limited to the planned + /// key capacity and re-encoded to the planned schema after slicing. This + /// preserves the existing one-time materialization semantics while keeping + /// repartition and coalescing schemas stable. + pub(super) fn next_partial_output_batch_inner( + &mut self, + materialize_accumulator_fn: MaterializeAccumulatorFn, + ) -> Result> { + let planned_schema = Arc::clone(&self.output_schema); + let configured_batch_size = self.batch_size; + + let mut output = + match std::mem::replace(&mut self.state, AggregateHashTableState::Done) { + AggregateHashTableState::Outputting(mut state) => { + if state.group_values.is_empty() { + return Ok(None); + } + + let emit_to = EmitTo::All; + let timer = self.group_by_metrics.emitting_time.timer(); + let mut columns = state.group_values.emit(emit_to)?; + let num_group_columns = columns.len(); + for acc in state.accumulators.iter_mut() { + columns.extend(materialize_accumulator_fn(acc, emit_to)?); + } + drop(timer); + + let materialized_schema = schema_with_group_values( + &planned_schema, + &columns[..num_group_columns], + ); + let batch = RecordBatch::try_new(materialized_schema, columns)?; + debug_assert!(batch.num_rows() > 0); + let batch_size = group_value_emit_batch_size( + &planned_schema, + num_group_columns, + configured_batch_size, + ); + MaterializedAggregateOutput::new_with_output_schema( + batch, + planned_schema, + num_group_columns, + batch_size, + ) + } + AggregateHashTableState::OutputtingMaterialized(output) => output, + AggregateHashTableState::Done => return Ok(None), + AggregateHashTableState::Building(_) => { + return internal_err!( + "next_output_batch must be called in the outputting state" + ); + } + }; + + let batch = output.next_batch()?; + if output.is_exhausted() { + self.state = AggregateHashTableState::Done; + } else { + self.state = AggregateHashTableState::OutputtingMaterialized(output); + } + Ok(batch) + } + /// Materializes the full output once, then returns it downstream incrementally /// by slicing it into `batch_size` chunks. /// @@ -233,14 +302,19 @@ impl AggregateHashTable { let emit_to = EmitTo::All; let timer = self.group_by_metrics.emitting_time.timer(); let mut columns = state.group_values.emit(emit_to)?; + let num_group_columns = columns.len(); for acc in state.accumulators.iter_mut() { columns.extend(materialize_accumulator_fn(acc, emit_to)?); } drop(timer); + let output_schema = schema_with_group_values( + &output_schema, + &columns[..num_group_columns], + ); let batch = RecordBatch::try_new(output_schema, columns)?; debug_assert!(batch.num_rows() > 0); - MaterializedAggregateOutput::new(batch) + MaterializedAggregateOutput::new(batch, batch_size) } AggregateHashTableState::OutputtingMaterialized(output) => output, AggregateHashTableState::Done => return Ok(None), @@ -251,7 +325,7 @@ impl AggregateHashTable { } }; - let batch = output.next_batch(batch_size); + let batch = output.next_batch()?; if output.is_exhausted() { self.state = AggregateHashTableState::Done; } else { @@ -413,7 +487,7 @@ pub(super) enum AggregateHashTableState { Building(AggregateHashTableBuffer), /// Emitting results directly from group keys and aggregate state. Outputting(AggregateHashTableBuffer), - /// Materialize all the output results, and then incrementally output in the `OutputtingMaterialized` state. + /// Materialized output that is being sliced into downstream batches. /// /// Note this is a temporary solution until the `GroupValues` issue is solved: /// Issue: @@ -423,29 +497,65 @@ pub(super) enum AggregateHashTableState { /// Fully evaluated aggregate output and the next row offset to emit. /// -/// Final aggregate evaluation consumes accumulator state, and partial terminal -/// output should not repeatedly renumber group values with `EmitTo::First`. -/// Materialize once and then slice to honor `batch_size` across output polls. +/// Partial output can use a narrower target schema than the materialized batch; +/// final output uses the materialized schema directly. pub(super) struct MaterializedAggregateOutput { batch: RecordBatch, + output_schema: SchemaRef, + num_group_columns: usize, + batch_size: usize, offset: usize, } impl MaterializedAggregateOutput { - pub(super) fn new(batch: RecordBatch) -> Self { - Self { batch, offset: 0 } + pub(super) fn new(batch: RecordBatch, batch_size: usize) -> Self { + Self { + output_schema: batch.schema(), + batch, + num_group_columns: 0, + batch_size, + offset: 0, + } + } + + pub(super) fn new_with_output_schema( + batch: RecordBatch, + output_schema: SchemaRef, + num_group_columns: usize, + batch_size: usize, + ) -> Self { + Self { + batch, + output_schema, + num_group_columns, + batch_size, + offset: 0, + } } - pub(super) fn next_batch(&mut self, batch_size: usize) -> Option { - debug_assert!(batch_size > 0); + pub(super) fn next_batch(&mut self) -> Result> { + debug_assert!(self.batch_size > 0); if self.is_exhausted() { - return None; + return Ok(None); } - let length = batch_size.min(self.batch.num_rows() - self.offset); + let length = self.batch_size.min(self.batch.num_rows() - self.offset); let batch = self.batch.slice(self.offset, length); self.offset += length; - Some(batch) + + if batch.schema() == self.output_schema { + return Ok(Some(batch)); + } + + let mut columns = batch.columns().to_vec(); + cast_group_values_to_schema( + &mut columns[..self.num_group_columns], + &self.output_schema, + )?; + Ok(Some(RecordBatch::try_new( + Arc::clone(&self.output_schema), + columns, + )?)) } pub(super) fn is_exhausted(&self) -> bool { @@ -630,12 +740,12 @@ mod tests { schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))], )?; - let mut output = MaterializedAggregateOutput::new(batch); + let mut output = MaterializedAggregateOutput::new(batch, 2); - assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![1, 2]); - assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![3, 4]); - assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![5]); - assert!(output.next_batch(2).is_none()); + assert_eq!(int32_values(&output.next_batch()?.unwrap(), 0), vec![1, 2]); + assert_eq!(int32_values(&output.next_batch()?.unwrap(), 0), vec![3, 4]); + assert_eq!(int32_values(&output.next_batch()?.unwrap(), 0), vec![5]); + assert!(output.next_batch()?.is_none()); assert!(output.is_exhausted()); Ok(()) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index c83303c51d6e8..fa77f7a4d6e4f 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -30,7 +30,10 @@ use datafusion_expr::EmitTo; use crate::InputOrderMode; use crate::PhysicalExpr; -use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::group_values::{ + GroupByMetrics, GroupValues, group_value_emit_batch_size, new_group_values, + schema_with_group_values, +}; use crate::aggregates::grouped_hash_stream::create_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ @@ -234,7 +237,7 @@ impl OrderedAggregateTable { + self.buffer.group_indices.allocated_size() } - /// Returns the [`EmitTo`], clamped to the specified batch size + /// Returns the [`EmitTo`], clamped to the specified maximum batch size. /// /// Returns `(emit_to, should_remove_groups)`, where `emit_to` is the number /// of groups to emit from `GroupValues` / accumulators, and @@ -244,16 +247,18 @@ impl OrderedAggregateTable { &self, group_count: usize, emit_to: EmitTo, + max_batch_size: usize, ) -> (EmitTo, bool) { match emit_to { - EmitTo::First(n) => (EmitTo::First(n.min(self.batch_size)), true), - EmitTo::All if group_count <= self.batch_size => (EmitTo::All, false), - EmitTo::All => (EmitTo::First(self.batch_size), false), + EmitTo::First(n) => (EmitTo::First(n.min(max_batch_size)), true), + EmitTo::All if group_count <= max_batch_size => (EmitTo::All, false), + EmitTo::All => (EmitTo::First(max_batch_size), false), } } + /// Aggregates one evaluated input batch. /// - /// This common utility is used by ordered partial and ordered final aggregation. + /// This common utility is used by ordered partial and final aggregation. /// /// # Argument: `is_final` /// @@ -308,12 +313,12 @@ impl OrderedAggregateTable { /// Emits groups allowed by `GroupOrdering`, leaving only the current /// unfinished ordered-key range buffered. /// - /// This common utility is used by ordered partial and ordered final aggregation. + /// This common utility is used by ordered partial and final aggregation. /// /// # Argument: `is_final` /// - /// - `true`: output final aggregate values. - /// - `false`: output partial accumulator states. + /// - `true`: output final aggregate values and permit dictionary key promotion. + /// - `false`: output partial accumulator states with the planned schema. pub(super) fn next_output_batch_for_mode( &mut self, is_final: bool, @@ -325,11 +330,21 @@ impl OrderedAggregateTable { let Some(emit_to) = self.buffer.group_ordering.emit_to() else { return Ok(None); }; + let max_batch_size = if is_final { + self.batch_size + } else { + group_value_emit_batch_size( + &self.output_schema, + self.buffer.group_by.num_group_exprs(), + self.batch_size, + ) + }; let (emit_to, should_remove_groups) = - self.clamp_emit_to(self.buffer.group_values.len(), emit_to); + self.clamp_emit_to(self.buffer.group_values.len(), emit_to, max_batch_size); let timer = self.group_by_metrics.emitting_time.timer(); let mut output = self.buffer.group_values.emit(emit_to)?; + let num_group_columns = output.len(); if should_remove_groups { match emit_to { EmitTo::First(n) => self.buffer.group_ordering.remove_groups(n), @@ -349,7 +364,12 @@ impl OrderedAggregateTable { } drop(timer); - let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; + let output_schema = if is_final { + schema_with_group_values(&self.output_schema, &output[..num_group_columns]) + } else { + Arc::clone(&self.output_schema) + }; + let batch = RecordBatch::try_new(output_schema, output)?; debug_assert!(batch.num_rows() > 0); Ok(Some(batch)) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs index 0d2495a1b556c..8e969c55a6a5b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs @@ -28,3 +28,6 @@ pub(super) use common::{ PartialSkipMarker, }; pub(super) use common_ordered::OrderedAggregateTable; + +#[cfg(test)] +mod tests; diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs index d8e92c5928b8a..d27df2dea933e 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs @@ -42,14 +42,15 @@ impl AggregateHashTable { /// Emits the next batch of aggregated group keys and aggregate states. /// - /// The output batch size is determined by `self.batch_size`. + /// The output batch size is bounded by both `self.batch_size` and any + /// dictionary group key capacity required to preserve the planned schema. /// /// Returns `Some(batch)` for each emitted batch, `None` when output is /// exhausted, and an internal error if polled in the `Building` state. pub(in crate::aggregates) fn next_output_batch( &mut self, ) -> Result> { - self.next_output_batch_inner(HashAggregateAccumulator::state) + self.next_partial_output_batch_inner(HashAggregateAccumulator::state) } /// Partial-reduce aggregation consumes partial aggregate states and merges diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index ffac42feaa3b3..81d29ba4ff240 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -58,14 +58,15 @@ impl AggregateHashTable { /// Emits the next batch of aggregated group keys and aggregate states. /// - /// The output batch size is determined by `self.batch_size`. + /// The output batch size is bounded by both `self.batch_size` and any + /// dictionary group key capacity required to preserve the planned schema. /// /// Returns `Some(batch)` for each emitted batch, `None` when output is /// exhausted, and an internal error if polled in the `Building` state. pub(in crate::aggregates) fn next_output_batch( &mut self, ) -> Result> { - self.next_output_batch_inner(HashAggregateAccumulator::state) + self.next_partial_output_batch_inner(HashAggregateAccumulator::state) } pub(in crate::aggregates) fn can_skip_aggregation(&self) -> bool { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs new file mode 100644 index 0000000000000..28b765f69a654 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs @@ -0,0 +1,152 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{ArrayRef, DictionaryArray, StringArray, StringDictionaryBuilder}; +use arrow::datatypes::{DataType, Field, Schema, UInt8Type, UInt16Type}; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_physical_expr::expressions::col; + +use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; +use crate::test::TestMemoryExec; + +use super::{AggregateHashTable, FinalMarker, PartialMarker}; + +#[test] +fn dictionary_groups_keep_partial_schema_and_promote_final_output() -> Result<()> { + let input_schema = Arc::new(Schema::new(vec![Field::new( + "group_col", + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), + false, + )])); + let input_batches = (0u32..257) + .map(|value| -> Result { + let mut builder = StringDictionaryBuilder::::new(); + builder.append_value(format!("group_{value}")); + let array: ArrayRef = Arc::new(builder.finish()); + Ok(RecordBatch::try_new( + Arc::clone(&input_schema), + vec![array], + )?) + }) + .collect::>>()?; + + let input = TestMemoryExec::try_new( + std::slice::from_ref(&input_batches), + Arc::clone(&input_schema), + None, + )?; + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + let group_by = PhysicalGroupBy::new_single(vec![( + col("group_col", &input_schema)?, + "group_col".to_string(), + )]); + let partial_exec = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![], + vec![], + input, + Arc::clone(&input_schema), + )?; + let partial_schema = Arc::clone(&partial_exec.schema); + let mut partial_table = AggregateHashTable::::new( + &partial_exec, + 0, + Arc::clone(&partial_schema), + 1024, + )?; + for batch in &input_batches { + partial_table.aggregate_batch(batch)?; + } + partial_table.start_output()?; + + let mut partial_batches = vec![]; + while let Some(batch) = partial_table.next_output_batch()? { + assert_eq!(batch.schema(), partial_schema); + assert_eq!( + batch.schema().field(0).data_type(), + &DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)) + ); + assert!( + batch + .column(0) + .as_any() + .downcast_ref::>() + .is_some() + ); + partial_batches.push(batch); + } + assert_eq!( + partial_batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![256, 1] + ); + + let partial_input = TestMemoryExec::try_new( + std::slice::from_ref(&partial_batches), + Arc::clone(&partial_schema), + None, + )?; + let partial_input = Arc::new(TestMemoryExec::update_cache(&Arc::new(partial_input))); + let final_exec = AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![], + vec![], + partial_input, + Arc::clone(&input_schema), + )?; + let mut final_table = AggregateHashTable::::new( + &final_exec, + 0, + Arc::clone(&final_exec.schema), + 1024, + )?; + for batch in &partial_batches { + final_table.aggregate_batch(batch)?; + } + final_table.start_output()?; + + let batch = final_table.next_output_batch()?.unwrap(); + assert_eq!(batch.num_rows(), 257); + assert_eq!( + batch.schema().field(0).data_type(), + &DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)) + ); + let dictionary = batch + .column(0) + .as_any() + .downcast_ref::>() + .unwrap(); + let values = dictionary + .values() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..257 { + let key = dictionary.keys().value(row) as usize; + assert_eq!(values.value(key), format!("group_{row}")); + } + assert!(final_table.next_output_batch()?.is_none()); + + Ok(()) +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index ee253e5d7afdd..ef9fd54776025 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -23,11 +23,153 @@ use arrow::array::types::{ TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, }; use arrow::array::{ArrayRef, downcast_primitive}; -use arrow::datatypes::{DataType, SchemaRef, TimeUnit}; +use arrow::compute::cast; +use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef, TimeUnit}; use datafusion_common::Result; +use std::sync::Arc; use datafusion_expr::EmitTo; +/// Returns `schema` with its leading group fields updated to the data types +/// actually emitted by [`GroupValues`]. Dictionary key types may grow at +/// runtime, while all field names, nullability, and metadata remain unchanged. +pub(crate) fn schema_with_group_values( + schema: &SchemaRef, + group_values: &[ArrayRef], +) -> SchemaRef { + if group_values + .iter() + .zip(schema.fields()) + .all(|(array, field)| array.data_type() == field.data_type()) + { + return Arc::clone(schema); + } + + let fields = schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| match group_values.get(index) { + Some(array) if array.data_type() != field.data_type() => Arc::new( + field + .as_ref() + .clone() + .with_data_type(array.data_type().clone()), + ), + _ => Arc::clone(field), + }) + .collect::>(); + + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) +} + +/// Returns the largest partial aggregate batch that can be emitted without +/// changing any top-level dictionary group key type in `schema`. +/// +/// Partial aggregate batches cross fixed-schema transport boundaries such as +/// repartition coalescing. Splitting them at the dictionary key capacity keeps +/// that transport schema stable; the final aggregate can still combine all +/// groups and promote its output key type dynamically. +pub(crate) fn group_value_emit_batch_size( + schema: &SchemaRef, + num_group_fields: usize, + batch_size: usize, +) -> usize { + schema + .fields() + .iter() + .take(num_group_fields) + .filter_map(|field| match field.data_type() { + DataType::Dictionary(key_type, _) => dictionary_key_capacity(key_type), + _ => None, + }) + .fold(batch_size, usize::min) + .max(1) +} + +fn dictionary_key_capacity(key_type: &DataType) -> Option { + let capacity = match key_type { + DataType::Int8 => 1u64 << 7, + DataType::Int16 => 1u64 << 15, + DataType::Int32 => 1u64 << 31, + DataType::UInt8 => 1u64 << 8, + DataType::UInt16 => 1u64 << 16, + DataType::UInt32 => 1u64 << 32, + DataType::Int64 | DataType::UInt64 => return Some(usize::MAX), + _ => return None, + }; + Some(usize::try_from(capacity).unwrap_or(usize::MAX)) +} + +/// Casts the leading group-value arrays to the corresponding schema fields. +/// Dictionary arrays are decoded and re-encoded when their key type changes so +/// sliced batches receive compact keys starting at zero. +pub(crate) fn cast_group_values_to_schema( + group_values: &mut [ArrayRef], + schema: &SchemaRef, +) -> Result<()> { + for (array, field) in group_values.iter_mut().zip(schema.fields()) { + if array.data_type() == field.data_type() { + continue; + } + + *array = match (array.data_type(), field.data_type()) { + (DataType::Dictionary(_, _), DataType::Dictionary(_, value_type)) => { + let values = cast(array.as_ref(), value_type.as_ref())?; + cast(values.as_ref(), field.data_type())? + } + _ => cast(array.as_ref(), field.data_type())?, + }; + } + Ok(()) +} + +/// Returns a stable spill schema whose top-level dictionary group keys use +/// their widest key type. This widening is internal to spill IPC and does not +/// change planned or non-spill aggregate output schemas. +pub(crate) fn group_value_spill_schema( + schema: &SchemaRef, + num_group_fields: usize, +) -> SchemaRef { + let fields = schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| { + if index < num_group_fields { + group_value_spill_field(field.as_ref()) + } else { + Arc::clone(field) + } + }) + .collect::>(); + + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) +} + +fn group_value_spill_field(field: &Field) -> FieldRef { + let data_type = match field.data_type() { + DataType::Dictionary(key_type, value_type) => { + let key_type = match key_type.as_ref() { + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => { + DataType::Int64 + } + DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => DataType::UInt64, + _ => key_type.as_ref().clone(), + }; + DataType::Dictionary( + Box::new(key_type), + Box::new(value_type.as_ref().clone()), + ) + } + _ => field.data_type().clone(), + }; + Arc::new(field.clone().with_data_type(data_type)) +} + pub mod multi_group_by; mod row; diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 4976a098ecee5..c688a40f85699 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -15,13 +15,14 @@ // specific language governing permissions and limitations // under the License. -use crate::aggregates::group_values::GroupValues; +use crate::aggregates::group_values::{GroupValues, schema_with_group_values}; use arrow::array::{ Array, ArrayRef, ListArray, PrimitiveArray, RunArray, StructArray, downcast_run_end_index, }; use arrow::compute::cast; -use arrow::datatypes::{DataType, SchemaRef}; +use arrow::datatypes::{DataType, Schema, SchemaRef}; +use arrow::error::ArrowError; use arrow::row::{RowConverter, Rows, SortField}; use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; @@ -43,9 +44,14 @@ use std::sync::Arc; /// It uses the arrow-rs [`Rows`] to store the group values, which is a row-wise /// representation. pub struct GroupValuesRows { - /// The output schema + /// The current output schema. Dictionary key types may grow, but never shrink. schema: SchemaRef, + /// Schema used by the row converter. Top-level dictionaries are materialized + /// to their value type so input batches with promoted dictionary keys remain + /// compatible with the same converter. + row_schema: SchemaRef, + /// Converter for the group values row_converter: RowConverter, @@ -87,8 +93,9 @@ impl GroupValuesRows { // Print a debugging message, so it is clear when the (slower) fallback // GroupValuesRows is used. debug!("Creating GroupValuesRows for schema: {schema}"); + let row_schema = materialized_row_schema(&schema); let row_converter = RowConverter::new( - schema + row_schema .fields() .iter() .map(|f| SortField::new(f.data_type().clone())) @@ -104,6 +111,7 @@ impl GroupValuesRows { row_converter.empty_rows(starting_rows_capacity, starting_data_capacity); Ok(Self { schema, + row_schema, row_converter, map, map_size: 0, @@ -117,11 +125,23 @@ impl GroupValuesRows { impl GroupValues for GroupValuesRows { fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { - // 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.iter().map(normalize_float_zero).collect(); + // Dictionary keys are an encoding detail and can be promoted between + // aggregate stages. Materialize top-level dictionaries before row + // conversion so the stored row representation remains stable. + // Also normalize -0.0 → +0.0 so RowConverter (IEEE 754 totalOrder) and + // primitive hashing both group ±0 together. + let normalized_cols = cols + .iter() + .zip(self.row_schema.fields()) + .map(|(col, field)| { + let col = if col.data_type() == field.data_type() { + Arc::clone(col) + } else { + cast(col.as_ref(), field.data_type())? + }; + Ok(normalize_float_zero(&col)) + }) + .collect::>>()?; let cols = normalized_cols.as_slice(); // Convert the group keys into the row format @@ -243,12 +263,13 @@ impl GroupValues for GroupValuesRows { } }; - // TODO: Materialize dictionaries in group keys - // https://github.com/apache/datafusion/issues/7647 - for (field, array) in self.schema.fields.iter().zip(&mut output) { - let expected = field.data_type(); - *array = dictionary_encode_if_necessary(array, expected)?; + // Re-encode top-level dictionary group keys using the current key type. + // If an emitted array no longer fits, grow to the next key width and + // retain that promoted type for all later emissions. + for (field, array) in self.schema.fields().iter().zip(&mut output) { + *array = dictionary_encode_with_promotion(array, field.data_type())?; } + self.schema = schema_with_group_values(&self.schema, &output); self.group_values = Some(group_values); Ok(output) @@ -267,6 +288,46 @@ impl GroupValues for GroupValuesRows { } } +fn materialized_row_schema(schema: &SchemaRef) -> SchemaRef { + let fields = schema + .fields() + .iter() + .map(|field| match field.data_type() { + DataType::Dictionary(_, value_type) => Arc::new( + field + .as_ref() + .clone() + .with_data_type(value_type.as_ref().clone()), + ), + _ => Arc::clone(field), + }) + .collect::>(); + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) +} + +fn dictionary_encode_with_promotion( + array: &ArrayRef, + expected: &DataType, +) -> Result { + if !matches!(expected, DataType::Dictionary(_, _)) { + return dictionary_encode_if_necessary(array, expected); + } + + let mut target = expected.clone(); + loop { + match cast(array.as_ref(), &target) { + Ok(array) => return Ok(array), + Err(ArrowError::DictionaryKeyOverflowError) => { + let Some(promoted) = promote_dictionary_type(&target) else { + return Err(ArrowError::DictionaryKeyOverflowError.into()); + }; + target = promoted; + } + Err(error) => return Err(error.into()), + } + } +} + fn dictionary_encode_if_necessary( array: &ArrayRef, expected: &DataType, @@ -332,3 +393,136 @@ fn dictionary_encode_if_necessary( (_, _) => Ok(Arc::::clone(array)), } } + +fn promote_dictionary_type(data_type: &DataType) -> Option { + let DataType::Dictionary(key_type, value_type) = data_type else { + return None; + }; + let key_type = match key_type.as_ref() { + DataType::Int8 => DataType::Int16, + DataType::Int16 => DataType::Int32, + DataType::Int32 => DataType::Int64, + DataType::UInt8 => DataType::UInt16, + DataType::UInt16 => DataType::UInt32, + DataType::UInt32 => DataType::UInt64, + DataType::Int64 | DataType::UInt64 => return None, + _ => return None, + }; + Some(DataType::Dictionary( + Box::new(key_type), + Box::new(value_type.as_ref().clone()), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{DictionaryArray, StringArray, StringDictionaryBuilder}; + use arrow::datatypes::{Field, Int8Type, Int16Type, Schema, UInt8Type, UInt16Type}; + + #[test] + fn dict_uint8_utf8_promotes_and_does_not_shrink() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "k", + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), + false, + )])); + + let mut group_values = GroupValuesRows::try_new(Arc::clone(&schema))?; + let mut groups = vec![]; + for i in 0u32..257 { + let mut builder = StringDictionaryBuilder::::new(); + builder.append_value(format!("group_{i}")); + let array: ArrayRef = Arc::new(builder.finish()); + group_values.intern(&[array], &mut groups)?; + } + + let arrays = group_values.emit(EmitTo::All)?; + assert_eq!( + arrays[0].data_type(), + &DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)) + ); + let dictionary = arrays[0] + .as_any() + .downcast_ref::>() + .unwrap(); + let values = dictionary + .values() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..257 { + let key = dictionary.keys().value(row) as usize; + assert_eq!(values.value(key), format!("group_{row}")); + } + + let mut builder = StringDictionaryBuilder::::new(); + builder.append_value("after_promotion"); + group_values.intern(&[Arc::new(builder.finish())], &mut groups)?; + let arrays = group_values.emit(EmitTo::All)?; + assert_eq!( + arrays[0].data_type(), + &DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)) + ); + + Ok(()) + } + + #[test] + fn dict_uint8_utf8_keeps_key_type_at_capacity() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "k", + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), + false, + )])); + let mut group_values = GroupValuesRows::try_new(schema)?; + let mut groups = vec![]; + for i in 0u32..256 { + let mut builder = StringDictionaryBuilder::::new(); + builder.append_value(format!("group_{i}")); + group_values.intern(&[Arc::new(builder.finish())], &mut groups)?; + } + + let arrays = group_values.emit(EmitTo::All)?; + assert_eq!( + arrays[0].data_type(), + &DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)) + ); + assert!( + arrays[0] + .as_any() + .downcast_ref::>() + .is_some() + ); + Ok(()) + } + + #[test] + fn dict_int8_utf8_promotes_after_capacity() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "k", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + false, + )])); + let mut group_values = GroupValuesRows::try_new(schema)?; + let mut groups = vec![]; + for i in 0u32..129 { + let mut builder = StringDictionaryBuilder::::new(); + builder.append_value(format!("group_{i}")); + group_values.intern(&[Arc::new(builder.finish())], &mut groups)?; + } + + let arrays = group_values.emit(EmitTo::All)?; + assert_eq!( + arrays[0].data_type(), + &DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)) + ); + assert!( + arrays[0] + .as_any() + .downcast_ref::>() + .is_some() + ); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 0d00e5c4d0d86..8cf5531b2507b 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -24,7 +24,11 @@ use std::vec; use super::order::GroupOrdering; use super::skip_partial::SkipAggregationProbe; use super::{AggregateExec, format_human_display}; -use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::group_values::{ + GroupByMetrics, GroupValues, cast_group_values_to_schema, + group_value_emit_batch_size, group_value_spill_schema, new_group_values, + schema_with_group_values, +}; use crate::aggregates::order::GroupOrderingFull; use crate::aggregates::{ AggregateInputMode, AggregateMode, AggregateOutputMode, PhysicalGroupBy, @@ -424,6 +428,7 @@ impl GroupedHashAggregateStream { .collect::>()?; let group_schema = agg_group_by.group_schema(&agg.input().schema())?; + let num_group_fields = group_schema.fields().len(); // fix https://github.com/apache/datafusion/issues/13949 // Builds a **partial aggregation** schema by combining the group columns and @@ -446,6 +451,7 @@ impl GroupedHashAggregateStream { &aggregate_exprs, AggregateMode::Partial, )?); + let spill_schema = group_value_spill_schema(&spill_schema, num_group_fields); // Need to update the GROUP BY expressions to point to the correct column after schema change let merging_group_by_expr = agg_group_by @@ -780,12 +786,18 @@ impl Stream for GroupedHashAggregateStream { (self.exec_state, output_batch) = if batch.num_rows() <= size { ( if self.input_done { - ExecutionState::Done + if self.group_values.is_empty() { + ExecutionState::Done + } else { + ExecutionState::ReadingInput + } } // In Partial aggregation, we also need to check - // if we should trigger partial skipping + // if we should trigger partial skipping after every + // accumulated group has been emitted. else if self.mode == AggregateMode::Partial && self.should_skip_aggregation() + && self.group_values.is_empty() { ExecutionState::SkippingAggregation } else { @@ -1023,17 +1035,45 @@ impl GroupedHashAggregateStream { /// Create an output RecordBatch with the group keys and /// accumulator states/values specified in emit_to fn emit(&mut self, emit_to: EmitTo, spilling: bool) -> Result> { - let schema = if spilling { - Arc::clone(&self.spill_state.spill_schema) - } else { - self.schema() - }; if self.group_values.is_empty() { return Ok(None); } + let emit_to = + if !spilling && self.mode.output_mode() == AggregateOutputMode::Partial { + let max_groups = group_value_emit_batch_size( + &self.schema, + self.group_by.num_group_exprs(), + self.batch_size, + ); + match emit_to { + EmitTo::First(n) => EmitTo::First(n.min(max_groups)), + EmitTo::All if self.group_values.len() > max_groups => { + EmitTo::First(max_groups) + } + EmitTo::All => EmitTo::All, + } + } else { + emit_to + }; + let timer = self.group_by_metrics.emitting_time.timer(); let mut output = self.group_values.emit(emit_to)?; + let num_group_columns = output.len(); + let schema = if spilling { + cast_group_values_to_schema( + &mut output[..num_group_columns], + &self.spill_state.spill_schema, + )?; + Arc::clone(&self.spill_state.spill_schema) + } else if self.mode.output_mode() == AggregateOutputMode::Partial { + Arc::clone(&self.schema) + } else { + let schema = + schema_with_group_values(&self.schema, &output[..num_group_columns]); + self.schema = Arc::clone(&schema); + schema + }; if let EmitTo::First(n) = emit_to { self.group_ordering.remove_groups(n); } @@ -1312,10 +1352,7 @@ impl GroupedHashAggregateStream { // in first-seen order, as required by `GroupOrderingFull`. // The pre-spill multi-column collector may use `vectorized_intern`, which // can assign new group ids out of input order under hash collisions. - let group_schema = self - .spill_state - .merging_group_by - .group_schema(&self.spill_state.spill_schema)?; + let group_schema = self.group_by.group_schema(&self.input_schema)?; if group_schema.fields().len() > 1 { self.group_values = new_group_values(group_schema, &self.group_ordering)?; } @@ -1338,7 +1375,7 @@ impl GroupedHashAggregateStream { fn update_skip_aggregation_probe(&mut self, input_rows: usize) { if let Some(probe) = self.skip_aggregation_probe.as_mut() { // Skip aggregation probe is not supported if stream has any spills, - // currently spilling is not supported for Partial aggregation + // currently spilling is not supported in Partial aggregation assert!(self.spill_state.spills.is_empty()); probe.update_state(input_rows, self.group_values.len()); }; @@ -1383,6 +1420,7 @@ impl GroupedHashAggregateStream { "group_values expected to have single element" ); let mut output = group_values.swap_remove(0); + cast_group_values_to_schema(&mut output, &self.schema)?; let iter = self .accumulators @@ -1406,13 +1444,68 @@ mod tests { use super::*; use crate::InputOrderMode; use crate::test::TestMemoryExec; - use arrow::array::{Int32Array, Int64Array}; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::array::{ + DictionaryArray, Int32Array, Int64Array, StringDictionaryBuilder, + }; + use arrow::datatypes::{DataType, Field, Schema, UInt8Type, UInt16Type}; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::count::count_udaf; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::col; + #[tokio::test] + async fn dictionary_group_key_promotes_runtime_schema() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "group_col", + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), + false, + )])); + let batches = (0u32..257) + .map(|value| -> Result { + let mut builder = StringDictionaryBuilder::::new(); + builder.append_value(format!("group_{value}")); + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(builder.finish())], + )?) + }) + .collect::>>()?; + let exec = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![( + col("group_col", &schema)?, + "group_col".to_string(), + )]), + vec![], + vec![], + exec, + Arc::clone(&schema), + )?; + + let mut stream = GroupedHashAggregateStream::new( + &aggregate_exec, + &Arc::new(TaskContext::default()), + 0, + )?; + let batch = stream.next().await.unwrap()?; + assert_eq!(batch.num_rows(), 257); + assert_eq!( + batch.schema().field(0).data_type(), + &DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)) + ); + assert!( + batch + .column(0) + .as_any() + .downcast_ref::>() + .is_some() + ); + assert!(stream.next().await.is_none()); + Ok(()) + } + // Migrated to PartialHashAggregateStream coverage in hash_stream.rs; // kept here for the legacy GroupedHashAggregateStream implementation. #[tokio::test] diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 62b92965030ae..323f90022175f 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -899,6 +899,7 @@ impl FinalHashAggregateStream { match result { Ok(Some(batch)) => { + self.schema = batch.schema(); let _ = self .reservation .try_resize(original_state.hash_table().memory_size()); diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 89653e05ab4c7..e7ae8e30568d3 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -169,6 +169,7 @@ impl OrderedFinalAggregateStream { // Some finalized groups can be emitted. Yield them, then // continue aggregating input in the current state. Ok(Some(batch)) => { + self.schema = batch.schema(); let next_state = OrderedFinalAggregateState::ReadingInput { table }; self.resize_reservation_for_state(&next_state); @@ -237,6 +238,7 @@ impl OrderedFinalAggregateStream { match result { Ok(Some(batch)) => { + self.schema = batch.schema(); let next_state = if table.is_empty() { OrderedFinalAggregateState::Done } else { diff --git a/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt b/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt new file mode 100644 index 0000000000000..ae7b6b577828c --- /dev/null +++ b/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt @@ -0,0 +1,75 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Each input dictionary batch fits in UInt8. With four target partitions and a +# batch size of 128, one partial aggregate receives 257 distinct groups. Its +# output must be split into schema-stable UInt8 batches before repartitioning. +statement ok +SET datafusion.execution.batch_size = 128; + +statement ok +SET datafusion.execution.target_partitions = 4; + +query I +SELECT COUNT(*) +FROM ( + SELECT k + FROM ( + SELECT arrow_cast(CAST(value AS VARCHAR), 'Dictionary(UInt8, Utf8)') AS k + FROM generate_series(0, 1024) AS series(value) + ) + GROUP BY k +); +---- +1025 + +# A single final aggregate sees all 257 groups and must promote the emitted +# dictionary key type from UInt8 to UInt16. +statement ok +SET datafusion.execution.target_partitions = 1; + +query I +SELECT COUNT(*) +FROM ( + SELECT k + FROM ( + SELECT arrow_cast(CAST(value AS VARCHAR), 'Dictionary(UInt8, Utf8)') AS k + FROM generate_series(0, 256) AS series(value) + ) + GROUP BY k +); +---- +257 + +query T rowsort +SELECT DISTINCT arrow_typeof(k) +FROM ( + SELECT k + FROM ( + SELECT arrow_cast(CAST(value AS VARCHAR), 'Dictionary(UInt8, Utf8)') AS k + FROM generate_series(0, 256) AS series(value) + ) + GROUP BY k +); +---- +Dictionary(UInt16, Utf8) + +statement ok +SET datafusion.execution.batch_size = 8192; + +statement ok +SET datafusion.execution.target_partitions = 4;