@@ -4114,3 +4114,114 @@ mod native_sort_spill_tests {
         }
     }
 }
+
+#[cfg(test)]
+mod aggregate_offset_overflow_tests {
+    use arrow::array::{Int32Array, StringArray};
+    use arrow::buffer::{Buffer, OffsetBuffer};
+    use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+    use arrow::record_batch::RecordBatch;
+    use datafusion::datasource::memory::MemorySourceConfig;
+    use datafusion::execution::TaskContext;
+    use datafusion::physical_expr::expressions::col;
+    use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
+    use datafusion::physical_plan::ExecutionPlan;
+    use datafusion::prelude::SessionConfig;
+    use futures::StreamExt;
+    use std::sync::Arc;
+
+    const KEY_BYTES: usize = 1 << 20;
+    const ROWS: usize = 1024;
+    const BATCHES: usize = 3;
+
+    /// `(day, product_id)` batches of distinct 1 MiB keys, 1 GiB per batch and 3 GiB in all,
+    /// sharing one buffer of pseudo-random letters so that the input itself stays at 1 GiB.
+    fn wide_keys(schema: &SchemaRef) -> Vec<RecordBatch> {
+        let mut seed = 0x2545_f491_4f6c_dd1du64;
+        let letters: Vec<u8> = (0..ROWS * KEY_BYTES + BATCHES)
+            .map(|_| {
+                seed ^= seed << 13;
+                seed ^= seed >> 7;
+                seed ^= seed << 17;
+                b'a' + (seed % 26) as u8
+            })
+            .collect();
+        let values = Buffer::from_vec(letters);
+        (0..BATCHES)
+            .map(|batch| {
+                let offsets = OffsetBuffer::new(
+                    (0..=ROWS)
+                        .map(|row| (batch + row * KEY_BYTES) as i32)
+                        .collect(),
+                );
+                RecordBatch::try_new(
+                    Arc::clone(schema),
+                    vec![
+                        Arc::new(Int32Array::from(vec![20_000; ROWS])),
+                        Arc::new(StringArray::try_new(offsets, values.clone(), None).unwrap()),
+                    ],
+                )
+                .unwrap()
+            })
+            .collect()
+    }
+
+    /// `SELECT DISTINCT` of `columns` in one aggregate of `mode`, with nothing but the
+    /// 32-bit offsets of its group keys to stop it from holding every group.
+    async fn distinct_wide_keys(
+        mode: AggregateMode,
+        columns: &[&str],
+    ) -> datafusion::common::Result<usize> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("day", DataType::Int32, false),
+            Field::new("product_id", DataType::Utf8, false),
+        ]));
+        let input =
+            MemorySourceConfig::try_new_exec(&[wide_keys(&schema)], Arc::clone(&schema), None)?;
+        let group_by = PhysicalGroupBy::new_single(
+            columns
+                .iter()
+                .map(|name| Ok((col(name, &schema)?, name.to_string())))
+                .collect::<datafusion::common::Result<_>>()?,
+        );
+        let aggregate = AggregateExec::try_new(mode, group_by, vec![], vec![], input, schema)?;
+        let ctx =
+            TaskContext::default().with_session_config(SessionConfig::new().with_batch_size(64));
+        let mut stream = aggregate.execute(0, Arc::new(ctx))?;
+        let mut rows = 0;
+        while let Some(batch) = stream.next().await {
+            rows += batch?.num_rows();
+        }
+        Ok(rows)
+    }
+
+    #[tokio::test]
+    async fn partial_aggregate_emits_before_group_keys_overflow_32_bit_offsets() {
+        assert_eq!(
+            distinct_wide_keys(AggregateMode::Partial, &["day", "product_id"])
+                .await
+                .unwrap(),
+            ROWS * BATCHES
+        );
+    }
+
+    #[tokio::test]
+    async fn final_aggregate_spills_before_group_keys_overflow_32_bit_offsets() {
+        assert_eq!(
+            distinct_wide_keys(AggregateMode::Final, &["day", "product_id"])
+                .await
+                .unwrap(),
+            ROWS * BATCHES
+        );
+    }
+
+    #[tokio::test]
+    async fn single_key_aggregate_emits_before_group_keys_overflow_32_bit_offsets() {
+        assert_eq!(
+            distinct_wide_keys(AggregateMode::Partial, &["product_id"])
+                .await
+                .unwrap(),
+            ROWS * BATCHES
+        );
+    }
+}
@@ -185,12 +185,27 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
     /// Each aggregation mode chooses a different `aggregate_fn` according to its
     /// semantics. For example, partial aggregation takes raw inputs, and update them
     /// into stored partial states, so [`GroupsAccumulator::update_batch`] is used.
