mz_timely_util/columnar/

batcher.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License in the LICENSE file at the
// root of this repository, or online at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Types for consolidating, merging, and extracting columnar update collections.

use std::collections::VecDeque;
use std::marker::PhantomData;

use crate::columnation::ColumnationStack;
use columnar::Container as _;
use columnar::Push as _;
use columnar::bytes::indexed;
use columnar::{BorrowedOf, Clear, Columnar, Index, Len};
use columnation::Columnation;
use differential_dataflow::difference::Semigroup;
use differential_dataflow::trace::implementations::merge_batcher::Merger;
use timely::Container;
use timely::PartialOrder;
use timely::container::{ContainerBuilder, PushInto};
use timely::progress::frontier::{Antichain, AntichainRef};

use crate::columnar::Column;
use crate::columnar::body::ColumnBody;

/// A chunker to transform input data into sorted columns.
#[derive(Default)]
pub struct Chunker<C> {
    /// Buffer into which we'll consolidate.
    ///
    /// Also the buffer where we'll stage responses to `extract` and `finish`.
    /// When these calls return, the buffer is available for reuse.
    target: C,
    /// Consolidated buffers ready to go.
    ready: VecDeque<C>,
}

impl<C: Container + Clone + 'static> ContainerBuilder for Chunker<C> {
    type Container = C;

    fn extract(&mut self) -> Option<&mut Self::Container> {
        if let Some(ready) = self.ready.pop_front() {
            self.target = ready;
            Some(&mut self.target)
        } else {
            None
        }
    }

    fn finish(&mut self) -> Option<&mut Self::Container> {
        self.extract()
    }
}

impl<'a, D, T, R> PushInto<&'a mut Column<(D, T, R)>> for Chunker<ColumnationStack<(D, T, R)>>
where
    D: Columnar + Columnation,
    for<'b> columnar::Ref<'b, D>: Ord + Copy,
    T: Columnar + Columnation,
    for<'b> columnar::Ref<'b, T>: Ord + Copy,
    R: Columnar + Columnation + Semigroup + for<'b> Semigroup<columnar::Ref<'b, R>>,
    for<'b> columnar::Ref<'b, R>: Ord,
{
    fn push_into(&mut self, container: &'a mut Column<(D, T, R)>) {
        // Sort input data
        // TODO: consider `Vec<usize>` that we retain, containing indexes.
        let borrowed = container.borrow();
        let mut permutation = Vec::with_capacity(borrowed.len());
        Extend::extend(&mut permutation, borrowed.into_index_iter());
        permutation.sort();

        self.target.clear();
        // Iterate over the data, accumulating diffs for like keys.
        let mut iter = permutation.drain(..);
        if let Some((data, time, diff)) = iter.next() {
            let mut owned_data = D::into_owned(data);
            let mut owned_time = T::into_owned(time);

            let mut prev_data = data;
            let mut prev_time = time;
            let mut prev_diff = <R as Columnar>::into_owned(diff);

            for (data, time, diff) in iter {
                if (&prev_data, &prev_time) == (&data, &time) {
                    prev_diff.plus_equals(&diff);
                } else {
                    if !prev_diff.is_zero() {
                        D::copy_from(&mut owned_data, prev_data);
                        T::copy_from(&mut owned_time, prev_time);
                        let tuple = (owned_data, owned_time, prev_diff);
                        self.target.push_into(&tuple);
                        (owned_data, owned_time, prev_diff) = tuple;
                    }
                    prev_data = data;
                    prev_time = time;
                    R::copy_from(&mut prev_diff, diff);
                }
            }

            if !prev_diff.is_zero() {
                D::copy_from(&mut owned_data, prev_data);
                T::copy_from(&mut owned_time, prev_time);
                let tuple = (owned_data, owned_time, prev_diff);
                self.target.push_into(&tuple);
            }
        }

        if !self.target.is_empty() {
            self.ready.push_back(std::mem::take(&mut self.target));
        }
    }
}

/// A chunker that consolidates `Column<(D, T, R)>` updates into sorted
/// [`ColumnBody`] chunks, without round-tripping through columnation.
///
/// Drop-in counterpart to [`Chunker`] for the merge-batcher path: same control
/// flow (sort borrowed refs, fold equal `(data, time)` runs, drop zero diffs).
/// This is where records leave the edge container: the input is whatever the
/// edge delivered, the output is a body the batcher chains.
pub struct ColumnChunker<U: Columnar> {
    /// Body we consolidate into and present to extract/finish callers.
    target: ColumnBody<U>,
    /// Sorted, consolidated chunks pending extraction.
    ready: VecDeque<ColumnBody<U>>,
}

// Manual impl rather than `#[derive(Default)]`: the derive would synthesize
// `impl<U: Columnar + Default>`, but `ColumnBody<U>: Default` only requires
// `U: Columnar`, and adding a spurious `U: Default` bound would propagate
// through every `ContainerBuilder for ColumnChunker<U>` impl.
impl<U: Columnar> Default for ColumnChunker<U> {
    fn default() -> Self {
        Self {
            target: ColumnBody::default(),
            ready: VecDeque::new(),
        }
    }
}

impl<U: Columnar> ContainerBuilder for ColumnChunker<U>
where
    U::Container: Clone + 'static,
{
    type Container = ColumnBody<U>;

    fn extract(&mut self) -> Option<&mut Self::Container> {
        if let Some(ready) = self.ready.pop_front() {
            self.target = ready;
            Some(&mut self.target)
        } else {
            None
        }
    }

    fn finish(&mut self) -> Option<&mut Self::Container> {
        self.extract()
    }
}

