Revert "Refactor emission handling and simplify logic"

This reverts commit baa8054.

Author: kosiew

datafusion/functions-aggregate/src/string_agg.rs

@@ -334,6 +334,8 @@ 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> {
@@ -420,6 +422,7 @@ impl StringAggGroupsAccumulator {
             total_data_bytes: 0,
             batches: Vec::new(),
             batch_entries: Vec::new(),
+            num_groups: 0,
         }
     }
 
@@ -430,32 +433,38 @@ 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;
-        (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();
+        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;
     }
 
     fn append_rows_typed<'a, A>(array: &A, group_indices: &[usize]) -> Vec<(u32, u32)>
@@ -520,12 +529,12 @@ impl StringAggGroupsAccumulator {
     ) where
         A: ArrayAccessor<Item = &'a str>,
     {
-        for (group_idx, row_idx) in entries
-            .iter()
-            .copied()
-            .filter(|(group_idx, _)| (*group_idx as usize) < emit_groups)
-        {
+        for &(group_idx, row_idx) in entries {
             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(
@@ -537,42 +546,31 @@ impl StringAggGroupsAccumulator {
         }
     }
 
-    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,
+    fn append_batch_values(
         values: &mut [Option<String>],
+        entries: &[(u32, u32)],
+        array: &ArrayRef,
+        delimiter: &str,
         emit_groups: usize,
     ) -> Result<()> {
-        for (batch, entries) in self.batches.iter().zip(&self.batch_entries) {
-            StringInputArray::try_new(batch)?.append_batch_values(
-                values,
-                entries,
-                &self.delimiter,
-                emit_groups,
-            );
-        }
-
+        StringInputArray::try_new(array)?.append_batch_values(
+            values,
+            entries,
+            delimiter,
+            emit_groups,
+        );
         Ok(())
     }
 
-    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),
-        }
+    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)
     }
 }
 
@@ -584,6 +582,7 @@ 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)?;
@@ -606,18 +605,31 @@ impl GroupsAccumulator for StringAggGroupsAccumulator {
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
-        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 mut to_emit = 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;
-        self.replay_deferred_batches(&mut to_emit, emit_groups)?;
-        self.finish_emit(emit_to, emit_groups);
+
+        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();
+            }
+        }
 
         Ok(Arc::new(LargeStringArray::from(to_emit)))
     }