Refactor emission handling and simplify logic

Remove redundant num_groups field and derive emission size from
values. Collapse deferred-state retention into a tighter
iterator/unzip flow. Eliminate the extra append_batch_values
forwarding helper. Split evaluate() into smaller private steps
with replay_deferred_batches and finish_emit. Simplify
should_defer and refine the deferred replay loop for clarity
and efficiency.

Author: kosiew

datafusion/functions-aggregate/src/string_agg.rs

@@ -334,8 +334,6 @@ struct StringAggGroupsAccumulator {
     batches: Vec<ArrayRef>,
     /// Per-batch `(group_idx, row_idx)` pairs for non-null rows.
     batch_entries: Vec<Vec<(u32, u32)>>,
-    /// Total number of groups tracked.
-    num_groups: usize,
 }
 
 enum StringInputArray<'a> {
@@ -422,7 +420,6 @@ impl StringAggGroupsAccumulator {
             total_data_bytes: 0,
             batches: Vec::new(),
             batch_entries: Vec::new(),
-            num_groups: 0,
         }
     }
 
@@ -433,38 +430,32 @@ impl StringAggGroupsAccumulator {
         self.total_data_bytes = 0;
         self.batches = Vec::new();
         self.batch_entries = Vec::new();
-        self.num_groups = 0;
     }
 
     fn retain_after_emit(&mut self, emit_groups: usize) {
         let emit_groups = emit_groups as u32;
-        let mut retained_batches = Vec::with_capacity(self.batches.len());
-        let mut retained_entries = Vec::with_capacity(self.batch_entries.len());
-
-        for (batch, entries) in self.batches.drain(..).zip(self.batch_entries.drain(..)) {
-            let entries: Vec<_> = entries
-                .into_iter()
-                .filter_map(|(group_idx, row_idx)| {
-                    if group_idx >= emit_groups {
-                        Some((group_idx - emit_groups, row_idx))
-                    } else {
-                        None
-                    }
-                })
-                .collect();
-            if entries.is_empty() {
-                continue;
-            }
-
-            // Keep the original arrays for this prototype and only renumber
-            // retained groups. SUB_ISSUE_04 will compact mixed batches so
-            // partially emitted batches no longer pin their full inputs.
-            retained_batches.push(batch);
-            retained_entries.push(entries);
-        }
-
-        self.batches = retained_batches;
-        self.batch_entries = retained_entries;
+        (self.batches, self.batch_entries) = self
+            .batches
+            .drain(..)
+            .zip(self.batch_entries.drain(..))
+            .filter_map(|(batch, entries)| {
+                let entries = entries
+                    .into_iter()
+                    .filter_map(|(group_idx, row_idx)| {
+                        if group_idx >= emit_groups {
+                            Some((group_idx - emit_groups, row_idx))
+                        } else {
+                            None
+                        }
+                    })
+                    .collect::<Vec<_>>();
+
+                // Keep the original arrays for this prototype and only renumber
+                // retained groups. SUB_ISSUE_04 will compact mixed batches so
+                // partially emitted batches no longer pin their full inputs.
+                (!entries.is_empty()).then_some((batch, entries))
+            })
+            .unzip();
     }
 
     fn append_rows_typed<'a, A>(array: &A, group_indices: &[usize]) -> Vec<(u32, u32)>
@@ -529,12 +520,12 @@ impl StringAggGroupsAccumulator {
     ) where
         A: ArrayAccessor<Item = &'a str>,
     {
-        for &(group_idx, row_idx) in entries {
+        for (group_idx, row_idx) in entries
+            .iter()
+            .copied()
+            .filter(|(group_idx, _)| (*group_idx as usize) < emit_groups)
+        {
             let group_idx = group_idx as usize;
-            if group_idx >= emit_groups {
-                continue;
-            }
-
             let row_idx = row_idx as usize;
             debug_assert!(!array.is_null(row_idx));
             let _ = Self::append_group_value(
@@ -546,31 +537,42 @@ impl StringAggGroupsAccumulator {
         }
     }
 
-    fn append_batch_values(
+    fn should_defer(
+        &self,
+        input: &StringInputArray<'_>,
+        total_num_groups: usize,
+    ) -> bool {
+        if total_num_groups < Self::DEFER_GROUP_THRESHOLD {
+            return false;
+        }
+
+        input
+            .sample_non_null_len()
+            .is_some_and(|len| len >= Self::DEFER_PAYLOAD_LEN_THRESHOLD)
+    }
+
+    fn replay_deferred_batches(
+        &self,
         values: &mut [Option<String>],
-        entries: &[(u32, u32)],
-        array: &ArrayRef,
-        delimiter: &str,
         emit_groups: usize,
     ) -> Result<()> {
-        StringInputArray::try_new(array)?.append_batch_values(
-            values,
-            entries,
-            delimiter,
-            emit_groups,
-        );
+        for (batch, entries) in self.batches.iter().zip(&self.batch_entries) {
+            StringInputArray::try_new(batch)?.append_batch_values(
+                values,
+                entries,
+                &self.delimiter,
+                emit_groups,
+            );
+        }
+
         Ok(())
     }
 
-    fn should_defer(
-        &self,
-        input: &StringInputArray<'_>,
-        total_num_groups: usize,
-    ) -> bool {
-        total_num_groups >= Self::DEFER_GROUP_THRESHOLD
-            && input
-                .sample_non_null_len()
-                .is_some_and(|len| len >= Self::DEFER_PAYLOAD_LEN_THRESHOLD)
+    fn finish_emit(&mut self, emit_to: EmitTo, emit_groups: usize) {
+        match emit_to {
+            EmitTo::All => self.clear_state(),
+            EmitTo::First(_) => self.retain_after_emit(emit_groups),
+        }
     }
 }
 
@@ -582,7 +584,6 @@ impl GroupsAccumulator for StringAggGroupsAccumulator {
         opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
-        self.num_groups = self.num_groups.max(total_num_groups);
         self.values.resize(total_num_groups, None);
         let array = apply_filter_as_nulls(&values[0], opt_filter)?;
         let input = StringInputArray::try_new(&array)?;
@@ -605,31 +606,18 @@ impl GroupsAccumulator for StringAggGroupsAccumulator {
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
-        let mut to_emit = emit_to.take_needed(&mut self.values);
+        let mut to_emit = match emit_to {
+            EmitTo::All => std::mem::take(&mut self.values),
+            EmitTo::First(_) => emit_to.take_needed(&mut self.values),
+        };
         let emit_groups = to_emit.len();
         let emitted_bytes: usize = to_emit
             .iter()
             .filter_map(|opt| opt.as_ref().map(|s| s.len()))
             .sum();
         self.total_data_bytes -= emitted_bytes;
-
-        for (batch, entries) in self.batches.iter().zip(&self.batch_entries) {
-            Self::append_batch_values(
-                &mut to_emit,
-                entries,
-                batch,
-                &self.delimiter,
-                emit_groups,
-            )?;
-        }
-
-        match emit_to {
-            EmitTo::All => self.clear_state(),
-            EmitTo::First(_) => {
-                self.retain_after_emit(emit_groups);
-                self.num_groups = self.values.len();
-            }
-        }
+        self.replay_deferred_batches(&mut to_emit, emit_groups)?;
+        self.finish_emit(emit_to, emit_groups);
 
         Ok(Arc::new(LargeStringArray::from(to_emit)))
     }