+    ///
+    /// If the batch's group keys could overflow the 32-bit offsets of the stored group
+    /// values, every group accumulated so far is first taken out as by
+    /// [`Self::take_state_batch`] and returned, for the caller to emit or spill.
     pub(super) fn aggregate_batch_inner(
         &mut self,
         batch: &RecordBatch,
         aggregate_fn: AggregateBatchFn,
-    ) -> Result<()> {
+    ) -> Result<Option<RecordBatch>> {
         let evaluated_batch = self.evaluate_batch(batch)?;
+        let state = self.state.building();
+        let taken = if !state.group_values.is_empty()
+            && !evaluated_batch
+                .grouping_set_args
+                .iter()
+                .all(|group_values| state.group_values.has_offset_room(group_values))
+        {
+            self.take_state_batch()?
+        } else {
+            None
+        };
         let state = self.state.building_mut();
 
         let _timer = self.group_by_metrics.aggregation_time.timer();
@@ -210,7 +225,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
             }
         }
 
-        Ok(())
+        Ok(taken)
     }
 
     /// Materializes the full output once, then returns it downstream incrementally
 
@@ -66,7 +66,7 @@ impl AggregateHashTable<FinalMarker> {
     pub(in crate::aggregates) fn aggregate_batch(
         &mut self,
         batch: &RecordBatch,
-    ) -> Result<()> {
+    ) -> Result<Option<RecordBatch>> {
         self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch)
     }
 
 
@@ -60,7 +60,7 @@ impl AggregateHashTable<PartialReduceMarker> {
     pub(in crate::aggregates) fn aggregate_batch(
         &mut self,
         batch: &RecordBatch,
-    ) -> Result<()> {
+    ) -> Result<Option<RecordBatch>> {
         self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch)
     }
 
 
