perf: reserve group-key storage across partial flushes

datafusion/common/src/config.rs

@@ -1136,6 +1136,28 @@ config_namespace! {
         /// aggregation ratio check and trying to switch to skipping aggregation mode
         pub skip_partial_aggregation_probe_rows_threshold: usize, default = 100_000
 
+        /// (experimental) Allocated-byte threshold for flushing unordered partial
+        /// hash aggregation tables. Checked after each input batch; the table can
+        /// exceed this size by one batch. Flushed groups count toward the existing
+        /// skip-partial probe. Repeated keys do not disable flushing. Aggregations
+        /// with a soft group limit or nested aggregate state are excluded.
+        /// Threshold flushes retain group-key hash tables and scratch buffers.
+        /// Key buffers transferred to output are reserved again using the previous
+        /// lengths, including primitive, string, and multi-column keys. StringView
+        /// payloads use contiguous buffers up to the view offset limit. Retained
+        /// capacity counts toward this threshold and memory reservations.
+        /// Memory-pressure emission still releases allocations.
+        /// Set to 0 to disable this threshold. If both flush thresholds are
+        /// enabled, reaching either triggers a flush. Final aggregation is unchanged.
+        pub partial_aggregation_flush_bytes: usize, default = 0
+
+        /// (experimental) Number of distinct group rows above which an unordered
+        /// partial hash aggregation table is flushed. This counts groups held in
+        /// the table, not input rows. Uses the same eligibility and skip-partial
+        /// accounting as partial_aggregation_flush_bytes. Set to 0 to disable this
+        /// threshold. If both thresholds are enabled, reaching either triggers a flush.
+        pub partial_aggregation_flush_rows: usize, default = 0
+
         /// Should DataFusion use row number estimates at the input to decide
         /// whether increasing parallelism is beneficial or not. By default,
         /// only exact row numbers (not estimates) are used for this decision.

datafusion/physical-expr-common/src/binary_map.rs

@@ -336,6 +336,25 @@ where
         new_self
     }
 
+    /// Emits keys while retaining the hash table and reserving replacement key buffers.
+    /// Output arrays own the previous key buffers, so subsequent input cannot change them.
+    pub fn take_state_reusing_allocation(&mut self) -> ArrayRef {
+        let payload_bytes = self.buffer.len();
+        let num_offsets = self.offsets.len();
+        let mut outgoing = Self::new_inner(self.output_type, 0, 0);
+        swap(self, &mut outgoing);
+        swap(&mut self.map, &mut outgoing.map);
+        self.map.clear();
+        swap(&mut self.hashes_buffer, &mut outgoing.hashes_buffer);
+        self.hashes_buffer.clear();
+        swap(&mut self.random_state, &mut outgoing.random_state);
+        self.initial_map_capacity = outgoing.initial_map_capacity;
+        self.initial_buffer_capacity = outgoing.initial_buffer_capacity;
+        self.buffer.reserve_exact(payload_bytes);
+        self.offsets.reserve_exact(num_offsets.saturating_sub(1));
+        outgoing.into_state()
+    }
+
     /// Empties this map and releases every allocation it holds, so
     /// [`Self::size`] drops to approximately zero.
     ///

datafusion/physical-expr-common/src/binary_view_map.rs

@@ -145,6 +145,8 @@ where
     views: Vec<u128>,
     /// In-progress buffer for out-of-line string data
     in_progress: Vec<u8>,
+    /// Maximum payload block size; extended after a reusable Partial flush.
+    payload_block_limit: usize,
     /// Completed buffers containing string data
     completed: Vec<Buffer>,
 
@@ -190,6 +192,7 @@ where
             initial_map_capacity: map_capacity,
             views: Vec::new(),
             in_progress: Vec::new(),
+            payload_block_limit: BYTE_VIEW_MAX_BLOCK_SIZE,
             completed: Vec::new(),
             random_state: RandomState::default(),
             hashes_buffer: vec![],
@@ -206,6 +209,33 @@ where
         new_self
     }
 
