From ad6241ebde1ff168f77c33acfc4e88808be35eeb Mon Sep 17 00:00:00 2001
From: Jamie Currier <247854506+StealthEyeLLC@users.noreply.github.com>
Date: Thu, 9 Jul 2026 04:57:15 +0000
Subject: [PATCH 01/28] fix: avoid dictionary key overflow in group value
 output

---
 .../src/aggregates/group_values/mod.rs        | 61 ++++++++++++++++++-
 .../src/aggregates/group_values/row.rs        | 57 ++++++++++++++++-
 .../physical-plan/src/aggregates/mod.rs       | 30 ++++++++-
 3 files changed, 144 insertions(+), 4 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
index ee253e5d7afdd..99823ae8b774b 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
@@ -23,11 +23,70 @@ use arrow::array::types::{
     TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
 };
 use arrow::array::{ArrayRef, downcast_primitive};
-use arrow::datatypes::{DataType, SchemaRef, TimeUnit};
+use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef, TimeUnit};
 use datafusion_common::Result;
+use std::sync::Arc;
 
 use datafusion_expr::EmitTo;
 
+pub(crate) fn group_value_output_schema(schema: &Schema) -> SchemaRef {
+    Arc::new(Schema::new(
+        schema
+            .fields()
+            .iter()
+            .map(|field| group_value_output_field(field.as_ref()))
+            .collect::<Vec<_>>(),
+    ))
+}
+
+pub(crate) fn group_value_output_field(field: &Field) -> FieldRef {
+    Arc::new(
+        Field::new(
+            field.name(),
+            group_value_output_data_type(field.data_type()),
+            field.is_nullable(),
+        )
+        .with_metadata(field.metadata().clone()),
+    )
+}
+
+fn group_value_output_data_type(data_type: &DataType) -> DataType {
+    match data_type {
+        DataType::Dictionary(_, value_type) => DataType::Dictionary(
+            Box::new(DataType::UInt64),
+            Box::new(group_value_output_data_type(value_type)),
+        ),
+        DataType::Struct(fields) => DataType::Struct(
+            fields
+                .iter()
+                .map(|field| group_value_output_field(field.as_ref()))
+                .collect::<Vec<_>>()
+                .into(),
+        ),
+        DataType::List(field) => DataType::List(group_value_output_field(field.as_ref())),
+        DataType::LargeList(field) => {
+            DataType::LargeList(group_value_output_field(field.as_ref()))
+        }
+        DataType::ListView(field) => {
+            DataType::ListView(group_value_output_field(field.as_ref()))
+        }
+        DataType::LargeListView(field) => {
+            DataType::LargeListView(group_value_output_field(field.as_ref()))
+        }
+        DataType::FixedSizeList(field, size) => {
+            DataType::FixedSizeList(group_value_output_field(field.as_ref()), *size)
+        }
+        DataType::Map(field, sorted) => {
+            DataType::Map(group_value_output_field(field.as_ref()), *sorted)
+        }
+        DataType::RunEndEncoded(run_ends, values) => DataType::RunEndEncoded(
+            Arc::clone(run_ends),
+            group_value_output_field(values.as_ref()),
+        ),
+        _ => data_type.clone(),
+    }
+}
+
 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..61cc21c801708 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/row.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs
@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::aggregates::group_values::GroupValues;
+use crate::aggregates::group_values::{GroupValues, group_value_output_schema};
 use arrow::array::{
     Array, ArrayRef, ListArray, PrimitiveArray, RunArray, StructArray,
     downcast_run_end_index,
@@ -83,6 +83,7 @@ pub struct GroupValuesRows {
 }
 
 impl GroupValuesRows {
+    #[expect(clippy::needless_pass_by_value)]
     pub fn try_new(schema: SchemaRef) -> Result<Self> {
         // Print a debugging message, so it is clear when the (slower) fallback
         // GroupValuesRows is used.
@@ -102,8 +103,9 @@ impl GroupValuesRows {
         let starting_data_capacity = 64 * starting_rows_capacity;
         let rows_buffer =
             row_converter.empty_rows(starting_rows_capacity, starting_data_capacity);
+        let output_schema = group_value_output_schema(schema.as_ref());
         Ok(Self {
-            schema,
+            schema: output_schema,
             row_converter,
             map,
             map_size: 0,
@@ -332,3 +334,54 @@ fn dictionary_encode_if_necessary(
         (_, _) => Ok(Arc::<dyn Array>::clone(array)),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{DictionaryArray, StringArray, StringDictionaryBuilder};
+    use arrow::datatypes::{Field, Schema, UInt8Type, UInt64Type};
+
+    #[test]
+    fn dict_uint8_utf8_257_groups_emit_promotes_key_type() -> 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::<UInt8Type>::new();
+            builder.append_value(format!("group_{i}"));
+            let array: ArrayRef = Arc::new(builder.finish());
+            group_values.intern(&[array], &mut groups)?;
+        }
+
+        assert_eq!(group_values.len(), 257);
+        let arrays = group_values.emit(EmitTo::All)?;
+        assert_eq!(arrays.len(), 1);
+        assert_eq!(arrays[0].len(), 257);
+        assert_eq!(
+            arrays[0].data_type(),
+            &DataType::Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8))
+        );
+
+        let dictionary = arrays[0]
+            .as_any()
+            .downcast_ref::<DictionaryArray<UInt64Type>>()
+            .unwrap();
+        let values = dictionary
+            .values()
+            .as_any()
+            .downcast_ref::<StringArray>()
+            .unwrap();
+
+        for row in 0..257 {
+            let key = dictionary.keys().value(row) as usize;
+            assert_eq!(values.value(key), format!("group_{row}"));
+        }
+
+        Ok(())
+    }
+}
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index d7c72253ecc0c..9ee06e92e3b11 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -465,7 +465,10 @@ impl PhysicalGroupBy {
     fn output_fields(&self, input_schema: &Schema) -> Result<Vec<FieldRef>> {
         let mut fields = self.group_fields(input_schema)?;
         fields.truncate(self.num_output_exprs());
-        Ok(fields)
+        Ok(fields
+            .iter()
+            .map(|field| group_values::group_value_output_field(field.as_ref()))
+            .collect())
     }
 
     /// Returns the `PhysicalGroupBy` for a final aggregation if `self` is used for a partial
@@ -7016,4 +7019,29 @@ mod tests {
         assert!(agg.with_dynamic_filter_expr(df).is_err());
         Ok(())
     }
+
+    #[test]
+    fn aggregate_output_schema_widens_dictionary_group_key_type() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "k",
+            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
+            false,
+        )]));
+        let group_by =
+            PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".to_string())]);
+
+        let group_schema = group_by.group_schema(&schema)?;
+        assert_eq!(
+            group_schema.field(0).data_type(),
+            &DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8))
+        );
+
+        let output_schema = create_schema(&schema, &group_by, &[], AggregateMode::Final)?;
+        assert_eq!(
+            output_schema.field(0).data_type(),
+            &DataType::Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8))
+        );
+
+        Ok(())
+    }
 }

From 711433c9cc6e0be0942f2fbcdb5202c4fe020d04 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 11:48:01 -0400
Subject: [PATCH 02/28] fix: support runtime dictionary key promotion

---
 .../src/aggregates/group_values/mod.rs        | 131 ++++++++++++++----
 1 file changed, 104 insertions(+), 27 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
index 99823ae8b774b..33985d7ca142a 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
@@ -23,65 +23,142 @@ use arrow::array::types::{
     TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
 };
 use arrow::array::{ArrayRef, downcast_primitive};
+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;
 
-pub(crate) fn group_value_output_schema(schema: &Schema) -> SchemaRef {
-    Arc::new(Schema::new(
-        schema
-            .fields()
-            .iter()
-            .map(|field| group_value_output_field(field.as_ref()))
-            .collect::<Vec<_>>(),
+/// 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::<Vec<_>>();
+
+    Arc::new(Schema::new_with_metadata(
+        fields,
+        schema.metadata().clone(),
     ))
 }
 
-pub(crate) fn group_value_output_field(field: &Field) -> FieldRef {
+/// Casts the leading group-value arrays to the corresponding schema fields.
+/// This is used by spill files, which require one stable IPC schema even though
+/// normal aggregate output may promote dictionary keys dynamically.
+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() {
+            *array = cast(array.as_ref(), field.data_type())?;
+        }
+    }
+    Ok(())
+}
+
+/// Returns a stable spill schema whose dictionary group keys use their widest
+/// key type. The 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::<Vec<_>>();
+
+    Arc::new(Schema::new_with_metadata(
+        fields,
+        schema.metadata().clone(),
+    ))
+}
+
+fn group_value_spill_field(field: &Field) -> FieldRef {
     Arc::new(
-        Field::new(
-            field.name(),
-            group_value_output_data_type(field.data_type()),
-            field.is_nullable(),
-        )
-        .with_metadata(field.metadata().clone()),
+        field
+            .clone()
+            .with_data_type(group_value_spill_data_type(field.data_type())),
     )
 }
 
