replace InternalMerger with ColumnationMerger for per-item capacity checks

The generic InternalMerger from differential-dataflow delegates extract()
to InternalMerge::extract which processes entire buffers without capacity
checks, leading to oversized output containers. This was introduced in
TimelyDataflow/differential-dataflow#689.

Replace it with ColumnationMerger, a direct Merger impl that iterates
items individually with capacity checks after each push — restoring the
pre-#689 behavior. The InternalMerge impl on ColumnationStack is retained
because DD's Builder requires it as a trait bound.

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

Author: antiguru

src/timely-util/src/columnation.rs

@@ -26,6 +26,7 @@ use columnation::{Columnation, Region};
 use differential_dataflow::consolidation::consolidate_updates;
 use differential_dataflow::difference::Semigroup;
 use differential_dataflow::lattice::Lattice;
+use differential_dataflow::trace::implementations::merge_batcher::Merger;
 use differential_dataflow::trace::implementations::merge_batcher::container::InternalMerge;
 use differential_dataflow::trace::implementations::{BatchContainer, BuilderInput};
 use timely::container::{ContainerBuilder, DrainContainer, PushInto, SizableContainer};
@@ -633,28 +634,222 @@ where
 
     fn extract(
         &mut self,
+        position: &mut usize,
         upper: AntichainRef<T>,
         frontier: &mut Antichain<T>,
         keep: &mut Self,
         ship: &mut Self,
     ) {
-        for (data, time, diff) in self.iter() {
+        let len = self[..].len();
+        while *position < len && !keep.at_capacity() && !ship.at_capacity() {
+            let (data, time, diff) = &self[*position];
             if upper.less_equal(time) {
                 frontier.insert_with(time, |time| time.clone());
                 keep.copy_destructured(data, time, diff);
             } else {
                 ship.copy_destructured(data, time, diff);
             }
+            *position += 1;
         }
     }
 }
 
 // ---------------------------------------------------------------------------
-// ColInternalMerger type alias
+// ColumnationMerger
 // ---------------------------------------------------------------------------
 