+    /// Emits all keys while retaining the hash allocation for a Partial flush.
+    /// The returned array owns the emitted strings. The empty map uses one payload
+    /// buffer, reserved from the previous payload length, up to the u32 offset
+    /// limit. Normal output and memory-pressure emission should use [`Self::take`].
+    pub fn take_state_reusing_allocation(&mut self) -> ArrayRef {
+        let payload_bytes = self
+            .completed
+            .iter()
+            .fold(self.in_progress.len(), |total, buffer| {
+                total.saturating_add(buffer.len())
+            });
+        let num_views = self.views.len();
+        let mut outgoing = Self::with_capacity(self.output_type, 0);
+        std::mem::swap(self, &mut outgoing);
+        std::mem::swap(&mut self.map, &mut outgoing.map);
+        self.map.clear();
+        std::mem::swap(&mut self.hashes_buffer, &mut outgoing.hashes_buffer);
+        self.hashes_buffer.clear();
+        self.views.reserve_exact(num_views);
+        self.initial_map_capacity = outgoing.initial_map_capacity;
+        std::mem::swap(&mut self.random_state, &mut outgoing.random_state);
+        self.payload_block_limit = u32::MAX as usize;
+        self.in_progress
+            .reserve_exact(payload_bytes.min(self.payload_block_limit));
+        outgoing.into_state()
+    }
+
     /// Empties this map and releases every allocation it holds, so
     /// [`Self::size`] drops to approximately zero.
     ///
