differential_dataflow/trace/chunk/

mod.rs

//! Sorted, consolidated runs of updates, and operators over sequences of them.
//!
//! A [`Chunk`] is a consolidated, sorted run of `(data, time, diff)` updates.
//! A sequence of chunks is also expected to be consolidated and sorted.
//!
//! The [`Chunk`] trait exposes whole-chunk operations, so that the implementor can internally divert to their best implementations, with amortized overhead.
//! Each operation is invoked as if "streaming", providing input and output queues.
//! An implementor is expected to drain as much as possible of the inputs, and any chunk written to the output is "committed" and likely to be shipped onward.
//!
//! # Wiring a `Chunk` into an arrangement
//!
//! Implementing [`Chunk`] for a type `C` is the only bespoke code needed; three aliases then expand into a full trace:
//!
//! * [`ChunkBatcher<Chu, C>`](ChunkBatcher) — the merge batcher.
//! * [`ChunkBuilder<C>`](ChunkBuilder) — the batch builder.
//! * [`ChunkSpine<C>`](ChunkSpine) — the trace, a spine of `Rc`-shared batches.
//!
//! These are the `Batcher` / `Builder` / `Spine` to hand to [`arrange_core`](crate::operators::arrange::arrangement::arrange_core), along with a chunker that
//! forms `C` from the input stream — typically [`ContainerChunker<C>`](crate::batcher::merge::chunker::ContainerChunker).
//! These aliases use [`ChunkBatch`], [`ChunkMerger`], [`ChunkBatchMerger`], and [`ChunkBatchBuilder`] for batch formation and trace maintenance.
//! The [`vec`](mod@vec) module provides a worked implementation and specializes the aliases for its chunk type.
//! The `chunks` example builds an arrangement from it.
//!
//! Reading chunks through cursors is an optional capability provided by the [`cursor`] module.
//!
//! # Bounded footprint
//!
//! There is a `TARGET` associated constant that signals the intended chunk size.
//! The constant should be chosen large enough to amortize overheads, but small enough that per-chunk work does not "stall" the system when invoked.
//! The implementor is trusted to make a reasonable choice here.
//!
//! The [`Chunk::settle`] method "settles" sequences of chunks, and is called as chunks are no longer expected to be needed in the near future.
//! The implementor should ensure the chunks are "graded", in that the sequence of chunks are all at most `TARGET` in size, any two in order sum to strictly more than `TARGET`.
//! This is also an opportunity to compress data, or spill to disk or cloud storage.
//!
//! Producers settle their committed output as they go, keeping the active, unsettled chunk set small.
//! Implementors must keep [`len`](Chunk::len) cheap even when a chunk's body is paged out.

use std::collections::VecDeque;

use timely::progress::Antichain;
use timely::progress::frontier::AntichainRef;
use crate::lattice::Lattice;
use crate::trace::implementations::spine_fueled::SpineBatch;

pub mod vec;
pub mod cursor;

pub use cursor::{ChunkBatchCursor, NavigableChunk};

/// A non-empty, bounded, consolidated, sorted sequence of `(data, time, diff)`.
///
/// An implementor gains access to types and trait implementations that provide
/// batch formation and trace maintenance with no additional effort.
///
/// The necessary implementations are either "data" or "metadata" operations.
/// The "data" operations transform lists of chunks, are expected to do roughly
/// "one chunk's worth" of work at a time; they can afford to compress and page.
/// The "metadata" operations provide chunk information, and should be lightweight.
///
/// The trait exposes only time, which trace maintenance needs.
/// Reading a chunk's contents is a separate, optional capability: see the [`cursor`] module.
pub trait Chunk: Sized + Clone {
    /// The timestamp type of the chunk's updates.
    ///
    /// Trace maintenance uses time to describe intervals and to advance and compact updates.
    type Time: Lattice + timely::progress::Timestamp;

    /// The intended maximum chunk size.
    const TARGET: usize;

