differential_dataflow/trace/implementations/

merge_batcher.rs

//! A `Batcher` implementation based on merge sort.
//!
//! The `MergeBatcher` requires support from two types, a "chunker" and a "merger".
//! The chunker receives input batches and consolidates them, producing sorted output
//! "chunks" that are fully consolidated (no adjacent updates can be accumulated).
//! The merger implements the [`Merger`] trait, and provides hooks for manipulating
//! sorted "chains" of chunks as needed by the merge batcher: merging chunks and also
//! splitting them apart based on time.
//!
//! Implementations of `MergeBatcher` can be instantiated through the choice of both
//! the chunker and the merger, provided their respective output and input types align.

use std::marker::PhantomData;

use timely::progress::frontier::AntichainRef;
use timely::progress::{frontier::Antichain, Timestamp};
use timely::container::{ContainerBuilder, PushInto};

use crate::logging::{BatcherEvent, Logger};
use crate::trace::{Batcher, Builder, Description};

/// Creates batches from containers of unordered tuples.
///
/// To implement `Batcher`, the container builder `C` must accept `&mut Input` as inputs,
/// and must produce outputs of type `M::Chunk`.
pub struct MergeBatcher<Input, C, M: Merger> {
    /// Transforms input streams to chunks of sorted, consolidated data.
    chunker: C,
    /// A sequence of power-of-two length lists of sorted, consolidated containers.
    ///
    /// Do not push/pop directly but use the corresponding functions ([`Self::chain_push`]/[`Self::chain_pop`]).
    chains: Vec<Vec<M::Chunk>>,
    /// Stash of empty chunks, recycled through the merging process.
    stash: Vec<M::Chunk>,
    /// Merges consolidated chunks, and extracts the subset of an update chain that lies in an interval of time.
    merger: M,
    /// Current lower frontier, we sealed up to here.
    lower: Antichain<M::Time>,
    /// The lower-bound frontier of the data, after the last call to seal.
    frontier: Antichain<M::Time>,
    /// Logger for size accounting.
    logger: Option<Logger>,
    /// Timely operator ID.
    operator_id: usize,
    /// The `Input` type needs to be called out as the type of container accepted, but it is not otherwise present.
    _marker: PhantomData<Input>,
}

impl<Input, C, M> Batcher for MergeBatcher<Input, C, M>
where
    C: ContainerBuilder<Container=M::Chunk> + for<'a> PushInto<&'a mut Input>,
    M: Merger<Time: Timestamp>,
{
    type Input = Input;
    type Time = M::Time;
    type Output = M::Chunk;

    fn new(logger: Option<Logger>, operator_id: usize) -> Self {
        Self {
            logger,
            operator_id,
            chunker: C::default(),
            merger: M::default(),
            chains: Vec::new(),
            stash: Vec::new(),
            frontier: Antichain::new(),
            lower: Antichain::from_elem(M::Time::minimum()),
            _marker: PhantomData,
        }
    }

    /// Push a container of data into this merge batcher. Updates the internal chain structure if
    /// needed.
    fn push_container(&mut self, container: &mut Input) {
        self.chunker.push_into(container);
        while let Some(chunk) = self.chunker.extract() {
            let chunk = std::mem::take(chunk);
            self.insert_chain(vec![chunk]);
        }
    }

    // Sealing a batch means finding those updates with times not greater or equal to any time
    // in `upper`. All updates must have time greater or equal to the previously used `upper`,
    // which we call `lower`, by assumption that after sealing a batcher we receive no more
    // updates with times not greater or equal to `upper`.
    fn seal<B: Builder<Input = Self::Output, Time = Self::Time>>(&mut self, upper: Antichain<M::Time>) -> B::Output {
        // Finish
        while let Some(chunk) = self.chunker.finish() {
            let chunk = std::mem::take(chunk);
            self.insert_chain(vec![chunk]);
        }

        // Merge all remaining chains into a single chain.
        while self.chains.len() > 1 {
            let list1 = self.chain_pop().unwrap();
            let list2 = self.chain_pop().unwrap();
            let merged = self.merge_by(list1, list2);
            self.chain_push(merged);
        }
        let merged = self.chain_pop().unwrap_or_default();

        // Extract readied data.
        let mut kept = Vec::new();
        let mut readied = Vec::new();
        self.frontier.clear();

        self.merger.extract(merged, upper.borrow(), &mut self.frontier, &mut readied, &mut kept, &mut self.stash);

        if !kept.is_empty() {
            self.chain_push(kept);
        }

        self.stash.clear();

        let description = Description::new(self.lower.clone(), upper.clone(), Antichain::from_elem(M::Time::minimum()));
        let seal = B::seal(&mut readied, description);
        self.lower = upper;
        seal
    }

    /// The frontier of elements remaining after the most recent call to `self.seal`.
    #[inline]
    fn frontier(&mut self) -> AntichainRef<'_, M::Time> {
        self.frontier.borrow()
    }
}