-fn group_value_output_data_type(data_type: &DataType) -> DataType {
+fn group_value_spill_data_type(data_type: &DataType) -> DataType {
     match data_type {
-        DataType::Dictionary(_, value_type) => DataType::Dictionary(
-            Box::new(DataType::UInt64),
-            Box::new(group_value_output_data_type(value_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(group_value_spill_data_type(value_type)),
+            )
+        }
         DataType::Struct(fields) => DataType::Struct(
             fields
                 .iter()
-                .map(|field| group_value_output_field(field.as_ref()))
+                .map(|field| group_value_spill_field(field.as_ref()))
                 .collect::<Vec<_>>()
                 .into(),
         ),
-        DataType::List(field) => DataType::List(group_value_output_field(field.as_ref())),
+        DataType::List(field) => DataType::List(group_value_spill_field(field.as_ref())),
         DataType::LargeList(field) => {
-            DataType::LargeList(group_value_output_field(field.as_ref()))
+            DataType::LargeList(group_value_spill_field(field.as_ref()))
         }
         DataType::ListView(field) => {
-            DataType::ListView(group_value_output_field(field.as_ref()))
+            DataType::ListView(group_value_spill_field(field.as_ref()))
         }
         DataType::LargeListView(field) => {
-            DataType::LargeListView(group_value_output_field(field.as_ref()))
+            DataType::LargeListView(group_value_spill_field(field.as_ref()))
         }
         DataType::FixedSizeList(field, size) => {
-            DataType::FixedSizeList(group_value_output_field(field.as_ref()), *size)
+            DataType::FixedSizeList(group_value_spill_field(field.as_ref()), *size)
         }
         DataType::Map(field, sorted) => {
-            DataType::Map(group_value_output_field(field.as_ref()), *sorted)
+            DataType::Map(group_value_spill_field(field.as_ref()), *sorted)
         }
         DataType::RunEndEncoded(run_ends, values) => DataType::RunEndEncoded(
             Arc::clone(run_ends),
-            group_value_output_field(values.as_ref()),
+            group_value_spill_field(values.as_ref()),
         ),
         _ => data_type.clone(),
     }

From 6a133e50e8073673ccf11ccc3a998f4e99207de7 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 11:50:58 -0400
Subject: [PATCH 03/28] fix: promote dictionary group keys on emit

---
 .../src/aggregates/group_values/row.rs        | 220 +++++++++++++++---
 1 file changed, 187 insertions(+), 33 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs
index 61cc21c801708..a149e46ca6bc4 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, group_value_output_schema};
+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, FieldRef, 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,
 
@@ -83,13 +89,13 @@ pub struct GroupValuesRows {
 }
 
 impl GroupValuesRows {
-    #[expect(clippy::needless_pass_by_value)]
     pub fn try_new(schema: SchemaRef) -> Result<Self> {
         // 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()))
@@ -103,9 +109,9 @@ impl GroupValuesRows {
         let starting_data_capacity = 64 * starting_rows_capacity;
         let rows_buffer =
             row_converter.empty_rows(starting_rows_capacity, starting_data_capacity);
-        let output_schema = group_value_output_schema(schema.as_ref());
         Ok(Self {
-            schema: output_schema,
+            schema,
+            row_schema,
             row_converter,
             map,
             map_size: 0,
@@ -119,11 +125,23 @@ impl GroupValuesRows {
 
 impl GroupValues for GroupValuesRows {
     fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> 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<ArrayRef> =
-            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::<Result<Vec<_>>>()?;
         let cols = normalized_cols.as_slice();
 
         // Convert the group keys into the row format
@@ -245,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 dictionary group keys using the current key type. If the
+        // emitted cardinality 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_if_necessary(array, field.data_type())?;
         }
+        self.schema = schema_with_group_values(&self.schema, &output);
 
         self.group_values = Some(group_values);
         Ok(output)
@@ -269,6 +288,26 @@ 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::<Vec<_>>();
+    Arc::new(Schema::new_with_metadata(
+        fields,
+        schema.metadata().clone(),
+    ))
+}
+
 fn dictionary_encode_if_necessary(
     array: &ArrayRef,
     expected: &DataType,
@@ -283,27 +322,49 @@ fn dictionary_encode_if_necessary(
                     dictionary_encode_if_necessary(column, expected_field.data_type())
                 })
                 .collect::<Result<Vec<_>>>()?;
+            let fields = expected_fields
+                .iter()
+                .zip(&arrays)
+                .map(|(field, array)| field_with_array_type(field, array))
+                .collect::<Vec<_>>()
+                .into();
 
             Ok(Arc::new(StructArray::try_new(
-                expected_fields.clone(),
+                fields,
                 arrays,
                 struct_array.nulls().cloned(),
             )?))
         }
         (DataType::List(expected_field), &DataType::List(_)) => {
             let list = array.as_any().downcast_ref::<ListArray>().unwrap();
+            let values = dictionary_encode_if_necessary(
+                list.values(),
+                expected_field.data_type(),
+            )?;
+            let field = field_with_array_type(expected_field, &values);
 
             Ok(Arc::new(ListArray::try_new(
-                Arc::<arrow::datatypes::Field>::clone(expected_field),
+                field,
                 list.offsets().clone(),
-                dictionary_encode_if_necessary(
-                    list.values(),
-                    expected_field.data_type(),
-                )?,
+                values,
                 list.nulls().cloned(),
             )?))
         }
-        (DataType::Dictionary(_, _), _) => Ok(cast(array.as_ref(), expected)?),
+        (DataType::Dictionary(_, _), _) => {
+            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()),
+                }
+            }
+        }
         (
             DataType::RunEndEncoded(run_ends_field, expected_values_field),
             &DataType::RunEndEncoded(_, _),
@@ -335,14 +396,49 @@ fn dictionary_encode_if_necessary(
     }
 }
 
+fn field_with_array_type(field: &FieldRef, array: &ArrayRef) -> FieldRef {
+    if field.data_type() == array.data_type() {
+        Arc::clone(field)
+    } else {
+        Arc::new(
+            field
+                .as_ref()
+                .clone()
+                .with_data_type(array.data_type().clone()),
+        )
+    }
+}
+
+fn promote_dictionary_type(data_type: &DataType) -> Option<DataType> {
+    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, Schema, UInt8Type, UInt64Type};
+    use arrow::datatypes::{
+        Field, Int8Type, Int16Type, Schema, UInt8Type, UInt16Type,
+    };
 
     #[test]
-    fn dict_uint8_utf8_257_groups_emit_promotes_key_type() -> Result<()> {
+    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)),
@@ -358,30 +454,88 @@ mod tests {
             group_values.intern(&[array], &mut groups)?;
         }
 
-        assert_eq!(group_values.len(), 257);
         let arrays = group_values.emit(EmitTo::All)?;
-        assert_eq!(arrays.len(), 1);
-        assert_eq!(arrays[0].len(), 257);
         assert_eq!(
             arrays[0].data_type(),
-            &DataType::Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8))
+            &DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8))
         );
-
         let dictionary = arrays[0]
             .as_any()
-            .downcast_ref::<DictionaryArray<UInt64Type>>()
+            .downcast_ref::<DictionaryArray<UInt16Type>>()
             .unwrap();
         let values = dictionary
             .values()
             .as_any()
             .downcast_ref::<StringArray>()
             .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::<UInt8Type>::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::<UInt8Type>::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::<DictionaryArray<UInt8Type>>()
+            .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::<Int8Type>::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::<DictionaryArray<Int16Type>>()
+            .is_some());
         Ok(())
     }
 }

From ca2a9e05e18ac18b11805019cca190966b2bd3f3 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 11:51:23 -0400
Subject: [PATCH 04/28] fix: keep planned aggregate dictionary schema unchanged

---
 .../physical-plan/src/aggregates/mod.rs       | 30 +------------------
 1 file changed, 1 insertion(+), 29 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index 9ee06e92e3b11..d7c72253ecc0c 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -465,10 +465,7 @@ impl PhysicalGroupBy {
     fn output_fields(&self, input_schema: &Schema) -> Result<Vec<FieldRef>> {
         let mut fields = self.group_fields(input_schema)?;
         fields.truncate(self.num_output_exprs());
-        Ok(fields
-            .iter()
-            .map(|field| group_values::group_value_output_field(field.as_ref()))
-            .collect())
+        Ok(fields)
     }
 
     /// Returns the `PhysicalGroupBy` for a final aggregation if `self` is used for a partial
@@ -7019,29 +7016,4 @@ mod tests {
         assert!(agg.with_dynamic_filter_expr(df).is_err());
         Ok(())
     }
-
-    #[test]
-    fn aggregate_output_schema_widens_dictionary_group_key_type() -> Result<()> {
-        let schema = Arc::new(Schema::new(vec![Field::new(
-            "k",
-            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
-            false,
-        )]));
-        let group_by =
-            PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".to_string())]);
-
-        let group_schema = group_by.group_schema(&schema)?;
-        assert_eq!(
-            group_schema.field(0).data_type(),
-            &DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8))
-        );
-
-        let output_schema = create_schema(&schema, &group_by, &[], AggregateMode::Final)?;
-        assert_eq!(
-            output_schema.field(0).data_type(),
-            &DataType::Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8))
-        );
-
-        Ok(())
-    }
 }

From 3b556fa653ea181dd6b347c5dfa71e832d0e2f61 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 11:54:12 -0400
Subject: [PATCH 05/28] fix: use emitted group schema in hash aggregate output

---
 .../src/aggregates/aggregate_hash_table/common.rs     | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
index e6e690c4d1e08..7c7c2cba836c3 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,9 @@ 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, new_group_values, schema_with_group_values,
+};
 use crate::aggregates::grouped_hash_stream::create_group_accumulator;
 use crate::aggregates::order::GroupOrdering;
 use crate::aggregates::{
@@ -233,11 +235,16 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
                     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)