impl<'a, D, T, R> PushInto<&'a mut Column<(D, T, R)>> for ColumnChunker<(D, T, R)>
where
    D: Columnar,
    for<'b> columnar::Ref<'b, D>: Copy + Ord,
    T: Columnar,
    for<'b> columnar::Ref<'b, T>: Copy + Ord,
    R: Columnar + Default + Semigroup + for<'b> Semigroup<columnar::Ref<'b, R>>,
    for<'b> columnar::Ref<'b, R>: Ord,
    for<'b> <D as Columnar>::Container: columnar::Push<columnar::Ref<'b, D>>,
    for<'b> <T as Columnar>::Container: columnar::Push<columnar::Ref<'b, T>>,
    for<'b> <R as Columnar>::Container: columnar::Push<&'b R>,
{
    fn push_into(&mut self, container: &'a mut Column<(D, T, R)>) {
        // Clearing keeps the allocations of a typed target, which is the
        // steady state: possibly a chunk just handed back via `extract`.
        self.target.clear();

        // Sort input by columnar ref order.
        let borrowed = container.borrow();
        let mut permutation = Vec::with_capacity(borrowed.len());
        Extend::extend(&mut permutation, borrowed.into_index_iter());
        permutation.sort();

        // Sweep sorted refs, accumulating diffs over equal `(data, time)`
        // pairs and pushing non-zero results to the target's leaves. Refs
        // from the input borrow are valid through the sweep, so D and T
        // are pushed directly via each leaf's `Push<Ref<_>>` impl. Only R
        // needs an owned scratch since it carries the consolidated sum.
        {
            let (target_d, target_t, target_r) = self.target.typed_mut();

            let mut iter = permutation.drain(..);
            if let Some((data, time, diff)) = iter.next() {
                let mut prev_data = data;
                let mut prev_time = time;
                let mut prev_diff = <R as Columnar>::into_owned(diff);

                for (data, time, diff) in iter {
                    if (&prev_data, &prev_time) == (&data, &time) {
                        prev_diff.plus_equals(&diff);
                    } else {
                        if !prev_diff.is_zero() {
                            target_d.push(prev_data);
                            target_t.push(prev_time);
                            target_r.push(&prev_diff);
                        }
                        prev_data = data;
                        prev_time = time;
                        R::copy_from(&mut prev_diff, diff);
                    }
                }

                if !prev_diff.is_zero() {
                    target_d.push(prev_data);
                    target_t.push(prev_time);
                    target_r.push(&prev_diff);
                }
            }
        }

        if !self.target.is_empty() {
            self.ready.push_back(std::mem::take(&mut self.target));
        }
    }
}

/// Advance `*lower` past every position in `[*lower, upper)` where `cmp`
/// returns true.
///
/// On return, `*lower` is the first index `>= initial *lower` where `cmp`
/// returns false, or `upper` if `cmp` holds through the end.
///
/// Takes the predicate as `FnMut(usize) -> bool` rather than a value-bearing
/// closure so callers can index whichever subset of the input columns they
/// actually need to compare — for the merger's `(d, t)`-keyed sort, this lets
/// each probe touch only the D and T leaf views, skipping the diff column.
///
/// Compared to a linear scan, this is `O(log K)` for a run of length `K`
/// satisfying `cmp` — useful when one side of a sorted merge has long runs
/// dominated by the other side.
pub(crate) fn gallop(upper: usize, lower: &mut usize, mut cmp: impl FnMut(usize) -> bool) {
    if *lower < upper && cmp(*lower) {
        let mut step = 1;
        while *lower + step < upper && cmp(*lower + step) {
            *lower += step;
            step <<= 1;
        }

        step >>= 1;
        while step > 0 {
            if *lower + step < upper && cmp(*lower + step) {
                *lower += step;
            }
            step >>= 1;
        }

        // `*lower` is the last index where `cmp` holds; step to the first
        // where it does not.
        *lower += 1;
    }
}

/// Counterpart to `ColInternalMerger` (which merges `ColumnationStack` chunks).
/// Drives the merge batcher with [`ColumnBody`] chunks, no columnation
/// detour, by way of the inherent `merge_from` / `extract` methods on
/// `ColumnBody<(D, T, R)>` below.
pub struct ColumnMerger<D, T, R> {
    _marker: PhantomData<(D, T, R)>,
}

impl<D, T, R> Default for ColumnMerger<D, T, R> {
    fn default() -> Self {
        Self {
            _marker: PhantomData,
        }
    }
}

/// A chunk the merge and extract bodies below write to: the merger's
/// [`ColumnBody`], and the column pager's [`Column`] until that pager is
/// retired.
pub trait MergeChunk<C: Columnar>: Default {
    /// The typed containers, for writing to.
    fn typed(&mut self) -> &mut C::Container;
    /// The chunk as a columnar view.
    fn view(&self) -> BorrowedOf<'_, C>;
    /// True when the chunk holds no records.
    fn is_empty(&self) -> bool;
}

impl<C: Columnar> MergeChunk<C> for ColumnBody<C> {
    fn typed(&mut self) -> &mut C::Container {
        self.typed_mut()
    }
    fn view(&self) -> BorrowedOf<'_, C> {
        self.borrow()
    }
    fn is_empty(&self) -> bool {
        self.is_empty()
    }
}

impl<C: Columnar> MergeChunk<C> for Column<C> {
    fn typed(&mut self) -> &mut C::Container {
        // The pager merges into recycled typed chunks, so this is a move. A
        // serialized chunk is copied, as `ColumnBody::typed_mut` does.
        if !matches!(self, Column::Typed(_)) {
            let view = self.borrow();
            let mut fresh = C::Container::default();
            fresh.extend_from_self(view, 0..view.len());
            *self = Column::Typed(fresh);
        }
        let Column::Typed(typed) = self else {
            unreachable!("a serialized chunk was materialized above");
        };
        typed
    }
    fn view(&self) -> BorrowedOf<'_, C> {
        self.borrow()
    }
    fn is_empty(&self) -> bool {
        self.is_empty()
    }
}

/// Per-chunk merge and extract for sorted [`ColumnBody`] chunks.
///
/// These are the building blocks that [`Merger for ColumnMerger`] orchestrates
/// over chains of chunks. They're inherent methods rather than a trait impl
/// so the merger can call them without going through any wrapper indirection.
/// The output side of each is written to, so a serialized output body is
/// materialized first (see [`ColumnBody::typed_mut`]).
impl<D, T, R> ColumnBody<(D, T, R)>
where
    D: Columnar,
    for<'a> columnar::Ref<'a, D>: Copy + Ord,
    T: Columnar + Default + Clone + PartialOrder,
    for<'a> columnar::Ref<'a, T>: Copy + Ord,
    R: Columnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
{
    /// Merge items from sorted inputs into `self`, advancing positions. See
    /// [`merge_from`].
    #[must_use]
    pub fn merge_from(&mut self, others: &mut [Self], positions: &mut [usize]) -> bool {
        merge_from(self, others, positions)
    }

    /// Partition records starting at `*position` into `keep` and `ship`. See
    /// [`extract`].
    pub fn extract(
        &mut self,
        position: &mut usize,
        upper: AntichainRef<T>,
        frontier: &mut Antichain<T>,
        keep: &mut Self,
        ship: &mut Self,
    ) {
        extract(self, position, upper, frontier, keep, ship)
    }
}

