Add preallocate to GroupsAccumulator trait and key implementations

Add `preallocate(total_num_groups)` method to the `GroupsAccumulator`
trait (default no-op) and implement it for:
- PrimitiveGroupsAccumulator (SUM, MIN, MAX, etc.)
- CountGroupsAccumulator
- VarianceGroupsAccumulator
- CorrelationGroupsAccumulator

Call preallocate on all accumulators in GroupedHashAggregateStream when
NDV capacity hint is available.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Dandandan

datafusion/expr-common/src/groups_accumulator.rs

@@ -247,6 +247,16 @@ pub trait GroupsAccumulator: Send + std::any::Any {
         false
     }
 
+    /// Pre-allocates internal storage for the given number of groups.
+    ///
+    /// This is an optional optimization hint. When statistics (such as NDV
+    /// from Parquet metadata) predict the number of distinct groups, calling
+    /// this before the first `update_batch` avoids repeated resizing of
+    /// internal vectors.
+    ///
+    /// The default implementation is a no-op.
+    fn preallocate(&mut self, _total_num_groups: usize) {}
+
     /// Amount of memory used to store the state of this accumulator,
     /// in bytes.
     ///

datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs

@@ -195,6 +195,10 @@ where
         true
     }
 
+    fn preallocate(&mut self, total_num_groups: usize) {
+        self.values.resize(total_num_groups, self.starting_value);
+    }
+
     fn size(&self) -> usize {
         self.values.capacity() * size_of::<T::Native>() + self.null_state.size()
     }

datafusion/functions-aggregate/src/correlation.rs

@@ -536,6 +536,15 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator {
         Ok(())
     }
 
+    fn preallocate(&mut self, total_num_groups: usize) {
+        self.count.resize(total_num_groups, 0);
+        self.sum_x.resize(total_num_groups, 0.0);
+        self.sum_y.resize(total_num_groups, 0.0);
+        self.sum_xy.resize(total_num_groups, 0.0);
+        self.sum_xx.resize(total_num_groups, 0.0);
+        self.sum_yy.resize(total_num_groups, 0.0);
+    }
+
     fn size(&self) -> usize {
         self.count.capacity() * size_of::<u64>()
             + self.sum_x.capacity() * size_of::<f64>()

datafusion/functions-aggregate/src/count.rs

@@ -712,6 +712,10 @@ impl GroupsAccumulator for CountGroupsAccumulator {
         true
     }
 
+    fn preallocate(&mut self, total_num_groups: usize) {
+        self.counts.resize(total_num_groups, 0);
+    }
+
     fn size(&self) -> usize {
         self.counts.capacity() * size_of::<usize>()
     }

datafusion/functions-aggregate/src/variance.rs

@@ -575,6 +575,12 @@ impl GroupsAccumulator for VarianceGroupsAccumulator {
         ])
     }
 
+    fn preallocate(&mut self, total_num_groups: usize) {
+        self.m2s.resize(total_num_groups, 0.0);
+        self.means.resize(total_num_groups, 0.0);
+        self.counts.resize(total_num_groups, 0);
+    }
+
     fn size(&self) -> usize {
         self.m2s.capacity() * size_of::<f64>()
             + self.means.capacity() * size_of::<f64>()

datafusion/physical-plan/src/aggregates/row_hash.rs

@@ -498,7 +498,7 @@ impl GroupedHashAggregateStream {
         };
 
         // Instantiate the accumulators
-        let accumulators: Vec<_> = aggregate_exprs
+        let mut accumulators: Vec<_> = aggregate_exprs
             .iter()
             .map(create_group_accumulator)
             .collect::<Result<_>>()?;
@@ -597,6 +597,12 @@ impl GroupedHashAggregateStream {
             .and_then(|stats| agg.compute_group_ndv(&stats))
             .map(|ndv: usize| ndv.min(MAX_NDV_CAPACITY));
 
+        if let Some(capacity) = capacity_hint {
+            for acc in &mut accumulators {
+                acc.preallocate(capacity);
+            }
+        }
+
         let group_values =
             new_group_values(group_schema, &group_ordering, capacity_hint)?;
         let reservation = MemoryConsumer::new(name)