@@ -336,7 +343,7 @@ pub(super) type AggregateAccumulator = HashAggregateAccumulator;
 /// Arguments:
 /// * accumulator to update.
 /// * accumulator's evaluated arguments and optional filter.
-/// * one group index per input row, mapping each row to its interned group.
+/// * group index per input row, mapping each row to its interned group.
 /// * total number of groups currently interned in that buffer, including newly
 ///   interned groups.
 pub(super) type AggregateBatchFn = fn(

From 2bf64ac3b876071c4910a004a39df76f5a307c0a Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 11:54:47 -0400
Subject: [PATCH 06/28] fix: use emitted group schema in ordered aggregate
 output

---
 .../aggregate_hash_table/common_ordered.rs        | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

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..ec6d0c066d701 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,9 @@ 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, new_group_values, schema_with_group_values,
+};
 use crate::aggregates::grouped_hash_stream::create_group_accumulator;
 use crate::aggregates::order::GroupOrdering;
 use crate::aggregates::{
@@ -253,7 +255,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
     }
     /// 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,7 +310,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
     /// 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`
     ///
@@ -330,6 +332,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
 
         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 +352,11 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         }
         drop(timer);
 
-        let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?;
+        let output_schema = schema_with_group_values(
+            &self.output_schema,
+            &output[..num_group_columns],
+        );
+        let batch = RecordBatch::try_new(output_schema, output)?;
         debug_assert!(batch.num_rows() > 0);
 
         Ok(Some(batch))

From aaf35effc0522c8684a2101602649aa5cdca706e Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 11:59:24 -0400
Subject: [PATCH 07/28] fix: propagate promoted dictionary schemas through
 legacy aggregation

---
 .../src/aggregates/grouped_hash_stream.rs     | 96 ++++++++++++++++---
 1 file changed, 81 insertions(+), 15 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