/// The same merge and extract for the column pager's [`Column`] chunks.
impl<D, T, R> Column<(D, T, R)>
where
    D: Columnar,
    for<'a> columnar::Ref<'a, D>: Copy + Ord,
    T: Columnar + Default + Clone + PartialOrder,
    for<'a> columnar::Ref<'a, T>: Copy + Ord,
    R: Columnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
{
    /// Merge items from sorted inputs into `self`, advancing positions. See
    /// [`merge_from`].
    #[must_use]
    pub fn merge_from(&mut self, others: &mut [Self], positions: &mut [usize]) -> bool {
        merge_from(self, others, positions)
    }

    /// Partition records starting at `*position` into `keep` and `ship`. See
    /// [`extract`].
    pub fn extract(
        &mut self,
        position: &mut usize,
        upper: AntichainRef<T>,
        frontier: &mut Antichain<T>,
        keep: &mut Self,
        ship: &mut Self,
    ) {
        extract(self, position, upper, frontier, keep, ship)
    }
}

/// Merge items from sorted inputs into `target`, advancing positions.
///
/// Mirrors the dispatch shape used by the merge-batcher framework:
/// - **0**: no-op
/// - **1**: bulk copy (or swap, if `target` is empty and `*pos == 0`)
/// - **2**: merge two sorted streams, with diff consolidation on equal
///   `(data, time)` keys and gallop bulk-copy of long single-side runs.
///
/// Returns `true` if the merge stopped because the amortized ship-threshold
/// check inside the inner loop fired (the caller should ship `target` before
/// the next call). Returns `false` if the merge stopped because at least
/// one input was exhausted at its position (the caller should refill that
/// side; `target` may still be at capacity from accumulation across short
/// calls and the caller should also check `at_capacity` in that case).
///
/// The 0- and 1-input dispatches always return `false`: 0 does no work,
/// 1 is a bulk copy or swap that runs to completion.
#[must_use]
pub fn merge_from<D, T, R, Ch>(target: &mut Ch, others: &mut [Ch], positions: &mut [usize]) -> bool
where
    D: Columnar,
    for<'a> columnar::Ref<'a, D>: Copy + Ord,
    T: Columnar + Default + Clone + PartialOrder,
    for<'a> columnar::Ref<'a, T>: Copy + Ord,
    R: Columnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
    Ch: MergeChunk<(D, T, R)>,
{
    match others.len() {
        0 => false,
        1 => {
            let other = &mut others[0];
            let pos = &mut positions[0];
            if target.is_empty() && *pos == 0 {
                std::mem::swap(target, other);
                return false;
            }
            let self_c = target.typed();
            let src_c = other.view();
            self_c.extend_from_self(src_c, *pos..other.view().len());
            *pos = other.view().len();
            false
        }
        2 => {
            let (left, right) = others.split_at(1);
            let (left_pos, right_pos) = positions.split_at_mut(1);
            let left_borrow = left[0].view();
            let right_borrow = right[0].view();

            let self_c = target.typed();

            // Split the input borrows into per-leaf views.
            //
            // A columnar tuple `Borrow::Ref` is recursive: indexing the
            // tuple borrow walks every leaf (and reconstructs the nested
            // ref tuple) regardless of which leaves the caller actually
            // reads. Indexing each leaf view directly cuts probe-path
            // work to the columns we consult — for the merge step
            // that's the `(D, T)` key, with the diff column read only
            // when we push.
            let l_d = left_borrow.0;
            let l_t = left_borrow.1;
            let l_r = left_borrow.2;
            let r_d = right_borrow.0;
            let r_t = right_borrow.1;
            let r_r = right_borrow.2;
            let upper_l = l_d.len();
            let upper_r = r_d.len();

            // Mirror the split on the output container. Tuple
            // containers split into per-leaf containers, which lets us
            // address each leaf independently — both gallop bulk-copies
            // and single-record pushes resolve to a primitive operation
            // per leaf. The leaves stay length-synchronized as long as
            // every record path pushes exactly one element to each.
            let (sd, st, sr) = self_c;

            // Pre-size each output leaf for the worst-case merge (no
            // consolidation). `reserve_for` walks each input's
            // `as_bytes`, which is accurate for variable-length leaves,
            // where reserving a record count wouldn't size the byte
            // buffer correctly. The reservation covers both whole
            // inputs, however little of them this call consumes, so
            // reserve only when the inputs fit within two ship-sized
            // chunks, which bounds it to 4 MiB.
            let input_bytes = 8
                * (indexed::length_in_words(&left_borrow)
                    + indexed::length_in_words(&right_borrow));
            if input_bytes <= 2 * crate::columnar::SHIP_WORDS * 8 {
                let inputs = [left_borrow, right_borrow];
                sd.reserve_for(inputs.iter().map(|b| b.0));
                st.reserve_for(inputs.iter().map(|b| b.1));
                sr.reserve_for(inputs.iter().map(|b| b.2));
            }

            let mut stash = R::default();

            // The size check walks every leaf slice, which is not cheap on
            // variable-length leaves, so it runs only after spending the
            // ship threshold's headroom at the inputs' average width.
            // That bounds uniform rows to just under the ship size. For
            // mixed widths it is a cadence hint: a stretch of wide rows
            // can overshoot. `ColumnChunk::settle` enforces the bound for
            // chunks that reach the pool. The column pager's merge-batcher
            // output has no size classes and ships as cut.
            let average_bytes = input_bytes / (upper_l + upper_r).max(1);
            let check_records =
                (crate::columnar::SHIP_WORDS * 8 / 10 / average_bytes.max(1)).clamp(1, 1024);
            let mut next_check = sd.len() + check_records;
            let mut yielded = false;

            while left_pos[0] < upper_l && right_pos[0] < upper_r {
                let d1 = l_d.get(left_pos[0]);
                let t1 = l_t.get(left_pos[0]);
                let d2 = r_d.get(right_pos[0]);
                let t2 = r_t.get(right_pos[0]);
                match (d1, t1).cmp(&(d2, t2)) {
                    std::cmp::Ordering::Less => {
                        // Common case (interleaved data): single-record
                        // advance. Skip the gallop call entirely — its
                        // setup plus the first cmp probe is more
                        // expensive than just pushing this record and
                        // re-entering the outer loop. Galloping is only
                        // worthwhile when there's an actual run, which
                        // we detect with the peek check below.
                        sd.push(d1);
                        st.push(t1);
                        sr.push(l_r.get(left_pos[0]));
                        left_pos[0] += 1;
                        // Long-run case: peek at the next record; if
                        // it's still strictly less than `(d2, t2)`,
                        // we have a run worth galloping (and bulk-
                        // copying).
                        if left_pos[0] < upper_l
                            && (l_d.get(left_pos[0]), l_t.get(left_pos[0])) < (d2, t2)
                        {
                            let start = left_pos[0];
                            gallop(
                                upper_l.min(left_pos[0] + next_check.saturating_sub(sd.len())),
                                &mut left_pos[0],
                                |i| (l_d.get(i), l_t.get(i)) < (d2, t2),
                            );
                            // Per-leaf bulk copy of the run: each call
                            // resolves to an `extend_from_slice` on its
                            // leaf (recursively for nested leaves).
                            sd.extend_from_self(l_d, start..left_pos[0]);
                            st.extend_from_self(l_t, start..left_pos[0]);
                            sr.extend_from_self(l_r, start..left_pos[0]);
                        }
                    }
                    std::cmp::Ordering::Greater => {
                        // Symmetric on the right side.
                        sd.push(d2);
                        st.push(t2);
                        sr.push(r_r.get(right_pos[0]));
                        right_pos[0] += 1;
                        if right_pos[0] < upper_r
                            && (r_d.get(right_pos[0]), r_t.get(right_pos[0])) < (d1, t1)
                        {
                            let start = right_pos[0];
                            gallop(
                                upper_r.min(right_pos[0] + next_check.saturating_sub(sd.len())),
                                &mut right_pos[0],
                                |i| (r_d.get(i), r_t.get(i)) < (d1, t1),
                            );
                            sd.extend_from_self(r_d, start..right_pos[0]);
                            st.extend_from_self(r_t, start..right_pos[0]);
                            sr.extend_from_self(r_r, start..right_pos[0]);
                        }
                    }
                    std::cmp::Ordering::Equal => {
                        let r1 = l_r.get(left_pos[0]);
                        let r2 = r_r.get(right_pos[0]);
                        R::copy_from(&mut stash, r1);
                        stash.plus_equals(&r2);
                        if !stash.is_zero() {
                            sd.push(d1);
                            st.push(t1);
                            sr.push(&stash);
                        }
                        left_pos[0] += 1;
                        right_pos[0] += 1;
                    }
                }

                if sd.len() >= next_check {
                    use columnar::Borrow as _;
                    if crate::columnar::at_serialized_capacity(&(
                        sd.borrow(),
                        st.borrow(),
                        sr.borrow(),
                    )) {
                        yielded = true;
                        break;
                    }
                    next_check = sd.len() + check_records;
                }
            }
            yielded
        }
        // `Merger::merge` only ever calls `merge_from` with 0/1/2-input
        // slices (k-way merge isn't part of the merge-batcher contract).
        n => unreachable!("merge_from called with {n} inputs; expected 0, 1, or 2"),
    }
}