@@ -105,7 +105,7 @@ impl AggregateHashTable<PartialMarker> {
     pub(in crate::aggregates) fn aggregate_batch(
         &mut self,
         batch: &RecordBatch,
-    ) -> Result<()> {
+    ) -> Result<Option<RecordBatch>> {
         self.aggregate_batch_inner(batch, HashAggregateAccumulator::update_batch)
     }
 
 
@@ -65,7 +65,7 @@ impl AggregateHashTable<SingleMarker> {
     pub(in crate::aggregates) fn aggregate_batch(
         &mut self,
         batch: &RecordBatch,
-    ) -> Result<()> {
+    ) -> Result<Option<RecordBatch>> {
         self.aggregate_batch_inner(batch, HashAggregateAccumulator::update_batch)
     }
 
 
@@ -22,7 +22,7 @@ use arrow::array::types::{
     Time64MicrosecondType, Time64NanosecondType, TimestampMicrosecondType,
     TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
 };
-use arrow::array::{ArrayRef, downcast_primitive};
+use arrow::array::{Array, ArrayRef, AsArray, downcast_primitive};
 use arrow::datatypes::{DataType, SchemaRef, TimeUnit};
 use datafusion_common::Result;
 
@@ -115,6 +115,106 @@ pub trait GroupValues: Send {
 
     /// Clear the contents and shrink the capacity to the size of the batch (free up memory usage)
     fn clear_shrink(&mut self, num_rows: usize);
+
+    /// Returns false if interning `cols` could take a 32-bit offset buffer of the stored
+    /// group values, or of the arrays they are emitted as, past `i32::MAX`.
+    fn has_offset_room(&self, _cols: &[ArrayRef]) -> bool {
+        true
+    }
+}
+
+/// The largest extent `array` addresses in any of its 32-bit offset buffers (`Utf8`,
+/// `Binary`, `List` and `Map` at any depth): bytes for the former, entries for the latter.
+pub(crate) fn max_offset_extent(array: &dyn Array) -> usize {
+    max_extent(array, 0, array.len())
+}
+
+fn max_extent(array: &dyn Array, start: usize, end: usize) -> usize {
+    match array.data_type() {
+        DataType::Utf8 => {
+            let offsets = array.as_string::<i32>().value_offsets();
+            (offsets[end] - offsets[start]) as usize
+        }
+        DataType::Binary => {
+            let offsets = array.as_binary::<i32>().value_offsets();
+            (offsets[end] - offsets[start]) as usize
+        }
+        DataType::List(_) => {
+            let list = array.as_list::<i32>();
+            let offsets = list.value_offsets();
+            let (start, end) = (offsets[start] as usize, offsets[end] as usize);
+            (end - start).max(max_extent(list.values().as_ref(), start, end))
+        }
+        DataType::Map(_, _) => {
+            let map = array.as_map();
+            let offsets = map.value_offsets();
+            let (start, end) = (offsets[start] as usize, offsets[end] as usize);
+            (end - start).max(max_extent(map.entries(), start, end))
+        }
+        DataType::LargeList(_) => {
+            let list = array.as_list::<i64>();
+            let offsets = list.value_offsets();
+            max_extent(
+                list.values().as_ref(),
+                offsets[start] as usize,
+                offsets[end] as usize,
+            )
+        }
+        DataType::FixedSizeList(_, size) => {
+            let size = *size as usize;
+            max_extent(
+                array.as_fixed_size_list().values().as_ref(),
+                start * size,
+                end * size,
+            )
+        }
+        DataType::Struct(_) => array
+            .as_struct()
+            .columns()
+            .iter()
+            .map(|column| max_extent(column.as_ref(), start, end))
+            .max()
+            .unwrap_or(0),
+        DataType::Dictionary(_, _) => {
+            let values = array.as_any_dictionary().values();
+            max_extent(values.as_ref(), 0, values.len()).saturating_mul(end - start)
+        }
+        _ => 0,
+    }
+}
+
+/// True if `data_type` has a 32-bit offset buffer at any depth.
+pub(crate) fn has_offset_buffer(data_type: &DataType) -> bool {
+    match data_type {
+        DataType::Utf8 | DataType::Binary | DataType::List(_) | DataType::Map(_, _) => {
+            true
+        }
+        DataType::LargeList(field)
+        | DataType::FixedSizeList(field, _)
+        | DataType::ListView(field)
+        | DataType::LargeListView(field) => has_offset_buffer(field.data_type()),
+        DataType::Struct(fields) => fields
+            .iter()
+            .any(|field| has_offset_buffer(field.data_type())),
+        DataType::Dictionary(_, values) => has_offset_buffer(values),
+        DataType::RunEndEncoded(_, values) => has_offset_buffer(values.data_type()),
+        DataType::Union(fields, _) => fields
+            .iter()
+            .any(|(_, field)| has_offset_buffer(field.data_type())),
+        _ => false,
+    }
+}
+
+/// Whether encoded group rows of `held_bytes`, after appending rows decoded from `cols`,
+/// still decode into arrays whose 32-bit offsets fit. Every row-format encoding takes at least
+/// one byte per value byte and per list entry, so the decoded extent of each offset buffer is
+/// at most the encoded size.
+pub(crate) fn rows_have_offset_room(held_bytes: usize, cols: &[ArrayRef]) -> bool {
+    let incoming = cols
+        .iter()
+        .map(|col| max_offset_extent(col.as_ref()))
+        .fold(0usize, usize::saturating_add);
+    held_bytes.saturating_add(incoming) <= i32::MAX as usize
 }
 
 /// Return a specialized implementation of [`GroupValues`] for the given schema.
 
@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::aggregates::group_values::max_offset_extent;
 use crate::aggregates::group_values::multi_group_by::{
     GroupColumn, Nulls, nulls_equal_to,
 };
@@ -219,6 +220,10 @@ where
     }
 }
 
+fn large_extent(offsets: &[i64]) -> usize {
+    (offsets[offsets.len() - 1] - offsets[0]) as usize
+}
+
 impl<O> GroupColumn for ByteGroupValueBuilder<O>
 where
     O: OffsetSizeTrait,
@@ -244,6 +249,19 @@ where
         }
     }
 