index 0d00e5c4d0d86..3df5c4aceb305 100644
--- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
@@ -24,7 +24,10 @@ 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_spill_schema, new_group_values, schema_with_group_values,
+};
 use crate::aggregates::order::GroupOrderingFull;
 use crate::aggregates::{
     AggregateInputMode, AggregateMode, AggregateOutputMode, PhysicalGroupBy,
@@ -424,6 +427,7 @@ impl GroupedHashAggregateStream {
             .collect::<Result<_>>()?;
 
         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 +450,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
@@ -760,7 +765,7 @@ impl Stream for GroupedHashAggregateStream {
                             if !self.group_values.is_empty() {
                                 return Poll::Ready(Some(internal_err!(
                                     "Switching from SkippingAggregation to Done with {} groups still in hash table. \
-                                    This is a bug - all groups should have been emitted before skip aggregation started.",
+This is a bug - all groups should have been emitted before skip aggregation started.",
                                     self.group_values.len()
                                 )));
                             }
@@ -819,7 +824,7 @@ impl Stream for GroupedHashAggregateStream {
                     if !self.group_values.is_empty() {
                         return Poll::Ready(Some(internal_err!(
                             "AggregateStream was in Done state with {} groups left in hash table. \
-                            This is a bug - all groups should have been emitted before entering Done state.",
+This is a bug - all groups should have been emitted before entering Done state.",
                             self.group_values.len()
                         )));
                     }
@@ -1023,17 +1028,27 @@ 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<Option<RecordBatch>> {
-        let schema = if spilling {
-            Arc::clone(&self.spill_state.spill_schema)
-        } else {
-            self.schema()
-        };
         if self.group_values.is_empty() {
             return Ok(None);
         }
 
         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 {
+            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 +1327,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 +1350,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 +1395,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 +1419,66 @@ 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| {
+                let mut builder = StringDictionaryBuilder::<UInt8Type>::new();
+                builder.append_value(format!("group_{value}"));
+                RecordBatch::try_new(
+                    Arc::clone(&schema),
+                    vec![Arc::new(builder.finish())],
+                )
+            })
+            .collect::<Result<Vec<_>>>()?;
+        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::<DictionaryArray<UInt16Type>>()
+            .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]

From 80f89614633ccf3ec0eb2f24f45dad42ff0ec723 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:22:12 -0400
Subject: [PATCH 08/28] fix: distinguish dynamic output from stable aggregate
 transport

---
 .../src/aggregates/group_values/mod.rs        | 160 +++++++++++-------
 1 file changed, 102 insertions(+), 58 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
index 33985d7ca142a..3da5fc9410f69 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
@@ -24,12 +24,15 @@ use arrow::array::types::{
 };
 use arrow::array::{ArrayRef, downcast_primitive};
 use arrow::compute::cast;
-use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef, TimeUnit};
-use datafusion_common::Result;
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
+use datafusion_common::{Result, internal_err};
 use std::sync::Arc;
 
 use datafusion_expr::EmitTo;
 
+const ORIGINAL_DICTIONARY_KEY_TYPE: &str =
+    "datafusion:group_value_original_dictionary_key_type";
+
 /// 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.
@@ -67,8 +70,8 @@ pub(crate) fn schema_with_group_values(
 }
 
 /// Casts the leading group-value arrays to the corresponding schema fields.
-/// This is used by spill files, which require one stable IPC schema even though
-/// normal aggregate output may promote dictionary keys dynamically.
+/// Stable partial/spill schemas use this before crossing fixed-schema transport
+/// boundaries such as repartition coalescing and Arrow IPC.
 pub(crate) fn cast_group_values_to_schema(
     group_values: &mut [ArrayRef],
     schema: &SchemaRef,
@@ -81,10 +84,23 @@ pub(crate) fn cast_group_values_to_schema(
     Ok(())
 }
 
-/// Returns a stable spill schema whose dictionary group keys use their widest
-/// key type. The widening is internal to spill IPC and does not change planned
-/// or non-spill aggregate output schemas.
-pub(crate) fn group_value_spill_schema(
+/// Returns true when the leading group fields use the stable transport encoding.
+pub(crate) fn has_group_value_transport_fields(
+    schema: &SchemaRef,
+    num_group_fields: usize,
+) -> bool {
+    schema
+        .fields()
+        .iter()
+        .take(num_group_fields)
+        .any(|field| field.metadata().contains_key(ORIGINAL_DICTIONARY_KEY_TYPE))
+}
+
+/// Returns a stable internal transport schema for partial aggregate and spill
+/// batches. Dictionary group keys use the widest key with the same signedness,
+/// while metadata records the original key type so final aggregate output can
+/// start from the planned minimum width and promote only when necessary.
+pub(crate) fn group_value_transport_schema(
     schema: &SchemaRef,
     num_group_fields: usize,
 ) -> SchemaRef {
@@ -94,7 +110,7 @@ pub(crate) fn group_value_spill_schema(
         .enumerate()
         .map(|(index, field)| {
             if index < num_group_fields {
-                group_value_spill_field(field.as_ref())
+                group_value_transport_field(field.as_ref())
             } else {
                 Arc::clone(field)
             }
@@ -107,60 +123,88 @@ pub(crate) fn group_value_spill_schema(
     ))
 }
 
-fn group_value_spill_field(field: &Field) -> FieldRef {
+/// Restores the planned minimum dictionary key type from an internal transport
+/// field. The transport marker is removed so it does not leak into final output.
+pub(crate) fn group_value_output_field(field: Field) -> Result<Field> {
+    let Some(original_key) = field
+        .metadata()
+        .get(ORIGINAL_DICTIONARY_KEY_TYPE)
+        .cloned()
+    else {
+        return Ok(field);
+    };
+
+    let DataType::Dictionary(_, value_type) = field.data_type() else {
+        return internal_err!(
+            "Group value transport marker found on non-dictionary field {}",
+            field.name()
+        );
+    };
+    let value_type = value_type.as_ref().clone();
+    let Some(key_type) = dictionary_key_type_from_name(&original_key) else {
+        return internal_err!(
+            "Invalid group value dictionary key type marker: {original_key}"
+        );
+    };
+
+    let mut metadata = field.metadata().clone();
+    metadata.remove(ORIGINAL_DICTIONARY_KEY_TYPE);
+    Ok(field
+        .with_data_type(DataType::Dictionary(
+            Box::new(key_type),
+            Box::new(value_type),
+        ))
+        .with_metadata(metadata))
+}
+
+fn group_value_transport_field(field: &Field) -> Arc<Field> {
+    let DataType::Dictionary(key_type, value_type) = field.data_type() else {
+        return Arc::new(field.clone());
+    };
+    let Some((original_key, transport_key)) = dictionary_transport_key(key_type) else {
+        return Arc::new(field.clone());
+    };
+
+    let mut metadata = field.metadata().clone();
+    metadata
+        .entry(ORIGINAL_DICTIONARY_KEY_TYPE.to_string())
+        .or_insert_with(|| original_key.to_string());
+
     Arc::new(
         field
             .clone()
-            .with_data_type(group_value_spill_data_type(field.data_type())),
+            .with_data_type(DataType::Dictionary(
+                Box::new(transport_key),
+                Box::new(value_type.as_ref().clone()),
+            ))
+            .with_metadata(metadata),
     )
 }
 
-fn group_value_spill_data_type(data_type: &DataType) -> DataType {
-    match 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(group_value_spill_data_type(value_type)),
-            )
-        }
-        DataType::Struct(fields) => DataType::Struct(
-            fields
-                .iter()
-                .map(|field| group_value_spill_field(field.as_ref()))
-                .collect::<Vec<_>>()
-                .into(),
-        ),
-        DataType::List(field) => DataType::List(group_value_spill_field(field.as_ref())),
-        DataType::LargeList(field) => {
-            DataType::LargeList(group_value_spill_field(field.as_ref()))
-        }
-        DataType::ListView(field) => {
-            DataType::ListView(group_value_spill_field(field.as_ref()))
-        }
-        DataType::LargeListView(field) => {
-            DataType::LargeListView(group_value_spill_field(field.as_ref()))
-        }
-        DataType::FixedSizeList(field, size) => {
-            DataType::FixedSizeList(group_value_spill_field(field.as_ref()), *size)
-        }
-        DataType::Map(field, sorted) => {
-            DataType::Map(group_value_spill_field(field.as_ref()), *sorted)
-        }
-        DataType::RunEndEncoded(run_ends, values) => DataType::RunEndEncoded(
-            Arc::clone(run_ends),
-            group_value_spill_field(values.as_ref()),
-        ),
-        _ => data_type.clone(),
+fn dictionary_transport_key(key_type: &DataType) -> Option<(&'static str, DataType)> {
+    match key_type {
+        DataType::Int8 => Some(("Int8", DataType::Int64)),
+        DataType::Int16 => Some(("Int16", DataType::Int64)),
+        DataType::Int32 => Some(("Int32", DataType::Int64)),
+        DataType::UInt8 => Some(("UInt8", DataType::UInt64)),
+        DataType::UInt16 => Some(("UInt16", DataType::UInt64)),
+        DataType::UInt32 => Some(("UInt32", DataType::UInt64)),
+        DataType::Int64 | DataType::UInt64 => None,
+        _ => None,
+    }
+}
+
+fn dictionary_key_type_from_name(name: &str) -> Option<DataType> {
+    match name {
+        "Int8" => Some(DataType::Int8),
+        "Int16" => Some(DataType::Int16),
+        "Int32" => Some(DataType::Int32),
+        "Int64" => Some(DataType::Int64),
+        "UInt8" => Some(DataType::UInt8),
+        "UInt16" => Some(DataType::UInt16),
+        "UInt32" => Some(DataType::UInt32),
+        "UInt64" => Some(DataType::UInt64),
+        _ => None,
     }
 }
 
@@ -265,7 +309,7 @@ pub trait GroupValues: Send {
 ///   - Otherwise, the general implementation `GroupValuesRows` will be chosen.
 ///
 /// `GroupColumn`:  crate::aggregates::group_values::multi_group_by::GroupColumn
-/// `GroupValuesColumn`: crate::aggregates::group_values::multi_group_by::GroupValuesColumn
+/// `GroupValuesColumn`: crate::aggregates::group_values::GroupValuesColumn
 /// `GroupValuesRows`: crate::aggregates::group_values::GroupValuesRows
 pub fn new_group_values(
     schema: SchemaRef,

From 2d8d9e8257a6bbc920597034e9689f37f3040695 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:24:39 -0400
Subject: [PATCH 09/28] fix: keep dynamic group schema helpers self-contained

---
 .../src/aggregates/group_values/mod.rs        | 158 +++++++-----------
 1 file changed, 57 insertions(+), 101 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
index 3da5fc9410f69..315952f43da22 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
@@ -24,15 +24,12 @@ use arrow::array::types::{
 };
 use arrow::array::{ArrayRef, downcast_primitive};
 use arrow::compute::cast;
-use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
-use datafusion_common::{Result, internal_err};
+use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef, TimeUnit};
+use datafusion_common::Result;
 use std::sync::Arc;
 
 use datafusion_expr::EmitTo;
 
-const ORIGINAL_DICTIONARY_KEY_TYPE: &str =
-    "datafusion:group_value_original_dictionary_key_type";
-
 /// 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.
@@ -70,8 +67,8 @@ pub(crate) fn schema_with_group_values(
 }
 
 /// Casts the leading group-value arrays to the corresponding schema fields.
-/// Stable partial/spill schemas use this before crossing fixed-schema transport
-/// boundaries such as repartition coalescing and Arrow IPC.
+/// This is used by spill files, which require one stable IPC schema even though
+/// normal aggregate output may promote dictionary keys dynamically.
 pub(crate) fn cast_group_values_to_schema(
     group_values: &mut [ArrayRef],
     schema: &SchemaRef,
@@ -84,23 +81,10 @@ pub(crate) fn cast_group_values_to_schema(
     Ok(())
 }
 
-/// Returns true when the leading group fields use the stable transport encoding.
-pub(crate) fn has_group_value_transport_fields(
-    schema: &SchemaRef,
-    num_group_fields: usize,
-) -> bool {
-    schema
-        .fields()
-        .iter()
-        .take(num_group_fields)
-        .any(|field| field.metadata().contains_key(ORIGINAL_DICTIONARY_KEY_TYPE))
-}
-
-/// Returns a stable internal transport schema for partial aggregate and spill
-/// batches. Dictionary group keys use the widest key with the same signedness,
-/// while metadata records the original key type so final aggregate output can
-/// start from the planned minimum width and promote only when necessary.
-pub(crate) fn group_value_transport_schema(
+/// Returns a stable spill schema whose dictionary group keys use their widest
+/// key type. The 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 {
@@ -110,7 +94,7 @@ pub(crate) fn group_value_transport_schema(
         .enumerate()
         .map(|(index, field)| {
             if index < num_group_fields {
-                group_value_transport_field(field.as_ref())
+                group_value_spill_field(field.as_ref())
             } else {
                 Arc::clone(field)
             }
@@ -123,88 +107,60 @@ pub(crate) fn group_value_transport_schema(
     ))
 }
 
-/// Restores the planned minimum dictionary key type from an internal transport
-/// field. The transport marker is removed so it does not leak into final output.
-pub(crate) fn group_value_output_field(field: Field) -> Result<Field> {
-    let Some(original_key) = field
-        .metadata()
-        .get(ORIGINAL_DICTIONARY_KEY_TYPE)
-        .cloned()
-    else {
-        return Ok(field);
-    };
-
-    let DataType::Dictionary(_, value_type) = field.data_type() else {
-        return internal_err!(
-            "Group value transport marker found on non-dictionary field {}",
-            field.name()
-        );
-    };
-    let value_type = value_type.as_ref().clone();
-    let Some(key_type) = dictionary_key_type_from_name(&original_key) else {
-        return internal_err!(
-            "Invalid group value dictionary key type marker: {original_key}"
-        );
-    };
-
-    let mut metadata = field.metadata().clone();
-    metadata.remove(ORIGINAL_DICTIONARY_KEY_TYPE);
-    Ok(field
-        .with_data_type(DataType::Dictionary(
-            Box::new(key_type),
-            Box::new(value_type),
-        ))
-        .with_metadata(metadata))
-}
-
-fn group_value_transport_field(field: &Field) -> Arc<Field> {
-    let DataType::Dictionary(key_type, value_type) = field.data_type() else {
-        return Arc::new(field.clone());
-    };
-    let Some((original_key, transport_key)) = dictionary_transport_key(key_type) else {
-        return Arc::new(field.clone());
-    };
-
-    let mut metadata = field.metadata().clone();
-    metadata
-        .entry(ORIGINAL_DICTIONARY_KEY_TYPE.to_string())
-        .or_insert_with(|| original_key.to_string());
-
+fn group_value_spill_field(field: &Field) -> FieldRef {
     Arc::new(
         field
             .clone()
-            .with_data_type(DataType::Dictionary(
-                Box::new(transport_key),
-                Box::new(value_type.as_ref().clone()),
-            ))
-            .with_metadata(metadata),
+            .with_data_type(group_value_spill_data_type(field.data_type())),
     )
 }
 
-fn dictionary_transport_key(key_type: &DataType) -> Option<(&'static str, DataType)> {
-    match key_type {
-        DataType::Int8 => Some(("Int8", DataType::Int64)),
-        DataType::Int16 => Some(("Int16", DataType::Int64)),
-        DataType::Int32 => Some(("Int32", DataType::Int64)),
-        DataType::UInt8 => Some(("UInt8", DataType::UInt64)),
-        DataType::UInt16 => Some(("UInt16", DataType::UInt64)),
-        DataType::UInt32 => Some(("UInt32", DataType::UInt64)),
-        DataType::Int64 | DataType::UInt64 => None,
-        _ => None,
-    }
-}
-
-fn dictionary_key_type_from_name(name: &str) -> Option<DataType> {
-    match name {
-        "Int8" => Some(DataType::Int8),
-        "Int16" => Some(DataType::Int16),
-        "Int32" => Some(DataType::Int32),
-        "Int64" => Some(DataType::Int64),
-        "UInt8" => Some(DataType::UInt8),
-        "UInt16" => Some(DataType::UInt16),
-        "UInt32" => Some(DataType::UInt32),
-        "UInt64" => Some(DataType::UInt64),
-        _ => None,
+fn group_value_spill_data_type(data_type: &DataType) -> DataType {
+    match 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(group_value_spill_data_type(value_type)),
+            )
+        }
+        DataType::Struct(fields) => DataType::Struct(
+            fields
+                .iter()
+                .map(|field| group_value_spill_field(field.as_ref()))
+                .collect::<Vec<_>>()
+                .into(),
+        ),
+        DataType::List(field) => DataType::List(group_value_spill_field(field.as_ref())),
+        DataType::LargeList(field) => {
+            DataType::LargeList(group_value_spill_field(field.as_ref()))
+        }
+        DataType::ListView(field) => {
+            DataType::ListView(group_value_spill_field(field.as_ref()))
+        }
+        DataType::LargeListView(field) => {
+            DataType::LargeListView(group_value_spill_field(field.as_ref()))
+        }
+        DataType::FixedSizeList(field, size) => {
+            DataType::FixedSizeList(group_value_spill_field(field.as_ref()), *size)
+        }
+        DataType::Map(field, sorted) => {
+            DataType::Map(group_value_spill_field(field.as_ref()), *sorted)
+        }
+        DataType::RunEndEncoded(run_ends, values) => DataType::RunEndEncoded(
+            Arc::clone(run_ends),
+            group_value_spill_field(values.as_ref()),
+        ),
+        _ => data_type.clone(),
     }
 }
 

From 1dff0a2bb1ee2dbc135ad39e68e5dfc3570e93eb Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:28:11 -0400
Subject: [PATCH 10/28] fix: bound partial dictionary group batches by key
 capacity

---
 .../src/aggregates/group_values/mod.rs        | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
index 315952f43da22..23d90b9c28de1 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
@@ -66,6 +66,44 @@ pub(crate) fn schema_with_group_values(
     ))
 }
 
+/// 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<usize> {
+    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.
 /// This is used by spill files, which require one stable IPC schema even though
 /// normal aggregate output may promote dictionary keys dynamically.

From fe1131a1bd1fe4976f631468891cb662aa63d514 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:29:19 -0400
Subject: [PATCH 11/28] fix: emit schema-stable partial dictionary batches

---
 .../aggregates/aggregate_hash_table/common.rs | 76 +++++++++++++++++--
 1 file changed, 69 insertions(+), 7 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
index 7c7c2cba836c3..f500c72b22e07 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
@@ -28,7 +28,8 @@ use datafusion_physical_expr::aggregate::AggregateFunctionExpr;
 
 use crate::PhysicalExpr;
 use crate::aggregates::group_values::{
-    GroupByMetrics, GroupValues, new_group_values, schema_with_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;
@@ -207,12 +208,74 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         Ok(())
     }
 
+    /// Emits one partial-state batch without changing the planned schema.
+    ///
+    /// Dictionary group columns are capped at the number of values representable
+    /// by their planned key type. Remaining groups stay in the table and are
+    /// emitted on later polls. This keeps repartition/coalescing schemas stable
+    /// while allowing a downstream final aggregate to combine all groups and
+    /// promote its final output key type when required.
+    pub(super) fn next_partial_output_batch_inner(
+        &mut self,
+        materialize_accumulator_fn: MaterializeAccumulatorFn,
+    ) -> Result<Option<RecordBatch>> {
+        let output_schema = Arc::clone(&self.output_schema);
+        let batch_size = self.batch_size;
+
+        let mut state =
+            match std::mem::replace(&mut self.state, AggregateHashTableState::Done) {
+                AggregateHashTableState::Outputting(state) => state,
+                AggregateHashTableState::Done => return Ok(None),
+                AggregateHashTableState::Building(_) => {
+                    return internal_err!(
+                        "next_output_batch must be called in the outputting state"
+                    );
+                }
+                AggregateHashTableState::OutputtingMaterialized(_) => {
+                    return internal_err!(
+                        "partial aggregate output must not be materialized"
+                    );
+                }
+            };
+
+        if state.group_values.is_empty() {
+            return Ok(None);
+        }
+
+        let max_groups = group_value_emit_batch_size(
+            &output_schema,
+            state.group_by.num_group_exprs(),
+            batch_size,
+        );
+        let num_groups = max_groups.min(state.group_values.len());
+        let emit_to = if num_groups == state.group_values.len() {
+            EmitTo::All
+        } else {
+            EmitTo::First(num_groups)
+        };
+
+        let timer = self.group_by_metrics.emitting_time.timer();
+        let mut columns = state.group_values.emit(emit_to)?;
+        for acc in state.accumulators.iter_mut() {
+            columns.extend(materialize_accumulator_fn(acc, emit_to)?);
+        }
+        drop(timer);
+
+        let batch = RecordBatch::try_new(output_schema, columns)?;
+        debug_assert!(batch.num_rows() > 0);
+        if state.group_values.is_empty() {
+            self.state = AggregateHashTableState::Done;
+        } else {
+            self.state = AggregateHashTableState::Outputting(state);
+        }
+        Ok(Some(batch))
+    }
+
     /// Materializes the full output once, then returns it downstream incrementally
     /// by slicing it into `batch_size` chunks.
     ///
-    /// Each aggregation mode chooses a different `materialize_accumulator_fn`
-    /// according to its semantics. For example, partial aggregation emits
-    /// partial states to feed the final stage, so it uses [`GroupsAccumulator::state`].
+    /// Final aggregate evaluation consumes accumulator state, so it is
+    /// materialized once and then sliced on later polls.
     ///
     /// This is a temporary solution until blocked state management is implemented:
     /// Issue: <https://github.com/apache/datafusion/issues/7065>
@@ -430,9 +493,8 @@ 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.
+/// Final aggregate evaluation consumes accumulator state, so it is materialized
+/// once and then sliced to honor `batch_size` across output polls.
 pub(super) struct MaterializedAggregateOutput {
     batch: RecordBatch,
     offset: usize,

From b541d9d548a83561ff809230bf91bba30266e5ac Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:29:43 -0400
Subject: [PATCH 12/28] fix: keep partial dictionary output schema stable

---
 .../src/aggregates/aggregate_hash_table/partial_table.rs     | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

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<PartialMarker> {
 
     /// 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<Option<RecordBatch>> {
-        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 {

From 063461e98cdf21d25b14e1dd3264843dc66234aa Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:29:57 -0400
Subject: [PATCH 13/28] fix: keep partial-reduce dictionary output schema
 stable

---
 .../aggregates/aggregate_hash_table/partial_reduce_table.rs  | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

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<PartialReduceMarker> {
 
     /// 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<Option<RecordBatch>> {
-        self.next_output_batch_inner(HashAggregateAccumulator::state)
+        self.next_partial_output_batch_inner(HashAggregateAccumulator::state)
     }
 
     /// Partial-reduce aggregation consumes partial aggregate states and merges

From 023a0fe6979fdda6cea34b4afeec8b80b77c0cb6 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:30:49 -0400
Subject: [PATCH 14/28] fix: keep ordered partial dictionary output schema
 stable

---
 .../aggregate_hash_table/common_ordered.rs    | 45 +++++++++++++------
 1 file changed, 32 insertions(+), 13 deletions(-)

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 ec6d0c066d701..855568fa39aa0 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
@@ -31,7 +31,8 @@ use datafusion_expr::EmitTo;
 use crate::InputOrderMode;
 use crate::PhysicalExpr;
 use crate::aggregates::group_values::{
-    GroupByMetrics, GroupValues, new_group_values, schema_with_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;
@@ -236,7 +237,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
             + 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
@@ -246,13 +247,15 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         &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 final aggregation.
@@ -314,8 +317,8 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
     ///
     /// # 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,
@@ -327,8 +330,20 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         let Some(emit_to) = self.buffer.group_ordering.emit_to() else {
             return Ok(None);
         };
-        let (emit_to, should_remove_groups) =
-            self.clamp_emit_to(self.buffer.group_values.len(), emit_to);
+        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,
+            max_batch_size,
+        );
 
         let timer = self.group_by_metrics.emitting_time.timer();
         let mut output = self.buffer.group_values.emit(emit_to)?;
@@ -352,10 +367,14 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         }
         drop(timer);
 
-        let output_schema = schema_with_group_values(
-            &self.output_schema,
-            &output[..num_group_columns],
-        );
+        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);
 

From 44829fb836468de56d724b3c29ae1f661731e00f Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:32:04 -0400
Subject: [PATCH 15/28] test: add aggregate hash table dictionary regressions

---
 .../physical-plan/src/aggregates/aggregate_hash_table/mod.rs   | 3 +++
 1 file changed, 3 insertions(+)

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;

From 60afd283226f48aa594a705eaee990084467cdeb Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:32:30 -0400
Subject: [PATCH 16/28] test: cover dictionary promotion across partial and
 final aggregation

---
 .../aggregates/aggregate_hash_table/tests.rs  | 149 ++++++++++++++++++
 1 file changed, 149 insertions(+)
 create mode 100644 datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs

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..9c455e1994f99
--- /dev/null
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs
@@ -0,0 +1,149 @@
+// 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| {
+            let mut builder = StringDictionaryBuilder::<UInt8Type>::new();
+            builder.append_value(format!("group_{value}"));
+            let array: ArrayRef = Arc::new(builder.finish());
+            RecordBatch::try_new(Arc::clone(&input_schema), vec![array])
+        })
+        .collect::<Result<Vec<_>>>()?;
+
+    let input = TestMemoryExec::try_new(
+        &[input_batches.clone()],
+        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::<PartialMarker>::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::<DictionaryArray<UInt8Type>>()
+            .is_some());
+        partial_batches.push(batch);
+    }
+    assert_eq!(
+        partial_batches
+            .iter()
+            .map(RecordBatch::num_rows)
+            .collect::<Vec<_>>(),
+        vec![256, 1]
+    );
+
+    let partial_input = TestMemoryExec::try_new(
+        &[partial_batches.clone()],
+        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::<FinalMarker>::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::<DictionaryArray<UInt16Type>>()
+        .unwrap();
+    let values = dictionary
+        .values()
+        .as_any()
+        .downcast_ref::<StringArray>()
+        .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(())
+}

From a9fc6be56862b6a77325195518fa6a3f91f351cd Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:33:28 -0400
Subject: [PATCH 17/28] test: use DataFusion result for dictionary batches

---
 .../src/aggregates/aggregate_hash_table/tests.rs           | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs
index 9c455e1994f99..98eeb6cdd7382 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs
@@ -38,11 +38,14 @@ fn dictionary_groups_keep_partial_schema_and_promote_final_output() -> Result<()
         false,
     )]));
     let input_batches = (0u32..257)
-        .map(|value| {
+        .map(|value| -> Result<RecordBatch> {
             let mut builder = StringDictionaryBuilder::<UInt8Type>::new();
             builder.append_value(format!("group_{value}"));
             let array: ArrayRef = Arc::new(builder.finish());
-            RecordBatch::try_new(Arc::clone(&input_schema), vec![array])
+            Ok(RecordBatch::try_new(
+                Arc::clone(&input_schema),
+                vec![array],
+            )?)
         })
         .collect::<Result<Vec<_>>>()?;
 

From 917aed8380fc0c459775fd277fcffbc5bb378f3d Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:35:03 -0400
Subject: [PATCH 18/28] test: cover dictionary key promotion through
 repartitioned aggregation

---
 .../dictionary_group_by_key_promotion.slt     | 50 +++++++++++++++++++
 1 file changed, 50 insertions(+)
 create mode 100644 datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt

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..a5d7c62cf3d85
--- /dev/null
+++ b/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt
@@ -0,0 +1,50 @@
+# 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 batch can be encoded with UInt8, but the complete GROUP BY result
+# has 257 distinct values and therefore requires UInt16 dictionary keys.
+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, 256)
+  )
+  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)
+  )
+  GROUP BY k
+);
+----
+Dictionary(UInt16, Utf8)