@@ -538,7 +568,7 @@ where
             make_view(value, 0, 0)
         } else {
             // Ensure buffer is big enough
-            if self.in_progress.len() + len > BYTE_VIEW_MAX_BLOCK_SIZE {
+            if self.in_progress.len() + len > self.payload_block_limit {
                 let flushed = std::mem::replace(
                     &mut self.in_progress,
                     Vec::with_capacity(BYTE_VIEW_MAX_BLOCK_SIZE),
@@ -890,6 +920,48 @@ mod tests {
         assert_eq!(lazy.map.capacity(), 0);
     }
 
+    #[test]
+    fn partial_flush_reuses_table_and_preserves_emitted_strings() {
+        let mut map = ArrowBytesViewMap::<usize>::new(OutputType::Utf8View);
+        let long_value = "x".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1);
+        let other_value = "y".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1);
+        let expected = vec![
+            Some("inline"),
+            None,
+            Some(long_value.as_str()),
+            Some(other_value.as_str()),
+        ];
+        let values: ArrayRef = Arc::new(StringViewArray::from(expected.clone()));
+        let mut emitted = Vec::new();
+        for _ in 0..3 {
+            let mut next_group = 0;
+            let mut groups = Vec::new();
+            map.insert_if_new(
+                &values,
+                |_| {
+                    let group = next_group;
+                    next_group += 1;
+                    group
+                },
+                |group| groups.push(group),
+            );
+            assert_eq!(groups, vec![0, 1, 2, 3]);
+            let capacity = map.map.capacity();
+            emitted.push(map.take_state_reusing_allocation());
+            assert!(map.is_empty());
+            assert_eq!(map.map.capacity(), capacity);
+            assert!(map.in_progress.capacity() >= long_value.len() + other_value.len());
+        }
+        for (index, array) in emitted.into_iter().enumerate() {
+            let array = array.as_string_view();
+            assert_eq!(array.iter().collect::<Vec<_>>(), expected);
+            assert_eq!(array.data_buffers().len(), if index == 0 { 2 } else { 1 });
+        }
+        map.clear_and_release();
+        assert_eq!(map.map.allocation_size(), 0);
+        assert_eq!(map.in_progress.capacity(), 0);
+    }
+
     #[test]
     fn clear_and_release_frees_the_preallocation_that_take_keeps() {
         let mut map = ArrowBytesViewMap::<()>::with_capacity(

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs

@@ -372,6 +372,20 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
     /// spilling without finalizing the same group more than once.
     pub(in crate::aggregates) fn take_state_batch(
         &mut self,
+    ) -> Result<Option<RecordBatch>> {
+        self.take_state_batch_inner(false)
+    }
+
+    /// Emits a threshold-triggered Partial batch while allowing key storage reuse.
+    pub(in crate::aggregates) fn take_state_batch_for_partial_flush(
+        &mut self,
+    ) -> Result<Option<RecordBatch>> {
+        self.take_state_batch_inner(true)
+    }
+
+    fn take_state_batch_inner(
+        &mut self,
+        is_partial_flush: bool,
     ) -> Result<Option<RecordBatch>> {
         let state_schema = Arc::clone(&self.state_schema);
         let accumulator_metrics = Arc::clone(&self.aggregate_accumulator_metrics);
@@ -382,7 +396,11 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         }
 
         let output = group_by_metrics.time_emitting(|| {
-            let mut output = state.group_values.emit(EmitTo::All)?;
+            let mut output = if is_partial_flush {
+                state.group_values.emit_for_partial_flush()?
+            } else {
+                state.group_values.emit(EmitTo::All)?
+            };
             for (idx, acc) in state.accumulators.iter_mut().enumerate() {
                 output.extend(accumulator_metrics.time(
                     idx,
@@ -399,9 +417,13 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the
         // key/index buffers too so the memory reservation can be released
         // before the batch is sorted for spilling.
-        state.group_values.clear_shrink(0);
+        if !is_partial_flush {
+            state.group_values.clear_shrink(0);
+        }
         state.batch_group_indices.clear();
-        state.batch_group_indices.shrink_to_fit();
+        if !is_partial_flush {
+            state.batch_group_indices.shrink_to_fit();
+        }
 
         Ok(Some(batch))
     }

datafusion/physical-plan/src/aggregates/group_values/mod.rs

@@ -111,6 +111,14 @@ pub trait GroupValues: Send {
     /// Emits the group values
     fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>>;
 
+    /// Emits all keys for a Partial threshold flush, preparing for more input.
+    /// Implementations may retain allocations when the table is filled again.
+    fn emit_for_partial_flush(&mut self) -> Result<Vec<ArrayRef>> {
+        let values = self.emit(EmitTo::All)?;
+        self.clear_shrink(0);
+        Ok(values)
+    }
+
     /// Materializes selected group values without changing the stored values or
     /// their group indices.
     ///
@@ -215,14 +223,110 @@ mod tests {
     use std::sync::Arc;
 
     use arrow::array::{
-        ArrayRef, AsArray, BooleanArray, Int32Array, StringArray, StringViewArray,
+        Array, ArrayRef, AsArray, BooleanArray, FixedSizeBinaryArray, Int32Array,
+        ListArray, StringArray, StringViewArray, StructArray, UInt32Array,
     };
+    use arrow::compute;
     use arrow::datatypes::{DataType, Field, Int32Type, Schema};
     use datafusion_expr::{EmitTo, GroupSelection};
 
     use super::new_group_values;
     use crate::aggregates::order::GroupOrdering;
 
+    #[test]
+    fn partial_flush_preserves_output_and_resets_all_key_types() {
+        let strings: ArrayRef = Arc::new(StringArray::from(vec![
+            Some("first long string value"),
+            None,
+            Some("second long string value"),
+        ]));
+        let mut arrays: Vec<ArrayRef> = vec![
+            Arc::new(Int32Array::from(vec![Some(11), None, Some(22)])),
+            Arc::new(BooleanArray::from(vec![Some(true), None, Some(false)])),
+            Arc::clone(&strings),
+            Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+                Some(vec![Some(1), None]),
+                None,
+                Some(vec![Some(2), Some(3)]),
+            ])),
+            Arc::new(
+                FixedSizeBinaryArray::try_from_sparse_iter_with_size(
+                    vec![Some(b"abc".as_slice()), None, Some(b"def".as_slice())]
+                        .into_iter(),
+                    3,
+                )
+                .unwrap(),
+            ),
+        ];
+        arrays.push(Arc::new(StructArray::from(vec![(
+            Arc::new(Field::new("nested", DataType::Int32, true)),
+            Arc::clone(&arrays[0]),
+        )])));
+        for data_type in [
+            DataType::LargeUtf8,
+            DataType::Binary,
+            DataType::LargeBinary,
+            DataType::Utf8View,
+            DataType::BinaryView,
+            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
+        ] {
+            arrays.push(compute::cast(&strings, &data_type).unwrap());
+        }
+        let mut cases: Vec<Vec<ArrayRef>> =
+            arrays.iter().map(|a| vec![Arc::clone(a)]).collect();
+        cases.push(arrays);
+        for columns in cases {
+            let schema = Arc::new(Schema::new(
+                columns
+                    .iter()
+                    .enumerate()
+                    .map(|(i, a)| {
+                        Field::new(format!("key{i}"), a.data_type().clone(), true)
+                    })
+                    .collect::<Vec<_>>(),
+            ));
+            let mut values = new_group_values(schema, &GroupOrdering::None).unwrap();
+            let mut retained_outputs = vec![];
+            let mut groups = vec![];
+            for indices in [[0, 1, 2, 0], [2, 0, 1, 2], [1, 2, 0, 1]] {
+                let input = columns
+                    .iter()
+                    .map(|a| {
+                        compute::take(
+                            a.as_ref(),
+                            &UInt32Array::from(indices.to_vec()),
+                            None,
+                        )
+                        .unwrap()
+                    })
+                    .collect::<Vec<_>>();
+                values.intern(&input, &mut groups).unwrap();
+                assert_eq!(groups, vec![0, 1, 2, 0]);
+                let output = values.emit_for_partial_flush().unwrap();
+                assert!(values.is_empty());
+                let expected = input.iter().map(|a| a.slice(0, 3)).collect::<Vec<_>>();
+                retained_outputs.push((output, expected));
+            }
+            // Earlier batches stay alive while the same builders are refilled.
+            for (output, expected) in retained_outputs {
+                for (actual, expected) in output.iter().zip(expected.iter()) {
+                    // Dictionary output may compact or reorder its dictionary values.
+                    let data_type = match expected.data_type() {
+                        DataType::Dictionary(_, value_type) => value_type.as_ref(),
+                        data_type => data_type,
+                    };
+                    let actual = compute::cast(actual, data_type).unwrap();
+                    let expected = compute::cast(expected, data_type).unwrap();
+                    assert_eq!(actual.as_ref(), expected.as_ref());
+                }
+            }
+            values.clear_shrink(0);
+            assert!(values.is_empty());
+            values.intern(&columns, &mut groups).unwrap();
+            assert_eq!(groups, vec![0, 1, 2]);
+        }
+    }
+
     #[test]
     fn preserving_values_keep_group_indices_valid() {
         let schema = Arc::new(Schema::new(vec![Field::new(

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs

@@ -419,6 +419,15 @@ where
         })
     }
 
+    fn take_for_partial_flush(&mut self) -> ArrayRef {
+        let num_values = self.len();
+        let payload_bytes = self.buffer.len();
+        let output = self.take_n(num_values);
+        self.offsets.reserve_exact(num_values);
+        self.buffer.reserve(payload_bytes);
+        output
+    }
+
     fn take_n(&mut self, n: usize) -> ArrayRef {
         debug_assert!(self.len() >= n);
         let null_buffer = self.nulls.take_n(n);

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs

@@ -652,6 +652,23 @@ impl<B: ByteViewType> GroupColumn for ByteViewGroupValueBuilder<B> {
         self.values_preserving_inner(selection)
     }
 
+    fn take_for_partial_flush(&mut self) -> ArrayRef {
+        let num_views = self.views.len();
+        let payload_bytes = self
+            .completed
+            .iter()
+            .fold(self.in_progress.len(), |total, buffer| {
+                total.saturating_add(buffer.len())
+            });
+        let mut fresh = Self::new().with_max_block_size(u32::MAX as usize);
+        fresh.views.reserve_exact(num_views);
+        fresh
+            .in_progress
+            .reserve_exact(payload_bytes.min(u32::MAX as usize));
+        fresh.nulls = NullBufferBuilder::new(num_views);
+        replace(self, fresh).build_inner()
+    }
+
     fn take_n(&mut self, n: usize) -> ArrayRef {
         self.take_n_inner(n)
     }

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs

@@ -549,6 +549,18 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
         Ok(Self::into_dict(values, &selected_groups, null_inner_slot))
     }
 
+    fn take_for_partial_flush(&mut self) -> ArrayRef {
+        let values = self.inner.take_for_partial_flush();
+        let output = Self::into_dict(values, &self.group_to_inner, self.null_inner_slot);
+        self.group_to_inner.clear();
+        self.value_dedup.clear();
+        self.null_inner_slot = None;
+        self.val_to_inner.clear();
+        self.val_hashes.clear();
+        self.cached_values = None;
+        output
+    }
+
     fn take_n(&mut self, n: usize) -> ArrayRef {
         let old_inner_len = self.inner.len();
         let all_inner_values = self.inner.take_n(old_inner_len);