+    fn has_offset_room(&self, column: &ArrayRef) -> bool {
+        let incoming = match column.data_type() {
+            DataType::LargeUtf8 => {
+                large_extent(column.as_string::<i64>().value_offsets())
+            }
+            DataType::LargeBinary => {
+                large_extent(column.as_binary::<i64>().value_offsets())
+            }
+            _ => max_offset_extent(column.as_ref()),
+        };
+        self.buffer.len().saturating_add(incoming) <= self.max_buffer_size
+    }
+
     fn append_val(&mut self, column: &ArrayRef, row: usize) -> Result<()> {
         // Sanity array type
         match self.output_type {
 
@@ -112,6 +112,12 @@ pub trait GroupColumn: Send + Sync {
     /// Builds a new array from the first `n` stored rows, shifting the
     /// remaining rows to the start of the builder
     fn take_n(&mut self, n: usize) -> ArrayRef;
+
+    /// Returns false if appending every row of `array` could take a 32-bit offset
+    /// buffer of this builder, or of the array it builds, past `i32::MAX`
+    fn has_offset_room(&self, _array: &ArrayRef) -> bool {
+        true
+    }
 }
 
 /// Determines if the nullability of the existing and new input array can be used
@@ -1165,6 +1171,13 @@ impl<const STREAMING: bool> GroupValues for GroupValuesColumn<STREAMING> {
         group_values_size + self.map_size + self.hashes_buffer.allocated_size()
     }
 
+    fn has_offset_room(&self, cols: &[ArrayRef]) -> bool {
+        self.group_values
+            .iter()
+            .zip(cols)
+            .all(|(group_column, col)| group_column.has_offset_room(col))
+    }
+
     fn is_empty(&self) -> bool {
         self.len() == 0
     }
 
@@ -48,6 +48,7 @@
 
 use crate::aggregates::group_values::multi_group_by::GroupColumn;
 use crate::aggregates::group_values::row::encode_array_if_necessary;
+use crate::aggregates::group_values::{has_offset_buffer, rows_have_offset_room};
 
 use arrow::array::{Array, ArrayRef, BooleanBufferBuilder};
 use arrow::datatypes::DataType;
@@ -290,6 +291,14 @@ impl GroupColumn for RowsGroupColumn {
         self.row_converter.size() + self.group_values.size()
     }
 
+    fn has_offset_room(&self, array: &ArrayRef) -> bool {
+        !has_offset_buffer(&self.output_type)
+            || rows_have_offset_room(
+                self.group_values.size(),
+                std::slice::from_ref(array),
+            )
+    }
+
     fn build(self: Box<Self>) -> ArrayRef {
         self.rows_to_array(&self.group_values)
     }
Original file line number	Diff line number	Diff line change
`@@ -66,7 +66,7 @@ impl AggregateHashTable<FinalMarker> {`
`66`	`66`	`pub(in crate::aggregates) fn aggregate_batch(`
`67`	`67`	`&mut self,`
`68`	`68`	`batch: &RecordBatch,`
`69`		`- ) -> Result<()> {`
	`69`	`+ ) -> Result<Option<RecordBatch>> {`
`70`	`70`	`self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch)`
`71`	`71`	`}`
`72`	`72`
Original file line number	Diff line number	Diff line change
`@@ -60,7 +60,7 @@ impl AggregateHashTable<PartialReduceMarker> {`
`60`	`60`	`pub(in crate::aggregates) fn aggregate_batch(`
`61`	`61`	`&mut self,`
`62`	`62`	`batch: &RecordBatch,`
`63`		`- ) -> Result<()> {`
	`63`	`+ ) -> Result<Option<RecordBatch>> {`
`64`	`64`	`self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch)`
`65`	`65`	`}`
`66`	`66`
Original file line number	Diff line number	Diff line change
`@@ -105,7 +105,7 @@ impl AggregateHashTable<PartialMarker> {`
`105`	`105`	`pub(in crate::aggregates) fn aggregate_batch(`
`106`	`106`	`&mut self,`
`107`	`107`	`batch: &RecordBatch,`
`108`		`- ) -> Result<()> {`
	`108`	`+ ) -> Result<Option<RecordBatch>> {`
`109`	`109`	`self.aggregate_batch_inner(batch, HashAggregateAccumulator::update_batch)`
`110`	`110`	`}`
`111`	`111`
Original file line number	Diff line number	Diff line change
`@@ -65,7 +65,7 @@ impl AggregateHashTable<SingleMarker> {`
`65`	`65`	`pub(in crate::aggregates) fn aggregate_batch(`
`66`	`66`	`&mut self,`
`67`	`67`	`batch: &RecordBatch,`
`68`		`- ) -> Result<()> {`
	`68`	`+ ) -> Result<Option<RecordBatch>> {`
`69`	`69`	`self.aggregate_batch_inner(batch, HashAggregateAccumulator::update_batch)`
`70`	`70`	`}`
`71`	`71`