From 192693f6e530c85144e833d05ab77ff949d29e15 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:40:05 -0400
Subject: [PATCH 19/28] fix: keep legacy partial dictionary output schema
 stable

---
 .../src/aggregates/grouped_hash_stream.rs     | 44 +++++++++++++++----
 1 file changed, 36 insertions(+), 8 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
index 3df5c4aceb305..83ab3f570a542 100644
--- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
@@ -26,7 +26,8 @@ use super::skip_partial::SkipAggregationProbe;
 use super::{AggregateExec, format_human_display};
 use crate::aggregates::group_values::{
     GroupByMetrics, GroupValues, cast_group_values_to_schema,
-    group_value_spill_schema, new_group_values, schema_with_group_values,
+    group_value_emit_batch_size, group_value_spill_schema, new_group_values,
+    schema_with_group_values,
 };
 use crate::aggregates::order::GroupOrderingFull;
 use crate::aggregates::{
@@ -765,7 +766,7 @@ impl Stream for GroupedHashAggregateStream {
                             if !self.group_values.is_empty() {
                                 return Poll::Ready(Some(internal_err!(
                                     "Switching from SkippingAggregation to Done with {} groups still in hash table. \
-This is a bug - all groups should have been emitted before skip aggregation started.",
+                                    This is a bug - all groups should have been emitted before skip aggregation started.",
                                     self.group_values.len()
                                 )));
                             }
@@ -785,12 +786,18 @@ This is a bug - all groups should have been emitted before skip aggregation star
                     (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 {
@@ -824,7 +831,7 @@ This is a bug - all groups should have been emitted before skip aggregation star
                     if !self.group_values.is_empty() {
                         return Poll::Ready(Some(internal_err!(
                             "AggregateStream was in Done state with {} groups left in hash table. \
-This is a bug - all groups should have been emitted before entering Done state.",
+                            This is a bug - all groups should have been emitted before entering Done state.",
                             self.group_values.len()
                         )));
                     }
@@ -1032,6 +1039,25 @@ impl GroupedHashAggregateStream {
             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();
@@ -1041,6 +1067,8 @@ impl GroupedHashAggregateStream {
                 &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,
@@ -1436,13 +1464,13 @@ mod tests {
             false,
         )]));
         let batches = (0u32..257)
-            .map(|value| {
+            .map(|value| -> Result<RecordBatch> {
                 let mut builder = StringDictionaryBuilder::<UInt8Type>::new();
                 builder.append_value(format!("group_{value}"));
-                RecordBatch::try_new(
+                Ok(RecordBatch::try_new(
                     Arc::clone(&schema),
                     vec![Arc::new(builder.finish())],
-                )
+                )?)
             })
             .collect::<Result<Vec<_>>>()?;
         let exec = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None)?;

From 47a4f3d5683d43fa25839da64532d2a6c5f0d5e1 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:42:24 -0400
Subject: [PATCH 20/28] test: make dictionary regression series column explicit

---
 .../test_files/dictionary_group_by_key_promotion.slt          | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt b/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt
index a5d7c62cf3d85..4a36e71670a0c 100644
--- a/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt
+++ b/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt
@@ -29,7 +29,7 @@ FROM (
   SELECT k
   FROM (
     SELECT arrow_cast(CAST(value AS VARCHAR), 'Dictionary(UInt8, Utf8)') AS k
-    FROM generate_series(0, 256)
+    FROM generate_series(0, 256) AS series(value)
   )
   GROUP BY k
 );
@@ -42,7 +42,7 @@ FROM (
   SELECT k
   FROM (
     SELECT arrow_cast(CAST(value AS VARCHAR), 'Dictionary(UInt8, Utf8)') AS k
-    FROM generate_series(0, 256)
+    FROM generate_series(0, 256) AS series(value)
   )
   GROUP BY k
 );

From eb8f4eebdb740ea609a631748b857172ae508ab2 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:43:28 -0400
Subject: [PATCH 21/28] test: separate partial transport and final promotion
 coverage

---
 .../dictionary_group_by_key_promotion.slt     | 23 +++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt b/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt
index 4a36e71670a0c..04e080d2e7d3d 100644
--- a/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt
+++ b/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt
@@ -15,14 +15,33 @@
 # specific language governing permissions and limitations
 # under the License.
 
-# Each input batch can be encoded with UInt8, but the complete GROUP BY result
-# has 257 distinct values and therefore requires UInt16 dictionary keys.
+# 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 (

From f575ad3ccfe9d88c820e87538f83357b6c1ebc54 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:47:21 -0400
Subject: [PATCH 22/28] fix: limit dictionary promotion to top-level group
 columns

---
 .../src/aggregates/group_values/row.rs        | 80 ++++++++-----------
 1 file changed, 34 insertions(+), 46 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs
index a149e46ca6bc4..2aecd4cd1d53a 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/row.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs
@@ -21,7 +21,7 @@ use arrow::array::{
     downcast_run_end_index,
 };
 use arrow::compute::cast;
-use arrow::datatypes::{DataType, FieldRef, Schema, SchemaRef};
+use arrow::datatypes::{DataType, Schema, SchemaRef};
 use arrow::error::ArrowError;
 use arrow::row::{RowConverter, Rows, SortField};
 use datafusion_common::Result;
@@ -263,11 +263,11 @@ impl GroupValues for GroupValuesRows {
             }
         };
 
-        // Re-encode dictionary group keys using the current key type. If the
-        // emitted cardinality no longer fits, grow to the next key width and
+        // 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_if_necessary(array, field.data_type())?;
+            *array = dictionary_encode_with_promotion(array, field.data_type())?;
         }
         self.schema = schema_with_group_values(&self.schema, &output);
 
@@ -308,6 +308,29 @@ fn materialized_row_schema(schema: &SchemaRef) -> SchemaRef {
     ))
 }
 
