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