impl<Input, C, M: Merger> MergeBatcher<Input, C, M> {
    /// Insert a chain and maintain chain properties: Chains are geometrically sized and ordered
    /// by decreasing length.
    fn insert_chain(&mut self, chain: Vec<M::Chunk>) {
        if !chain.is_empty() {
            self.chain_push(chain);
            while self.chains.len() > 1 && (self.chains[self.chains.len() - 1].len() >= self.chains[self.chains.len() - 2].len() / 2) {
                let list1 = self.chain_pop().unwrap();
                let list2 = self.chain_pop().unwrap();
                let merged = self.merge_by(list1, list2);
                self.chain_push(merged);
            }
        }
    }

    // merges two sorted input lists into one sorted output list.
    fn merge_by(&mut self, list1: Vec<M::Chunk>, list2: Vec<M::Chunk>) -> Vec<M::Chunk> {
        // TODO: `list1` and `list2` get dropped; would be better to reuse?
        let mut output = Vec::with_capacity(list1.len() + list2.len());
        self.merger.merge(list1, list2, &mut output, &mut self.stash);

        output
    }

    /// Pop a chain and account size changes.
    #[inline]
    fn chain_pop(&mut self) -> Option<Vec<M::Chunk>> {
        let chain = self.chains.pop();
        self.account(chain.iter().flatten().map(M::account), -1);
        chain
    }

    /// Push a chain and account size changes.
    #[inline]
    fn chain_push(&mut self, chain: Vec<M::Chunk>) {
        self.account(chain.iter().map(M::account), 1);
        self.chains.push(chain);
    }

    /// Account size changes. Only performs work if a logger exists.
    ///
    /// Calculate the size based on the iterator passed along, with each attribute
    /// multiplied by `diff`. Usually, one wants to pass 1 or -1 as the diff.
    #[inline]
    fn account<I: IntoIterator<Item = (usize, usize, usize, usize)>>(&self, items: I, diff: isize) {
        if let Some(logger) = &self.logger {
            let (mut records, mut size, mut capacity, mut allocations) = (0isize, 0isize, 0isize, 0isize);
            for (records_, size_, capacity_, allocations_) in items {
                records = records.saturating_add_unsigned(records_);
                size = size.saturating_add_unsigned(size_);
                capacity = capacity.saturating_add_unsigned(capacity_);
                allocations = allocations.saturating_add_unsigned(allocations_);
            }
            logger.log(BatcherEvent {
                operator: self.operator_id,
                records_diff: records * diff,
                size_diff: size * diff,
                capacity_diff: capacity * diff,
                allocations_diff: allocations * diff,
            })
        }
    }
}

impl<Input, C, M: Merger> Drop for MergeBatcher<Input, C, M> {
    fn drop(&mut self) {
        // Cleanup chain to retract accounting information.
        while self.chain_pop().is_some() {}
    }
}