+fn dictionary_encode_with_promotion(
+    array: &ArrayRef,
+    expected: &DataType,
+) -> Result<ArrayRef> {
+    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,
@@ -322,49 +345,27 @@ fn dictionary_encode_if_necessary(
                     dictionary_encode_if_necessary(column, expected_field.data_type())
                 })
                 .collect::<Result<Vec<_>>>()?;
-            let fields = expected_fields
-                .iter()
-                .zip(&arrays)
-                .map(|(field, array)| field_with_array_type(field, array))
-                .collect::<Vec<_>>()
-                .into();
 
             Ok(Arc::new(StructArray::try_new(
-                fields,
+                expected_fields.clone(),
                 arrays,
                 struct_array.nulls().cloned(),
             )?))
         }
         (DataType::List(expected_field), &DataType::List(_)) => {
             let list = array.as_any().downcast_ref::<ListArray>().unwrap();
-            let values = dictionary_encode_if_necessary(
-                list.values(),
-                expected_field.data_type(),
-            )?;
-            let field = field_with_array_type(expected_field, &values);
 
             Ok(Arc::new(ListArray::try_new(
-                field,
+                Arc::<arrow::datatypes::Field>::clone(expected_field),
                 list.offsets().clone(),
-                values,
+                dictionary_encode_if_necessary(
+                    list.values(),
+                    expected_field.data_type(),
+                )?,
                 list.nulls().cloned(),
             )?))
         }
-        (DataType::Dictionary(_, _), _) => {
-            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()),
-                }
-            }
-        }
+        (DataType::Dictionary(_, _), _) => Ok(cast(array.as_ref(), expected)?),
         (
             DataType::RunEndEncoded(run_ends_field, expected_values_field),
             &DataType::RunEndEncoded(_, _),
@@ -396,19 +397,6 @@ fn dictionary_encode_if_necessary(
     }
 }
 