    /// The number of updates in the chunk.
    fn len(&self) -> usize;

    /// Merge the fronts of two input deques through their shared horizon.
    ///
    /// Both deques are non-empty (the caller guarantees it). The two queues are both
    /// the heads of lists of chunks, and the implementor should only merge through the
    /// least last `(key, val, time)` update, or risk emitting an unconsolidated
    /// output chunk.
    ///
    /// When a chunk cannot be completely retired, perhaps it had the larger last update,
    /// it should be rewritten as a new chunk and pushed back to the front of the queue.
    /// The invocation is expected to consume at least one of its inputs, and the harness
    /// may continually re-invoke if this doesn't happen.
    ///
    /// A merge concludes when the harness sees that either input is now empty, at which
    /// point it appends the queue to the output without the method's assistance.
    fn merge(in1: &mut VecDeque<Self>, in2: &mut VecDeque<Self>, out: &mut VecDeque<Self>);

    /// Partition `input` updates into `keep` (greater or equal `frontier`) or not (`ship`).
    ///
    /// An implementation should yield with some frequency to allow the output to "settle".
    /// The harness may guard against this, but it prefers to provide as much context as it
    /// can in order to allow broader chunk fusion where needed.
    fn extract(
        input: &mut VecDeque<Self>,
        frontier: AntichainRef<Self::Time>,
        residual: &mut Antichain<Self::Time>,
        keep: &mut VecDeque<Self>,
        ship: &mut VecDeque<Self>,
    );

    /// Advance times by `frontier` producing consolidated chunks.
    ///
    /// An output for `(key, val)` should generally not be produced until a later pair
    /// is observed, or `done` is set, to ensure the output chunks are consolidated.
    /// Incomplete work can be pushed back to the front of `input`.
    ///
    /// On `done` a single `(key, val)` group may span the whole input; advancing and
    /// consolidating it should cost time linear in its size, not quadratic.
    fn advance(
        input: &mut VecDeque<Self>,
        frontier: AntichainRef<Self::Time>,
        done: bool,
        out: &mut VecDeque<Self>,
    );

    /// Reshape `input` to a sequence that maintains the "grading" structural invariant.
    ///
    /// Specifically, the chunks in `output` should have a maximum size of `TARGET` and
    /// each adjacent pair should have lengths that sum to more than `TARGET`.
    /// This is also a good moment to consider compression or paging out the contents.
    /// When `done` is set the input must be moved to the output.
    ///
    /// This method may be called on already settled data, and should be efficient then.
    ///
    /// Implementors that want the standard maximal packing can delegate to the
    /// [`pack`] helper, supplying their layout's coalesce / split / commit closures.
    fn settle(input: &mut VecDeque<Self>, done: bool, out: &mut VecDeque<Self>);

}