/// Partition records of `source` starting at `*position` into `keep` (times
/// beyond `upper`, retained for the next round) and `ship` (times not beyond
/// `upper`, sealed into the output batch). Updates `frontier` with the
/// times of kept records.
///
/// The caller invokes `extract` repeatedly until `*position >= source.len()`,
/// swapping out a full output buffer between calls. This shape exists
/// because the framework only checks `at_capacity()` between calls, so
/// without an inner-loop yield a single call could quietly produce
/// oversized output chunks.
pub fn extract<D, T, R, Ch>(
    source: &Ch,
    position: &mut usize,
    upper: AntichainRef<T>,
    frontier: &mut Antichain<T>,
    keep: &mut Ch,
    ship: &mut Ch,
) where
    D: Columnar,
    for<'a> columnar::Ref<'a, D>: Copy + Ord,
    T: Columnar + Default + Clone + PartialOrder,
    for<'a> columnar::Ref<'a, T>: Copy + Ord,
    R: Columnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
    Ch: MergeChunk<(D, T, R)>,
{
    let keep_c = keep.typed();
    let ship_c = ship.typed();

    let self_view = source.view();
    let len = self_view.len();

    use columnar::Borrow as _;
    let mut owned_t = T::default();
    while *position < len
        && !crate::columnar::at_serialized_capacity(&keep_c.borrow())
        && !crate::columnar::at_serialized_capacity(&ship_c.borrow())
    {
        let (_, time, _) = self_view.get(*position);
        T::copy_from(&mut owned_t, time);
        if upper.less_equal(&owned_t) {
            // `insert_with` only clones when the time isn't already
            // present in the antichain.
            frontier.insert_with(&owned_t, |t| t.clone());
            keep_c.extend_from_self(self_view, *position..*position + 1);
        } else {
            ship_c.extend_from_self(self_view, *position..*position + 1);
        }
        *position += 1;
    }
}