-fn field_with_array_type(field: &FieldRef, array: &ArrayRef) -> FieldRef {
-    if field.data_type() == array.data_type() {
-        Arc::clone(field)
-    } else {
-        Arc::new(
-            field
-                .as_ref()
-                .clone()
-                .with_data_type(array.data_type().clone()),
-        )
-    }
-}
-
 fn promote_dictionary_type(data_type: &DataType) -> Option<DataType> {
     let DataType::Dictionary(key_type, value_type) = data_type else {
         return None;

From 385383bfff22db3819b2eb2afe40d1be78125e7a Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:48:05 -0400
Subject: [PATCH 23/28] fix: keep dictionary schema helpers top-level and
 scoped

---
 .../src/aggregates/group_values/mod.rs        | 55 +++----------------
 1 file changed, 9 insertions(+), 46 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
index 23d90b9c28de1..97409dfb69a0b 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
@@ -119,9 +119,9 @@ pub(crate) fn cast_group_values_to_schema(
     Ok(())
 }
 
-/// Returns a stable spill schema whose dictionary group keys use their widest
-/// key type. The widening is internal to spill IPC and does not change planned
-/// or non-spill aggregate output schemas.
+/// 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,
@@ -146,15 +146,7 @@ pub(crate) fn group_value_spill_schema(
 }
 
 fn group_value_spill_field(field: &Field) -> FieldRef {
-    Arc::new(
-        field
-            .clone()
-            .with_data_type(group_value_spill_data_type(field.data_type())),
-    )
-}
-
-fn group_value_spill_data_type(data_type: &DataType) -> DataType {
-    match data_type {
+    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 => {
@@ -166,40 +158,11 @@ fn group_value_spill_data_type(data_type: &DataType) -> DataType {
                 | DataType::UInt64 => DataType::UInt64,
                 _ => key_type.as_ref().clone(),
             };
-            DataType::Dictionary(
-                Box::new(key_type),
-                Box::new(group_value_spill_data_type(value_type)),
-            )
-        }
-        DataType::Struct(fields) => DataType::Struct(
-            fields
-                .iter()
-                .map(|field| group_value_spill_field(field.as_ref()))
-                .collect::<Vec<_>>()
-                .into(),
-        ),
-        DataType::List(field) => DataType::List(group_value_spill_field(field.as_ref())),
-        DataType::LargeList(field) => {
-            DataType::LargeList(group_value_spill_field(field.as_ref()))
-        }
-        DataType::ListView(field) => {
-            DataType::ListView(group_value_spill_field(field.as_ref()))
-        }
-        DataType::LargeListView(field) => {
-            DataType::LargeListView(group_value_spill_field(field.as_ref()))
+            DataType::Dictionary(Box::new(key_type), Box::new(value_type.as_ref().clone()))
         }
-        DataType::FixedSizeList(field, size) => {
-            DataType::FixedSizeList(group_value_spill_field(field.as_ref()), *size)
-        }
-        DataType::Map(field, sorted) => {
-            DataType::Map(group_value_spill_field(field.as_ref()), *sorted)
-        }
-        DataType::RunEndEncoded(run_ends, values) => DataType::RunEndEncoded(
-            Arc::clone(run_ends),
-            group_value_spill_field(values.as_ref()),
-        ),
-        _ => data_type.clone(),
-    }
+        _ => field.data_type().clone(),
+    };
+    Arc::new(field.clone().with_data_type(data_type))
 }
 
 pub mod multi_group_by;
@@ -303,7 +266,7 @@ pub trait GroupValues: Send {
 ///   - Otherwise, the general implementation `GroupValuesRows` will be chosen.
 ///
 /// `GroupColumn`:  crate::aggregates::group_values::multi_group_by::GroupColumn
-/// `GroupValuesColumn`: crate::aggregates::group_values::GroupValuesColumn
+/// `GroupValuesColumn`: crate::aggregates::group_values::multi_group_by::GroupValuesColumn
 /// `GroupValuesRows`: crate::aggregates::group_values::GroupValuesRows
 pub fn new_group_values(
     schema: SchemaRef,

From 8db894fbce0fdb6f699b77fce9c12bd31fcd7755 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:52:32 -0400
Subject: [PATCH 24/28] fix: slice materialized partial dictionary output
 safely

---
 .../aggregates/aggregate_hash_table/common.rs | 171 +++++++++++-------
 1 file changed, 106 insertions(+), 65 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
index f500c72b22e07..45db47f978923 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
@@ -28,8 +28,8 @@ use datafusion_physical_expr::aggregate::AggregateFunctionExpr;
 
 use crate::PhysicalExpr;
 use crate::aggregates::group_values::{
-    GroupByMetrics, GroupValues, group_value_emit_batch_size, new_group_values,
-    schema_with_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;
@@ -208,74 +208,78 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         Ok(())
     }
 
-    /// Emits one partial-state batch without changing the planned schema.
+    /// Materializes all partial state once, then returns schema-stable slices.
     ///
-    /// Dictionary group columns are capped at the number of values representable
-    /// by their planned key type. Remaining groups stay in the table and are
-    /// emitted on later polls. This keeps repartition/coalescing schemas stable
-    /// while allowing a downstream final aggregate to combine all groups and
-    /// promote its final output key type when required.
+    /// 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<Option<RecordBatch>> {
-        let output_schema = Arc::clone(&self.output_schema);
-        let batch_size = self.batch_size;
+        let planned_schema = Arc::clone(&self.output_schema);
+        let configured_batch_size = self.batch_size;
 
-        let mut state =
+        let mut output =
             match std::mem::replace(&mut self.state, AggregateHashTableState::Done) {
-                AggregateHashTableState::Outputting(state) => state,
+                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"
                     );
                 }
-                AggregateHashTableState::OutputtingMaterialized(_) => {
-                    return internal_err!(
-                        "partial aggregate output must not be materialized"
-                    );
-                }
             };
 
-        if state.group_values.is_empty() {
-            return Ok(None);
-        }
-
-        let max_groups = group_value_emit_batch_size(
-            &output_schema,
-            state.group_by.num_group_exprs(),
-            batch_size,
-        );
-        let num_groups = max_groups.min(state.group_values.len());
-        let emit_to = if num_groups == state.group_values.len() {
-            EmitTo::All
-        } else {
-            EmitTo::First(num_groups)
-        };
-
-        let timer = self.group_by_metrics.emitting_time.timer();
-        let mut columns = state.group_values.emit(emit_to)?;
-        for acc in state.accumulators.iter_mut() {
-            columns.extend(materialize_accumulator_fn(acc, emit_to)?);
-        }
-        drop(timer);
-
-        let batch = RecordBatch::try_new(output_schema, columns)?;
-        debug_assert!(batch.num_rows() > 0);
-        if state.group_values.is_empty() {
+        let batch = output.next_batch()?;
+        if output.is_exhausted() {
             self.state = AggregateHashTableState::Done;
         } else {
-            self.state = AggregateHashTableState::Outputting(state);
+            self.state = AggregateHashTableState::OutputtingMaterialized(output);
         }
-        Ok(Some(batch))
+        Ok(batch)
     }
 
     /// Materializes the full output once, then returns it downstream incrementally
     /// by slicing it into `batch_size` chunks.
     ///
-    /// Final aggregate evaluation consumes accumulator state, so it is
-    /// materialized once and then sliced on later polls.
+    /// Each aggregation mode chooses a different `materialize_accumulator_fn`
+    /// according to its semantics. For example, partial aggregation emits
+    /// partial states to feed the final stage, so it uses [`GroupsAccumulator::state`].
     ///
     /// This is a temporary solution until blocked state management is implemented:
     /// Issue: <https://github.com/apache/datafusion/issues/7065>
@@ -310,7 +314,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
                     );
                     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),
@@ -321,7 +325,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
                 }
             };
 