/// Maximal-packing driver an implementor's [`Chunk::settle`] may delegate to.
///
/// Holds a `carry` run of chunks, grown until it reaches `TARGET` (then emitted as
/// one chunk) and emitted early when the next chunk can't be absorbed without
/// exceeding `TARGET`; over-sized chunks are peeled with `split`. Each committed
/// chunk is passed through `seal` (the compress / spill hook — use the identity
/// closure when there's nothing to do). The closures are the only layout-specific
/// pieces:
///
/// * `combine(run)` — one chunk from a run of at least two (caller guarantees their
///   lengths sum to at most `TARGET`, and that they form one sorted, consolidated
///   chain). A run is combined once, when emitted, so packing stays linear even
///   when combining copies its inputs.
/// * `split(chunk, n)` — the first `n` updates and the remaining `len - n`.
/// * `seal(chunk)` — commit a chunk (e.g. compress or spill); identity to keep it.
pub fn pack<C: Chunk>(
    input: &mut VecDeque<C>,
    done: bool,
    out: &mut VecDeque<C>,
    mut combine: impl FnMut(Vec<C>) -> C,
    mut split: impl FnMut(C, usize) -> (C, C),
    mut seal: impl FnMut(C) -> C,
) {
    let (mut carry, mut carried) = (Vec::new(), 0);
    let mut fuse = |carry: &mut Vec<C>| if carry.len() == 1 { carry.pop().unwrap() } else { combine(std::mem::take(carry)) };
    while let Some(chunk) = input.pop_front() {
        if carried + chunk.len() <= C::TARGET {
            // Combines into one legal chunk; absorb it.
            carried += chunk.len();
            carry.push(chunk);
            if carried == C::TARGET { out.push_back(seal(fuse(&mut carry))); carried = 0; }
        } else if !carry.is_empty() {
            // The carry is maximal against this neighbour; emit it and absorb afresh.
            out.push_back(seal(fuse(&mut carry)));
            carried = 0;
            input.push_front(chunk);
        } else {
            // Peel `TARGET`-sized pieces off an over-sized chunk and absorb the remainder.
            let mut rest = chunk;
            while rest.len() > C::TARGET {
                let (head, tail) = split(rest, C::TARGET);
                out.push_back(seal(head));
                rest = tail;
            }
            input.push_front(rest);
        }
    }
    if done {
        if !carry.is_empty() { out.push_back(seal(fuse(&mut carry))); }
    } else {
        // Hand the run back uncombined, to be combined once when it is emitted.
        for chunk in carry.into_iter().rev() { input.push_front(chunk); }
    }
}

/// A batch: an ordered [`Chunk`] sequence whose concatenation is its updates.
pub struct ChunkBatch<C: Chunk> {
    /// Ordered, consolidated chunks; their concatenation is the batch.
    pub chunks: Vec<C>,
}

impl<C: Chunk> ChunkBatch<C> {
    /// Assemble a batch from ordered chunks.
    pub fn new(chunks: Vec<C>) -> Self {
        for chunk in &chunks {
            assert!(chunk.len() > 0, "ChunkBatch chunks must be non-empty");
        }
        ChunkBatch { chunks }
    }
}

impl<C: Chunk + 'static> SpineBatch for ChunkBatch<C>
where
    C::Time: timely::progress::Timestamp + Lattice + Ord,
{
    type Time = C::Time;
    type Merger = ChunkBatchMerger<C>;
    fn len(&self) -> usize { self.chunks.iter().map(C::len).sum() }
}

/// A merge-batcher [`Merger`](crate::batcher::merge::Merger)
/// over chains of [`Chunk`]s.
///
/// `merge` runs the whole-chain binary merger; `extract` splits by the seal frontier
/// using [`Chunk::extract`]. The batcher consolidates equal `(data, time)` updates
/// but does *not* advance times — time advancement is advance's job, handled later in
/// the trace. Both settle their output, since the batcher's chains want to be graded.
pub type ChunkBatcher<Chu, C> = crate::batcher::merge::MergeBatcher<Chu, ChunkMerger<C>, ChunkBuilder<C>>;

/// A spine of `Rc`-shared [`ChunkBatch`]s of type `C`: the trace type for `arrange`.
pub type ChunkSpine<C> = crate::trace::implementations::spine_fueled::Spine<std::rc::Rc<ChunkBatch<C>>>;

/// A [`ChunkBatch`] builder over chunks of type `C`, emitting `Rc`-shared updates.
pub type ChunkBuilder<C> = ChunkBatchBuilder<C>;

/// A merge-batcher [`Merger`](crate::batcher::merge::Merger)
/// over chains of [`Chunk`]s.
///
/// `merge` runs the whole-chain binary merger; `extract` splits by the seal frontier
/// using [`Chunk::extract`]. The batcher consolidates equal `(data, time)` updates
/// but does *not* advance times — time advancement is advance's job, handled later in
/// the trace. Both settle their output, since the batcher's chains want to be graded.
pub struct ChunkMerger<C> {
    _marker: std::marker::PhantomData<C>,
}