/// A trait to describe interesting moments in a merge batcher.
pub trait Merger: Default {
    /// The internal representation of chunks of data.
    type Chunk: Default;
    /// The type of time in frontiers to extract updates.
    type Time;
    /// Merge chains into an output chain.
    fn merge(&mut self, list1: Vec<Self::Chunk>, list2: Vec<Self::Chunk>, output: &mut Vec<Self::Chunk>, stash: &mut Vec<Self::Chunk>);
    /// Extract ready updates based on the `upper` frontier.
    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>,
    );

    /// Account size and allocation changes. Returns a tuple of (records, size, capacity, allocations).
    fn account(chunk: &Self::Chunk) -> (usize, usize, usize, usize);
}

pub use container::InternalMerger;

pub mod container {

    //! Merger implementations for the merge batcher.
    //!
    //! The `InternalMerge` trait allows containers to merge sorted, consolidated
    //! data using internal iteration. The `InternalMerger` type implements the
    //! `Merger` trait using `InternalMerge`, and is the standard merger for all
    //! container types.

    use std::marker::PhantomData;
    use timely::container::SizableContainer;
    use timely::progress::frontier::{Antichain, AntichainRef};
    use timely::{Accountable, PartialOrder};
    use crate::trace::implementations::merge_batcher::Merger;

    /// A container that supports the operations needed by the merge batcher:
    /// merging sorted chains and extracting updates by time.
    pub trait InternalMerge: Accountable + SizableContainer + Default {
        /// The owned time type, for maintaining antichains.
        type TimeOwned;

        /// The number of items in this container.
        fn len(&self) -> usize;

        /// Clear the container for reuse.
        fn clear(&mut self);

        /// Account the allocations behind the chunk.
        fn account(&self) -> (usize, usize, usize, usize) {
            let (size, capacity, allocations) = (0, 0, 0);
            (usize::try_from(self.record_count()).unwrap(), size, capacity, allocations)
        }

        /// Merge items from sorted inputs into `self`, advancing positions.
        /// Merges until `self` is at capacity or all inputs are exhausted.
        ///
        /// Dispatches based on the number of inputs:
        /// - **0**: no-op
        /// - **1**: bulk copy (may swap the input into `self`)
        /// - **2**: merge two sorted streams
        fn merge_from(
            &mut self,
            others: &mut [Self],
            positions: &mut [usize],
        );

        /// Extract updates from `self` into `ship` (times not beyond `upper`)
        /// and `keep` (times beyond `upper`), updating `frontier` with kept times.
        fn extract(
            &mut self,
            upper: AntichainRef<Self::TimeOwned>,
            frontier: &mut Antichain<Self::TimeOwned>,
            keep: &mut Self,
            ship: &mut Self,
        );
    }

    /// A `Merger` using internal iteration for `Vec` containers.
    pub type VecInternalMerger<D, T, R> = InternalMerger<Vec<(D, T, R)>>;
    /// A `Merger` using internal iteration for `TimelyStack` containers.
    pub type ColInternalMerger<D, T, R> = InternalMerger<crate::containers::TimelyStack<(D, T, R)>>;

    /// A merger that uses internal iteration via [`InternalMerge`].
    pub struct InternalMerger<MC> {
        _marker: PhantomData<MC>,
    }

    impl<MC> Default for InternalMerger<MC> {
        fn default() -> Self { Self { _marker: PhantomData } }
    }

    impl<MC> InternalMerger<MC> where MC: InternalMerge {
        #[inline]
        fn empty(&self, stash: &mut Vec<MC>) -> MC {
            stash.pop().unwrap_or_else(|| {
                let mut container = MC::default();
                container.ensure_capacity(&mut None);
                container
            })
        }
        #[inline]
        fn recycle(&self, mut chunk: MC, stash: &mut Vec<MC>) {
            chunk.clear();
            stash.push(chunk);
        }
        /// Drain remaining items from one side into `result`/`output`.
        fn drain_side(
            &self,
            head: &mut MC,
            pos: &mut usize,
            list: &mut std::vec::IntoIter<MC>,
            result: &mut MC,
            output: &mut Vec<MC>,
            stash: &mut Vec<MC>,
        ) {
            while *pos < head.len() {
                result.merge_from(
                    std::slice::from_mut(head),
                    std::slice::from_mut(pos),
                );
                if *pos >= head.len() {
                    let old = std::mem::replace(head, list.next().unwrap_or_default());
                    self.recycle(old, stash);
                    *pos = 0;
                }
                if result.at_capacity() {
                    output.push(std::mem::take(result));
                    *result = self.empty(stash);
                }
            }
        }
    }