-        let batch = output.next_batch(batch_size);
+        let batch = output.next_batch()?;
         if output.is_exhausted() {
             self.state = AggregateHashTableState::Done;
         } else {
@@ -406,7 +410,7 @@ pub(super) type AggregateAccumulator = HashAggregateAccumulator;
 /// Arguments:
 /// * accumulator to update.
 /// * accumulator's evaluated arguments and optional filter.
-/// * group index per input row, mapping each row to its interned group.
+/// * one group index per input row, mapping each row to its interned group.
 /// * total number of groups currently interned in that buffer, including newly
 ///   interned groups.
 pub(super) type AggregateBatchFn = fn(
@@ -483,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: <https://github.com/apache/datafusion/issues/23178>
@@ -493,28 +497,65 @@ pub(super) enum AggregateHashTableState {
 
 /// Fully evaluated aggregate output and the next row offset to emit.
 ///
-/// Final aggregate evaluation consumes accumulator state, so it is materialized
-/// once and then sliced 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 next_batch(&mut self, batch_size: usize) -> Option<RecordBatch> {
-        debug_assert!(batch_size > 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) -> Result<Option<RecordBatch>> {
+        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 {
@@ -699,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(())

From fe31614203b7c68b4bc0436b44bb64967fa94ea3 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 12:53:47 -0400
Subject: [PATCH 25/28] fix: re-encode sliced dictionary transport batches

---
 .../src/aggregates/group_values/mod.rs           | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
index 97409dfb69a0b..6888f841e7deb 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
@@ -105,16 +105,24 @@ fn dictionary_key_capacity(key_type: &DataType) -> Option<usize> {
 }
 
 /// Casts the leading group-value arrays to the corresponding schema fields.
-/// This is used by spill files, which require one stable IPC schema even though
-/// normal aggregate output may promote dictionary keys dynamically.
+/// 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() {
-            *array = cast(array.as_ref(), field.data_type())?;
+        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(())
 }

From 46a9983b1ef6de7b6c2df195d9f9cf7b8eb120f7 Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 13:09:35 -0400
Subject: [PATCH 26/28] style: apply rustfmt

---
 .../aggregate_hash_table/common_ordered.rs    | 12 ++---
 .../aggregates/aggregate_hash_table/tests.rs  | 16 +++---
 .../src/aggregates/group_values/mod.rs        | 15 +++---
 .../src/aggregates/group_values/row.rs        | 29 +++++------
 .../src/aggregates/grouped_hash_stream.rs     | 51 +++++++++----------
 5 files changed, 56 insertions(+), 67 deletions(-)

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 855568fa39aa0..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
@@ -339,11 +339,8 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
                 self.batch_size,
             )
         };
-        let (emit_to, should_remove_groups) = self.clamp_emit_to(
-            self.buffer.group_values.len(),
-            emit_to,
-            max_batch_size,
-        );
+        let (emit_to, should_remove_groups) =
+            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)?;
@@ -368,10 +365,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
         drop(timer);
 
         let output_schema = if is_final {
-            schema_with_group_values(
-                &self.output_schema,
-                &output[..num_group_columns],
-            )
+            schema_with_group_values(&self.output_schema, &output[..num_group_columns])
         } else {
             Arc::clone(&self.output_schema)
         };
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs
index 98eeb6cdd7382..5e22f538a6233 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs
@@ -17,9 +17,7 @@
 
 use std::sync::Arc;
 
-use arrow::array::{
-    ArrayRef, DictionaryArray, StringArray, StringDictionaryBuilder,
-};
+use arrow::array::{ArrayRef, DictionaryArray, StringArray, StringDictionaryBuilder};
 use arrow::datatypes::{DataType, Field, Schema, UInt8Type, UInt16Type};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
@@ -86,11 +84,13 @@ fn dictionary_groups_keep_partial_schema_and_promote_final_output() -> Result<()
             batch.schema().field(0).data_type(),
             &DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8))
         );
-        assert!(batch
-            .column(0)
-            .as_any()
-            .downcast_ref::<DictionaryArray<UInt8Type>>()
-            .is_some());
+        assert!(
+            batch
+                .column(0)
+                .as_any()
+                .downcast_ref::<DictionaryArray<UInt8Type>>()
+                .is_some()
+        );
         partial_batches.push(batch);
     }
     assert_eq!(
diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
index 6888f841e7deb..ef9fd54776025 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
@@ -60,10 +60,7 @@ pub(crate) fn schema_with_group_values(
         })
         .collect::<Vec<_>>();
 
-    Arc::new(Schema::new_with_metadata(
-        fields,
-        schema.metadata().clone(),
-    ))
+    Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))
 }
 
 /// Returns the largest partial aggregate batch that can be emitted without
@@ -147,10 +144,7 @@ pub(crate) fn group_value_spill_schema(
         })
         .collect::<Vec<_>>();
 
-    Arc::new(Schema::new_with_metadata(
-        fields,
-        schema.metadata().clone(),
-    ))
+    Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))
 }
 
 fn group_value_spill_field(field: &Field) -> FieldRef {
@@ -166,7 +160,10 @@ fn group_value_spill_field(field: &Field) -> FieldRef {
                 | DataType::UInt64 => DataType::UInt64,
                 _ => key_type.as_ref().clone(),
             };
-            DataType::Dictionary(Box::new(key_type), Box::new(value_type.as_ref().clone()))
+            DataType::Dictionary(
+                Box::new(key_type),
+                Box::new(value_type.as_ref().clone()),
+            )
         }
         _ => field.data_type().clone(),
     };
diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs
index 2aecd4cd1d53a..c688a40f85699 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/row.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs
@@ -302,10 +302,7 @@ fn materialized_row_schema(schema: &SchemaRef) -> SchemaRef {
             _ => Arc::clone(field),
         })
         .collect::<Vec<_>>();
-    Arc::new(Schema::new_with_metadata(
-        fields,
-        schema.metadata().clone(),
-    ))
+    Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))
 }
 
 fn dictionary_encode_with_promotion(
@@ -421,9 +418,7 @@ fn promote_dictionary_type(data_type: &DataType) -> Option<DataType> {
 mod tests {
     use super::*;
     use arrow::array::{DictionaryArray, StringArray, StringDictionaryBuilder};
-    use arrow::datatypes::{
-        Field, Int8Type, Int16Type, Schema, UInt8Type, UInt16Type,
-    };
+    use arrow::datatypes::{Field, Int8Type, Int16Type, Schema, UInt8Type, UInt16Type};
 
     #[test]
     fn dict_uint8_utf8_promotes_and_does_not_shrink() -> Result<()> {
@@ -493,10 +488,12 @@ mod tests {
             arrays[0].data_type(),
             &DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8))
         );
-        assert!(arrays[0]
-            .as_any()
-            .downcast_ref::<DictionaryArray<UInt8Type>>()
-            .is_some());
+        assert!(
+            arrays[0]
+                .as_any()
+                .downcast_ref::<DictionaryArray<UInt8Type>>()
+                .is_some()
+        );
         Ok(())
     }
 
@@ -520,10 +517,12 @@ mod tests {
             arrays[0].data_type(),
             &DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8))
         );
-        assert!(arrays[0]
-            .as_any()
-            .downcast_ref::<DictionaryArray<Int16Type>>()
-            .is_some());
+        assert!(
+            arrays[0]
+                .as_any()
+                .downcast_ref::<DictionaryArray<Int16Type>>()
+                .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 83ab3f570a542..8cf5531b2507b 100644
--- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
@@ -1039,24 +1039,23 @@ impl GroupedHashAggregateStream {
             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)
+        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,
                 }
-                EmitTo::All => EmitTo::All,
-            }
-        } else {
-            emit_to
-        };
+            } else {
+                emit_to
+            };
 
         let timer = self.group_by_metrics.emitting_time.timer();
         let mut output = self.group_values.emit(emit_to)?;
@@ -1070,10 +1069,8 @@ impl GroupedHashAggregateStream {
         } 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],
-            );
+            let schema =
+                schema_with_group_values(&self.schema, &output[..num_group_columns]);
             self.schema = Arc::clone(&schema);
             schema
         };
@@ -1498,11 +1495,13 @@ mod tests {
             batch.schema().field(0).data_type(),
             &DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8))
         );
-        assert!(batch
-            .column(0)
-            .as_any()
-            .downcast_ref::<DictionaryArray<UInt16Type>>()
-            .is_some());
+        assert!(
+            batch
+                .column(0)
+                .as_any()
+                .downcast_ref::<DictionaryArray<UInt16Type>>()
+                .is_some()
+        );
         assert!(stream.next().await.is_none());
         Ok(())
     }

From c86c39d7ab31e31108daddc1b4787773ffdd5f50 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]"
 <41898282+github-actions[bot]@users.noreply.github.com>
Date: Fri, 10 Jul 2026 17:39:28 +0000
Subject: [PATCH 27/28] fix: keep promoted aggregate stream schemas current

---
 datafusion/physical-plan/src/aggregates/hash_stream.rs      | 1 +
 .../physical-plan/src/aggregates/ordered_final_stream.rs    | 2 ++
 .../test_files/dictionary_group_by_key_promotion.slt        | 6 ++++++
 3 files changed, 9 insertions(+)

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
index 04e080d2e7d3d..ae7b6b577828c 100644
--- a/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt
+++ b/datafusion/sqllogictest/test_files/dictionary_group_by_key_promotion.slt
@@ -67,3 +67,9 @@ FROM (
 );
 ----
 Dictionary(UInt16, Utf8)
+
+statement ok
+SET datafusion.execution.batch_size = 8192;
+
+statement ok
+SET datafusion.execution.target_partitions = 4;

From 23d6d1ef60fb643bd1d38b0d3b9e3eda6d65a0ed Mon Sep 17 00:00:00 2001
From: StealthEyeLLC <5gcrth85jw@privaterelay.appleid.com>
Date: Fri, 10 Jul 2026 13:58:55 -0400
Subject: [PATCH 28/28] test: avoid cloned slices in dictionary aggregation
 regression

---
 .../src/aggregates/aggregate_hash_table/tests.rs              | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs
index 5e22f538a6233..28b765f69a654 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/tests.rs
@@ -48,7 +48,7 @@ fn dictionary_groups_keep_partial_schema_and_promote_final_output() -> Result<()
         .collect::<Result<Vec<_>>>()?;
 
     let input = TestMemoryExec::try_new(
-        &[input_batches.clone()],
+        std::slice::from_ref(&input_batches),
         Arc::clone(&input_schema),
         None,
     )?;
@@ -102,7 +102,7 @@ fn dictionary_groups_keep_partial_schema_and_promote_final_output() -> Result<()
     );
 
     let partial_input = TestMemoryExec::try_new(
-        &[partial_batches.clone()],
+        std::slice::from_ref(&partial_batches),
         Arc::clone(&partial_schema),
         None,
     )?;