/// `Merger` impl driving [`MergeBatcher`] over [`ColumnBody`] chunks,
/// built on the inherent `ColumnBody::merge_from` and `ColumnBody::extract`
/// methods.
/// Exhausted input chunks are recycled through `stash`, and remaining full
/// chunks on a drained side move to the output directly, with no per-element
/// copy.
///
/// [`MergeBatcher`]: differential_dataflow::trace::implementations::merge_batcher::MergeBatcher
impl<D, T, R> Merger for ColumnMerger<D, T, R>
where
    D: Columnar,
    for<'a> columnar::Ref<'a, D>: Copy + Ord,
    T: Columnar + Default + Clone + Ord + PartialOrder,
    for<'a> columnar::Ref<'a, T>: Copy + Ord,
    R: Columnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
{
    type Time = T;
    type Chunk = ColumnBody<(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>,
    ) {
        let mut list1 = list1.into_iter();
        let mut list2 = list2.into_iter();

        let mut heads = [
            list1.next().unwrap_or_default(),
            list2.next().unwrap_or_default(),
        ];
        let mut positions = [0usize, 0usize];

        let mut result = empty_chunk(stash);

        // Main merge loop: both sides have data.
        loop {
            let upper_l = heads[0].borrow().len();
            let upper_r = heads[1].borrow().len();
            if positions[0] >= upper_l || positions[1] >= upper_r {
                break;
            }

            // Whole-chunk passthrough fast path: when one head's tail (from
            // its current position) sorts entirely before the other head's
            // current record, the head moves to `output` wholesale, with two
            // probe records in place of per-record compares and per-leaf
            // byte copies. Restricted to `positions[i] == 0` so the head can
            // be handed off intact; a partial tail is what gallop already
            // handles inside the merge loop.
            let lhs_passthrough = positions[0] == 0 && upper_l > 0 && {
                let lhs = heads[0].borrow();
                let rhs = heads[1].borrow();
                let last_l = (lhs.0.get(upper_l - 1), lhs.1.get(upper_l - 1));
                let cur_r = (rhs.0.get(positions[1]), rhs.1.get(positions[1]));
                last_l < cur_r
            };
            if lhs_passthrough {
                if !result.is_empty() {
                    output.push(std::mem::take(&mut result));
                    result = empty_chunk(stash);
                }
                let head = std::mem::replace(&mut heads[0], list1.next().unwrap_or_default());
                output.push(head);
                positions[0] = 0;
                continue;
            }

            let rhs_passthrough = positions[1] == 0 && upper_r > 0 && {
                let lhs = heads[0].borrow();
                let rhs = heads[1].borrow();
                let last_r = (rhs.0.get(upper_r - 1), rhs.1.get(upper_r - 1));
                let cur_l = (lhs.0.get(positions[0]), lhs.1.get(positions[0]));
                last_r < cur_l
            };
            if rhs_passthrough {
                if !result.is_empty() {
                    output.push(std::mem::take(&mut result));
                    result = empty_chunk(stash);
                }
                let head = std::mem::replace(&mut heads[1], list2.next().unwrap_or_default());
                output.push(head);
                positions[1] = 0;
                continue;
            }

            // Per-record merge. `merge_from` returns `true` when its inner
            // amortized ship-threshold check fires — short-circuit the
            // outer `at_capacity` walk in that case.
            let yielded = result.merge_from(&mut heads, &mut positions);

            if positions[0] >= heads[0].borrow().len() {
                let old = std::mem::replace(&mut heads[0], list1.next().unwrap_or_default());
                recycle_chunk(old, stash);
                positions[0] = 0;
            }
            if positions[1] >= heads[1].borrow().len() {
                let old = std::mem::replace(&mut heads[1], list2.next().unwrap_or_default());
                recycle_chunk(old, stash);
                positions[1] = 0;
            }
            if yielded || result.at_capacity() {
                output.push(std::mem::take(&mut result));
                result = empty_chunk(stash);
            }
        }

        // Drain remaining from each side: copy partial head, then append
        // full chunks directly to output (no per-element copy).
        drain_side(
            &mut heads[0],
            &mut positions[0],
            &mut list1,
            &mut result,
            output,
            stash,
        );
        drain_side(
            &mut heads[1],
            &mut positions[1],
            &mut list2,
            &mut result,
            output,
            stash,
        );
        if !result.is_empty() {
            output.push(result);
        }
    }

    fn extract(
        &mut self,
        merged: Vec<Self::Chunk>,
        upper: AntichainRef<Self::Time>,
        frontier: &mut Antichain<Self::Time>,
        ship: &mut Vec<Self::Chunk>,
        kept: &mut Vec<Self::Chunk>,
        stash: &mut Vec<Self::Chunk>,
    ) {
        let mut keep = empty_chunk(stash);
        let mut ready = empty_chunk(stash);

        for mut buffer in merged {
            let mut position = 0;
            let len = buffer.borrow().len();
            while position < len {
                buffer.extract(&mut position, upper, frontier, &mut keep, &mut ready);
                if keep.at_capacity() {
                    kept.push(std::mem::take(&mut keep));
                    keep = empty_chunk(stash);
                }
                if ready.at_capacity() {
                    ship.push(std::mem::take(&mut ready));
                    ready = empty_chunk(stash);
                }
            }
            recycle_chunk(buffer, stash);
        }
        if !keep.is_empty() {
            kept.push(keep);
        }
        if !ready.is_empty() {
            ship.push(ready);
        }
    }

    fn len(chunk: &Self::Chunk) -> usize {
        chunk.len()
    }

    fn allocation(chunk: &Self::Chunk) -> (usize, usize, usize) {
        // Serialized footprint stands in for both `size` and `capacity`: the
        // chunk owns one logical allocation worth of leaf storage, and we
        // ship/recycle the whole thing rather than tracking per-leaf
        // capacities. Treating `size == capacity` matches how the framework
        // accounts already-shipped chunks (no slack to absorb).
        let bytes = chunk.length_in_bytes();
        (bytes, bytes, 1)
    }
}

/// Pop a chunk from `stash` or allocate a fresh one. Stashed chunks are
/// already cleared via `recycle_chunk`, so they're ready for push.
#[inline]
pub(crate) fn empty_chunk<C: Columnar>(stash: &mut Vec<ColumnBody<C>>) -> ColumnBody<C> {
    stash.pop().unwrap_or_default()
}

/// Reset `chunk` to an empty `Typed` and push it to `stash` for reuse.
///
/// Chunks recycled here come from the merger and chunker, both of which
/// produce `Typed`; only the typed allocations are worth caching for reuse.
/// A serialized chunk has no typed-side allocation to preserve, so we simply
/// drop it: `empty_chunk` will produce a fresh default just as cheaply, and
/// pushing it onto `stash` would only displace useful recycled allocations.
#[inline]
pub(crate) fn recycle_chunk<C: Columnar>(mut chunk: ColumnBody<C>, stash: &mut Vec<ColumnBody<C>>) {
    if let ColumnBody::Typed(c) = &mut chunk {
        c.clear();
        stash.push(chunk);
    }
}

/// Drain remaining items from one side into `result` / `output`.
///
/// Copies the partially-consumed head into `result` via `merge_from`'s 1-input
/// path, then appends remaining full chunks directly to `output` without
/// per-element copy.
fn drain_side<D, T, R>(
    head: &mut ColumnBody<(D, T, R)>,
    pos: &mut usize,
    list: &mut std::vec::IntoIter<ColumnBody<(D, T, R)>>,
    result: &mut ColumnBody<(D, T, R)>,
    output: &mut Vec<ColumnBody<(D, T, R)>>,
    stash: &mut Vec<ColumnBody<(D, T, R)>>,
) where
    D: Columnar,
    for<'a> columnar::Ref<'a, D>: Copy + Ord,
    T: Columnar + Default + Clone + PartialOrder,
    for<'a> columnar::Ref<'a, T>: Copy + Ord,
    R: Columnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
{
    if *pos < head.borrow().len() {
        // 1-input dispatch — bulk copy that runs to completion; the yield
        // signal is unused.
        let _ = result.merge_from(std::slice::from_mut(head), std::slice::from_mut(pos));
    }
    if !result.is_empty() {
        output.push(std::mem::take(result));
        *result = empty_chunk(stash);
    }
    Extend::extend(output, list);
}

#[cfg(test)]
mod tests {
    use timely::dataflow::channels::ContainerBytes;

    use super::*;

    /// Re-encode `column` as the `Align` variant, to drive the chunker over serialized
    /// input.
    ///
    /// `into_bytes` writes whole `u64` words, read back in native byte order.
    fn serialize<C: Columnar>(column: &Column<C>) -> Column<C> {
        let mut bytes: Vec<u8> = Vec::new();
        column.into_bytes(&mut bytes);
        assert_eq!(bytes.len() % 8, 0);
        let words = bytes
            .chunks_exact(8)
            .map(|w| u64::from_ne_bytes(w.try_into().expect("chunk is 8 bytes")))
            .collect();
        Column::Align(words)
    }