    impl<MC> Merger for InternalMerger<MC>
    where
        MC: InternalMerge<TimeOwned: Ord + PartialOrder + Clone + 'static> + 'static,
    {
        type Time = MC::TimeOwned;
        type Chunk = MC;

        fn merge(&mut self, list1: Vec<MC>, list2: Vec<MC>, output: &mut Vec<MC>, stash: &mut Vec<MC>) {
            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 = self.empty(stash);

            // Main merge loop: both sides have data.
            while positions[0] < heads[0].len() && positions[1] < heads[1].len() {
                result.merge_from(&mut heads, &mut positions);

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

            // Drain remaining from side 0.
            self.drain_side(&mut heads[0], &mut positions[0], &mut list1, &mut result, output, stash);
            if !result.is_empty() {
                output.push(std::mem::take(&mut result));
                result = self.empty(stash);
            }
            output.extend(list1);

            // Drain remaining from side 1.
            self.drain_side(&mut heads[1], &mut positions[1], &mut list2, &mut result, output, stash);
            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>,
            ship: &mut Vec<Self::Chunk>,
            kept: &mut Vec<Self::Chunk>,
            stash: &mut Vec<Self::Chunk>,
        ) {
            let mut keep = self.empty(stash);
            let mut ready = self.empty(stash);

            for mut buffer in merged {
                buffer.extract(upper, frontier, &mut keep, &mut ready);
                self.recycle(buffer, stash);
                if keep.at_capacity() {
                    kept.push(std::mem::take(&mut keep));
                    keep = self.empty(stash);
                }
                if ready.at_capacity() {
                    ship.push(std::mem::take(&mut ready));
                    ready = self.empty(stash);
                }
            }
            if !keep.is_empty() {
                kept.push(keep);
            }
            if !ready.is_empty() {
                ship.push(ready);
            }
        }

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

    /// Implementation of `InternalMerge` for `Vec<(D, T, R)>`.
    pub mod vec_internal {
        use std::cmp::Ordering;
        use timely::PartialOrder;
        use timely::container::SizableContainer;
        use timely::progress::frontier::{Antichain, AntichainRef};
        use crate::difference::Semigroup;
        use super::InternalMerge;

        impl<D: Ord + Clone + 'static, T: Ord + Clone + PartialOrder + 'static, R: Semigroup + Clone + 'static> InternalMerge for Vec<(D, T, R)> {
            type TimeOwned = T;

            fn len(&self) -> usize { Vec::len(self) }
            fn clear(&mut self) { Vec::clear(self) }

            fn merge_from(
                &mut self,
                others: &mut [Self],
                positions: &mut [usize],
            ) {
                match others.len() {
                    0 => {},
                    1 => {
                        let other = &mut others[0];
                        let pos = &mut positions[0];
                        if self.is_empty() && *pos == 0 {
                            std::mem::swap(self, other);
                            return;
                        }
                        self.extend_from_slice(&other[*pos ..]);
                        *pos = other.len();
                    },
                    2 => {
                        let (left, right) = others.split_at_mut(1);
                        let other1 = &mut left[0];
                        let other2 = &mut right[0];

                        while positions[0] < other1.len() && positions[1] < other2.len() && !self.at_capacity() {
                            let (d1, t1, _) = &other1[positions[0]];
                            let (d2, t2, _) = &other2[positions[1]];
                            match (d1, t1).cmp(&(d2, t2)) {
                                Ordering::Less => {
                                    self.push(other1[positions[0]].clone());
                                    positions[0] += 1;
                                }
                                Ordering::Greater => {
                                    self.push(other2[positions[1]].clone());
                                    positions[1] += 1;
                                }
                                Ordering::Equal => {
                                    let (d, t, mut r1) = other1[positions[0]].clone();
                                    let (_, _, ref r2) = other2[positions[1]];
                                    r1.plus_equals(r2);
                                    if !r1.is_zero() {
                                        self.push((d, t, r1));
                                    }
                                    positions[0] += 1;
                                    positions[1] += 1;
                                }
                            }
                        }
                    },
                    n => unimplemented!("{n}-way merge not yet supported"),
                }
            }

            fn extract(
                &mut self,
                upper: AntichainRef<T>,
                frontier: &mut Antichain<T>,
                keep: &mut Self,
                ship: &mut Self,
            ) {
                for (data, time, diff) in self.drain(..) {
                    if upper.less_equal(&time) {
                        frontier.insert_with(&time, |time| time.clone());
                        keep.push((data, time, diff));
                    } else {
                        ship.push((data, time, diff));
                    }
                }
            }
        }
    }

    /// Implementation of `InternalMerge` for `TimelyStack<(D, T, R)>`.
    pub mod columnation_internal {
        use std::cmp::Ordering;
        use columnation::Columnation;
        use timely::PartialOrder;
        use timely::container::SizableContainer;
        use timely::progress::frontier::{Antichain, AntichainRef};
        use crate::containers::TimelyStack;
        use crate::difference::Semigroup;
        use super::InternalMerge;

        impl<D, T, R> InternalMerge for TimelyStack<(D, T, R)>
        where
            D: Ord + Columnation + Clone + 'static,
            T: Ord + Columnation + Clone + PartialOrder + 'static,
            R: Default + Semigroup + Columnation + Clone + 'static,
        {
            type TimeOwned = T;

            fn len(&self) -> usize { self[..].len() }
            fn clear(&mut self) { TimelyStack::clear(self) }

            fn account(&self) -> (usize, usize, usize, usize) {
                let (mut size, mut capacity, mut allocations) = (0, 0, 0);
                let cb = |siz, cap| {
                    size += siz;
                    capacity += cap;
                    allocations += 1;
                };
                self.heap_size(cb);
                (self.len(), size, capacity, allocations)
            }

            fn merge_from(
                &mut self,
                others: &mut [Self],
                positions: &mut [usize],
            ) {
                match others.len() {
                    0 => {},
                    1 => {
                        let other = &mut others[0];
                        let pos = &mut positions[0];
                        if self[..].is_empty() && *pos == 0 {
                            std::mem::swap(self, other);
                            return;
                        }
                        for i in *pos .. other[..].len() {
                            self.copy(&other[i]);
                        }
                        *pos = other[..].len();
                    },
                    2 => {
                        let (left, right) = others.split_at_mut(1);
                        let other1 = &left[0];
                        let other2 = &right[0];

                        let mut stash = R::default();

                        while positions[0] < other1[..].len() && positions[1] < other2[..].len() && !self.at_capacity() {
                            let (d1, t1, _) = &other1[positions[0]];
                            let (d2, t2, _) = &other2[positions[1]];
                            match (d1, t1).cmp(&(d2, t2)) {
                                Ordering::Less => {
                                    self.copy(&other1[positions[0]]);
                                    positions[0] += 1;
                                }
                                Ordering::Greater => {
                                    self.copy(&other2[positions[1]]);
                                    positions[1] += 1;
                                }
                                Ordering::Equal => {
                                    let (_, _, r1) = &other1[positions[0]];
                                    let (_, _, r2) = &other2[positions[1]];
                                    stash.clone_from(r1);
                                    stash.plus_equals(r2);
                                    if !stash.is_zero() {
                                        let (d, t, _) = &other1[positions[0]];
                                        self.copy_destructured(d, t, &stash);
                                    }
                                    positions[0] += 1;
                                    positions[1] += 1;
                                }
                            }
                        }
                    },
                    n => unimplemented!("{n}-way merge not yet supported"),
                }
            }

            fn extract(
                &mut self,
                upper: AntichainRef<T>,
                frontier: &mut Antichain<T>,
                keep: &mut Self,
                ship: &mut Self,
            ) {
                for (data, time, diff) in self.iter() {
                    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);
                    }
                }
            }
        }
    }
}