impl<C> Default for ChunkMerger<C> {
    fn default() -> Self { Self { _marker: std::marker::PhantomData } }
}

impl<C> crate::batcher::merge::Merger for ChunkMerger<C>
where
    C: Chunk + Default + 'static,
    C::Time: Clone + timely::PartialOrder + 'static,
{
    type Chunk = C;

    fn merge(
        &mut self,
        list1: Vec<C>,
        list2: Vec<C>,
        output: &mut Vec<C>,
        _stash: &mut Vec<C>,
    ) {
        // Settle the output after each merge, to maintain bounded active chunks.
        let mut in1: VecDeque<C> = list1.into();
        let mut in2: VecDeque<C> = list2.into();
        let (mut staged, mut settled) = (VecDeque::new(), VecDeque::new());
        while !in1.is_empty() && !in2.is_empty() {
            C::merge(&mut in1, &mut in2, &mut staged);
            C::settle(&mut staged, false, &mut settled);
        }
        // Append the non-empty tail from either input, settle as we go.
        for tail in in1.drain(..).chain(in2.drain(..)) {
            staged.push_back(tail);
            C::settle(&mut staged, false, &mut settled);
        }
        C::settle(&mut staged, true, &mut settled);
        output.extend(settled);
    }

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

impl<C> crate::batcher::merge::Extractor for ChunkMerger<C>
where
    C: Chunk + Default + 'static,
    C::Time: Clone + timely::PartialOrder + 'static,
{
    type Time = C::Time;

    fn extract(
        &mut self,
        merged: Vec<C>,
        upper: AntichainRef<C::Time>,
        frontier: &mut Antichain<C::Time>,
        ship: &mut Vec<C>,
        kept: &mut Vec<C>,
        _stash: &mut Vec<C>,
    ) {
        // `extract` keeps updates greater-or-equal `upper` and ships the rest, folding
        // the lower envelope of kept times into `frontier`. Drive it a bounded amount
        // per call (≈ one input chunk) and `settle` each side as it accumulates, so
        // neither `keep` (retained across yields) nor `ship` (handed to the builder)
        // builds up unsettled in core. `settle` may withhold a sub-`TARGET` carry
        // between calls; the final `settle(done)` flushes it.
        let mut input: VecDeque<C> = merged.into();
        let (mut keep, mut shipped) = (VecDeque::new(), VecDeque::new());
        let (mut kept_q, mut shipped_q) = (VecDeque::new(), VecDeque::new());
        while !input.is_empty() {
            C::extract(&mut input, upper, frontier, &mut keep, &mut shipped);
            C::settle(&mut keep, false, &mut kept_q);
            C::settle(&mut shipped, false, &mut shipped_q);
        }
        C::settle(&mut keep, true, &mut kept_q);
        C::settle(&mut shipped, true, &mut shipped_q);
        kept.extend(kept_q);
        ship.extend(shipped_q);
    }
}

/// The resumable [`SpineBatch::Merger`] for [`ChunkBatch`]: merges two updates and advances
/// their times to the compaction frontier, a fuel-bounded step at a time.
///
/// Each step pipelines [`merge`](Chunk::merge) → [`advance`](Chunk::advance) →
/// [`settle`](Chunk::settle) and settles its output, so a suspended merge holds only
/// graded chunks. The sources are read by cloning (a cheap refcount bump) and must be
/// supplied unchanged on every call.
pub struct ChunkBatchMerger<C: Chunk> {
    /// Compaction frontier supplied at construction.
    frontier: Antichain<C::Time>,
    /// Input deques, refilled from the sources (clones) head-of-list at a time.
    in1: VecDeque<C>,
    in2: VecDeque<C>,
    /// Next source chunk to clone into `in1` / `in2`.
    idx1: usize,
    idx2: usize,
    /// `advance`'s input: the merge output plus advance's withheld carry at the front.
    merged: VecDeque<C>,
    /// `advance`'s output and `settle`'s input: merged-and-advanced chunks, with
    /// settle's withheld sub-`TARGET` carry at the front.
    advanced: VecDeque<C>,
    /// `settle`'s output: the committed, graded result, grown by `work`. Graded at
    /// every yield, so a suspended merge holds well-formed (spillable) chunk state.
    settled: VecDeque<C>,
    /// Set once both sources are drained and advance's and settle's final flushes ran.
    complete: bool,
}

impl<C> crate::trace::implementations::spine_fueled::Merger<ChunkBatch<C>> for ChunkBatchMerger<C>
where
    C: Chunk + 'static,
    C::Time: timely::progress::Timestamp + Lattice + Ord + 'static,
{
    fn new(_source1: &ChunkBatch<C>, _source2: &ChunkBatch<C>, frontier: AntichainRef<C::Time>) -> Self {
        Self {
            frontier: frontier.to_owned(),
            in1: VecDeque::new(),
            in2: VecDeque::new(),
            idx1: 0,
            idx2: 0,
            merged: VecDeque::new(),
            advanced: VecDeque::new(),
            settled: VecDeque::new(),
            complete: false,
        }
    }

    fn work(&mut self, source1: &ChunkBatch<C>, source2: &ChunkBatch<C>, fuel: &mut isize) {

        // TODO: The logic is a bit tortured here, and should be improved.

        if self.complete { return; }

        while *fuel > 0 {
            // Refill each input deque up to a burst of source chunks (clones).
            // The constant trades away fuel precision for overhead amortization.
            const BURST: usize = 8;
            while self.in1.len() < BURST && self.idx1 < source1.chunks.len() {
                self.in1.push_back(source1.chunks[self.idx1].clone());
                self.idx1 += 1;
            }
            while self.in2.len() < BURST && self.idx2 < source2.chunks.len() {
                self.in2.push_back(source2.chunks[self.idx2].clone());
                self.idx2 += 1;
            }

            // Merge's per-tick output (a burst's worth, or one tail chunk), measured
            // for fuel before it joins the carry already in `merged`.
            let mut produced = VecDeque::new();
            if !self.in1.is_empty() && !self.in2.is_empty() {
                // Both sides have data: drain the loaded burst.
                C::merge(&mut self.in1, &mut self.in2, &mut produced);
            } else if let Some(chunk) = self.in1.pop_front().or_else(|| self.in2.pop_front()) {
                // Exactly one side has data: flush its verbatim tail, one chunk a step.
                produced.push_back(chunk);
            } else {
                // Both sources drained: final flush of advance's and settle's carries.
                C::advance(&mut self.merged, self.frontier.borrow(), true, &mut self.advanced);
                C::settle(&mut self.advanced, true, &mut self.settled);
                self.complete = true;
                break;
            }

            let work: usize = produced.iter().map(C::len).sum();
            self.merged.extend(produced);
            C::advance(&mut self.merged, self.frontier.borrow(), false, &mut self.advanced);
            // Maintain grading at the yield boundary: this step may exhaust `fuel` and
            // suspend with `advanced` held, and held chunk state must be graded.
            C::settle(&mut self.advanced, false, &mut self.settled);
            *fuel -= work as isize;
        }
    }

    fn done(self) -> Option<ChunkBatch<C>> {
        debug_assert!(self.merged.is_empty() && self.advanced.is_empty());
        wrap(self.settled.into())
    }
}

/// A [`Builder`](crate::trace::Builder) that collects a chunk sequence into a [`ChunkBatch`].
pub struct ChunkBatchBuilder<C: Chunk> {
    /// Pushed chunks awaiting settling; holds settle's sub-`TARGET` carry at the front.
    input: VecDeque<C>,
    /// The graded chunks emitted so far.
    output: VecDeque<C>,
}

impl<C: Chunk> Default for ChunkBatchBuilder<C> {
    fn default() -> Self {
        Self { input: VecDeque::new(), output: VecDeque::new() }
    }
}

impl<C> crate::trace::Builder for ChunkBatchBuilder<C>
where
    C: Chunk + Default + 'static,
    C::Time: timely::progress::Timestamp,
{
    type Input = C;
    type Time = C::Time;
    type Output = ChunkBatch<C>;

    fn push(&mut self, chunk: &mut C) {
        let chunk = std::mem::take(chunk);
        if chunk.len() > 0 {
            self.input.push_back(chunk);
            C::settle(&mut self.input, false, &mut self.output);
        }
    }

    fn done(self) -> Option<Self::Output> {
        let ChunkBatchBuilder { mut input, mut output } = self;
        C::settle(&mut input, true, &mut output);
        let chunks: Vec<C> = output.into();
        wrap(chunks)
    }

}

impl<C> crate::batcher::merge::Sealer<C> for ChunkBatchBuilder<C>
where
    C: Chunk + Default + 'static,
    C::Time: timely::progress::Timestamp,
{
    type Output = ChunkBatch<C>;

    fn seal(chain: &mut Vec<C>) -> Option<Self::Output> {
        // We settle the chain because we are not guaranteed to received pre-settled data.
        // This should be efficient on pre-settled data.
        wrap(settle_all(std::mem::take(chain)))
    }
}

/// Wraps settled chunks as a batch, absent when there are no chunks.
fn wrap<C: Chunk>(chunks: Vec<C>) -> Option<ChunkBatch<C>> {
    (!chunks.is_empty()).then(|| ChunkBatch::new(chunks))
}

/// Whether `chunks` satisfy the [`Chunk::TARGET`] grading invariant: every chunk
/// at most `TARGET`, and every adjacent pair summing to more than `TARGET` (so no
/// two neighbours could be combined into one legal chunk — a *maximal packing*).
///
/// This is the post-[`settle`](Chunk::settle) shape; useful as a test/debug check.
pub fn is_graded<C: Chunk>(chunks: &[C]) -> bool {
    chunks.iter().all(|c| c.len() <= C::TARGET)
        && chunks.windows(2).all(|w| w[0].len() + w[1].len() > C::TARGET)
}

/// Settle `input` to completion into a fresh graded `Vec` (see [`Chunk::settle`]).
///
/// A convenience for the one-shot callers (batch sealing, the batcher's merge and
/// extract) that have a whole sequence in hand and want it graded; the streaming
/// callers drive [`Chunk::settle`] directly across ticks.
pub fn settle_all<C: Chunk>(input: impl IntoIterator<Item = C>) -> Vec<C> {
    let mut input: VecDeque<C> = input.into_iter().collect();
    let mut out = VecDeque::new();
    C::settle(&mut input, true, &mut out);
    debug_assert!(input.is_empty());
    out.into()
}

/// Merge two full chains of chunks into one, to completion, appending to `out`.
///
/// The plain whole-chain driver: ticks [`Chunk::merge`] until one deque empties, then
/// appends the other's remainder (the verbatim tail). Output is near-graded, not
/// settled. The batcher's `merge` runs the same loop but settles after each push (the
/// bounded-footprint discipline) and so does not use this; it stays as the simplest way
/// to drive [`Chunk::merge`] to completion.
pub fn merge_chains<C: Chunk>(
    chain1: Vec<C>,
    chain2: Vec<C>,
    out: &mut VecDeque<C>,
) {
    let mut in1: VecDeque<C> = chain1.into();
    let mut in2: VecDeque<C> = chain2.into();
    while !in1.is_empty() && !in2.is_empty() {
        C::merge(&mut in1, &mut in2, out);
    }
    // One deque is empty; the other's remainder is all greater than everything merged.
    out.extend(in1.drain(..));
    out.extend(in2.drain(..));
}