-/// A `Merger` using internal iteration for `ColumnationStack` containers.
-pub type ColInternalMerger<D, T, R> =
-    differential_dataflow::trace::implementations::merge_batcher::InternalMerger<
-        ColumnationStack<(D, T, R)>,
-    >;
+/// A `Merger` for `ColumnationStack` containers that performs per-item
+/// merge and extract with capacity checks.
+///
+/// This replaces differential-dataflow's generic `InternalMerger<MC>` which
+/// delegates to `InternalMerge::extract` and processes entire buffers without
+/// capacity checks, leading to oversized output containers.
+pub struct ColumnationMerger<D, T, R> {
+    _marker: std::marker::PhantomData<(D, T, R)>,
+}
+
+impl<D, T, R> Default for ColumnationMerger<D, T, R> {
+    fn default() -> Self {
+        Self {
+            _marker: std::marker::PhantomData,
+        }
+    }
+}
+
+type CS<D, T, R> = ColumnationStack<(D, T, R)>;
+
+impl<D, T, R> ColumnationMerger<D, T, R>
+where
+    D: Columnation,
+    T: Columnation,
+    R: Columnation,
+{
+    #[inline]
+    fn empty(&self, stash: &mut Vec<CS<D, T, R>>) -> CS<D, T, R> {
+        stash.pop().unwrap_or_else(|| {
+            let mut container = CS::<D, T, R>::default();
+            container.ensure_capacity(&mut None);
+            container
+        })
+    }
+
+    #[inline]
+    fn recycle(&self, mut chunk: CS<D, T, R>, stash: &mut Vec<CS<D, T, R>>) {
+        chunk.clear();
+        stash.push(chunk);
+    }
+}
+
+impl<D, T, R> Merger for ColumnationMerger<D, T, R>
+where
+    D: Ord + Columnation + Clone + 'static,
+    T: Ord + Columnation + Clone + PartialOrder + 'static,
+    R: Default + Semigroup + Columnation + Clone + 'static,
+{
+    type Time = T;
+    type Chunk = CS<D, T, R>;
+
+    fn merge(
+        &mut self,
+        list1: Vec<Self::Chunk>,
+        list2: Vec<Self::Chunk>,
+        output: &mut Vec<Self::Chunk>,
+        stash: &mut Vec<Self::Chunk>,
+    ) {
+        use std::cmp::Ordering;
+
+        let mut list1 = list1.into_iter();
+        let mut list2 = list2.into_iter();
+
+        let mut head1 = list1.next().unwrap_or_default();
+        let mut head2 = list2.next().unwrap_or_default();
+        let mut pos1 = 0usize;
+        let mut pos2 = 0usize;
+
+        let mut result = self.empty(stash);
+        let mut diff_stash = R::default();
+
+        // Main merge loop: both sides have data.
+        while pos1 < head1[..].len() && pos2 < head2[..].len() {
+            while pos1 < head1[..].len() && pos2 < head2[..].len() && !result.at_capacity() {
+                let (d1, t1, _) = &head1[pos1];
+                let (d2, t2, _) = &head2[pos2];
+                match (d1, t1).cmp(&(d2, t2)) {
+                    Ordering::Less => {
+                        result.copy(&head1[pos1]);
+                        pos1 += 1;
+                    }
+                    Ordering::Greater => {
+                        result.copy(&head2[pos2]);
+                        pos2 += 1;
+                    }
+                    Ordering::Equal => {
+                        let (_, _, r1) = &head1[pos1];
+                        let (_, _, r2) = &head2[pos2];
+                        diff_stash.clone_from(r1);
+                        diff_stash.plus_equals(r2);
+                        if !diff_stash.is_zero() {
+                            let (d, t, _) = &head1[pos1];
+                            result.copy_destructured(d, t, &diff_stash);
+                        }
+                        pos1 += 1;
+                        pos2 += 1;
+                    }
+                }
+            }
+
+            // Advance exhausted heads.
+            if pos1 >= head1[..].len() {
+                self.recycle(std::mem::take(&mut head1), stash);
+                head1 = list1.next().unwrap_or_default();
+                pos1 = 0;
+            }
+            if pos2 >= head2[..].len() {
+                self.recycle(std::mem::take(&mut head2), stash);
+                head2 = list2.next().unwrap_or_default();
+                pos2 = 0;
+            }
+            if result.at_capacity() {
+                output.push(std::mem::take(&mut result));
+                result = self.empty(stash);
+            }
+        }
+
+        // Drain remaining from head1.
+        if pos1 < head1[..].len() {
+            for i in pos1..head1[..].len() {
+                result.copy(&head1[i]);
+            }
+        }
+        if !result.is_empty() {
+            output.push(std::mem::take(&mut result));
+            result = self.empty(stash);
+        }
+        output.extend(list1);
+
+        // Drain remaining from head2.
+        if pos2 < head2[..].len() {
+            for i in pos2..head2[..].len() {
+                result.copy(&head2[i]);
+            }
+        }
+        if !result.is_empty() {
+            output.push(std::mem::take(&mut result));
+        }
+        output.extend(list2);
+    }
+
+    fn extract(
+        &mut self,
+        merged: Vec<Self::Chunk>,
+        upper: AntichainRef<Self::Time>,
+        frontier: &mut Antichain<Self::Time>,
+        readied: &mut Vec<Self::Chunk>,
+        kept: &mut Vec<Self::Chunk>,
+        stash: &mut Vec<Self::Chunk>,
+    ) {
+        let mut keep = self.empty(stash);
+        let mut ship = self.empty(stash);
+
+        for buffer in merged {
+            for (data, time, diff) in buffer.iter() {
+                if upper.less_equal(time) {
+                    frontier.insert_with(time, |time| time.clone());
+                    keep.copy_destructured(data, time, diff);
+                    if keep.at_capacity() {
+                        kept.push(std::mem::take(&mut keep));
+                        keep = self.empty(stash);
+                    }
+                } else {
+                    ship.copy_destructured(data, time, diff);
+                    if ship.at_capacity() {
+                        readied.push(std::mem::take(&mut ship));
+                        ship = self.empty(stash);
+                    }
+                }
+            }
+            self.recycle(buffer, stash);
+        }
+
+        if !keep.is_empty() {
+            kept.push(keep);
+        }
+        if !ship.is_empty() {
+            readied.push(ship);
+        }
+    }
+
+    fn account(chunk: &Self::Chunk) -> (usize, usize, usize, usize) {
+        let (mut size, mut capacity, mut allocations) = (0, 0, 0);
+        let cb = |siz, cap| {
+            size += siz;
+            capacity += cap;
+            allocations += 1;
+        };
+        chunk.heap_size(cb);
+        (chunk.len(), size, capacity, allocations)
+    }
+}
+
+/// A `Merger` for `ColumnationStack` containers with per-item capacity checks.
+pub type ColInternalMerger<D, T, R> = ColumnationMerger<D, T, R>;