Skip to main content

differential_dataflow/columnar/trace/
chunk.rs

1//! A worked [`Chunk`]: the columnar `UpdatesTyped<U>` trie, resident or paged.
2//!
3//! Where [`vec`](crate::trace::chunk::vec) backs a chunk with a flat `Vec<((K,V),T,R)>`, this
4//! backs it with the column-oriented trie from [`crate::columnar::updates`] —
5//! deduplicated keys, per-key val runs, per-val `(time, diff)` runs. It is a
6//! *retargeting* of the columnar trace pile at the [`Chunk`] abstraction: the
7//! storage (`UpdatesTyped`) and the trie-native merge (`trie_merger`) are reused
8//! verbatim, and the four transducers delegate to them. The harness
9//! ([`ChunkBatch`](crate::trace::chunk::ChunkBatch), the straddle cursor, the batcher/builder/
10//! spine aliases) is shared with `vec`.
11//!
12//! This makes columnar trace merges trie-native (the old `OrdValBatch`-backed
13//! trace ran them through ord_neu's row-oriented merger).
14//!
15//! # Resident vs paged
16//!
17//! A [`ColChunk`] is either [`Resident`](ColChunk::Resident) (the trie in memory)
18//! or [`Paged`](ColChunk::Paged) (resident bounds + a byte handle). [`Chunk::settle`]
19//! is the spill point: it pages committed chunks out via
20//! [`spill`] when a worker has installed a spiller. Reads
21//! fetch a paged chunk's trie back, caching it in a [`OnceCell`] so repeated
22//! cursor access pays the fetch once; [`len`](Chunk::len) and [`bounds`](crate::trace::chunk::NavigableChunk::bounds)
23//! read the resident metadata and never fetch.
24//!
25//! [`Chunk::advance`] is trie-native: it melds its input in place and rewrites
26//! only the time column (keys/vals/diffs stay columnar refs), then re-consolidates
27//! the advanced time runs (a global pass a future per-`(key,val)` streaming
28//! consolidation could narrow).
29
30use std::cell::OnceCell;
31use std::collections::VecDeque;
32use std::marker::PhantomData;
33use std::rc::Rc;
34
35use columnar::{Borrow, Columnar, Container, ContainerOf, Index, Len, Push};
36use timely::Accountable;
37use timely::container::{PushInto, SizableContainer};
38use timely::progress::Antichain;
39use timely::progress::frontier::AntichainRef;
40
41use crate::consolidation::Consolidate;
42use crate::lattice::Lattice;
43use crate::trace::Navigable;
44use crate::trace::cursor::Cursor;
45use crate::trace::implementations::{BatchContainer, Layout};
46
47use crate::columnar::layout::{ColumnarLayout, ColumnarUpdate, Coltainer};
48use crate::columnar::updates::{child_range, UpdatesBuilder, UpdatesTyped};
49use crate::columnar::trie_merger;
50
51use super::spill::{self, BytesSource};
52
53use crate::trace::chunk::Chunk;
54
55/// The chunk size: the [`Chunk::TARGET`] grading value.
56///
57/// Smaller than `crate::columnar::LINK_TARGET` so multi-chunk batches — and the
58/// straddle cursor — are exercised; the trie operations are happy at any size.
59const TARGET: usize = 8192;
60
61/// Resident bounds for a paged chunk: the first and last `(key, val, time)` as
62/// single-element columnar containers (so [`crate::trace::chunk::NavigableChunk::bounds`] returns refs without
63/// fetching), plus the record count.
64pub struct ChunkMeta<U: ColumnarUpdate> {
65    fk: ContainerOf<U::Key>, fv: ContainerOf<U::Val>, ft: ContainerOf<U::Time>,
66    lk: ContainerOf<U::Key>, lv: ContainerOf<U::Val>, lt: ContainerOf<U::Time>,
67    len: usize,
68}
69
70/// A paged chunk: resident `meta`, a byte handle to fetch the trie, and a cache
71/// populated on first read. Opaque — held inside [`ColChunk::Paged`].
72pub struct PagedChunk<U: ColumnarUpdate> {
73    meta: ChunkMeta<U>,
74    source: Box<dyn BytesSource>,
75    cache: OnceCell<Rc<UpdatesTyped<U>>>,
76}
77
78/// A single-element container holding `borrowed[i]`.
79fn singleton<C: Container + Default>(borrowed: <C as Borrow>::Borrowed<'_>, i: usize) -> C {
80    let mut out = C::default();
81    out.push(borrowed.get(i));
82    out
83}
84
85/// Snapshot a trie's first/last `(key, val, time)` and length as resident metadata.
86fn meta_of<U: ColumnarUpdate>(t: &UpdatesTyped<U>) -> ChunkMeta<U> {
87    let v = t.view();
88    let (nk, nv, nt) = (v.keys.values.len(), v.vals.values.len(), v.times.values.len());
89    ChunkMeta {
90        fk: singleton(v.keys.values, 0),  lk: singleton(v.keys.values, nk - 1),
91        fv: singleton(v.vals.values, 0),  lv: singleton(v.vals.values, nv - 1),
92        ft: singleton(v.times.values, 0), lt: singleton(v.times.values, nt - 1),
93        len: t.len(),
94    }
95}
96
97/// A sorted, consolidated columnar trie of `((key, val), time, diff)`: resident, or
98/// paged to backing storage.
99pub enum ColChunk<U: ColumnarUpdate> {
100    /// The trie in memory, shared via `Rc`.
101    Resident(Rc<UpdatesTyped<U>>),
102    /// Spilled out: resident bounds + a fetch handle (and a fill-once cache).
103    Paged(Rc<PagedChunk<U>>),
104}
105
106impl<U: ColumnarUpdate> Clone for ColChunk<U> {
107    fn clone(&self) -> Self {
108        match self {
109            ColChunk::Resident(rc) => ColChunk::Resident(Rc::clone(rc)),
110            ColChunk::Paged(p) => ColChunk::Paged(Rc::clone(p)),
111        }
112    }
113}
114impl<U: ColumnarUpdate> Default for ColChunk<U> {
115    fn default() -> Self { ColChunk::Resident(Rc::new(UpdatesTyped::default())) }
116}
117
118impl<U: ColumnarUpdate> ColChunk<U> {
119    /// Wrap an already sorted, consolidated trie as a resident chunk. The chunker
120    /// uses this to hand its melded `UpdatesTyped` blobs to the `Chunk` harness.
121    pub fn from_trie(updates: UpdatesTyped<U>) -> Self { ColChunk::Resident(Rc::new(updates)) }
122
123    /// Borrow the chunk's trie, fetching (and caching) it if paged. Reads tie to
124    /// `&self`, so cursor refs remain valid for the borrow.
125    fn trie(&self) -> &UpdatesTyped<U> {
126        match self {
127            ColChunk::Resident(rc) => &**rc,
128            ColChunk::Paged(p) => &**p.cache.get_or_init(|| {
129                spill::note_fetched();
130                Rc::new(spill::decode::<U>(&*p.source))
131            }),
132        }
133    }
134
135    /// Mutable access to the backing trie (materializing if paged, cloning if the
136    /// `Rc` is shared), for builder closures that populate a chunk in place — e.g.
137    /// `reduce_abelian`.
138    pub fn updates_mut(&mut self) -> &mut UpdatesTyped<U> {
139        if let ColChunk::Paged(_) = self {
140            *self = ColChunk::Resident(Rc::new(into_trie(std::mem::take(self))));
141        }
142        match self {
143            ColChunk::Resident(rc) => Rc::make_mut(rc),
144            ColChunk::Paged(_) => unreachable!(),
145        }
146    }
147}
148
149/// Take a chunk's trie by value, fetching it if paged (and notifying the spiller
150/// for budget accounting).
151fn into_trie<U: ColumnarUpdate>(chunk: ColChunk<U>) -> UpdatesTyped<U> {
152    match chunk {
153        ColChunk::Resident(rc) => Rc::try_unwrap(rc).unwrap_or_else(|rc| (*rc).clone()),
154        ColChunk::Paged(p) => match p.cache.get() {
155            Some(rc) => (**rc).clone(),
156            None => { spill::note_fetched(); spill::decode::<U>(&*p.source) }
157        },
158    }
159}
160
161// --- Container traits (batcher side, via `ContainerChunker<ColChunk>`) ---
162
163impl<U: ColumnarUpdate> Accountable for ColChunk<U> {
164    fn record_count(&self) -> i64 { Chunk::len(self) as i64 }
165}
166
167impl<U: ColumnarUpdate> SizableContainer for ColChunk<U> {
168    // Absorb at `TARGET`, the grading size, so the chunker emits pre-graded chunks.
169    fn at_capacity(&self) -> bool { Chunk::len(self) >= TARGET }
170    // The trie grows as updates are pushed; nothing to pre-size.
171    fn ensure_capacity(&mut self, _stash: &mut Option<Self>) { }
172}
173
174impl<U: ColumnarUpdate> Consolidate for ColChunk<U> {
175    fn len(&self) -> usize { Chunk::len(self) }
176    fn clear(&mut self) { *self.updates_mut() = UpdatesTyped::default(); }
177    fn consolidate_into(&mut self, target: &mut Self) {
178        let taken = std::mem::take(self.updates_mut());
179        *target.updates_mut() = taken.consolidate();
180    }
181}
182
183impl<U: ColumnarUpdate> PushInto<((U::Key, U::Val), U::Time, U::Diff)> for ColChunk<U> {
184    fn push_into(&mut self, item: ((U::Key, U::Val), U::Time, U::Diff)) {
185        self.updates_mut().push_into(item);
186    }
187}
188
189// --- Cursor (trace side), navigating the trie directly (cf. `OrdValCursor`) ---
190
191/// A cursor over a [`ColChunk`], tracking the current key and value as absolute
192/// indices into the trie's flat `keys.values` / `vals.values` columns.
193pub struct ColChunkCursor<U: ColumnarUpdate> {
194    key_cursor: usize,
195    val_cursor: usize,
196    phantom: PhantomData<U>,
197}
198
199impl<U: ColumnarUpdate> Cursor for ColChunkCursor<U> {
200    type Storage = ColChunk<U>;
201
202    type KeyContainer = <ColumnarLayout<U> as Layout>::KeyContainer;
203    type Key<'a> = <<ColumnarLayout<U> as Layout>::KeyContainer as BatchContainer>::ReadItem<'a>;
204    type ValContainer = <ColumnarLayout<U> as Layout>::ValContainer;
205    type Val<'a> = <<ColumnarLayout<U> as Layout>::ValContainer as BatchContainer>::ReadItem<'a>;
206    type ValOwn = <<ColumnarLayout<U> as Layout>::ValContainer as BatchContainer>::Owned;
207    type TimeContainer = <ColumnarLayout<U> as Layout>::TimeContainer;
208    type TimeGat<'a> = <<ColumnarLayout<U> as Layout>::TimeContainer as BatchContainer>::ReadItem<'a>;
209    type Time = <<ColumnarLayout<U> as Layout>::TimeContainer as BatchContainer>::Owned;
210    type DiffContainer = <ColumnarLayout<U> as Layout>::DiffContainer;
211    type DiffGat<'a> = <<ColumnarLayout<U> as Layout>::DiffContainer as BatchContainer>::ReadItem<'a>;
212    type Diff = <<ColumnarLayout<U> as Layout>::DiffContainer as BatchContainer>::Owned;
213
214    fn key_valid(&self, s: &Self::Storage) -> bool { self.key_cursor < s.trie().view().keys.values.len() }
215    fn val_valid(&self, s: &Self::Storage) -> bool {
216        let view = s.trie().view();
217        self.key_cursor < view.keys.values.len()
218            && self.val_cursor < child_range(view.vals.bounds, self.key_cursor).end
219    }
220    fn key<'a>(&self, s: &'a Self::Storage) -> Self::Key<'a> { s.trie().view().keys.values.get(self.key_cursor) }
221    fn val<'a>(&self, s: &'a Self::Storage) -> Self::Val<'a> { s.trie().view().vals.values.get(self.val_cursor) }
222    fn get_key<'a>(&self, s: &'a Self::Storage) -> Option<Self::Key<'a>> {
223        if self.key_valid(s) { Some(self.key(s)) } else { None }
224    }
225    fn get_val<'a>(&self, s: &'a Self::Storage) -> Option<Self::Val<'a>> {
226        if self.val_valid(s) { Some(self.val(s)) } else { None }
227    }
228    fn map_times<L: FnMut(Self::TimeGat<'_>, Self::DiffGat<'_>)>(&mut self, s: &Self::Storage, mut logic: L) {
229        if !self.val_valid(s) { return; }
230        let view = s.trie().view();
231        for t in child_range(view.times.bounds, self.val_cursor) {
232            let time = view.times.values.get(t);
233            for d in child_range(view.diffs.bounds, t) {
234                logic(time, view.diffs.values.get(d));
235            }
236        }
237    }
238    fn step_key(&mut self, s: &Self::Storage) {
239        self.key_cursor += 1;
240        if self.key_valid(s) { self.rewind_vals(s); }
241        else { self.key_cursor = s.trie().view().keys.values.len(); }
242    }
243    fn seek_key(&mut self, s: &Self::Storage, key: Self::Key<'_>) {
244        let view = s.trie().view();
245        let n = view.keys.values.len();
246        let mut lo = self.key_cursor;
247        trie_merger::gallop(view.keys.values, &mut lo, n, |x|
248            <Coltainer<U::Key> as BatchContainer>::reborrow(x).lt(&<Coltainer<U::Key> as BatchContainer>::reborrow(key)));
249        self.key_cursor = lo;
250        if self.key_valid(s) { self.rewind_vals(s); }
251    }
252    fn step_val(&mut self, s: &Self::Storage) {
253        self.val_cursor += 1;
254        if !self.val_valid(s) {
255            self.val_cursor = child_range(s.trie().view().vals.bounds, self.key_cursor).end;
256        }
257    }
258    fn seek_val(&mut self, s: &Self::Storage, val: Self::Val<'_>) {
259        if !self.key_valid(s) { return; }
260        let view = s.trie().view();
261        let upper = child_range(view.vals.bounds, self.key_cursor).end;
262        let mut lo = self.val_cursor;
263        trie_merger::gallop(view.vals.values, &mut lo, upper, |x|
264            <Coltainer<U::Val> as BatchContainer>::reborrow(x).lt(&<Coltainer<U::Val> as BatchContainer>::reborrow(val)));
265        self.val_cursor = lo;
266    }
267    fn rewind_keys(&mut self, s: &Self::Storage) { self.key_cursor = 0; self.rewind_vals(s); }
268    fn rewind_vals(&mut self, s: &Self::Storage) {
269        if self.key_valid(s) {
270            self.val_cursor = child_range(s.trie().view().vals.bounds, self.key_cursor).start;
271        }
272    }
273}
274
275/// Wrap a non-empty resident trie as a chunk and append it to `out`.
276fn emit<U: ColumnarUpdate>(updates: UpdatesTyped<U>, out: &mut VecDeque<ColChunk<U>>) {
277    if updates.len() > 0 { out.push_back(ColChunk::Resident(Rc::new(updates))); }
278}
279
280/// Drop the empty diff lists `merge_pair`'s `write_diffs` leaves where updates
281/// cancel — but only when some actually cancelled. With no cancellation every
282/// time keeps its singleton diff, so `diffs.values.len() == times.values.len()`
283/// and we skip [`UpdatesTyped::filter_zero`]'s rebuild entirely.
284fn consolidated<U: ColumnarUpdate>(merged: UpdatesTyped<U>) -> UpdatesTyped<U> {
285    if merged.diffs.values.len() == merged.times.values.len() { merged } else { merged.filter_zero() }
286}
287
288impl<U: ColumnarUpdate> Navigable for ColChunk<U>
289where U::Time: 'static {
290    type Cursor = ColChunkCursor<U>;
291
292    fn cursor(&self) -> Self::Cursor {
293        ColChunkCursor { key_cursor: 0, val_cursor: 0, phantom: PhantomData }
294    }
295}
296
297impl<U: ColumnarUpdate> crate::trace::chunk::NavigableChunk for ColChunk<U>
298where U::Time: 'static {
299    fn bounds(&self) -> (
300        (<Self::Cursor as Cursor>::Key<'_>, <Self::Cursor as Cursor>::Val<'_>, <Self::Cursor as Cursor>::TimeGat<'_>),
301        (<Self::Cursor as Cursor>::Key<'_>, <Self::Cursor as Cursor>::Val<'_>, <Self::Cursor as Cursor>::TimeGat<'_>),
302    ) {
303        match self {
304            ColChunk::Resident(rc) => {
305                let view = rc.view();
306                let (nk, nv, nt) = (view.keys.values.len(), view.vals.values.len(), view.times.values.len());
307                // Sorted trie: the first/last of each flat column are the bounding triples.
308                ((view.keys.values.get(0), view.vals.values.get(0), view.times.values.get(0)),
309                 (view.keys.values.get(nk - 1), view.vals.values.get(nv - 1), view.times.values.get(nt - 1)))
310            }
311            ColChunk::Paged(p) => {
312                // Read the resident metadata — no fetch.
313                let m = &p.meta;
314                ((m.fk.borrow().get(0), m.fv.borrow().get(0), m.ft.borrow().get(0)),
315                 (m.lk.borrow().get(0), m.lv.borrow().get(0), m.lt.borrow().get(0)))
316            }
317        }
318    }
319}
320
321impl<U: ColumnarUpdate> Chunk for ColChunk<U>
322where U::Time: 'static {
323    type Time = U::Time;
324
325    const TARGET: usize = TARGET;
326
327    fn len(&self) -> usize {
328        match self {
329            ColChunk::Resident(rc) => rc.len(),
330            ColChunk::Paged(p) => p.meta.len,
331        }
332    }
333
334    /// Trie-native binary merge of the two deques' front chunks through their
335    /// shared horizon, via [`trie_merger::merge_pair`]: it merges one input fully
336    /// and a prefix of the other (the survey/`write_layer` bulk range-copy),
337    /// leaving the survivor's suffix as a cursor we materialize with
338    /// [`trie_merger::suffix_chunk`] and push back. The harness re-invokes.
339    ///
340    /// One pair per call (not a resumable drain): `merge_pair` surveys each batch
341    /// from index 0, so re-passing a large survivor would re-survey its whole body
342    /// every call (quadratic). Re-materializing the suffix as a standalone chunk
343    /// keeps each survey bounded by the shrinking remainder instead.
344    fn merge(in1: &mut VecDeque<Self>, in2: &mut VecDeque<Self>, out: &mut VecDeque<Self>) {
345        let mut cursor1 = Some(((0, 0, 0), into_trie(in1.pop_front().unwrap())));
346        let mut cursor2 = Some(((0, 0, 0), into_trie(in2.pop_front().unwrap())));
347        emit(consolidated(trie_merger::merge_pair(&mut cursor1, &mut cursor2)), out);
348        // Push the survivor's unconsumed suffix back to the front of its deque.
349        if let Some((cursor, batch)) = cursor1 {
350            let suffix = trie_merger::suffix_chunk(cursor, &batch);
351            if suffix.len() > 0 { in1.push_front(ColChunk::Resident(Rc::new(suffix))); }
352        }
353        if let Some((cursor, batch)) = cursor2 {
354            let suffix = trie_merger::suffix_chunk(cursor, &batch);
355            if suffix.len() > 0 { in2.push_front(ColChunk::Resident(Rc::new(suffix))); }
356        }
357    }
358
359    /// Partition the front chunk by `frontier` (keep `>=`, ship `<`), folding kept
360    /// times into `residual`, via [`trie_merger::extract`]. One chunk per call.
361    fn extract(
362        input: &mut VecDeque<Self>,
363        frontier: AntichainRef<U::Time>,
364        residual: &mut Antichain<U::Time>,
365        keep: &mut VecDeque<Self>,
366        ship: &mut VecDeque<Self>,
367    ) {
368        let Some(chunk) = input.pop_front() else { return };
369        trie_merger::extract(
370            std::iter::once(into_trie(chunk)),
371            frontier,
372            residual,
373            |c| emit(c, ship),
374            |c| emit(c, keep),
375        );
376    }
377
378    /// Advance times by `frontier`, consolidating each complete `(key, val)` group
379    /// and withholding the last unless `done`.
380    ///
381    /// Streams the input one chunk at a time, holding at most a single chunk plus the
382    /// withheld trailing `(key, val)` group (the `carry`) — it never pulls the whole
383    /// chain into core. Cross-chunk consolidation happens *only* where a `(key, val)`
384    /// straddles a boundary: the (un-advanced) carry is melded onto the next chunk's
385    /// head, a sorted-chain meld that just stitches that one group. `advance_trie`
386    /// then advances and consolidates the time runs *within* each chunk (keys / vals /
387    /// diffs stay columnar refs, no global re-sort). Grading the emitted chunks is
388    /// left to [`Chunk::settle`].
389    fn advance(
390        input: &mut VecDeque<Self>,
391        frontier: AntichainRef<U::Time>,
392        done: bool,
393        out: &mut VecDeque<Self>,
394    ) {
395        // The withheld trailing `(key, val)` group, kept un-advanced so it melds
396        // cleanly onto the next chunk's head if the group straddles the boundary.
397        let mut carry: Option<UpdatesTyped<U>> = None;
398        while let Some(chunk) = input.pop_front() {
399            // Stitch the carry onto this chunk via a sorted-chain meld (touches only
400            // the straddling group); otherwise take the chunk as-is.
401            let combined = match carry.take() {
402                None => into_trie(chunk),
403                Some(c) => {
404                    let mut builder = UpdatesBuilder::new_from(c);
405                    builder.meld(&into_trie(chunk));
406                    builder.done()
407                }
408            };
409            if combined.len() == 0 { continue; }
410
411            // Withhold the trailing `(key, val)` group (it may continue in the next
412            // chunk). Its size is read from the resident bounds, no body scan.
413            let tail = {
414                let v = combined.view();
415                let last_val = child_range(v.vals.bounds, v.keys.values.len() - 1).end - 1;
416                let times = child_range(v.times.bounds, last_val);
417                child_range(v.diffs.bounds, times.end - 1).end
418                    - child_range(v.diffs.bounds, times.start).start
419            };
420            if tail == combined.len() {
421                // A single `(key, val)` spans the chunk; hold it all as the carry.
422                carry = Some(combined);
423                continue;
424            }
425            let split = combined.len() - tail;
426            let (keep, rest) = trie_merger::split_at(combined, split);
427            carry = Some(rest);
428            // Advance the now-complete groups and emit one chunk; `settle` grades.
429            emit(advance_trie(keep, frontier), out);
430        }
431        // Flush the final carry: advance and emit if `done`, else hold for next call.
432        if let Some(c) = carry {
433            if done { emit(advance_trie(c, frontier), out); }
434            else { input.push_front(ColChunk::Resident(Rc::new(c))); }
435        }
436    }
437
438    /// Maximal packing via the harness [`pack`](crate::trace::chunk::pack): coalesce by melding
439    /// each trie of the run onto the first (adjacent chunks of a sorted, consolidated
440    /// chain, so meld's "strictly greater first triple" precondition holds), split
441    /// with [`trie_merger::split_at`], and seal through `seal_chunk` (the spill
442    /// point — pages a committed chunk when a spiller is installed).
443    fn settle(input: &mut VecDeque<Self>, done: bool, out: &mut VecDeque<Self>) {
444        crate::trace::chunk::pack(
445            input, done, out,
446            |run| {
447                let mut run = run.into_iter();
448                let mut build = UpdatesBuilder::new_from(into_trie(run.next().unwrap()));
449                for next in run { build.meld(&into_trie(next)); }
450                ColChunk::Resident(Rc::new(build.done()))
451            },
452            |chunk, n| {
453                let (first, rest) = trie_merger::split_at(into_trie(chunk), n);
454                (ColChunk::Resident(Rc::new(first)), ColChunk::Resident(Rc::new(rest)))
455            },
456            seal_chunk,
457        );
458    }
459}
460
461/// The columnar spill point: when a spiller is installed and over the high-water
462/// mark, page a committed chunk out (serialize via [`spill::try_page`], keep
463/// resident bounds + a fetch handle). Otherwise keep it resident.
464fn seal_chunk<U: ColumnarUpdate>(chunk: ColChunk<U>) -> ColChunk<U> {
465    let ColChunk::Resident(rc) = chunk else { return chunk };
466    if !spill::active() { return ColChunk::Resident(rc); }
467    let updates = Rc::try_unwrap(rc).unwrap_or_else(|rc| (*rc).clone());
468    let meta = meta_of(&updates);
469    match spill::try_page(updates) {
470        Ok(source) => ColChunk::Paged(Rc::new(PagedChunk { meta, source, cache: OnceCell::new() })),
471        Err(updates) => ColChunk::Resident(Rc::new(updates)),
472    }
473}
474
475/// Advance the times in `keep` by `frontier` and rebuild the trie, consolidating
476/// each `(key, val)`'s time run *in isolation*. Advancing preserves `(key, val)`
477/// order, so only the per-group time runs can reorder — they are re-sorted and
478/// merged locally (a no-op sort when advancing is monotone, e.g. total-order
479/// times), never the whole trie. Keys / vals / diffs are copied by columnar ref;
480/// a `(key, val)` whose diffs fully cancel is dropped.
481fn advance_trie<U: ColumnarUpdate>(
482    keep: UpdatesTyped<U>,
483    frontier: AntichainRef<U::Time>,
484) -> UpdatesTyped<U> {
485    use crate::difference::{IsZero, Semigroup};
486    let view = keep.view();
487    let mut out = UpdatesTyped::<U>::default();
488    let mut run: Vec<(U::Time, U::Diff)> = Vec::new();
489    let mut time = U::Time::default();
490    let mut any_key = false;
491
492    for key_idx in 0..view.keys.values.len() {
493        let mut key_emitted = false;
494        for val_idx in child_range(view.vals.bounds, key_idx) {
495            // Gather (advanced time, summed diff) for this `(key, val)`.
496            run.clear();
497            for time_idx in child_range(view.times.bounds, val_idx) {
498                <U::Time as Columnar>::copy_from(&mut time, view.times.values.get(time_idx));
499                time.advance_by(frontier);
500                let mut diff = U::Diff::default();
501                for diff_idx in child_range(view.diffs.bounds, time_idx) {
502                    diff.plus_equals(&<U::Diff as Columnar>::into_owned(view.diffs.values.get(diff_idx)));
503                }
504                run.push((time.clone(), diff));
505            }
506            // Re-sort the run by advanced time and merge equal times, dropping zeros.
507            run.sort_by(|a, b| a.0.cmp(&b.0));
508            let (mut w, mut r) = (0, 0);
509            while r < run.len() {
510                let t = run[r].0.clone();
511                let mut d = run[r].1.clone();
512                r += 1;
513                while r < run.len() && run[r].0 == t { d.plus_equals(&run[r].1); r += 1; }
514                if !d.is_zero() { run[w] = (t, d); w += 1; }
515            }
516            run.truncate(w);
517            if run.is_empty() { continue; }
518
519            // Emit the surviving group, sealing bounds level by level (cf. `form`).
520            if !key_emitted {
521                out.keys.values.push(view.keys.values.get(key_idx));
522                key_emitted = true;
523            }
524            out.vals.values.push(view.vals.values.get(val_idx));
525            for (t, d) in &run {
526                out.times.values.push(t);
527                out.diffs.values.push(d);
528                out.diffs.bounds.push(out.diffs.values.len() as u64);
529            }
530            out.times.bounds.push(out.times.values.len() as u64);
531        }
532        if key_emitted {
533            out.vals.bounds.push(out.vals.values.len() as u64);
534            any_key = true;
535        }
536    }
537    if any_key { out.keys.bounds.push(out.keys.values.len() as u64); }
538    out
539}
540
541#[cfg(test)]
542mod test {
543    use timely::progress::Antichain;
544    use std::collections::VecDeque;
545    use columnar::Push;
546    use super::{ColChunk, Chunk};
547    use crate::columnar::updates::UpdatesTyped;
548    use crate::trace::chunk::merge_chains;
549    use crate::trace::Navigable;
550
551    type Upd = (u64, u64, u64, i64);
552
553    // A sorted, consolidated columnar chunk from raw updates.
554    fn chunk(updates: Vec<Upd>) -> ColChunk<Upd> {
555        let mut u = UpdatesTyped::<Upd>::default();
556        for (k, v, t, d) in updates { u.push((&k, &v, &t, &d)); }
557        ColChunk::from_trie(u.consolidate())
558    }
559
560    // Flatten a chunk sequence back to its update stream.
561    fn flat<I: IntoIterator<Item = ColChunk<Upd>>>(chunks: I) -> Vec<Upd> {
562        chunks.into_iter().flat_map(|c| c.trie().iter().map(|(k, v, t, d)| (*k, *v, *t, *d)).collect::<Vec<_>>()).collect()
563    }
564
565    // Cut a consolidated set into a chain of small chunks, so groups straddle boundaries.
566    fn chain(updates: &[Upd], sz: usize) -> Vec<ColChunk<Upd>> {
567        updates.chunks(sz).map(|c| chunk(c.to_vec())).collect()
568    }
569
570    // Property test: merging two multi-chunk chains (driven through `merge` by
571    // `merge_chains`) reproduces the union of all updates, consolidated. Tiny
572    // chunks force `(key, val)` groups to straddle chunk boundaries on both
573    // sides, exercising the survey/`merge_pair` horizon and suffix push-back.
574    #[test]
575    fn merge_matches_reference() {
576        use crate::consolidation::consolidate_updates;
577
578        let mut seed = 0x2545F4914F6CDD1Du64;
579        let mut rng = move || { seed ^= seed << 13; seed ^= seed >> 7; seed ^= seed << 17; seed };
580
581        fn gen(rng: &mut impl FnMut() -> u64, n: usize) -> Vec<Upd> {
582            let mut v: Vec<Upd> = (0..n).map(|_| {
583                let k = rng() % 20; let val = rng() % 3; let t = rng() % 8;
584                let d = if rng() % 4 == 0 { -1 } else { 1 };
585                (k, val, t, d)
586            }).collect();
587            let mut rows: Vec<((u64, u64), u64, i64)> = v.drain(..).map(|(k, val, t, d)| ((k, val), t, d)).collect();
588            consolidate_updates(&mut rows);
589            rows.into_iter().map(|((k, val), t, d)| (k, val, t, d)).collect()
590        }
591
592        for _ in 0..300 {
593            let (n1, n2) = ((rng() as usize % 60) + 1, (rng() as usize % 60) + 1);
594            let u1 = gen(&mut rng, n1);
595            let u2 = gen(&mut rng, n2);
596            if u1.is_empty() || u2.is_empty() { continue; }
597            let sz = (rng() as usize % 5) + 1;
598
599            let mut out = VecDeque::new();
600            merge_chains(chain(&u1, sz), chain(&u2, sz), &mut out);
601            let merged = flat(out);
602
603            let mut reference: Vec<((u64, u64), u64, i64)> =
604                u1.iter().chain(u2.iter()).map(|&(k, v, t, d)| ((k, v), t, d)).collect();
605            consolidate_updates(&mut reference);
606            let reference: Vec<Upd> = reference.into_iter().map(|((k, v), t, d)| (k, v, t, d)).collect();
607
608            assert_eq!(merged, reference, "chunk size {sz}\n  u1={u1:?}\n  u2={u2:?}");
609        }
610    }
611
612    // `settle` produces a maximal packing: chunks `<= TARGET`, adjacent pairs
613    // summing past `TARGET`, contents preserved exactly.
614    #[test]
615    fn settle_maximal_packing() {
616        use super::TARGET;
617        use crate::trace::chunk::is_graded;
618
619        let t = TARGET;
620        let sizes = [t / 3, t / 3, t / 3, t, t / 2, t / 2, t, 1, t - 1];
621        let total: usize = sizes.iter().sum();
622        let mut key = 0u64;
623        let mut input = VecDeque::new();
624        let mut output = VecDeque::new();
625        for &s in &sizes {
626            let updates: Vec<Upd> = (0..s).map(|_| { let k = key; key += 1; (k, 0, 0, 1) }).collect();
627            input.push_back(chunk(updates));
628            ColChunk::settle(&mut input, false, &mut output);
629        }
630        ColChunk::settle(&mut input, true, &mut output);
631        let chunks: Vec<_> = output.into();
632
633        assert!(is_graded(&chunks), "not graded: {:?}", chunks.iter().map(Chunk::len).collect::<Vec<_>>());
634        let got = flat(chunks);
635        assert_eq!(got.len(), total);
636        assert!(got.windows(2).all(|w| w[0].0 < w[1].0));
637    }
638
639    // The straddle-aware `ChunkBatch` cursor reconstructs the same grouped
640    // updates as a flat reference, even when a key — and a `(key, val)`'s times —
641    // span a chunk boundary.
642    #[test]
643    fn cursor_handles_straddle() {
644        use crate::trace::cursor::Cursor;
645        use crate::trace::chunk::ChunkBatch;
646
647        let chunks = vec![
648            chunk(vec![(0, 0, 0, 1), (1, 0, 0, 1), (1, 1, 0, 1)]),
649            chunk(vec![(1, 1, 1, 1), (1, 2, 0, 1)]),
650            chunk(vec![(2, 0, 0, 1)]),
651        ];
652        let batch = ChunkBatch::new(chunks);
653
654        let mut cursor = batch.cursor();
655        let got = cursor.to_vec(&batch, |k| *k, |v| *v);
656        let want = vec![
657            ((0u64, 0u64), vec![(0u64, 1i64)]),
658            ((1, 0), vec![(0, 1)]),
659            ((1, 1), vec![(0, 1), (1, 1)]),
660            ((1, 2), vec![(0, 1)]),
661            ((2, 0), vec![(0, 1)]),
662        ];
663        assert_eq!(got, want);
664    }
665
666    // Driving `ChunkBatchMerger` to completion with tiny `fuel` (so it suspends
667    // and settles on nearly every tick) yields the same advanced-and-consolidated
668    // batch as a one-shot reference, and that batch is graded. Exercises the
669    // resumable merge -> advance -> settle pipeline end to end.
670    #[test]
671    fn batch_merger_resumable_matches_reference() {
672        use crate::trace::implementations::spine_fueled::Merger;
673        use crate::trace::chunk::{ChunkBatch, ChunkBatchMerger, is_graded};
674        use crate::trace::cursor::Cursor;
675        use crate::consolidation::consolidate_updates;
676
677        let mut seed = 0x9E3779B97F4A7C15u64;
678        let mut rng = move || { seed ^= seed << 13; seed ^= seed >> 7; seed ^= seed << 17; seed };
679
680        fn gen(rng: &mut impl FnMut() -> u64) -> Vec<Upd> {
681            let n = rng() as usize % 40 + 1;
682            let mut rows: Vec<((u64, u64), u64, i64)> = (0..n).map(|_| {
683                let k = rng() % 10; let val = rng() % 3; let t = rng() % 6;
684                let d = if rng() % 4 == 0 { -1 } else { 1 };
685                ((k, val), t, d)
686            }).collect();
687            consolidate_updates(&mut rows);
688            rows.into_iter().map(|((k, v), t, d)| (k, v, t, d)).collect()
689        }
690        fn batch(updates: &[Upd], sz: usize) -> ChunkBatch<ColChunk<Upd>> {
691            let chunks: Vec<_> = updates.chunks(sz).map(|c| chunk(c.to_vec())).collect();
692            ChunkBatch::new(chunks)
693        }
694        fn read(b: &ChunkBatch<ColChunk<Upd>>) -> Vec<Upd> {
695            let mut out = Vec::new();
696            let mut c = b.cursor();
697            while c.key_valid(b) {
698                let k = *c.key(b);
699                while c.val_valid(b) {
700                    let v = *c.val(b);
701                    c.map_times(b, |t, d| out.push(((k, v), *t, *d)));
702                    c.step_val(b);
703                }
704                c.step_key(b);
705            }
706            consolidate_updates(&mut out);
707            out.into_iter().map(|((k, v), t, d)| (k, v, t, d)).collect()
708        }
709
710        for _ in 0..200 {
711            let u1 = gen(&mut rng);
712            let u2 = gen(&mut rng);
713            if u1.is_empty() || u2.is_empty() { continue; }
714            let sz = (rng() as usize % 4) + 1;
715            let f = rng() % 6;
716            let (s1, s2) = (batch(&u1, sz), batch(&u2, sz));
717            let frontier = Antichain::from_elem(f);
718
719            let mut merger = ChunkBatchMerger::new(&s1, &s2, frontier.borrow());
720            loop {
721                let mut fuel = 1isize;
722                merger.work(&s1, &s2, &mut fuel);
723                if fuel > 0 { break; }
724            }
725            let result = merger.done();
726
727            let chunks: &[ColChunk<Upd>] = result.as_ref().map_or(&[], |b| &b.chunks[..]);
728            assert!(is_graded(chunks), "ungraded result: {:?}",
729                chunks.iter().map(Chunk::len).collect::<Vec<_>>());
730            let got = result.as_ref().map_or_else(Vec::new, read);
731            let mut want: Vec<((u64, u64), u64, i64)> =
732                u1.iter().chain(u2.iter()).map(|&(k, v, t, d)| ((k, v), t.max(f), d)).collect();
733            consolidate_updates(&mut want);
734            let want: Vec<Upd> = want.into_iter().map(|((k, v), t, d)| (k, v, t, d)).collect();
735            // A merge that cancels to nothing reports an absent updates.
736            assert_eq!(result.is_none(), want.is_empty(), "absence must track emptiness\n  u1={u1:?}\n  u2={u2:?}\n  f={f}");
737            assert_eq!(got, want, "fuel-driven merge mismatch\n  u1={u1:?}\n  u2={u2:?}\n  f={f}");
738        }
739    }
740
741    // With a spiller installed and a tiny budget, `settle` pages its committed
742    // chunks out; reads must transparently fetch them back and reproduce the
743    // exact contents (exercises the trie byte codec + the OnceCell materialization).
744    #[test]
745    fn settle_pages_and_round_trips() {
746        use std::cell::RefCell;
747        use std::rc::Rc;
748        use std::sync::Arc;
749        use std::sync::atomic::Ordering::Relaxed;
750        use crate::columnar::trace::spill::{self, BytesSource, BytesStore, SpillStats};
751
752        // In-memory backing store: an arena of byte blobs.
753        struct MemStore(Rc<RefCell<Vec<Vec<u8>>>>);
754        struct MemSource(Rc<RefCell<Vec<Vec<u8>>>>, usize);
755        impl BytesStore for MemStore {
756            fn store(&mut self, bytes: &[u8]) -> Box<dyn BytesSource> {
757                let mut a = self.0.borrow_mut();
758                let id = a.len();
759                a.push(bytes.to_vec());
760                Box::new(MemSource(self.0.clone(), id))
761            }
762        }
763        impl BytesSource for MemSource { fn load(&self) -> Vec<u8> { self.0.borrow()[self.1].clone() } }
764
765        let arena = Rc::new(RefCell::new(Vec::new()));
766        let stats = Arc::new(SpillStats::default());
767        spill::install(1, Box::new(MemStore(arena)), stats.clone()); // budget 1 record → page everything
768
769        // Many single-update chunks with distinct keys; settle coalesces and pages.
770        let n = 5 * super::TARGET as u64;
771        let mut input: VecDeque<_> = (0..n).map(|k| chunk(vec![(k, 0, 0, 1)])).collect();
772        let mut out = VecDeque::new();
773        ColChunk::settle(&mut input, true, &mut out);
774
775        assert!(stats.spilled_chunks.load(Relaxed) > 0, "nothing was paged");
776        assert!(out.iter().any(|c| matches!(c, ColChunk::Paged(_))), "no paged chunk in output");
777        // Contents survive the disk round-trip exactly.
778        let got = flat(out);
779        let want: Vec<Upd> = (0..n).map(|k| (k, 0, 0, 1)).collect();
780        assert_eq!(got, want);
781        assert!(stats.fetched_chunks.load(Relaxed) > 0, "nothing was fetched back");
782
783        spill::uninstall();
784    }
785
786    // A single `(key, val)` spanning every pushed chunk: `advance` makes no
787    // progress until `done`, accumulating in the carry, and must still produce
788    // the right advanced-and-consolidated result.
789    #[test]
790    fn advance_single_key_spanning_pushes() {
791        let frontier = Antichain::from_elem(100u64);
792        let n = 50u64;
793        let mut q = VecDeque::new();
794        let mut out = VecDeque::new();
795        for t in 0..n {
796            q.push_back(chunk(vec![(7, 0, t, 1)]));
797            ColChunk::advance(&mut q, frontier.borrow(), false, &mut out);
798        }
799        ColChunk::advance(&mut q, frontier.borrow(), true, &mut out);
800        assert_eq!(flat(out), vec![(7, 0, 100, n as i64)]);
801    }
802
803    // `advance` advances and consolidates complete `(key, val)` groups eagerly,
804    // withholding the (possibly-growing) last group as the carry when not `done`.
805    #[test]
806    fn advance_emits_complete_groups_eagerly() {
807        let frontier = Antichain::from_elem(5u64);
808        // Group (0,0) is complete within this chunk; group (1,0) might still grow.
809        let mut q = VecDeque::from([chunk(vec![(0, 0, 0, 1), (0, 0, 1, 1), (1, 0, 0, 1)])]);
810        let mut out = VecDeque::new();
811        ColChunk::advance(&mut q, frontier.borrow(), false, &mut out);
812        // The trailing group (1,0) is withheld as the carry at the front of `input`.
813        assert_eq!(q.len(), 1);
814        assert_eq!(Chunk::len(&q[0]), 1);
815        // Group (0,0)'s times {0,1} advanced to 5 and consolidated, emitted now.
816        assert_eq!(flat(out), vec![(0, 0, 5, 2)]);
817    }
818
819    // Streaming the input one chunk at a time must yield exactly what a single
820    // all-at-once flush does — the resumable path is the one-shot path cut at
821    // group boundaries.
822    #[test]
823    fn advance_resumable_matches_oneshot() {
824        let frontier = Antichain::from_elem(3u64);
825        // Groups span chunk boundaries and carry several times each.
826        let input = || vec![
827            chunk(vec![(0, 0, 0, 1), (0, 0, 1, 1), (1, 0, 0, 1)]),
828            chunk(vec![(1, 0, 5, 1), (1, 1, 0, 1), (2, 0, 0, 1)]),
829            chunk(vec![(2, 0, 2, 1), (2, 0, 9, 1)]),
830        ];
831        let oneshot = {
832            let mut q: VecDeque<_> = input().into();
833            let mut out = VecDeque::new();
834            ColChunk::advance(&mut q, frontier.borrow(), false, &mut out);
835            ColChunk::advance(&mut q, frontier.borrow(), true, &mut out);
836            flat(out)
837        };
838        let incremental = {
839            let mut q = VecDeque::new();
840            let mut out = VecDeque::new();
841            for c in input() { q.push_back(c); ColChunk::advance(&mut q, frontier.borrow(), false, &mut out); }
842            ColChunk::advance(&mut q, frontier.borrow(), true, &mut out);
843            flat(out)
844        };
845        assert_eq!(oneshot, incremental);
846        for u in &oneshot { assert!(u.2 >= 3); }
847    }
848
849    // Property test: driving `advance` resumably over many tiny chunks must match a
850    // row oracle (advance each time to `max(t, frontier)`, then consolidate). Tiny
851    // chunks force `(key, val)` groups — with several times each — to straddle
852    // boundaries, exercising the meld / withhold / split path.
853    #[test]
854    fn advance_matches_row_reference() {
855        use crate::consolidation::consolidate_updates;
856
857        let mut seed = 0x2545F4914F6CDD1Du64;
858        let mut rng = move || { seed ^= seed << 13; seed ^= seed >> 7; seed ^= seed << 17; seed };
859
860        for _ in 0..200 {
861            // A sorted, consolidated set over a small space, so groups collide and
862            // a `(key, val)` carries several times.
863            let n = rng() as usize % 60 + 1;
864            let mut rows: Vec<((u64, u64), u64, i64)> = (0..n).map(|_| {
865                let k = rng() % 8; let v = rng() % 3; let t = rng() % 6;
866                let d = if rng() % 4 == 0 { -1 } else { 1 };
867                ((k, v), t, d)
868            }).collect();
869            consolidate_updates(&mut rows);
870            if rows.is_empty() { continue; }
871            let f = rng() % 6;
872            let frontier = Antichain::from_elem(f);
873            let sz = rng() as usize % 5 + 1; // tiny chunks → heavy straddling
874
875            // Drive resumably: push one chunk, advance(false), … then advance(true).
876            let mut q = VecDeque::new();
877            let mut out = VecDeque::new();
878            for c in rows.chunks(sz) {
879                q.push_back(chunk(c.iter().map(|&((k, v), t, d)| (k, v, t, d)).collect()));
880                ColChunk::advance(&mut q, frontier.borrow(), false, &mut out);
881            }
882            ColChunk::advance(&mut q, frontier.borrow(), true, &mut out);
883            let got = flat(out);
884
885            // Oracle: advance each time (u64 total order → `max(t, f)`), consolidate.
886            let mut want: Vec<((u64, u64), u64, i64)> =
887                rows.iter().map(|&((k, v), t, d)| ((k, v), t.max(f), d)).collect();
888            consolidate_updates(&mut want);
889            let want: Vec<Upd> = want.into_iter().map(|((k, v), t, d)| (k, v, t, d)).collect();
890
891            assert_eq!(got, want, "frontier {f}, chunk size {sz}, rows {rows:?}");
892        }
893    }
894}