    /// Drive a single `push_into` call with `inputs` and collect the
    /// consolidated output (if any) as owned tuples.
    fn run_chunker<D, T, R>(inputs: &[(D, T, R)]) -> Vec<(D, T, R)>
    where
        D: Columnar + Clone,
        for<'a> columnar::Ref<'a, D>: Copy + Ord,
        T: Columnar + Clone,
        for<'a> columnar::Ref<'a, T>: Copy + Ord,
        R: Columnar + Clone + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
        for<'a> columnar::Ref<'a, R>: Ord,
        <(D, T, R) as Columnar>::Container: Clone,
        for<'a> <(D, T, R) as Columnar>::Container: columnar::Push<&'a (D, T, R)>,
        <(D, T, R) as Columnar>::Container: columnar::Push<(D, T, R)>,
        for<'a> <D as Columnar>::Container: columnar::Push<columnar::Ref<'a, D>>,
        for<'a> <T as Columnar>::Container: columnar::Push<columnar::Ref<'a, T>>,
        for<'a> <R as Columnar>::Container: columnar::Push<&'a R>,
    {
        let mut input: Column<(D, T, R)> = Default::default();
        for tuple in inputs.iter().cloned() {
            input.push_into(tuple);
        }

        let mut chunker: ColumnChunker<(D, T, R)> = Default::default();
        chunker.push_into(&mut input);

        let mut out = Vec::new();
        while let Some(chunk) = chunker.extract() {
            for (d, t, r) in chunk.borrow().into_index_iter() {
                out.push((D::into_owned(d), T::into_owned(t), R::into_owned(r)));
            }
        }
        out
    }

    #[mz_ore::test]
    fn empty_input_yields_no_chunk() {
        let mut chunker: ColumnChunker<(u64, u64, i64)> = Default::default();
        let mut input: Column<(u64, u64, i64)> = Default::default();
        chunker.push_into(&mut input);
        assert!(chunker.extract().is_none());
        assert!(chunker.finish().is_none());
    }

    #[mz_ore::test]
    fn unsorted_input_is_sorted() {
        let out = run_chunker(&[(3u64, 0u64, 1i64), (1u64, 0u64, 1i64), (2u64, 0u64, 1i64)]);
        assert_eq!(out, vec![(1, 0, 1), (2, 0, 1), (3, 0, 1)]);
    }

    #[mz_ore::test]
    fn serialized_input_is_consolidated() {
        // Callers hand the chunker whatever the upstream edge delivered, which is
        // serialized once the data crossed an exchange.
        let mut input: Column<(u64, u64, i64)> = Default::default();
        for tuple in [
            (2u64, 0u64, 1i64),
            (1, 0, 1),
            (2, 0, 2),
            (3, 0, 1),
            (3, 0, -1),
        ] {
            input.push_into(tuple);
        }
        let mut input = serialize(&input);

        let mut chunker: ColumnChunker<(u64, u64, i64)> = Default::default();
        chunker.push_into(&mut input);

        let mut out = Vec::new();
        while let Some(chunk) = chunker.extract() {
            for (d, t, r) in chunk.borrow().into_index_iter() {
                out.push((*d, *t, *r));
            }
        }
        assert_eq!(out, vec![(1, 0, 1), (2, 0, 3)]);
    }

    #[mz_ore::test]
    fn duplicate_keys_consolidate() {
        let out = run_chunker(&[(1u64, 0u64, 1i64), (1u64, 0u64, 2i64), (1u64, 0u64, -1i64)]);
        assert_eq!(out, vec![(1, 0, 2)]);
    }

    #[mz_ore::test]
    fn diffs_summing_to_zero_are_dropped() {
        let out = run_chunker(&[(1u64, 0u64, 1i64), (1u64, 0u64, -1i64)]);
        assert!(out.is_empty());
    }

    #[mz_ore::test]
    fn mixed_consolidation() {
        // (1, 0): 1 + 2 + (-3) = 0  -> dropped
        // (2, 0): 1            = 1  -> kept
        // (1, 1): 5            = 5  -> kept (different time from the (1, 0) group)
        let out = run_chunker(&[
            (1u64, 0u64, 1i64),
            (2u64, 0u64, 1i64),
            (1u64, 0u64, 2i64),
            (1u64, 1u64, 5i64),
            (1u64, 0u64, -3i64),
        ]);
        assert_eq!(out, vec![(1, 1, 5), (2, 0, 1)]);
    }

    #[mz_ore::test]
    fn key_val_tuple_data() {
        // Exercise the actual val-batcher shape: `D = (K, V)`.
        let out = run_chunker(&[
            ((1u64, 10u64), 0u64, 1i64),
            ((1u64, 10u64), 0u64, 1i64),
            ((1u64, 11u64), 0u64, 1i64),
            ((2u64, 10u64), 0u64, 1i64),
        ]);
        assert_eq!(
            out,
            vec![((1, 10), 0, 2), ((1, 11), 0, 1), ((2, 10), 0, 1),]
        );
    }

    #[mz_ore::test]
    fn buffer_reuse_across_calls() {
        // Two sequential push_into calls; second runs after extract returned
        // the first chunk, exercising the in-place clear path.
        let mut input1: Column<(u64, u64, i64)> = Default::default();
        input1.push_into((1u64, 0u64, 1i64));
        input1.push_into((2u64, 0u64, 1i64));

        let mut input2: Column<(u64, u64, i64)> = Default::default();
        input2.push_into((3u64, 0u64, 1i64));
        input2.push_into((1u64, 0u64, 1i64));

        let mut chunker: ColumnChunker<(u64, u64, i64)> = Default::default();
        chunker.push_into(&mut input1);

        // Hand back the first chunk via extract, simulating the merge batcher
        // taking ownership of the &mut and then returning.
        {
            let _ = chunker.extract().expect("first chunk");
        }

        chunker.push_into(&mut input2);

        let chunk = chunker.extract().expect("second chunk");
        let collected: Vec<_> = chunk
            .borrow()
            .into_index_iter()
            .map(|(d, t, r)| (u64::into_owned(d), u64::into_owned(t), i64::into_owned(r)))
            .collect();
        assert_eq!(collected, vec![(1, 0, 1), (3, 0, 1)]);
    }

    /// Build a `ColumnBody<((u64, u64), u64, i64)>` from a slice of tuples.
    fn col(rows: &[((u64, u64), u64, i64)]) -> ColumnBody<((u64, u64), u64, i64)> {
        let mut c: ColumnBody<((u64, u64), u64, i64)> = Default::default();
        for &t in rows {
            c.push_into(t);
        }
        c
    }

    fn collect_chunks(
        chunks: &[ColumnBody<((u64, u64), u64, i64)>],
    ) -> Vec<((u64, u64), u64, i64)> {
        chunks
            .iter()
            .flat_map(|c| {
                c.borrow().into_index_iter().map(|((k, v), t, r)| {
                    (
                        (u64::into_owned(k), u64::into_owned(v)),
                        u64::into_owned(t),
                        i64::into_owned(r),
                    )
                })
            })
            .collect()
    }

    /// Disjoint-range chains exercise the whole-chunk passthrough fast path:
    /// every chunk in chain1 is sortable-before every chunk in chain2, so
    /// each outer-loop iteration should hand a chunk straight to `output`
    /// without recursing through the per-record merge.
    #[mz_ore::test]
    fn merger_disjoint_chains_passthrough() {
        let chain1 = vec![
            col(&[((0, 0), 0, 1), ((1, 0), 0, 1)]),
            col(&[((2, 0), 0, 1), ((3, 0), 0, 1)]),
        ];
        let chain2 = vec![
            col(&[((10, 0), 0, 1), ((11, 0), 0, 1)]),
            col(&[((12, 0), 0, 1), ((13, 0), 0, 1)]),
        ];

        let mut merger: ColumnMerger<(u64, u64), u64, i64> = Default::default();
        let mut output = Vec::new();
        let mut stash = Vec::new();
        Merger::merge(&mut merger, chain1, chain2, &mut output, &mut stash);

        let collected = collect_chunks(&output);
        let expected: Vec<_> = (0..4u64)
            .map(|d| ((d, 0u64), 0u64, 1i64))
            .chain((10..14u64).map(|d| ((d, 0u64), 0u64, 1i64)))
            .collect();
        assert_eq!(collected, expected);
    }

    /// Interleaved chains never satisfy the passthrough condition; each
    /// outer iteration falls through to `merge_from`. Same correctness
    /// expectation, exercises the non-passthrough path under
    /// `Merger::merge`.
    #[mz_ore::test]
    fn merger_interleaved_chains() {
        // Even keys on one chain, odd on the other; chunks alternate so the
        // per-record path is the only viable route.
        let chain1 = vec![
            col(&[((0, 0), 0, 1), ((2, 0), 0, 1)]),
            col(&[((4, 0), 0, 1), ((6, 0), 0, 1)]),
        ];
        let chain2 = vec![
            col(&[((1, 0), 0, 1), ((3, 0), 0, 1)]),
            col(&[((5, 0), 0, 1), ((7, 0), 0, 1)]),
        ];

        let mut merger: ColumnMerger<(u64, u64), u64, i64> = Default::default();
        let mut output = Vec::new();
        let mut stash = Vec::new();
        Merger::merge(&mut merger, chain1, chain2, &mut output, &mut stash);

        let collected = collect_chunks(&output);
        let expected: Vec<_> = (0..8u64).map(|d| ((d, 0u64), 0u64, 1i64)).collect();
        assert_eq!(collected, expected);
    }

    /// Passthrough must consolidate adjacent equal keys at chunk
    /// boundaries — i.e., must NOT fire when `chain1`'s last record's
    /// `(d, t)` equals `chain2`'s first.
    #[mz_ore::test]
    fn merger_passthrough_respects_equal_boundary() {
        // chain1's last == chain2's first key: equal-key consolidation
        // must kick in (sum of diffs would be 2). If passthrough fired
        // erroneously, both records would land in different output chunks
        // unconsolidated.
        let chain1 = vec![col(&[((0, 0), 0, 1), ((5, 0), 0, 1)])];
        let chain2 = vec![col(&[((5, 0), 0, 1), ((10, 0), 0, 1)])];

        let mut merger: ColumnMerger<(u64, u64), u64, i64> = Default::default();
        let mut output = Vec::new();
        let mut stash = Vec::new();
        Merger::merge(&mut merger, chain1, chain2, &mut output, &mut stash);

        let collected = collect_chunks(&output);
        assert_eq!(
            collected,
            vec![((0, 0), 0, 1), ((5, 0), 0, 2), ((10, 0), 0, 1)]
        );
    }
}

#[cfg(test)]
mod proptests {
    //! Property tests for `ColumnBody::merge_from` and `ColumnBody::extract`.
    //!
    //! Strategy: generate sorted+consolidated inputs (the merger's input
    //! contract), drive `merge_from` / `extract` the same way the framework
    //! would, and compare against a brute-force reference impl.
    //!
    //! Test types are `D = (u64, u64)`, `T = u64`, `R = i64` drawn from small
    //! ranges so that equal-key collisions are common and the consolidation
    //! path actually runs.
    use super::*;
    use mz_ore::cast::CastFrom;
    use proptest::prelude::*;
    use timely::progress::frontier::Antichain;

    type Tuple = ((u64, u64), u64, i64);

    /// Reference consolidation: sort by `(data, time)`, sum diffs over equal
    /// pairs, drop zeros.
    fn consolidate(mut v: Vec<Tuple>) -> Vec<Tuple> {
        v.sort();
        let mut out: Vec<Tuple> = Vec::new();
        for (d, t, r) in v {
            if let Some(last) = out.last_mut() {
                if last.0 == d && last.1 == t {
                    last.2 += r;
                    continue;
                }
            }
            out.push((d, t, r));
        }
        out.retain(|x| x.2 != 0);
        out
    }

    /// Strategy for sorted+consolidated input lists. Ranges are small to
    /// encourage equal-key collisions.
    fn arb_consolidated() -> impl Strategy<Value = Vec<Tuple>> {
        prop::collection::vec(((0u64..5, 0u64..5), 0u64..3, -3i64..=3i64), 0..30)
            .prop_map(consolidate)
    }

    fn build_column(v: &[Tuple]) -> ColumnBody<Tuple> {
        let mut col: ColumnBody<Tuple> = Default::default();
        for tup in v {
            col.push_into(*tup);
        }
        col
    }

    /// The body in `ColumnBody::Words` form when `words` is set, as a store hands it back.
    fn in_form(col: ColumnBody<Tuple>, words: bool) -> ColumnBody<Tuple> {
        if !words {
            return col;
        }
        let mut bytes = Vec::new();
        col.write_into(&mut bytes).expect("vec writes");
        ColumnBody::Words(bytemuck::allocation::pod_collect_to_vec(&bytes))
    }

    fn collect_column(col: &ColumnBody<Tuple>) -> Vec<Tuple> {
        col.borrow()
            .into_index_iter()
            .map(|((k, v), t, r)| {
                (
                    (u64::into_owned(k), u64::into_owned(v)),
                    u64::into_owned(t),
                    i64::into_owned(r),
                )
            })
            .collect()
    }

    /// Drive a 2-way merge the same way `Merger::merge` would: a 2-input
    /// call until one side exhausts, then a 1-input drain for whichever
    /// side still has data.
    fn drive_merge(left: ColumnBody<Tuple>, right: ColumnBody<Tuple>) -> ColumnBody<Tuple> {
        let mut self_col: ColumnBody<Tuple> = Default::default();
        let mut others = [left, right];
        let mut positions = [0usize, 0];
        let _ = self_col.merge_from(&mut others, &mut positions);

        let [left_done, right_done] = others;
        let [left_pos, right_pos] = positions;

        if left_pos < left_done.borrow().len() {
            let mut tail = [left_done];
            let mut p = [left_pos];
            let _ = self_col.merge_from(&mut tail, &mut p);
        } else if right_pos < right_done.borrow().len() {
            let mut tail = [right_done];
            let mut p = [right_pos];
            let _ = self_col.merge_from(&mut tail, &mut p);
        }

        self_col
    }

    proptest! {
        /// `merge_from` with two sorted+consolidated inputs equals the
        /// reference consolidate(union).
        #[mz_ore::test]
        #[cfg_attr(miri, ignore)]
        fn merge_from_equals_consolidated_union(
            a in arb_consolidated(),
            b in arb_consolidated(),
            a_words in any::<bool>(),
            b_words in any::<bool>(),
        ) {
            let merged = drive_merge(
                in_form(build_column(&a), a_words),
                in_form(build_column(&b), b_words),
            );

            let mut union = a.clone();
            Extend::extend(&mut union, b.iter().copied());
            let expected = consolidate(union);

            prop_assert_eq!(collect_column(&merged), expected);
        }

        /// `merge_from` 1-input bulk-copy from a non-zero position equals
        /// `other[*pos..]`.
        #[mz_ore::test]
        #[cfg_attr(miri, ignore)]
        fn merge_from_one_input_drains_tail(
            data in arb_consolidated(),
            pos_frac in 0u32..=100,
            self_words in any::<bool>(),
            other_words in any::<bool>(),
        ) {
            // Cap at len so we always have a valid position.
            let len = data.len();
            let start_pos = if len == 0 { 0 } else {
                (usize::cast_from(pos_frac) * len) / 101
            };

            // Self starts non-empty so we exercise the bulk-copy path, not the
            // empty-self swap shortcut.
            let mut self_col: ColumnBody<Tuple> = Default::default();
            let sentinel: Tuple = ((u64::MAX, u64::MAX), 0, 1);
            self_col.push_into(sentinel);
            // A serialized target is materialized before the copy appends to it.
            let mut self_col = in_form(self_col, self_words);

            let mut others = [in_form(build_column(&data), other_words)];
            let mut positions = [start_pos];
            let _ = self_col.merge_from(&mut others, &mut positions);

            let mut expected = vec![sentinel];
            Extend::extend(&mut expected, data[start_pos..].iter().copied());

            prop_assert_eq!(collect_column(&self_col), expected);
            prop_assert_eq!(positions[0], len);
        }

        /// `merge_from` 1-input swap shortcut: empty self + pos=0 should
        /// produce a column equal to the input.
        #[mz_ore::test]
        #[cfg_attr(miri, ignore)]
        fn merge_from_empty_self_swap(data in arb_consolidated(), words in any::<bool>()) {
            let mut self_col: ColumnBody<Tuple> = Default::default();
            let mut others = [in_form(build_column(&data), words)];
            let mut positions = [0usize];
            let _ = self_col.merge_from(&mut others, &mut positions);

            prop_assert_eq!(collect_column(&self_col), data);
        }

        /// `extract` partitions correctly:
        ///   - keep ∪ ship multiset-equals self
        ///   - upper.less_equal(t) for every kept time
        ///   - !upper.less_equal(t) for every shipped time
        ///   - frontier covers every kept time
        #[mz_ore::test]
        #[cfg_attr(miri, ignore)]
        fn extract_partitions_by_frontier(
            data in arb_consolidated(),
            upper_time in 0u64..=4,
            words in any::<bool>(),
        ) {
            let mut self_col = in_form(build_column(&data), words);
            let upper = Antichain::from_elem(upper_time);
            let mut frontier: Antichain<u64> = Antichain::new();
            let mut keep: ColumnBody<Tuple> = Default::default();
            let mut ship: ColumnBody<Tuple> = Default::default();
            let mut position = 0;

            self_col.extract(
                &mut position,
                upper.borrow(),
                &mut frontier,
                &mut keep,
                &mut ship,
            );

            // Single call drains the input (we removed the at_capacity yield).
            prop_assert_eq!(position, data.len());

            let kept = collect_column(&keep);
            let shipped = collect_column(&ship);

            // Partition predicate: kept times >= upper, shipped times < upper.
            for (_, t, _) in &kept {
                prop_assert!(
                    upper.borrow().less_equal(t),
                    "kept time {} should satisfy upper.less_equal", t,
                );
            }
            for (_, t, _) in &shipped {
                prop_assert!(
                    !upper.borrow().less_equal(t),
                    "shipped time {} should NOT satisfy upper.less_equal", t,
                );
            }

            // Union (multiset) equals input.
            let mut union = kept.clone();
            Extend::extend(&mut union, shipped.iter().copied());
            union.sort();
            let mut expected_sorted = data.clone();
            expected_sorted.sort();
            prop_assert_eq!(union, expected_sorted);

            // Frontier dominates every kept time.
            for (_, t, _) in &kept {
                prop_assert!(
                    frontier.less_equal(t),
                    "frontier should dominate kept time {}", t,
                );
            }
        }

        /// Empty input → no work, frontier untouched, position = 0.
        #[mz_ore::test]
        #[cfg_attr(miri, ignore)]
        fn extract_empty_input(upper_time in 0u64..=4) {
            let mut self_col: ColumnBody<Tuple> = Default::default();
            let upper = Antichain::from_elem(upper_time);
            let mut frontier: Antichain<u64> = Antichain::new();
            let mut keep: ColumnBody<Tuple> = Default::default();
            let mut ship: ColumnBody<Tuple> = Default::default();
            let mut position = 0;

            self_col.extract(
                &mut position,
                upper.borrow(),
                &mut frontier,
                &mut keep,
                &mut ship,
            );

            prop_assert_eq!(position, 0);
            prop_assert!(collect_column(&keep).is_empty());
            prop_assert!(collect_column(&ship).is_empty());
            prop_assert!(frontier.elements().is_empty());
        }
    }
}