Skip to main content

mz_timely_util/columnar/
chunk.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! [`ColumnChunk`]: differential's [`Chunk`] over [`ColumnBody`] updates.
11//!
12//! A chunk is a sorted, consolidated run of `(D, T, R)` updates in the flat
13//! columnar layout, in one of two homes:
14//!
15//! * **Resident**: an `Rc`-shared [`ColumnBody`] on the heap. Fresh input,
16//!   merge output, and small tails live here.
17//! * **Spilled**: the serialized body in the process [`Pool`], with the record
18//!   count and the first and last data items resident. The pool owns residency
19//!   from there, with slots under a memory budget and compression and device
20//!   pageout under pressure, and a body that dies before pressure reaches it
21//!   is freed without I/O.
22//!
23//! Reads of a spilled body are copy-out and scoped to the call that needs
24//! them, the contract that lets the pool evict with no reader accounting
25//! (see [`mz_ore::pool`]).
26//!
27//! Spilling happens in [`Chunk::settle`], the trait's designated commit point:
28//! chunks moved to settled output are handed to the pool while the chunk
29//! type's [`SpillGate`] is open. Grading is by serialized
30//! bytes, the ship size [`Column`] already targets, rather than by the
31//! record-count `TARGET`, since record count does not bound bytes for
32//! variable-width data.
33//!
34//! Chunks whose data is a `(key, val)` pair additionally implement
35//! [`UnloadChunk`], the bulk-read capability: sorted probe keys in, matching
36//! updates appended to caller-owned staging, with `locate` answered from the
37//! resident fence metadata so a probe set faults only the chunk bodies it
38//! actually touches.
39
40use std::borrow::Cow;
41#[cfg(test)]
42use std::cell::Cell;
43use std::cell::RefCell;
44use std::collections::VecDeque;
45use std::marker::PhantomData;
46use std::rc::Rc;
47use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
48use std::sync::{Arc, OnceLock};
49
50use columnar::bytes::indexed;
51use columnar::{Borrow, BorrowedOf, Columnar, Container as _, Index, Len, Push as _};
52use differential_dataflow::difference::Semigroup;
53use differential_dataflow::lattice::Lattice;
54use differential_dataflow::trace::chunk::Chunk;
55use mz_ore::cast::CastFrom;
56use mz_ore::pool::{ChunkHandle, ChunkHints, ExtentCodec, IDENTITY_CODEC, Pool, TypedBody};
57use smallvec::SmallVec;
58use timely::Accountable;
59use timely::PartialOrder;
60use timely::container::{ContainerBuilder, PushInto};
61use timely::progress::Timestamp;
62use timely::progress::frontier::{Antichain, AntichainRef};
63
64use crate::columnar::batcher::{ColumnChunker, gallop};
65use crate::columnar::body::{ColumnBody, borrow_words};
66use crate::columnar::unload::UnloadChunk;
67use crate::columnar::{Column, at_serialized_capacity};
68
69pub mod metrics;
70
71/// The gate of [`ComputeSpill`] chunks. See [`set_compute_spill_enabled`].
72static COMPUTE_SPILL_ENABLED: AtomicBool = AtomicBool::new(false);
73
74/// The gate of [`StorageSpill`] chunks. See [`set_storage_spill_enabled`].
75static STORAGE_SPILL_ENABLED: AtomicBool = AtomicBool::new(false);
76
77/// The gate for bodies spilled through [`try_spill_ref`]. See
78/// [`set_sink_spill_enabled`].
79static SINK_SPILL_ENABLED: AtomicBool = AtomicBool::new(false);
80
81thread_local! {
82    /// A thread-scoped pool override, taking precedence over the global
83    /// enable flag and pool. Lets tests and benches spill through a private
84    /// pool without touching process-global state.
85    static SPILL_OVERRIDE: RefCell<Option<Pool>> = const { RefCell::new(None) };
86
87    /// A thread-scoped depth-floor override, taking precedence over the
88    /// global value. Lets tests pin the floor without racing concurrently
89    /// running tests on the process-global state.
90    #[cfg(test)]
91    static COMPRESS_MIN_DEPTH_OVERRIDE: Cell<Option<u8>> = const { Cell::new(None) };
92
93    /// Reusable staging for call-scoped reads of spilled bodies.
94    static READ_SCRATCH: RefCell<Vec<u64>> = const { RefCell::new(Vec::new()) };
95}
96
97/// Which process-wide gate decides whether a [`ColumnChunk`] spills.
98///
99/// Compute and storage dataflows share a process and its worker threads, so
100/// the gate is part of the chunk type: each consumer names the gate of the
101/// subsystem that owns it, and one subsystem's flag cannot spill another's
102/// chunks. Every gate takes effect at the next commit, and only once
103/// `apply_pool_config` has installed the pool. With no pool installed chunks
104/// stay resident whatever the gate says.
105pub trait SpillGate: 'static {
106    /// Whether chunks of this gate spill at commit.
107    fn enabled() -> bool;
108}
109
110/// The gate of compute's chunk batchers. See [`set_compute_spill_enabled`].
111pub enum ComputeSpill {}
112
113impl SpillGate for ComputeSpill {
114    fn enabled() -> bool {
115        COMPUTE_SPILL_ENABLED.load(Ordering::Relaxed)
116    }
117}
118
119/// The gate of storage's upsert chunk batchers. See
120/// [`set_storage_spill_enabled`].
121pub enum StorageSpill {}
122
123impl SpillGate for StorageSpill {
124    fn enabled() -> bool {
125        STORAGE_SPILL_ENABLED.load(Ordering::Relaxed)
126    }
127}
128
129/// Enable or disable spilling of [`ComputeSpill`] chunks. Already-spilled
130/// chunks are unaffected either way.
131pub fn set_compute_spill_enabled(enabled: bool) {
132    COMPUTE_SPILL_ENABLED.store(enabled, Ordering::Relaxed);
133}
134
135/// Enable or disable spilling of [`StorageSpill`] chunks. Already-spilled
136/// chunks are unaffected either way.
137pub fn set_storage_spill_enabled(enabled: bool) {
138    STORAGE_SPILL_ENABLED.store(enabled, Ordering::Relaxed);
139}
140
141/// Enable or disable spilling of bodies offered through [`try_spill_ref`],
142/// which is how compute's MV sink correction buffer spills. Like the
143/// [`SpillGate`]s, it takes effect only once `apply_pool_config` has
144/// installed the pool.
145pub fn set_sink_spill_enabled(enabled: bool) {
146    SINK_SPILL_ENABLED.store(enabled, Ordering::Relaxed);
147}
148
149/// Set or unset the pool through which this thread's chunk spills are
150/// routed, taking precedence over the gates and the process pool. `None`
151/// restores the global resolution.
152pub fn set_spill_override(pool: Option<Pool>) {
153    SPILL_OVERRIDE.with(|cell| *cell.borrow_mut() = pool);
154}
155
156/// The youngest generational depth whose spilled bodies are compressed. See
157/// [`set_compress_min_depth`].
158static COMPRESS_MIN_DEPTH: AtomicU8 = AtomicU8::new(DEFAULT_COMPRESS_MIN_DEPTH);
159
160/// Set the youngest generational depth whose spilled bodies are compressed.
161///
162/// A chunk at depth `d` is rewritten (merged, extracted, advanced) with
163/// frequency proportional to `2^-d` under geometric merging, so compressing
164/// a shallow chunk buys a short stay in the pool at the cost of a guaranteed
165/// near-term codec round-trip: the body is encoded only to be read back and
166/// decoded by the next rewrite. Generations below the floor spill under the
167/// identity codec instead: still budgeted and swap-backed like every extent,
168/// but encode and decode are copies. The floor never exempts a body from the
169/// pool, so it cannot grow unbudgeted resident state.
170///
171/// The floor cannot strand a long-lived body uncompressed: a chunk that a
172/// merge carries forward untouched also ages a generation, and its body is
173/// re-spilled under the compressing codec when it crosses the floor (see
174/// `survive_merge`).
175///
176/// `0` compresses every spilled body. Consulted at every commit, so changes
177/// apply to running dataflows.
178pub fn set_compress_min_depth(depth: u8) {
179    COMPRESS_MIN_DEPTH.store(depth, Ordering::Relaxed);
180}
181
182/// Set or unset a thread-scoped depth-floor override, taking precedence over
183/// [`set_compress_min_depth`]. Tests run concurrently and must not race on
184/// the process-global floor.
185#[cfg(test)]
186pub fn set_compress_min_depth_override(depth: Option<u8>) {
187    COMPRESS_MIN_DEPTH_OVERRIDE.with(|cell| cell.set(depth));
188}
189
190/// The depth floor in effect for this thread's commits.
191fn compress_min_depth() -> u8 {
192    #[cfg(test)]
193    if let Some(depth) = COMPRESS_MIN_DEPTH_OVERRIDE.with(|cell| cell.get()) {
194        return depth;
195    }
196    COMPRESS_MIN_DEPTH.load(Ordering::Relaxed)
197}
198
199/// The codec a body at `depth` stores under, identity below the compression
200/// floor and lz4 at and past it, paired with whether that codec compresses.
201/// One read of the floor, so the pair cannot disagree with itself when the
202/// floor moves under a concurrent commit.
203fn codec_for_depth(depth: u8) -> (&'static dyn ExtentCodec, bool) {
204    if depth < compress_min_depth() {
205        (&IDENTITY_CODEC, false)
206    } else {
207        (&LZ4_CODEC, true)
208    }
209}
210
211/// The pool a body of `len_bytes` spills into while `enabled`, or `None` when
212/// it stays resident.
213///
214/// The one place the spill decision lives: the thread override, the gate, the
215/// installed pool, and the size floor.
216fn spill_target(enabled: bool, len_bytes: usize) -> Option<Pool> {
217    resolve_pool(enabled).filter(|_| len_bytes >= SPILL_MIN_BYTES)
218}
219
220/// The thread's override pool, else the installed pool while `enabled`.
221fn resolve_pool(enabled: bool) -> Option<Pool> {
222    if let Some(pool) = SPILL_OVERRIDE.with(|cell| cell.borrow().clone()) {
223        return Some(pool);
224    }
225    if enabled {
226        crate::pool_config::active_pool()
227    } else {
228        None
229    }
230}
231
232/// Scratch capacity retained across reads, in words. A read larger than this
233/// releases the buffer afterward, so a thread's scratch does not ratchet to
234/// the largest body it ever carried (heap no pool gauge can see).
235const SCRATCH_RETAIN_WORDS: usize = 1 << 18;
236
237/// Run `f` with this thread's read scratch, cleared of any previous use.
238fn with_scratch<Out>(f: impl FnOnce(&mut Vec<u64>) -> Out) -> Out {
239    READ_SCRATCH.with(|cell| {
240        let mut scratch = cell.take();
241        scratch.clear();
242        let out = f(&mut scratch);
243        if scratch.capacity() > SCRATCH_RETAIN_WORDS {
244            scratch.clear();
245            scratch.shrink_to_fit();
246        }
247        cell.replace(scratch);
248        out
249    })
250}
251
252/// The serialized-byte size committed chunks aim for, matching the ship size
253/// of the columnar merge machinery.
254const COMMIT_BYTES: usize = 2 << 20;
255
256/// Bodies smaller than this stay resident: the pool's smallest size class is
257/// 64 KiB, so spilling below it trades no meaningful memory for slot waste.
258///
259/// Sub-floor bodies are invisible to the pool's budget, which is safe only
260/// while they are rare. `settle` coalesces toward `COMMIT_BYTES` before
261/// committing, so in the harness only a final `done` tail commits below the
262/// floor. A caller that commits many small chunks directly accumulates
263/// unbudgeted heap, and no accounting here would catch it.
264const SPILL_MIN_BYTES: usize = 64 << 10;
265
266/// The default compression depth floor: fresh (depth 0) bodies spill
267/// uncompressed.
268///
269/// A fresh chunk is consumed by its first merge with certainty, so
270/// compressing it can never save pool bytes for longer than one merge
271/// cadence and always costs a full encode plus decode. Depth 1 and beyond
272/// have survived a merge and wait geometrically longer for the next, so
273/// their compression amortizes.
274const DEFAULT_COMPRESS_MIN_DEPTH: u8 = 1;
275
276/// Records of a body with `len` records in `bytes` that fit in `space`
277/// bytes at the body's average width, keeping 5% of `COMMIT_BYTES` as a
278/// margin for per-piece framing. Uniform rows cut this way fill their size
279/// class and leave one short remainder that `settle` can coalesce. Uneven
280/// widths can still overflow, so callers check the cut piece.
281fn cut_records(len: usize, bytes: usize, space: usize) -> usize {
282    let space = space.saturating_sub(COMMIT_BYTES / 20);
283    (len.saturating_mul(space) / bytes.max(1)).min(len)
284}
285
286/// Whether a body is big enough to commit on its own.
287fn at_commit_size<C: Columnar>(body: &ColumnBody<C>) -> bool {
288    body.length_in_bytes() >= COMMIT_BYTES - COMMIT_BYTES / 10
289}
290
291/// Narrow a columnar ref to a shorter lifetime, so refs from different
292/// borrows, such as a probe column and a chunk's own columns, can be compared
293/// (the refs are lifetime-invariant).
294#[inline(always)]
295fn rr<'b, 'a: 'b, C: Columnar>(item: columnar::Ref<'a, C>) -> columnar::Ref<'b, C> {
296    columnar::ContainerOf::<C>::reborrow_ref(item)
297}
298
299/// A spilled chunk body: the serialized column in the pool, plus the resident
300/// metadata every [`Chunk`] must answer without fetching. That metadata is
301/// the record count, the first and last data items (the fence entries
302/// [`UnloadChunk::locate`] consults), and the time bounds `extract` consults
303/// to pass frontier-disjoint chunks through without loading them.
304pub struct SpilledBody<D: Columnar, T, G> {
305    /// Number of updates in the body.
306    records: usize,
307    /// The first and last data items, as a two-element container. One
308    /// container rather than two singletons, so the leaf allocations are not
309    /// duplicated per fence.
310    fences: D::Container,
311    /// The minimal and maximal times in the body, shared with the pooled
312    /// body and filled on first use (see [`PooledBody`]). The minimal times
313    /// fold into `extract`'s residual frontier when the chunk is kept whole.
314    /// Some contained time is greater-or-equal to a frontier exactly when
315    /// some maximal time is, which is `extract`'s ship-whole test.
316    time_bounds: Arc<OnceLock<TimeBounds<T>>>,
317    /// Whether the body was inserted under the compressing codec. The pool
318    /// stores the codec itself and reads decode through it, so this is the
319    /// only handle chunk code has on what a body is stored as, and it is
320    /// what `survive_merge` consults to decide a body wants migrating.
321    /// Deriving that from depth instead would tie it to a single transition
322    /// and miss every path that skips it.
323    compressed: bool,
324    /// The body's serialized size, as it was before the pool's codec saw
325    /// it. Retained because the pool reports no per-chunk figure, and a
326    /// chunk that reported nothing would drop its operator's share of the
327    /// batcher's memory out of the introspection tables.
328    len_bytes: usize,
329    /// The pool chunk holding the serialized column.
330    handle: ChunkHandle,
331    /// Carries the chunk's [`SpillGate`], which the resident variant has no
332    /// field to hold.
333    gate: PhantomData<G>,
334}
335
336/// A sorted, consolidated run of `(D, T, R)` updates, resident or spilled.
337///
338/// Every chunk carries a generational depth counting the merge cadences it
339/// has lived through: fresh chunks are depth 0, a merge output is one
340/// generation past its deepest input (saturating at `u8::MAX`, where
341/// remerged long-lived chunks stay), a chunk a merge carries forward
342/// untouched also gains a generation (see `survive_merge`), and rewrites
343/// within a generation (extract, advance, settle coalescing) preserve
344/// depth.
345///
346/// Depth belongs to the chunk, not to the body: a body outlives the chunks
347/// that share it, and aging must not depend on whether a caller happens to
348/// hold the only reference. At spill time the depth becomes the pool's
349/// [`ChunkHints`], so repeatedly merged (older, colder) data lands in deeper
350/// eviction bands. Hints are fixed at insert, so a chunk aged without a
351/// re-spill keeps the band it spilled into.
352///
353/// `G` is the [`SpillGate`] that decides whether commits spill.
354pub enum ColumnChunk<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> {
355    /// Body on the heap, shared via `Rc`, with its generational depth.
356    Resident(Rc<ColumnBody<(D, T, R)>>, u8),
357    /// Body in the pool, with its generational depth. See [`SpilledBody`].
358    Spilled(Rc<SpilledBody<D, T, G>>, u8),
359}
360
361impl<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> Clone for ColumnChunk<D, T, R, G> {
362    fn clone(&self) -> Self {
363        match self {
364            ColumnChunk::Resident(col, depth) => ColumnChunk::Resident(Rc::clone(col), *depth),
365            ColumnChunk::Spilled(body, depth) => ColumnChunk::Spilled(Rc::clone(body), *depth),
366        }
367    }
368}
369
370impl<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> Default for ColumnChunk<D, T, R, G> {
371    fn default() -> Self {
372        ColumnChunk::Resident(Rc::new(ColumnBody::default()), 0)
373    }
374}
375
376impl<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> Accountable for ColumnChunk<D, T, R, G> {
377    fn record_count(&self) -> i64 {
378        i64::try_from(self.records()).expect("record count fits i64")
379    }
380}
381
382impl<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> ColumnChunk<D, T, R, G> {
383    /// Wrap a sorted, consolidated, non-empty body as a resident chunk of
384    /// the youngest generation.
385    pub fn from_body(body: ColumnBody<(D, T, R)>) -> Self {
386        mz_ore::soft_assert_no_log!(!body.is_empty(), "chunks must be non-empty");
387        ColumnChunk::Resident(Rc::new(body), 0)
388    }
389
390    /// Append pieces of at most `COMMIT_BYTES`, preserving order and
391    /// generational depth. A single update may exceed the bound because it
392    /// cannot be split.
393    fn push_bounded(body: ColumnBody<(D, T, R)>, depth: u8, out: &mut VecDeque<Self>) {
394        let len = body.len();
395        let bytes = body.length_in_bytes();
396        if len <= 1 || bytes <= COMMIT_BYTES {
397            if len > 0 {
398                out.push_back(Self::Resident(Rc::new(body), depth));
399            }
400            return;
401        }
402        Self::push_cuts(
403            &body,
404            cut_records(len, bytes, COMMIT_BYTES).max(1),
405            depth,
406            out,
407        );
408    }
409
410    /// Append `body` in consecutive pieces of `records` records. A piece
411    /// that still exceeds `COMMIT_BYTES`, from uneven widths, is halved until
412    /// it fits or holds one record. Halving bounds the recursion depth by
413    /// log2 of `records`, where cutting again at the piece's own average
414    /// width can shrink a piece holding one wide record by only a few
415    /// percent per level.
416    fn push_cuts(
417        body: &ColumnBody<(D, T, R)>,
418        records: usize,
419        depth: u8,
420        out: &mut VecDeque<Self>,
421    ) {
422        let view = body.borrow();
423        Self::push_range(view, 0..view.len(), records, depth, out);
424    }
425
426    /// [`Self::push_cuts`] over `range` of `view`. An oversized piece is
427    /// dropped before descending, and the halves are cut again from `view`,
428    /// so transient memory is one piece rather than one copy of a wide
429    /// update per level.
430    fn push_range(
431        view: BorrowedOf<'_, (D, T, R)>,
432        range: std::ops::Range<usize>,
433        records: usize,
434        depth: u8,
435        out: &mut VecDeque<Self>,
436    ) {
437        let mut start = range.start;
438        while start < range.end {
439            let end = range.end.min(start + records);
440            let mut part = <(D, T, R) as Columnar>::Container::default();
441            part.extend_from_self(view, start..end);
442            let part = ColumnBody::Typed(part);
443            if end - start > 1 && part.length_in_bytes() > COMMIT_BYTES {
444                drop(part);
445                Self::push_range(view, start..end, (end - start).div_ceil(2), depth, out);
446            } else {
447                out.push_back(Self::Resident(Rc::new(part), depth));
448            }
449            start = end;
450        }
451    }
452
453    /// The body, owned. A spilled body is copied out of the pool within this
454    /// call. A shared resident body is copied in its current form.
455    pub fn into_body(self) -> ColumnBody<(D, T, R)> {
456        match self {
457            ColumnChunk::Resident(col, _) => {
458                Rc::try_unwrap(col).unwrap_or_else(|shared| shared.duplicate())
459            }
460            ColumnChunk::Spilled(body, _) => match Rc::try_unwrap(body) {
461                Ok(spilled) => match spilled.handle.take_typed() {
462                    Ok(typed) => {
463                        typed
464                            .into_any()
465                            .downcast::<PooledBody<D, T, R>>()
466                            .expect("chunk bodies are pooled as `PooledBody`")
467                            .body
468                    }
469                    Err(handle) => read_words(&handle),
470                },
471                // A shared body is duplicated in its current form, as a
472                // shared resident body is.
473                Err(shared) => shared.handle.with_typed(|typed| match typed {
474                    Some(typed) => pooled::<D, T, R>(typed).body.duplicate(),
475                    None => read_words(&shared.handle),
476                }),
477            },
478        }
479    }
480
481    /// The body for the duration of `f`: a resident body is borrowed, a
482    /// spilled body is loaded from the pool for the call and dropped after it.
483    pub fn with_body<F, X>(&self, f: F) -> X
484    where
485        F: FnOnce(&ColumnBody<(D, T, R)>) -> X,
486    {
487        match self {
488            ColumnChunk::Resident(col, _) => f(col),
489            ColumnChunk::Spilled(body, _) => {
490                let mut f = Some(f);
491                let borrowed = body.handle.with_typed(|typed| {
492                    typed.map(|typed| f.take().expect("unused")(&pooled::<D, T, R>(typed).body))
493                });
494                match borrowed {
495                    Some(out) => out,
496                    None => f.take().expect("unused")(&read_words(&body.handle)),
497                }
498            }
499        }
500    }
501
502    /// True when the body lives in the pool.
503    pub fn is_spilled(&self) -> bool {
504        matches!(self, ColumnChunk::Spilled(_, _))
505    }
506
507    /// The number of updates, from resident state only.
508    fn records(&self) -> usize {
509        match self {
510            ColumnChunk::Resident(col, _) => col.len(),
511            ColumnChunk::Spilled(body, _) => body.records,
512        }
513    }
514
515    /// The generational depth, from resident state only.
516    fn depth(&self) -> u8 {
517        match self {
518            ColumnChunk::Resident(_, depth) | ColumnChunk::Spilled(_, depth) => *depth,
519        }
520    }
521
522    /// The first and last data items, from resident state only.
523    fn data_span(&self) -> (columnar::Ref<'_, D>, columnar::Ref<'_, D>) {
524        match self {
525            ColumnChunk::Resident(col, _) => {
526                let data = col.borrow().0;
527                (data.get(0), data.get(data.len() - 1))
528            }
529            ColumnChunk::Spilled(body, _) => {
530                let fences = body.fences.borrow();
531                (fences.get(0), fences.get(1))
532            }
533        }
534    }
535
536    /// Commit a non-empty body at the given generational depth: spill it to
537    /// the pool when `G` is open and the body is worth a slot, else keep it
538    /// resident.
539    fn commit(body: ColumnBody<(D, T, R)>, depth: u8) -> Self
540    where
541        D: PoolColumnar,
542        T: Timestamp + PoolColumnar,
543        R: PoolColumnar,
544    {
545        let len_bytes = body.length_in_bytes();
546        metrics::record(metrics::Stage::Commit, body.len(), len_bytes);
547        mz_ore::soft_assert_no_log!(!body.is_empty(), "chunks must be non-empty");
548        match spill_target(G::enabled(), len_bytes) {
549            Some(pool) => Self::spill_body(body, &pool, depth),
550            None => ColumnChunk::Resident(Rc::new(body), depth),
551        }
552    }
553
554    /// Spill a non-empty body into `pool` unconditionally, capturing the
555    /// resident fence metadata.
556    ///
557    /// Generations below the compression depth floor store under the
558    /// identity codec: rewritten too soon for compression to amortize, they
559    /// stay budgeted and swap-backed while encode and decode reduce to
560    /// copies.
561    fn spill_body(body: ColumnBody<(D, T, R)>, pool: &Pool, depth: u8) -> Self
562    where
563        D: PoolColumnar,
564        T: Timestamp + PoolColumnar,
565        R: PoolColumnar,
566    {
567        let (codec, compressed) = codec_for_depth(depth);
568        let len_bytes = body.length_in_bytes();
569        let view = body.borrow();
570        let records = view.len();
571        let mut fences = D::Container::default();
572        fences.push(view.0.get(0));
573        fences.push(view.0.get(records - 1));
574        let time_bounds = Arc::new(OnceLock::new());
575        let pooled = PooledBody {
576            body,
577            bounds: Arc::clone(&time_bounds),
578            compute: |body| {
579                let (lower, upper) = Self::time_bounds(body);
580                (lower, upper.into())
581            },
582        };
583        let handle = pool.insert_typed(Box::new(pooled), ChunkHints { depth }, codec);
584        ColumnChunk::Spilled(
585            Rc::new(SpilledBody {
586                records,
587                fences,
588                time_bounds,
589                compressed,
590                len_bytes,
591                handle,
592                gate: PhantomData,
593            }),
594            depth,
595        )
596    }
597
598    /// Age a chunk that a merge carried forward untouched by one generation.
599    /// Depth counts merge cadences lived through, not rewrites, so a
600    /// pass-through survivor ages like a merged chunk; the bump rides on the
601    /// chunk, so it is free whether or not the body is shared.
602    ///
603    /// An identity-coded body at or past the floor wants migrating, so a
604    /// sole owner re-spills it compressed; without the re-spill,
605    /// key-disjoint input would keep its whole spilled backlog
606    /// identity-coded for as long as it lived. The test is the body's stored
607    /// codec against the floor, not a depth transition, so a migration that
608    /// cannot happen now (a shared body, no pool installed, a floor lowered
609    /// long after the spill) is retried at the next survival rather than
610    /// consumed. A shared body is skipped because re-spilling this reference
611    /// cannot change what the other holder stores, and the compaction merger
612    /// that shares bodies rewrites its clones immediately.
613    fn survive_merge(self) -> Self
614    where
615        D: PoolColumnar,
616        T: Timestamp + PoolColumnar,
617        R: PoolColumnar,
618    {
619        let depth = self.depth().saturating_add(1);
620        match self {
621            ColumnChunk::Resident(col, _) => ColumnChunk::Resident(col, depth),
622            ColumnChunk::Spilled(body, was) => {
623                let migrate = !body.compressed && depth >= compress_min_depth();
624                if !migrate || Rc::strong_count(&body) > 1 {
625                    return ColumnChunk::Spilled(body, depth);
626                }
627                match resolve_pool(G::enabled()) {
628                    Some(pool) => {
629                        let body = ColumnChunk::Spilled(body, was).into_body();
630                        Self::spill_body(body, &pool, depth)
631                    }
632                    None => ColumnChunk::Spilled(body, depth),
633                }
634            }
635        }
636    }
637
638    /// The chunk's time bounds: borrowed from the stored metadata for
639    /// spilled bodies, computed by a time-column scan for resident ones. The
640    /// scan costs less than the copy it lets `extract` avoid when the chunk
641    /// passes through whole.
642    fn chunk_time_bounds(&self) -> (Cow<'_, Antichain<T>>, Cow<'_, [T]>)
643    where
644        T: Timestamp,
645    {
646        match self {
647            ColumnChunk::Resident(col, _) => {
648                let (lower, upper) = Self::time_bounds(col);
649                (Cow::Owned(lower), Cow::Owned(upper))
650            }
651            ColumnChunk::Spilled(body, _) => {
652                // Eviction serializes the body under the chunk's state lock
653                // and fills the bounds cell from there, so the lock order is
654                // state lock, then cell. Taking the state lock from inside
655                // the cell's initializer deadlocks against an evictor.
656                if body.time_bounds.get().is_none() {
657                    let filled = body.handle.with_typed(|typed| {
658                        typed
659                            .map(|typed| pooled::<D, T, R>(typed).bounds())
660                            .is_some()
661                    });
662                    // Leaving the typed form fills the bounds, so this is
663                    // reached only by bodies spilled some other way.
664                    if !filled {
665                        let (lower, upper) = Self::time_bounds(&read_words(&body.handle));
666                        let _ = body.time_bounds.set((lower, upper.into()));
667                    }
668                }
669                let (lower, upper) = body.time_bounds.get().expect("time bounds filled above");
670                (Cow::Borrowed(lower), Cow::Borrowed(&upper[..]))
671            }
672        }
673    }
674
675    /// The time bounds of a non-empty body: the antichain of minimal times
676    /// (every contained time is greater-or-equal to some element) and the
677    /// set of maximal times (some contained time is greater-or-equal to a
678    /// frontier exactly when some maximal one is).
679    fn time_bounds(body: &ColumnBody<(D, T, R)>) -> (Antichain<T>, Vec<T>)
680    where
681        T: Timestamp,
682    {
683        let (_, times, _) = body.borrow();
684        let mut lower = Antichain::new();
685        let mut upper: Vec<T> = Vec::new();
686        // One owned time reused across the scan, so times with owned
687        // allocations do not allocate per element; the bound sets clone only
688        // the elements they retain.
689        let mut time = T::minimum();
690        for i in 0..times.len() {
691            time.copy_from(rr::<T>(times.get(i)));
692            if !upper.iter().any(|u| PartialOrder::less_equal(&time, u)) {
693                upper.retain(|u| !PartialOrder::less_equal(u, &time));
694                upper.push(time.clone());
695            }
696            lower.insert_ref(&time);
697        }
698        (lower, upper)
699    }
700}
701
702/// The chunk-side [`ExtentCodec`]: a little-endian `u32` body-length prefix
703/// followed by one lz4 block, the framing
704/// `lz4_flex::block::compress_prepend_size` produces. Every chunk consumer
705/// passes [`LZ4_CODEC`] at insert; the pool itself has no codec opinion.
706#[derive(Debug)]
707pub struct Lz4Codec;
708
709/// The [`Lz4Codec`] instance chunk consumers pass to
710/// [`Pool::insert_with`].
711pub static LZ4_CODEC: Lz4Codec = Lz4Codec;
712
713impl ExtentCodec for Lz4Codec {
714    fn encode(&self, body: &[u8], out: &mut Vec<u8>) {
715        let max_out = lz4_flex::block::get_maximum_output_size(body.len());
716        out.resize(4 + max_out, 0);
717        let len = u32::try_from(body.len()).expect("chunk bodies are bounded by the size classes");
718        out[..4].copy_from_slice(&len.to_le_bytes());
719        let compressed = lz4_flex::block::compress_into(body, &mut out[4..])
720            .expect("output sized to the maximum");
721        out.truncate(4 + compressed);
722    }
723
724    fn decode(&self, stored: &[u8], body: &mut [u8]) {
725        let prefix: [u8; 4] = stored[..4].try_into().expect("prefix length");
726        let len = usize::try_from(u32::from_le_bytes(prefix)).expect("length fits usize");
727        assert_eq!(
728            len,
729            body.len(),
730            "destination must match the encoded body length"
731        );
732        let written = lz4_flex::block::decompress_into(&stored[4..], body)
733            .expect("stored bytes hold a valid lz4 block");
734        assert_eq!(written, body.len(), "decoded length mismatch");
735    }
736}
737
738/// A spilled chunk's time bounds: the antichain of minimal times and the set
739/// of maximal times (see `ColumnChunk::time_bounds`).
740type TimeBounds<T> = (Antichain<T>, SmallVec<[T; 1]>);
741
742/// A columnar type a spilled chunk can hold in its typed form. Spill threads serialize and drop
743/// typed bodies, so the containers must be `Send`. The pooled body shares its time bounds with
744/// the chunk through an `Arc`, so the values must also be `Sync`.
745pub trait PoolColumnar: Columnar<Container: Send> + Send + Sync {}
746
747impl<C: Columnar<Container: Send> + Send + Sync> PoolColumnar for C {}
748
749/// A [`ColumnBody`] as the pool holds it: in its typed form until the pool
750/// has to serialize it.
751///
752/// The time bounds are computed on first use and shared with the chunk's
753/// [`SpilledBody`]. Serialization fills them before the typed form can be
754/// dropped, so a chunk that has left its typed form never needs a decode to
755/// answer `extract`'s pass-through test.
756struct PooledBody<D: Columnar, T: Columnar, R: Columnar> {
757    body: ColumnBody<(D, T, R)>,
758    bounds: Arc<OnceLock<TimeBounds<T>>>,
759    compute: fn(&ColumnBody<(D, T, R)>) -> TimeBounds<T>,
760}
761
762impl<D: Columnar, T: Columnar, R: Columnar> PooledBody<D, T, R> {
763    fn bounds(&self) -> &TimeBounds<T> {
764        self.bounds.get_or_init(|| (self.compute)(&self.body))
765    }
766}
767
768impl<D: Columnar, T: Columnar, R: Columnar> std::fmt::Debug for PooledBody<D, T, R> {
769    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
770        f.debug_struct("PooledBody")
771            .field("len_bytes", &self.body.length_in_bytes())
772            .finish()
773    }
774}
775
776impl<D: PoolColumnar + 'static, T: PoolColumnar + 'static, R: PoolColumnar + 'static> TypedBody
777    for PooledBody<D, T, R>
778{
779    fn words(&self) -> usize {
780        self.body.length_in_bytes() / 8
781    }
782
783    fn serialize_into(&self, dst: &mut [u64]) {
784        self.bounds();
785        let len_bytes = self.body.length_in_bytes();
786        let bytes: &mut [u8] = bytemuck::cast_slice_mut(dst);
787        let mut cursor = std::io::Cursor::new(bytes);
788        self.body
789            .write_into(&mut cursor)
790            .expect("the destination is sized to the serialized body");
791        assert_eq!(
792            usize::try_from(cursor.position()).expect("usize position"),
793            len_bytes,
794            "serialized body must fill the destination exactly",
795        );
796    }
797
798    fn as_any(&self) -> &dyn std::any::Any {
799        self
800    }
801
802    fn into_any(self: Box<Self>) -> Box<dyn std::any::Any> {
803        self
804    }
805}
806
807/// A copy of a spilled body, in serialized form.
808fn read_words<C: Columnar>(handle: &ChunkHandle) -> ColumnBody<C> {
809    let mut words = Vec::new();
810    handle.read_into(&mut words);
811    ColumnBody::Words(words)
812}
813
814/// The [`PooledBody`] a chunk's typed body is.
815fn pooled<D: Columnar + 'static, T: Columnar + 'static, R: Columnar + 'static>(
816    typed: &dyn TypedBody,
817) -> &PooledBody<D, T, R> {
818    typed
819        .as_any()
820        .downcast_ref::<PooledBody<D, T, R>>()
821        .expect("chunk bodies are pooled as `PooledBody`")
822}
823
824/// Serialize a body into a pool slot, writing its serialized form through a cursor over the
825/// slot memory. A serialized body's encoding is one copy. Sizing is exact, so a short or
826/// overlong write is a contract violation and panics.
827fn spill_serialized<C: Columnar>(
828    body: &ColumnBody<C>,
829    pool: &Pool,
830    len_bytes: usize,
831    hints: ChunkHints,
832    codec: &'static dyn ExtentCodec,
833) -> ChunkHandle {
834    mz_ore::soft_assert_eq_no_log!(len_bytes % 8, 0);
835    pool.insert_with(len_bytes / 8, hints, codec, |dst| {
836        let bytes: &mut [u8] = bytemuck::cast_slice_mut(dst);
837        let mut cursor = std::io::Cursor::new(bytes);
838        body.write_into(&mut cursor)
839            .expect("the slot is sized to the serialized body");
840        assert_eq!(
841            usize::try_from(cursor.position()).expect("usize position"),
842            len_bytes,
843            "serialized body must fill the chunk exactly",
844        );
845    })
846}
847
848/// Spill a serialized copy of `body` into the process pool, leaving `body` untouched, or
849/// `None` when the body stays resident.
850///
851/// A body stays resident when the gate [`set_sink_spill_enabled`] sets is off, when no pool is
852/// installed, or when it is smaller than `SPILL_MIN_BYTES`. `depth` is the body's generational
853/// depth: it selects the stored codec and the pool's eviction band.
854///
855/// For consumers that keep their own chunk representation and so cannot use [`ColumnChunk`];
856/// the two share the spill decision through `spill_target`. The encode writes straight into
857/// the pool slot, so a spilled body costs no intermediate allocation and the caller can
858/// [`ColumnBody::clear`] and keep its allocation. A caller that gets `None` back still owns
859/// the body and must keep it resident itself.
860pub fn try_spill_ref<C: Columnar>(body: &ColumnBody<C>, depth: u8) -> Option<ChunkHandle> {
861    let len_bytes = body.length_in_bytes();
862    let pool = spill_target(SINK_SPILL_ENABLED.load(Ordering::Relaxed), len_bytes)?;
863    let (codec, _) = codec_for_depth(depth);
864    Some(spill_serialized(
865        body,
866        &pool,
867        len_bytes,
868        ChunkHints { depth },
869        codec,
870    ))
871}
872
873impl<D, T, R, G> Chunk for ColumnChunk<D, T, R, G>
874where
875    G: SpillGate,
876    D: PoolColumnar,
877    for<'a> columnar::Ref<'a, D>: Copy + Ord,
878    T: PoolColumnar + Default + Timestamp + Lattice + Ord,
879    for<'a> columnar::Ref<'a, T>: Copy + Ord,
880    R: PoolColumnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
881{
882    type Time = T;
883
884    /// A nominal record count for the harness's fuel and ladder accounting,
885    /// not a bound. Actual chunk sizing is by serialized bytes: `merge` and
886    /// `extract` cut output at the [`Column`] ship threshold, and `settle`
887    /// grades by `at_commit_size`, so a chunk of narrow records can hold more
888    /// records than this and nothing here consults it.
889    const TARGET: usize = 65536;
890
891    fn len(&self) -> usize {
892        self.records()
893    }
894
895    /// [`ColumnBody::merge_from`] does the work: gallop bulk-copies for disjoint
896    /// runs, semigroup consolidation on equal `(data, time)`, output cut at
897    /// the ship threshold.
898    ///
899    /// Fronts whose data ranges are disjoint never load at all: the resident
900    /// fence entries decide, and the lower front moves to the output whole.
901    fn merge(in1: &mut VecDeque<Self>, in2: &mut VecDeque<Self>, out: &mut VecDeque<Self>) {
902        // Disjoint fast path: when one front lies strictly below the other's
903        // first data item (equal boundary data could still interleave on
904        // time), the merged prefix through the shared horizon is exactly that
905        // front, unchanged.
906        let (a_first, a_last) = in1
907            .front()
908            .expect("caller guarantees non-empty input")
909            .data_span();
910        let (b_first, b_last) = in2
911            .front()
912            .expect("caller guarantees non-empty input")
913            .data_span();
914        let a_low = rr::<D>(a_last) < rr::<D>(b_first);
915        let b_low = rr::<D>(b_last) < rr::<D>(a_first);
916        if a_low {
917            let chunk = in1.pop_front().expect("front observed above");
918            out.push_back(chunk.survive_merge());
919            return;
920        }
921        if b_low {
922            let chunk = in2.pop_front().expect("front observed above");
923            out.push_back(chunk.survive_merge());
924            return;
925        }
926
927        let a = in1.pop_front().expect("caller guarantees non-empty input");
928        let b = in2.pop_front().expect("caller guarantees non-empty input");
929        // Merged output is one generation past its deepest input. A survivor
930        // (untouched or rewritten from its remainder) keeps its own depth.
931        let depths = [a.depth(), b.depth()];
932        let out_depth = depths[0].max(depths[1]).saturating_add(1);
933        // A spilled side keeps its body so an untouched survivor returns
934        // without a rewrite. A typed-resident body is cheaper to take out
935        // and re-insert than to duplicate, which keeping it shared would
936        // force on `into_body`.
937        let keep = |chunk: &Self| match chunk {
938            ColumnChunk::Spilled(body, _) if !body.handle.is_typed() => Some(Rc::clone(body)),
939            _ => None,
940        };
941        let respill = [a.is_spilled(), b.is_spilled()];
942        let mut spill_a = keep(&a);
943        let mut spill_b = keep(&b);
944        let mut cols = [a.into_body(), b.into_body()];
945        let mut positions = [0usize, 0usize];
946        loop {
947            let mut result: ColumnBody<(D, T, R)> = ColumnBody::default();
948            let yielded = result.merge_from(&mut cols, &mut positions);
949            if !result.is_empty() {
950                out.push_back(ColumnChunk::Resident(Rc::new(result), out_depth));
951            }
952            if !yielded {
953                break;
954            }
955        }
956        // The disjoint fast paths above move fronts without reading a row, so
957        // only this path counts as merge work.
958        metrics::record(metrics::Stage::Merge, positions[0] + positions[1], 0);
959        let [col_a, col_b] = &mut cols;
960        // Per input side: the loaded column and the merge's consumed position
961        // within it, the side's pre-merge depth, its original spilled body
962        // when it had one, and the deque a survivor returns to.
963        for (col, pos, depth, spilled, respill, queue) in [
964            (
965                col_a,
966                positions[0],
967                depths[0],
968                &mut spill_a,
969                respill[0],
970                in1,
971            ),
972            (
973                col_b,
974                positions[1],
975                depths[1],
976                &mut spill_b,
977                respill[1],
978                in2,
979            ),
980        ] {
981            let len = col.len();
982            if pos == 0 && len > 0 {
983                // Untouched survivor: restore it as it was (the loaded copy
984                // is dropped), aged one generation by its survival.
985                let chunk = match spilled.take() {
986                    Some(body) => ColumnChunk::Spilled(body, depth).survive_merge(),
987                    // A body taken out of the pool goes back in.
988                    None if respill => Self::commit(std::mem::take(col), depth.saturating_add(1)),
989                    None => {
990                        ColumnChunk::Resident(Rc::new(std::mem::take(col)), depth).survive_merge()
991                    }
992                };
993                queue.push_front(chunk);
994            } else if pos < len {
995                let view = col.borrow();
996                let mut rest = <(D, T, R) as Columnar>::Container::default();
997                rest.extend_from_self(view, pos..len);
998                queue.push_front(ColumnChunk::Resident(
999                    Rc::new(ColumnBody::Typed(rest)),
1000                    depth,
1001                ));
1002            }
1003        }
1004    }
1005
1006    /// Partition one front chunk by `frontier`, folding kept times into
1007    /// `residual`. One chunk per call, so the harness settles both sides
1008    /// between chunks. Output is cut at the ship threshold.
1009    fn extract(
1010        input: &mut VecDeque<Self>,
1011        frontier: AntichainRef<T>,
1012        residual: &mut timely::progress::Antichain<T>,
1013        keep: &mut VecDeque<Self>,
1014        ship: &mut VecDeque<Self>,
1015    ) {
1016        let Some(chunk) = input.pop_front() else {
1017            return;
1018        };
1019        // Whole-chunk pass-through from the resident time bounds: a chunk
1020        // the frontier is entirely past ships unchanged, one entirely at or
1021        // past the frontier keeps unchanged. Spilled bodies pass through
1022        // without a load, a re-commit, or any codec work; only chunks the
1023        // frontier actually splits are loaded below.
1024        let (time_lower, time_upper) = chunk.chunk_time_bounds();
1025        if time_upper.iter().all(|t| !frontier.less_equal(t)) {
1026            ship.push_back(chunk);
1027            return;
1028        }
1029        if time_lower.elements().iter().all(|m| frontier.less_equal(m)) {
1030            // The residual must lower-bound every kept time, which is the
1031            // chunk's lower bound antichain by construction.
1032            for m in time_lower.elements() {
1033                residual.insert_ref(m);
1034            }
1035            keep.push_back(chunk);
1036            return;
1037        }
1038        // Partitioning rewrites within a generation, so both sides keep the
1039        // input chunk's depth.
1040        let depth = chunk.depth();
1041        let mut col = chunk.into_body();
1042        let len = col.len();
1043        let mut pos = 0;
1044        let mut keep_col: ColumnBody<(D, T, R)> = ColumnBody::default();
1045        let mut ship_col: ColumnBody<(D, T, R)> = ColumnBody::default();
1046        // TODO: rewrite the underlying `ColumnBody::extract` as two passes, the
1047        // time column first to find run boundaries, then bulk per-range
1048        // copies of the remaining leaves.
1049        // Move a side's accumulation to its queue, at the ship threshold
1050        // mid-loop, or any non-empty remainder at the end.
1051        let cut = |col: &mut ColumnBody<(D, T, R)>, queue: &mut VecDeque<Self>, force: bool| {
1052            if !col.is_empty() && (force || at_serialized_capacity(&col.borrow())) {
1053                queue.push_back(ColumnChunk::Resident(Rc::new(std::mem::take(col)), depth));
1054            }
1055        };
1056        while pos < len {
1057            col.extract(&mut pos, frontier, residual, &mut keep_col, &mut ship_col);
1058            if pos < len {
1059                cut(&mut keep_col, keep, false);
1060                cut(&mut ship_col, ship, false);
1061            }
1062        }
1063        cut(&mut keep_col, keep, true);
1064        cut(&mut ship_col, ship, true);
1065    }
1066
1067    /// Advance times by `frontier` and consolidate, withholding the trailing
1068    /// `D` group as the carry unless `done` (its updates may continue in input
1069    /// this call has not seen).
1070    ///
1071    /// The input concatenates into the carry's container, so a group that
1072    /// grows across many calls is appended to, not rebuilt. Each record is
1073    /// copied once on arrival, keeping the run linear. Advancing is
1074    /// lattice-monotone but not order-monotone, so each group's advanced
1075    /// times are re-sorted before adjacent equal times fold.
1076    fn advance(
1077        input: &mut VecDeque<Self>,
1078        frontier: AntichainRef<T>,
1079        done: bool,
1080        out: &mut VecDeque<Self>,
1081    ) {
1082        let Some(front) = input.pop_front() else {
1083            return;
1084        };
1085        // Advancing rewrites within a generation, so output and carry keep
1086        // the deepest input depth. Only merges increment.
1087        let mut depth = front.depth();
1088        // Concatenate the input into one body, reusing the front chunk's
1089        // storage when it is exclusively owned (the usual case: it is last
1090        // call's carry).
1091        let mut base = front.into_body();
1092        {
1093            let base_c = base.typed_mut();
1094            for chunk in input.drain(..) {
1095                depth = depth.max(chunk.depth());
1096                let col = chunk.into_body();
1097                let view = col.borrow();
1098                base_c.extend_from_self(view, 0..view.len());
1099            }
1100        }
1101        let view = base.borrow();
1102        let total = view.len();
1103        if total == 0 {
1104            return;
1105        }
1106        let data = view.0;
1107
1108        // Giant-group early-out: if the whole input is one `D` group, nothing
1109        // is provably complete. Unless `done`, push it all back as the carry.
1110        if !done && data.get(0) == data.get(total - 1) {
1111            input.push_front(ColumnChunk::Resident(Rc::new(base), depth));
1112            return;
1113        }
1114
1115        // The processing bound: everything, or everything before the trailing
1116        // `D` group when it must be withheld.
1117        let end = if done {
1118            total
1119        } else {
1120            let last = data.get(total - 1);
1121            let mut end = total - 1;
1122            while end > 0 && data.get(end - 1) == last {
1123                end -= 1;
1124            }
1125            end
1126        };
1127        // Count only the rows this call advances. The withheld carry is
1128        // counted by the call that finally processes it, so a group spanning
1129        // many calls counts once.
1130        metrics::record(metrics::Stage::Advance, end, 0);
1131
1132        let mut result = <(D, T, R) as Columnar>::Container::default();
1133        // Per-group scratch: advanced owned times with owned diffs.
1134        let mut scratch: Vec<(T, R)> = Vec::new();
1135        let mut index = 0;
1136        // Cut output at the commit size, checked amortized by emitted records
1137        // (the size test walks the container's leaves, so probing it per
1138        // record would be quadratic). Records, not groups: a single group may
1139        // carry arbitrarily many advanced times, and a cut is legal anywhere
1140        // in the sorted sequence, so bounding by records keeps the largest
1141        // possible output chunk within one check period of the target. It
1142        // must not outgrow the pool's largest size class, past which a body
1143        // degrades to a permanently resident heap chunk.
1144        const CUT_CHECK_RECORDS: usize = 1024;
1145        let mut records_since_check = 0usize;
1146        // TODO: the output leaves are addressed independently, so a group
1147        // that folds nothing (no time collisions, no zeroed diffs) could bulk
1148        // `extend_from_self` the D leaf over the whole group range and push
1149        // only the advanced times and diffs per record, and a singleton group
1150        // (the common case for mostly-unique D) could skip the scratch and
1151        // sort round trip entirely.
1152        while index < end {
1153            let group_d = data.get(index);
1154            scratch.clear();
1155            while index < end && data.get(index) == group_d {
1156                let (_, t, r) = view.get(index);
1157                let mut owned_t = T::into_owned(t);
1158                owned_t.advance_by(frontier);
1159                scratch.push((owned_t, R::into_owned(r)));
1160                index += 1;
1161            }
1162            scratch.sort_by(|a, b| a.0.cmp(&b.0));
1163            let mut run = scratch.drain(..).peekable();
1164            while let Some((t, mut r)) = run.next() {
1165                while run.peek().is_some_and(|(t2, _)| *t2 == t) {
1166                    let (_, r2) = run.next().expect("peeked");
1167                    r.plus_equals(&r2);
1168                }
1169                if !r.is_zero() {
1170                    result.0.push(group_d);
1171                    result.1.push(&t);
1172                    result.2.push(&r);
1173                    records_since_check += 1;
1174                    if records_since_check >= CUT_CHECK_RECORDS {
1175                        records_since_check = 0;
1176                        if u64::cast_from(indexed::length_in_words(&result.borrow()))
1177                            >= u64::cast_from(COMMIT_BYTES / 8)
1178                        {
1179                            out.push_back(ColumnChunk::Resident(
1180                                Rc::new(ColumnBody::Typed(std::mem::take(&mut result))),
1181                                depth,
1182                            ));
1183                        }
1184                    }
1185                }
1186            }
1187        }
1188        if !result.is_empty() {
1189            out.push_back(ColumnChunk::Resident(
1190                Rc::new(ColumnBody::Typed(result)),
1191                depth,
1192            ));
1193        }
1194
1195        // Rebuild the withheld trailing group as the carry.
1196        if end < total {
1197            let mut carry = <(D, T, R) as Columnar>::Container::default();
1198            carry.extend_from_self(view, end..total);
1199            input.push_front(ColumnChunk::Resident(
1200                Rc::new(ColumnBody::Typed(carry)),
1201                depth,
1202            ));
1203        }
1204    }
1205
1206    /// Grade by serialized bytes and commit: spilled chunks pass through
1207    /// untouched, resident chunks at the commit size commit as they are, and
1208    /// smaller neighbors coalesce until the accumulation reaches it.
1209    /// Committing is the spill hook (see `ColumnChunk::commit`). Unless
1210    /// `done`, a tail below the commit size returns to the front of `input`.
1211    fn settle(input: &mut VecDeque<Self>, done: bool, out: &mut VecDeque<Self>) {
1212        Self::settle_graded(input, done, out, true)
1213    }
1214}
1215
1216impl<D, T, R, G> ColumnChunk<D, T, R, G>
1217where
1218    G: SpillGate,
1219    D: PoolColumnar,
1220    T: PoolColumnar + Timestamp,
1221    R: PoolColumnar,
1222    for<'a> columnar::Ref<'a, D>: Ord,
1223    for<'a> columnar::Ref<'a, T>: Ord,
1224{
1225    /// [`Chunk::settle`], with `commit` choosing what a graded chunk becomes.
1226    ///
1227    /// A caller that reads its output back immediately passes `false`: the
1228    /// chunk is graded resident, so the body is not serialized into a pool
1229    /// slot only to be copied straight back out, and the pool's budget is
1230    /// not charged for a body whose lifetime ends at the next read.
1231    fn settle_graded(
1232        input: &mut VecDeque<Self>,
1233        done: bool,
1234        out: &mut VecDeque<Self>,
1235        commit: bool,
1236    ) {
1237        // Resident chunks below the commit size accumulate into `acc` until
1238        // it reaches that size. Its depth is the deepest among the chunks it
1239        // holds, since coalescing rewrites within a generation.
1240        let mut acc: Option<(ColumnBody<(D, T, R)>, u8)> = None;
1241        while let Some(chunk) = input.pop_front() {
1242            let (rc, depth) = match chunk {
1243                spilled @ ColumnChunk::Spilled(_, _) => {
1244                    if let Some((body, depth)) = acc.take() {
1245                        out.push_back(Self::grade(body, depth, commit));
1246                    }
1247                    out.push_back(spilled);
1248                    continue;
1249                }
1250                ColumnChunk::Resident(rc, depth) => (rc, depth),
1251            };
1252            if rc.length_in_bytes() > COMMIT_BYTES && rc.len() > 1 {
1253                // Cut by borrow, so a shared body is copied once, into its
1254                // pieces.
1255                let records = cut_records(rc.len(), rc.length_in_bytes(), COMMIT_BYTES).max(1);
1256                let mut pieces = VecDeque::new();
1257                Self::push_cuts(&rc, records, depth, &mut pieces);
1258                for piece in pieces.into_iter().rev() {
1259                    input.push_front(piece);
1260                }
1261                continue;
1262            }
1263            let full = at_commit_size(&rc);
1264            // A chunk is copied into `acc` from its borrow, so a shared body
1265            // is never unwrapped.
1266            let fits = acc.as_ref().is_some_and(|(body, _)| {
1267                body.length_in_bytes().saturating_add(rc.length_in_bytes()) <= COMMIT_BYTES
1268            });
1269            if !full
1270                && fits
1271                && let Some((mut body, acc_depth)) = acc.take()
1272            {
1273                let view = rc.borrow();
1274                body.typed_mut().extend_from_self(view, 0..view.len());
1275                let acc_depth = acc_depth.max(depth);
1276                if at_commit_size(&body) {
1277                    out.push_back(Self::grade(body, acc_depth, commit));
1278                } else {
1279                    acc = Some((body, acc_depth));
1280                }
1281                continue;
1282            }
1283            // The chunk does not fit in `acc`. If the chunk is below the
1284            // commit size, as many of its records as fit move into `acc` and
1285            // the rest return to the front of `input`. Without that move, a
1286            // stream of chunks just over half the bound would commit each one
1287            // at about half the bound. Either way `acc` then goes to `out`.
1288            if let Some((mut body, acc_depth)) = acc.take() {
1289                let len = rc.len();
1290                let prefix = if full {
1291                    0
1292                } else {
1293                    let space = COMMIT_BYTES.saturating_sub(body.length_in_bytes());
1294                    cut_records(len, rc.length_in_bytes(), space)
1295                };
1296                if prefix == 0 || prefix >= len {
1297                    out.push_back(Self::grade(body, acc_depth, commit));
1298                } else {
1299                    let view = rc.borrow();
1300                    body.typed_mut().extend_from_self(view, 0..prefix);
1301                    let mut rest = <(D, T, R) as Columnar>::Container::default();
1302                    rest.extend_from_self(view, prefix..len);
1303                    input.push_front(ColumnChunk::Resident(
1304                        Rc::new(ColumnBody::Typed(rest)),
1305                        depth,
1306                    ));
1307                    let acc_depth = acc_depth.max(depth);
1308                    if body.length_in_bytes() > COMMIT_BYTES {
1309                        // Uneven record widths pushed `acc` past the bound.
1310                        let mut pieces = VecDeque::new();
1311                        Self::push_bounded(body, acc_depth, &mut pieces);
1312                        for piece in pieces {
1313                            let ColumnChunk::Resident(piece, piece_depth) = piece else {
1314                                unreachable!("push_bounded emits resident pieces");
1315                            };
1316                            let piece =
1317                                Rc::try_unwrap(piece).unwrap_or_else(|piece| piece.duplicate());
1318                            out.push_back(Self::grade(piece, piece_depth, commit));
1319                        }
1320                    } else {
1321                        out.push_back(Self::grade(body, acc_depth, commit));
1322                    }
1323                    continue;
1324                }
1325            }
1326            let body = Rc::try_unwrap(rc).unwrap_or_else(|rc| rc.duplicate());
1327            if full {
1328                out.push_back(Self::grade(body, depth, commit));
1329            } else {
1330                acc = Some((body, depth));
1331            }
1332        }
1333        if let Some((body, depth)) = acc {
1334            if done {
1335                out.push_back(Self::grade(body, depth, commit));
1336            } else {
1337                input.push_front(ColumnChunk::Resident(Rc::new(body), depth));
1338            }
1339        }
1340    }
1341
1342    /// A graded body as a chunk: committed, which may spill it, or left
1343    /// resident.
1344    fn grade(body: ColumnBody<(D, T, R)>, depth: u8, commit: bool) -> Self {
1345        if commit {
1346            Self::commit(body, depth)
1347        } else {
1348            ColumnChunk::Resident(Rc::new(body), depth)
1349        }
1350    }
1351}
1352
1353/// Append every update in `view` whose key matches a probe at or after
1354/// `*probe_index` into `staging`, per the [`UnloadChunk`] consume-index
1355/// protocol: probes strictly below the view's last key are consumed, a probe
1356/// equal to it is extracted but left for the next chunk.
1357fn extract_view_into<'v, 'p, K, V, T, R>(
1358    view: BorrowedOf<'v, ((K, V), T, R)>,
1359    probes: BorrowedOf<'p, K>,
1360    probe_index: &mut usize,
1361    staging: &mut <((K, V), T, R) as Columnar>::Container,
1362) where
1363    K: Columnar,
1364    V: Columnar,
1365    T: Columnar,
1366    R: Columnar,
1367    for<'b> columnar::Ref<'b, K>: Copy + Ord,
1368{
1369    let keys = view.0.0;
1370    let len = keys.len();
1371    let last = keys.get(len - 1);
1372    let count = probes.len();
1373    let mut pos = 0;
1374    while *probe_index < count {
1375        let probe = probes.get(*probe_index);
1376        mz_ore::soft_assert_no_log!(
1377            *probe_index == 0 || rr::<K>(probes.get(*probe_index - 1)) < rr::<K>(probe),
1378            "probe keys must be sorted and deduplicated"
1379        );
1380        if rr::<K>(probe) > rr::<K>(last) {
1381            return;
1382        }
1383        gallop(len, &mut pos, |i| rr::<K>(keys.get(i)) < rr::<K>(probe));
1384        let start = pos;
1385        while pos < len && rr::<K>(keys.get(pos)) == rr::<K>(probe) {
1386            pos += 1;
1387        }
1388        staging.extend_from_self(view, start..pos);
1389        if rr::<K>(probe) == rr::<K>(last) {
1390            return;
1391        }
1392        *probe_index += 1;
1393    }
1394}
1395
1396impl<K, V, T, R, G> UnloadChunk for ColumnChunk<(K, V), T, R, G>
1397where
1398    G: SpillGate,
1399    K: PoolColumnar,
1400    for<'a> columnar::Ref<'a, K>: Copy + Ord,
1401    V: PoolColumnar,
1402    for<'a> columnar::Ref<'a, V>: Copy + Ord,
1403    T: PoolColumnar + Default + Timestamp + Lattice + Ord,
1404    for<'a> columnar::Ref<'a, T>: Copy + Ord,
1405    R: PoolColumnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
1406{
1407    /// The flat columnar accumulation. Appends are bulk column-range copies,
1408    /// and a group straddling chunks stitches by plain concatenation.
1409    type Staging = <((K, V), T, R) as Columnar>::Container;
1410
1411    /// A borrowed key column, e.g. of a `Column<K>` the consumer assembled
1412    /// from its sorted, deduplicated probe keys.
1413    type Probes<'a> = BorrowedOf<'a, K>;
1414
1415    fn probe_count(probes: Self::Probes<'_>) -> usize {
1416        probes.len()
1417    }
1418
1419    fn locate(&self, probes: Self::Probes<'_>, probe_index: usize) -> std::cmp::Ordering {
1420        let probe = probes.get(probe_index);
1421        // A data ref is a `(key ref, val ref)` tuple, so the key fences are a
1422        // projection of the data fences.
1423        let (first, last) = self.data_span();
1424        let (first, last) = (first.0, last.0);
1425        if rr::<K>(probe) < rr::<K>(first) {
1426            std::cmp::Ordering::Less
1427        } else if rr::<K>(probe) > rr::<K>(last) {
1428            std::cmp::Ordering::Greater
1429        } else {
1430            std::cmp::Ordering::Equal
1431        }
1432    }
1433
1434    fn extract_into(
1435        &self,
1436        probes: Self::Probes<'_>,
1437        probe_index: &mut usize,
1438        staging: &mut Self::Staging,
1439    ) {
1440        match self {
1441            ColumnChunk::Resident(col, _) => {
1442                extract_view_into::<K, V, T, R>(col.borrow(), probes, probe_index, staging);
1443            }
1444            ColumnChunk::Spilled(body, _) => body.handle.with_typed(|typed| match typed {
1445                Some(typed) => {
1446                    let view = pooled::<(K, V), T, R>(typed).body.borrow();
1447                    extract_view_into::<K, V, T, R>(view, probes, probe_index, staging);
1448                }
1449                None => with_scratch(|scratch| {
1450                    // NOTE: deliberately the non-admitting read. One probe
1451                    // set touching a chunk is weak evidence it will be
1452                    // touched again, and probing a spilled trace must not
1453                    // accrete it back into residency. The cost is a full
1454                    // decode per probe set against an evicted chunk.
1455                    body.handle.read_into(scratch);
1456                    let view = borrow_words::<((K, V), T, R)>(scratch);
1457                    extract_view_into::<K, V, T, R>(view, probes, probe_index, staging);
1458                }),
1459            }),
1460        }
1461    }
1462
1463    fn fetch_into(&self, staging: &mut Self::Staging) {
1464        match self {
1465            ColumnChunk::Resident(col, _) => {
1466                let view = col.borrow();
1467                staging.extend_from_self(view, 0..view.len());
1468            }
1469            ColumnChunk::Spilled(body, _) => body.handle.with_typed(|typed| match typed {
1470                Some(typed) => {
1471                    let view = pooled::<(K, V), T, R>(typed).body.borrow();
1472                    staging.extend_from_self(view, 0..view.len());
1473                }
1474                None => with_scratch(|scratch| {
1475                    body.handle.read_into(scratch);
1476                    let view = borrow_words::<((K, V), T, R)>(scratch);
1477                    staging.extend_from_self(view, 0..view.len());
1478                }),
1479            }),
1480        }
1481    }
1482}
1483
1484/// A batch builder over [`ColumnChunk`] input that delegates to a builder
1485/// over [`ColumnBody`] input, loading each chunk's body as it is pushed.
1486///
1487/// This is the adapter that lets a [`ChunkBatcher`] feed the body-input batch
1488/// builders (and through them the existing spine layouts):
1489/// the batcher's chains carry pool-spillable chunks, and bodies are read back
1490/// copy-out only at the seal, one chunk at a time.
1491///
1492/// [`ChunkBatcher`]: differential_dataflow::trace::chunk::ChunkBatcher
1493pub struct UnchunkBuilder<Bu, D: Columnar, T: Columnar, R: Columnar, G: SpillGate> {
1494    inner: Bu,
1495    _marker: std::marker::PhantomData<(D, T, R, G)>,
1496}
1497
1498impl<Bu: Default, D: Columnar, T: Columnar, R: Columnar, G: SpillGate> Default
1499    for UnchunkBuilder<Bu, D, T, R, G>
1500{
1501    fn default() -> Self {
1502        Self {
1503            inner: Bu::default(),
1504            _marker: std::marker::PhantomData,
1505        }
1506    }
1507}
1508
1509impl<Bu, D, T, R, G> differential_dataflow::trace::Builder for UnchunkBuilder<Bu, D, T, R, G>
1510where
1511    Bu: differential_dataflow::trace::Builder<Input = ColumnBody<(D, T, R)>> + ChainState,
1512    D: Columnar + 'static,
1513    T: Columnar + 'static,
1514    R: Columnar + 'static,
1515    G: SpillGate,
1516{
1517    type Input = ColumnChunk<D, T, R, G>;
1518    type Time = Bu::Time;
1519    type Output = Bu::Output;
1520
1521    fn push(&mut self, chunk: &mut Self::Input) {
1522        let mut body = std::mem::take(chunk).into_body();
1523        self.inner.push(&mut body);
1524    }
1525
1526    fn done(self) -> Option<Self::Output> {
1527        self.inner.done()
1528    }
1529}
1530
1531impl<Bu, D, T, R, G>
1532    differential_dataflow::trace::implementations::merge_batcher::Sealer<ColumnChunk<D, T, R, G>>
1533    for UnchunkBuilder<Bu, D, T, R, G>
1534where
1535    Bu: differential_dataflow::trace::Builder<Input = ColumnBody<(D, T, R)>> + ChainState,
1536    D: Columnar + 'static,
1537    T: Columnar + 'static,
1538    R: Columnar + 'static,
1539    G: SpillGate,
1540{
1541    type Output = Bu::Output;
1542
1543    fn seal(chain: &mut Vec<ColumnChunk<D, T, R, G>>) -> Option<Self::Output> {
1544        use differential_dataflow::trace::Builder;
1545
1546        let mut state = Bu::State::default();
1547        for chunk in chain.iter() {
1548            // A resident body folds for the price of a borrow, so it goes to
1549            // `observe` whatever the builder asked for: the state gets exact
1550            // figures for free. Folding a spilled body costs a second pool
1551            // read on top of the push below, so it waits on `wants_bodies`.
1552            if !chunk.is_spilled() || Bu::wants_bodies() {
1553                chunk.with_body(|body| Bu::observe(&mut state, body));
1554            } else {
1555                Bu::observe_records(&mut state, chunk.records());
1556            }
1557        }
1558        let mut builder = Self {
1559            inner: Bu::from_state(state),
1560            _marker: std::marker::PhantomData,
1561        };
1562        // One chunk at a time through `push`, so peak transient memory is a
1563        // single loaded body rather than the whole chain at once.
1564        for chunk in chain.iter_mut() {
1565            builder.push(chunk);
1566        }
1567        chain.clear();
1568        builder.done()
1569    }
1570}
1571
1572/// A builder whose batches carry state derived from a whole chain, computed
1573/// before the chain's first push.
1574///
1575/// [`Sealer::seal`] receives the chain at once, which a chain of pool-backed
1576/// chunks cannot supply without holding every body resident at the same time.
1577/// An implementor splits the derivation instead: a caller folds the chain into
1578/// [`State`](Self::State) one entry at a time, then builds from it.
1579///
1580/// A caller folds each chain entry exactly once, through
1581/// [`observe`](Self::observe) when it holds the entry's contents and through
1582/// [`observe_records`](Self::observe_records) when it does not, so an
1583/// implementor may accumulate the same figure in both without double
1584/// counting. Which one an entry takes is the caller's choice, so an
1585/// implementor reads no meaning into it beyond what each carries.
1586///
1587/// [`Sealer::seal`]: differential_dataflow::trace::implementations::merge_batcher::Sealer::seal
1588pub trait ChainState: differential_dataflow::trace::Builder {
1589    /// State accumulated across a chain.
1590    type State: Default;
1591
1592    /// Whether the state is worth reading a chain entry's contents for when
1593    /// they are not already at hand. A caller that holds them passes them on
1594    /// regardless.
1595    fn wants_bodies() -> bool;
1596
1597    /// Fold one chain entry's contents into `state`.
1598    fn observe(state: &mut Self::State, input: &Self::Input);
1599
1600    /// Fold one chain entry's update count into `state`, for a caller that
1601    /// reads no bodies.
1602    fn observe_records(state: &mut Self::State, records: usize);
1603
1604    /// Build with `state` installed.
1605    fn from_state(state: Self::State) -> Self;
1606}
1607
1608/// A chunker for `arrange_core` over [`ColumnChunk`]s: sorts and consolidates
1609/// raw input columns through a [`ColumnChunker`] and wraps its output chunks.
1610pub struct ChunkChunker<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> {
1611    inner: ColumnChunker<(D, T, R)>,
1612    ready: VecDeque<ColumnChunk<D, T, R, G>>,
1613    staged: ColumnChunk<D, T, R, G>,
1614}
1615
1616impl<D, T, R, G> Default for ChunkChunker<D, T, R, G>
1617where
1618    D: Columnar,
1619    T: Columnar,
1620    R: Columnar,
1621    G: SpillGate,
1622    ColumnChunker<(D, T, R)>: Default,
1623{
1624    fn default() -> Self {
1625        Self {
1626            inner: Default::default(),
1627            ready: VecDeque::new(),
1628            staged: Default::default(),
1629        }
1630    }
1631}
1632
1633impl<'a, D, T, R, G> PushInto<&'a mut Column<(D, T, R)>> for ChunkChunker<D, T, R, G>
1634where
1635    D: Columnar,
1636    T: Columnar,
1637    R: Columnar,
1638    G: SpillGate,
1639    ColumnChunker<(D, T, R)>: PushInto<&'a mut Column<(D, T, R)>>,
1640{
1641    fn push_into(&mut self, item: &'a mut Column<(D, T, R)>) {
1642        self.inner.push_into(item);
1643    }
1644}
1645
1646impl<D, T, R, G> ContainerBuilder for ChunkChunker<D, T, R, G>
1647where
1648    D: Columnar + 'static,
1649    T: Columnar + 'static,
1650    R: Columnar + 'static,
1651    G: SpillGate,
1652    ColumnChunker<(D, T, R)>: ContainerBuilder<Container = ColumnBody<(D, T, R)>>,
1653{
1654    type Container = ColumnChunk<D, T, R, G>;
1655
1656    fn extract(&mut self) -> Option<&mut Self::Container> {
1657        if self.ready.is_empty() {
1658            let body = self.inner.extract()?;
1659            ColumnChunk::push_bounded(std::mem::take(body), 0, &mut self.ready);
1660        }
1661        self.staged = self.ready.pop_front()?;
1662        Some(&mut self.staged)
1663    }
1664
1665    fn finish(&mut self) -> Option<&mut Self::Container> {
1666        if self.ready.is_empty() {
1667            let body = self.inner.finish()?;
1668            ColumnChunk::push_bounded(std::mem::take(body), 0, &mut self.ready);
1669        }
1670        self.staged = self.ready.pop_front()?;
1671        Some(&mut self.staged)
1672    }
1673}
1674
1675/// The [`ChunkBatcher`] of a chunk chain, reporting resident bytes to the
1676/// batcher size logger.
1677///
1678/// `S` seals each extracted chain into a batch. A caller that wants the chain
1679/// itself takes [`AccountedChunkChainBatcher`] instead.
1680///
1681/// [`ChunkBatcher`]: differential_dataflow::trace::chunk::ChunkBatcher
1682pub type AccountedChunkBatcher<D, T, R, G, S> =
1683    differential_dataflow::trace::implementations::merge_batcher::MergeBatcher<
1684        ChunkChunker<D, T, R, G>,
1685        AccountedChunkMerger<D, T, R, G>,
1686        S,
1687    >;
1688
1689/// The [`AccountedChunkBatcher`] of a caller that wants the merged chain rather
1690/// than a batch built from it, such as the temporal-bucket store.
1691pub type AccountedChunkChainBatcher<D, T, R, G> = crate::operator::ConsolidatingBatcher<
1692    ChunkChunker<D, T, R, G>,
1693    AccountedChunkMerger<D, T, R, G>,
1694>;
1695
1696/// The chunk merger of [`AccountedChunkBatcher`]: differential's merger with
1697/// a resident readied side, plus the [`Merger::allocation`] figures the
1698/// `mz_arrangement_batcher_*_raw` introspection tables are assembled from.
1699///
1700/// A chunk reports the bytes of the body it holds whether that body sits on
1701/// the heap or in the pool. A spilled body is not free: the pool keeps it in
1702/// a slot, or compressed in an extent, or on the device, and only the last
1703/// of those is off the process's books. Which tier holds it is the pool's
1704/// own business and its metrics report that split, so attributing the body
1705/// to the operator that owns it is what keeps these tables whole.
1706///
1707/// [`Merger::allocation`]: differential_dataflow::trace::implementations::merge_batcher::Merger::allocation
1708pub struct AccountedChunkMerger<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> {
1709    inner: differential_dataflow::trace::chunk::ChunkMerger<ColumnChunk<D, T, R, G>>,
1710}
1711
1712impl<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> Default
1713    for AccountedChunkMerger<D, T, R, G>
1714{
1715    fn default() -> Self {
1716        Self {
1717            inner: Default::default(),
1718        }
1719    }
1720}
1721
1722impl<D, T, R, G> differential_dataflow::trace::implementations::merge_batcher::Merger
1723    for AccountedChunkMerger<D, T, R, G>
1724where
1725    D: PoolColumnar + 'static,
1726    T: PoolColumnar + Timestamp + 'static,
1727    R: PoolColumnar + 'static,
1728    G: SpillGate,
1729    for<'a> columnar::Ref<'a, D>: Ord,
1730    for<'a> columnar::Ref<'a, T>: Ord,
1731    ColumnChunk<D, T, R, G>: Chunk,
1732    <ColumnChunk<D, T, R, G> as Chunk>::Time: Clone + PartialOrder + 'static,
1733{
1734    type Chunk = ColumnChunk<D, T, R, G>;
1735
1736    fn merge(
1737        &mut self,
1738        list1: Vec<Self::Chunk>,
1739        list2: Vec<Self::Chunk>,
1740        output: &mut Vec<Self::Chunk>,
1741        stash: &mut Vec<Self::Chunk>,
1742    ) {
1743        self.inner.merge(list1, list2, output, stash)
1744    }
1745
1746    fn len(chunk: &Self::Chunk) -> usize {
1747        Chunk::len(chunk)
1748    }
1749
1750    fn allocation(chunk: &Self::Chunk) -> (usize, usize, usize) {
1751        let bytes = match chunk {
1752            ColumnChunk::Resident(col, _) => col.length_in_bytes(),
1753            ColumnChunk::Spilled(body, _) => body.len_bytes,
1754        };
1755        (bytes, bytes, 1)
1756    }
1757}
1758
1759impl<D, T, R, G> differential_dataflow::trace::implementations::merge_batcher::Extractor
1760    for AccountedChunkMerger<D, T, R, G>
1761where
1762    D: PoolColumnar + 'static,
1763    T: PoolColumnar + Timestamp + 'static,
1764    R: PoolColumnar + 'static,
1765    G: SpillGate,
1766    for<'a> columnar::Ref<'a, D>: Ord,
1767    for<'a> columnar::Ref<'a, T>: Ord,
1768    ColumnChunk<D, T, R, G>: Chunk,
1769    <ColumnChunk<D, T, R, G> as Chunk>::Time: Clone + PartialOrder + 'static,
1770{
1771    type Time = <ColumnChunk<D, T, R, G> as Chunk>::Time;
1772
1773    fn extract(
1774        &mut self,
1775        merged: Vec<Self::Chunk>,
1776        upper: AntichainRef<Self::Time>,
1777        frontier: &mut Antichain<Self::Time>,
1778        readied: &mut Vec<Self::Chunk>,
1779        kept: &mut Vec<Self::Chunk>,
1780        _stash: &mut Vec<Self::Chunk>,
1781    ) {
1782        // Differential's merger settles both sides with committing settles.
1783        // The readied side is handed to a builder that loads every chunk
1784        // back, so committing it spills bodies the next call copies out
1785        // again, and charges the pool's budget for them in between. Drive
1786        // the loop here instead and grade that side resident. The kept side
1787        // stays committed: it is the side that lives on, and spilling it is
1788        // the point.
1789        let mut input: VecDeque<Self::Chunk> = merged.into();
1790        let (mut keep, mut shipped) = (VecDeque::new(), VecDeque::new());
1791        let (mut kept_q, mut shipped_q) = (VecDeque::new(), VecDeque::new());
1792        while !input.is_empty() {
1793            Chunk::extract(&mut input, upper, frontier, &mut keep, &mut shipped);
1794            Chunk::settle(&mut keep, false, &mut kept_q);
1795            ColumnChunk::settle_graded(&mut shipped, false, &mut shipped_q, false);
1796        }
1797        Chunk::settle(&mut keep, true, &mut kept_q);
1798        ColumnChunk::settle_graded(&mut shipped, true, &mut shipped_q, false);
1799        Extend::extend(kept, kept_q);
1800        Extend::extend(readied, shipped_q);
1801    }
1802}
1803
1804#[cfg(test)]
1805mod tests {
1806    //! Property tests for the [`Chunk`] and [`UnloadChunk`] contracts on
1807    //! [`ColumnChunk`].
1808    //!
1809    //! Strategy: generate sorted+consolidated inputs (the chunk invariant),
1810    //! drive the trait methods the way the differential harness does, and
1811    //! compare against brute-force references on owned tuples. Test types are
1812    //! `D = (u64, u64)`, `T = u64`, `R = i64` from small ranges so equal-key
1813    //! collisions are common and consolidation actually runs.
1814
1815    use differential_dataflow::batcher::Extract;
1816    use differential_dataflow::trace::chunk::ChunkBatch;
1817    use mz_ore::pool::Pool;
1818    use proptest::prelude::*;
1819    use timely::container::PushInto;
1820    use timely::progress::Antichain;
1821
1822    use crate::columnar::unload::UnloadBatch;
1823
1824    use super::*;
1825
1826    type Tuple = ((u64, u64), u64, i64);
1827    type TestChunk = ColumnChunk<(u64, u64), u64, i64, ComputeSpill>;
1828    /// A batcher over [`TestChunk`]s that hands its chain back rather than sealing it.
1829    type TestChunkBatcher = crate::operator::ConsolidatingBatcher<
1830        ChunkChunker<(u64, u64), u64, i64, ComputeSpill>,
1831        differential_dataflow::trace::chunk::ChunkMerger<TestChunk>,
1832    >;
1833
1834    /// The delegated codec's stored form is byte-identical to the extent
1835    /// store's previous hard-coded framing: a little-endian `u32`
1836    /// body-length prefix followed by one lz4 block, which is exactly what
1837    /// `compress_prepend_size` produces.
1838    #[mz_ore::test]
1839    fn lz4_codec_matches_the_previous_extent_framing() {
1840        let body: Vec<u8> = (0..100_000u32).flat_map(|i| i.to_le_bytes()).collect();
1841        let mut stored = Vec::new();
1842        LZ4_CODEC.encode(&body, &mut stored);
1843        assert_eq!(stored, lz4_flex::block::compress_prepend_size(&body));
1844        let mut round = vec![0u8; body.len()];
1845        LZ4_CODEC.decode(&stored, &mut round);
1846        assert_eq!(round, body);
1847    }
1848
1849    #[mz_ore::test]
1850    #[should_panic(expected = "destination must match")]
1851    fn lz4_codec_decode_length_mismatch_panics() {
1852        let mut stored = Vec::new();
1853        LZ4_CODEC.encode(&[7u8; 64], &mut stored);
1854        let mut short = vec![0u8; 32];
1855        LZ4_CODEC.decode(&stored, &mut short);
1856    }
1857
1858    /// Reference consolidation: sort by `(data, time)`, sum diffs over equal
1859    /// pairs, drop zeros.
1860    fn consolidate(mut v: Vec<Tuple>) -> Vec<Tuple> {
1861        v.sort();
1862        let mut out: Vec<Tuple> = Vec::new();
1863        for (d, t, r) in v {
1864            if let Some(last) = out.last_mut() {
1865                if last.0 == d && last.1 == t {
1866                    last.2 += r;
1867                    continue;
1868                }
1869            }
1870            out.push((d, t, r));
1871        }
1872        out.retain(|x| x.2 != 0);
1873        out
1874    }
1875
1876    fn arb_consolidated() -> impl Strategy<Value = Vec<Tuple>> {
1877        prop::collection::vec(((0u64..5, 0u64..5), 0u64..4, -3i64..=3i64), 0..40)
1878            .prop_map(consolidate)
1879    }
1880
1881    fn build_column(v: &[Tuple]) -> ColumnBody<Tuple> {
1882        let mut col: ColumnBody<Tuple> = Default::default();
1883        for tup in v {
1884            col.push_into(*tup);
1885        }
1886        col
1887    }
1888
1889    fn collect_column(col: &ColumnBody<Tuple>) -> Vec<Tuple> {
1890        col.borrow()
1891            .into_index_iter()
1892            .map(|((k, v), t, r)| {
1893                (
1894                    (u64::into_owned(k), u64::into_owned(v)),
1895                    u64::into_owned(t),
1896                    i64::into_owned(r),
1897                )
1898            })
1899            .collect()
1900    }
1901
1902    fn collect_chunks(chunks: impl IntoIterator<Item = TestChunk>) -> Vec<Tuple> {
1903        chunks
1904            .into_iter()
1905            .flat_map(|chunk| collect_column(&chunk.into_body()))
1906            .collect()
1907    }
1908
1909    fn collect_staging(staging: &<Tuple as Columnar>::Container) -> Vec<Tuple> {
1910        staging
1911            .borrow()
1912            .into_index_iter()
1913            .map(|((k, v), t, r)| {
1914                (
1915                    (u64::into_owned(k), u64::into_owned(v)),
1916                    u64::into_owned(t),
1917                    i64::into_owned(r),
1918                )
1919            })
1920            .collect()
1921    }
1922
1923    /// Cut consolidated data into non-empty chunks at the given points.
1924    fn chunked(data: &[Tuple], cuts: &[usize]) -> VecDeque<TestChunk> {
1925        let mut chunks = VecDeque::new();
1926        let mut start = 0;
1927        for cut in cuts {
1928            let end = (start + 1 + cut % 7).min(data.len());
1929            if end > start {
1930                chunks.push_back(ColumnChunk::from_body(build_column(&data[start..end])));
1931                start = end;
1932            }
1933        }
1934        if start < data.len() {
1935            chunks.push_back(ColumnChunk::from_body(build_column(&data[start..])));
1936        }
1937        chunks
1938    }
1939
1940    /// The chunked cut, with every chunk force-spilled through a private pool
1941    /// (bounds captured, bodies in the pool) regardless of size thresholds.
1942    fn chunked_spilled(data: &[Tuple], cuts: &[usize], pool: &Pool) -> VecDeque<TestChunk> {
1943        chunked(data, cuts)
1944            .into_iter()
1945            .map(|chunk| force_spill(chunk, pool))
1946            .collect()
1947    }
1948
1949    /// Whether a spilled chunk's body is stored compressed. Deliberately
1950    /// does not retain the body: an `Rc` held across a `survive_merge` would
1951    /// itself make the body shared and suppress the migration under test.
1952    fn body_compressed(chunk: &TestChunk) -> bool {
1953        match chunk {
1954            ColumnChunk::Spilled(body, _) => body.compressed,
1955            ColumnChunk::Resident(_, _) => panic!("chunk must be spilled"),
1956        }
1957    }
1958
1959    /// Spill one chunk through `pool`, bypassing the size threshold and
1960    /// keeping the chunk's depth.
1961    fn force_spill(chunk: TestChunk, pool: &Pool) -> TestChunk {
1962        let depth = chunk.depth();
1963        TestChunk::spill_body(chunk.into_body(), pool, depth)
1964    }
1965
1966    /// The batcher size logger reports a body's bytes wherever the body
1967    /// lives, so an operator's share does not vanish when a chunk spills.
1968    #[mz_ore::test]
1969    #[cfg_attr(miri, ignore)]
1970    fn allocation_reports_a_body_wherever_it_lives() {
1971        use differential_dataflow::trace::implementations::merge_batcher::Merger;
1972
1973        let data: Vec<Tuple> = (0..64u64).map(|i| ((i, i), 0, 1)).collect();
1974        let resident = TestChunk::from_body(build_column(&data));
1975        let (size, capacity, allocations) =
1976            AccountedChunkMerger::<(u64, u64), u64, i64, ComputeSpill>::allocation(&resident);
1977        assert!(size > 0, "a resident body reports its bytes");
1978        assert_eq!((capacity, allocations), (size, 1));
1979
1980        let spilled = force_spill(resident, &test_pool());
1981        assert_eq!(
1982            AccountedChunkMerger::<(u64, u64), u64, i64, ComputeSpill>::allocation(&spilled),
1983            (size, capacity, allocations),
1984            "spilling a body does not change what its owner holds"
1985        );
1986    }
1987
1988    /// A chunk that reaches the seal resident is handed back resident, even
1989    /// with a pool installed. The readied side goes straight to a builder
1990    /// that loads every chunk, so committing it there would serialize a body
1991    /// the next call copies out again.
1992    ///
1993    /// A body the batcher spilled earlier, while merging chains it was
1994    /// holding, passes through as it is: that spill is the one worth paying
1995    /// for, and undoing it would be another copy.
1996    #[mz_ore::test]
1997    #[cfg_attr(miri, ignore)]
1998    fn seal_readies_resident_chunks() {
1999        // One push, so nothing merges before the seal and the chunk reaches
2000        // `extract` resident. Comfortably past the 64 KiB spill floor at 32
2001        // bytes per record, so a committing settle would spill it.
2002        let data: Vec<Tuple> = (0..20_000u64).map(|i| ((i, i), 0, 1)).collect();
2003
2004        set_spill_override(Some(test_pool()));
2005        let mut batcher =
2006            AccountedChunkChainBatcher::<(u64, u64), u64, i64, ComputeSpill>::new(None, 0);
2007        batcher.push_chunk(TestChunk::from_body(build_column(&data)));
2008        let (chain, _frontier) = batcher.extract(Antichain::new().borrow());
2009        let chain = chain.unwrap_or_default();
2010        set_spill_override(None);
2011
2012        assert!(!chain.is_empty(), "the seal ships what was pushed");
2013        assert!(
2014            chain.iter().all(|chunk| !chunk.is_spilled()),
2015            "a chunk that arrived resident is readied resident"
2016        );
2017    }
2018
2019    /// A single pool shared by every test in the module. A pool reserves a
2020    /// large slab of address space, so one per test (let alone per proptest
2021    /// case) exhausts the VM map under parallel test threads.
2022    fn test_pool() -> Pool {
2023        static POOL: std::sync::OnceLock<Pool> = std::sync::OnceLock::new();
2024        POOL.get_or_init(|| Pool::new().expect("pool creation"))
2025            .clone()
2026    }
2027
2028    proptest! {
2029        /// A full batcher round trip: push chunked inputs, seal everything,
2030        /// and compare with the reference consolidation of the union.
2031        #[mz_ore::test]
2032        #[cfg_attr(miri, ignore)]
2033        fn batcher_round_trip(
2034            inputs in prop::collection::vec(arb_consolidated(), 1..6),
2035            cuts in prop::collection::vec(0usize..7, 0..8),
2036        ) {
2037            let mut batcher = TestChunkBatcher::new(None, 0);
2038            let mut union = Vec::new();
2039            for input in &inputs {
2040                Extend::extend(&mut union, input.iter().copied());
2041                for chunk in chunked(input, &cuts) {
2042                    batcher.push_chunk(chunk);
2043                }
2044            }
2045            // An empty upper ships everything.
2046            let (sealed, _frontier) = batcher.extract(Antichain::new().borrow());
2047            let sealed = sealed.unwrap_or_default();
2048            prop_assert_eq!(collect_chunks(sealed), consolidate(union));
2049        }
2050
2051        /// The same round trip over force-spilled inputs: merge and extract
2052        /// read bodies back from the pool call-scoped.
2053        #[mz_ore::test]
2054        #[cfg_attr(miri, ignore)]
2055        fn batcher_round_trip_spilled(
2056            inputs in prop::collection::vec(arb_consolidated(), 1..4),
2057            cuts in prop::collection::vec(0usize..7, 0..6),
2058        ) {
2059            let pool = test_pool();
2060            let mut batcher = TestChunkBatcher::new(None, 0);
2061            let mut union = Vec::new();
2062            for input in &inputs {
2063                Extend::extend(&mut union, input.iter().copied());
2064                for chunk in chunked_spilled(input, &cuts, &pool) {
2065                    batcher.push_chunk(chunk);
2066                }
2067            }
2068            let (sealed, _frontier) = batcher.extract(Antichain::new().borrow());
2069            let sealed = sealed.unwrap_or_default();
2070            prop_assert_eq!(collect_chunks(sealed), consolidate(union));
2071        }
2072
2073        /// Sealing at an intermediate upper partitions by time and reports
2074        /// the kept lower envelope as the frontier.
2075        #[mz_ore::test]
2076        #[cfg_attr(miri, ignore)]
2077        fn seal_partitions_by_time(
2078            input in arb_consolidated(),
2079            cuts in prop::collection::vec(0usize..7, 0..8),
2080            upper in 0u64..5,
2081        ) {
2082            let mut batcher = TestChunkBatcher::new(None, 0);
2083            for chunk in chunked(&input, &cuts) {
2084                batcher.push_chunk(chunk);
2085            }
2086            let (shipped, frontier) = batcher.extract(Antichain::from_elem(upper).borrow());
2087            let shipped = shipped.unwrap_or_default();
2088            let frontier = frontier.to_owned();
2089            let expected_shipped: Vec<Tuple> =
2090                input.iter().copied().filter(|(_, t, _)| *t < upper).collect();
2091            prop_assert_eq!(collect_chunks(shipped), consolidate(expected_shipped));
2092
2093            let kept_min = input.iter().filter(|(_, t, _)| *t >= upper).map(|(_, t, _)| *t).min();
2094            prop_assert_eq!(frontier.elements().first().copied(), kept_min);
2095
2096            let (rest, _frontier) = batcher.extract(Antichain::new().borrow());
2097            let rest = rest.unwrap_or_default();
2098            let expected_rest: Vec<Tuple> =
2099                input.iter().copied().filter(|(_, t, _)| *t >= upper).collect();
2100            prop_assert_eq!(collect_chunks(rest), consolidate(expected_rest));
2101        }
2102
2103        /// The intermediate-upper partition of `seal_partitions_by_time`, over
2104        /// force-spilled inputs: bodies read back from the pool and split by
2105        /// time in one seal.
2106        #[mz_ore::test]
2107        #[cfg_attr(miri, ignore)]
2108        fn seal_partitions_by_time_spilled(
2109            input in arb_consolidated(),
2110            cuts in prop::collection::vec(0usize..7, 0..8),
2111            upper in 0u64..5,
2112        ) {
2113            let pool = test_pool();
2114            let mut batcher = TestChunkBatcher::new(None, 0);
2115            for chunk in chunked_spilled(&input, &cuts, &pool) {
2116                batcher.push_chunk(chunk);
2117            }
2118            let (shipped, frontier) = batcher.extract(Antichain::from_elem(upper).borrow());
2119            let shipped = shipped.unwrap_or_default();
2120            let frontier = frontier.to_owned();
2121            let expected_shipped: Vec<Tuple> =
2122                input.iter().copied().filter(|(_, t, _)| *t < upper).collect();
2123            prop_assert_eq!(collect_chunks(shipped), consolidate(expected_shipped));
2124
2125            let kept_min = input.iter().filter(|(_, t, _)| *t >= upper).map(|(_, t, _)| *t).min();
2126            prop_assert_eq!(frontier.elements().first().copied(), kept_min);
2127
2128            let (rest, _frontier) = batcher.extract(Antichain::new().borrow());
2129            let rest = rest.unwrap_or_default();
2130            let expected_rest: Vec<Tuple> =
2131                input.iter().copied().filter(|(_, t, _)| *t >= upper).collect();
2132            prop_assert_eq!(collect_chunks(rest), consolidate(expected_rest));
2133        }
2134
2135        /// `advance` equals per-record time advancement plus reference
2136        /// consolidation, including across a `done = false` carry.
2137        #[mz_ore::test]
2138        #[cfg_attr(miri, ignore)]
2139        fn advance_matches_reference(
2140            input in arb_consolidated(),
2141            cuts in prop::collection::vec(0usize..7, 0..8),
2142            frontier_elem in 0u64..5,
2143        ) {
2144            let frontier = Antichain::from_elem(frontier_elem);
2145            let mut chunks = chunked(&input, &cuts);
2146            let mut out = VecDeque::new();
2147            TestChunk::advance(&mut chunks, frontier.borrow(), false, &mut out);
2148            TestChunk::advance(&mut chunks, frontier.borrow(), true, &mut out);
2149            prop_assert!(chunks.is_empty());
2150
2151            let expected = consolidate(
2152                input
2153                    .iter()
2154                    .map(|&(d, mut t, r)| {
2155                        t.advance_by(frontier.borrow());
2156                        (d, t, r)
2157                    })
2158                    .collect(),
2159            );
2160            prop_assert_eq!(collect_chunks(out), expected);
2161        }
2162
2163        /// `settle` preserves contents and order, moves everything on `done`,
2164        /// and coalesces small neighbors.
2165        #[mz_ore::test]
2166        #[cfg_attr(miri, ignore)]
2167        fn settle_preserves_and_packs(
2168            input in arb_consolidated(),
2169            cuts in prop::collection::vec(0usize..7, 1..8),
2170        ) {
2171            let mut chunks = chunked(&input, &cuts);
2172            let mut out = VecDeque::new();
2173            TestChunk::settle(&mut chunks, true, &mut out);
2174            prop_assert!(chunks.is_empty());
2175            // Test chunks are far below the byte threshold, so maximal
2176            // packing coalesces everything into a single chunk.
2177            prop_assert!(out.len() <= 1);
2178            prop_assert_eq!(collect_chunks(out), input);
2179        }
2180
2181        /// `ChunkBatch::extract_into` over sorted, deduplicated probe keys
2182        /// equals the reference filter, resident and spilled alike, straddled
2183        /// keys included.
2184        #[mz_ore::test]
2185        #[cfg_attr(miri, ignore)]
2186        fn unload_extract_matches_filter(
2187            input in arb_consolidated(),
2188            cuts in prop::collection::vec(0usize..7, 0..8),
2189            probe_keys in prop::collection::btree_set(0u64..6, 0..6),
2190            spill in any::<bool>(),
2191        ) {
2192            prop_assume!(!input.is_empty());
2193            let pool = test_pool();
2194            let chunks: Vec<TestChunk> = if spill {
2195                chunked_spilled(&input, &cuts, &pool).into()
2196            } else {
2197                chunked(&input, &cuts).into()
2198            };
2199            let batch = ChunkBatch::new(chunks);
2200
2201            let mut probe_col = <u64 as Columnar>::Container::default();
2202            for key in &probe_keys {
2203                probe_col.push(*key);
2204            }
2205            let mut staging = <Tuple as Columnar>::Container::default();
2206            batch.extract_into(probe_col.borrow(), &mut staging);
2207
2208            let expected: Vec<Tuple> = input
2209                .iter()
2210                .copied()
2211                .filter(|((k, _), _, _)| probe_keys.contains(k))
2212                .collect();
2213            prop_assert_eq!(collect_staging(&staging), expected);
2214
2215            // The scan path reproduces the batch exactly, resident and
2216            // spilled alike.
2217            let mut staging = <Tuple as Columnar>::Container::default();
2218            batch.fetch_into(&mut staging);
2219            prop_assert_eq!(collect_staging(&staging), input);
2220        }
2221    }
2222
2223    /// `locate` answers the three-way span comparison for every probe
2224    /// placement: below, within, and past the chunk's keys.
2225    #[mz_ore::test]
2226    fn locate_spans_keys() {
2227        let chunk = TestChunk::from_body(build_column(&[
2228            ((2, 0), 0, 1),
2229            ((4, 0), 0, 1),
2230            ((6, 0), 0, 1),
2231        ]));
2232        let mut probe_col = <u64 as Columnar>::Container::default();
2233        for key in [0u64, 2, 3, 6, 9] {
2234            probe_col.push(key);
2235        }
2236        let probes = probe_col.borrow();
2237        use std::cmp::Ordering::*;
2238        let expected = [Less, Equal, Equal, Equal, Greater];
2239        for (index, expected) in expected.iter().enumerate() {
2240            assert_eq!(chunk.locate(probes, index), *expected, "probe {index}");
2241        }
2242    }
2243
2244    /// Collect chunk contents while asserting each chunk's serialized size
2245    /// stays within `bound` bytes.
2246    fn collect_bounded(chunks: impl IntoIterator<Item = TestChunk>, bound: usize) -> Vec<Tuple> {
2247        let mut collected = Vec::new();
2248        for chunk in chunks {
2249            let col = chunk.into_body();
2250            let bytes = col.length_in_bytes();
2251            assert!(bytes <= bound, "chunk of {bytes} bytes exceeds {bound}");
2252            Extend::extend(&mut collected, collect_column(&col));
2253        }
2254        collected
2255    }
2256
2257    type WideUpdate = ((u64, String), u64, i64);
2258    type WideChunk = ColumnChunk<(u64, String), u64, i64, ComputeSpill>;
2259
2260    fn wide_column(keys: impl Iterator<Item = u64>, bytes: usize) -> ColumnBody<WideUpdate> {
2261        let mut column = ColumnBody::default();
2262        for key in keys {
2263            column.push_into(&((key, "x".repeat(bytes)), 0, 1));
2264        }
2265        column
2266    }
2267
2268    fn assert_wide_byte_bound(chunks: VecDeque<WideChunk>, expected: usize, payload_bytes: usize) {
2269        let mut keys = Vec::new();
2270        for chunk in chunks {
2271            let column = chunk.into_body();
2272            assert!(
2273                column.length_in_bytes() <= COMMIT_BYTES || column.len() == 1,
2274                "{} bytes in a {}-record chunk",
2275                column.length_in_bytes(),
2276                column.len(),
2277            );
2278            let view = column.borrow();
2279            for index in 0..view.len() {
2280                let ((key, payload), time, diff) = view.get(index);
2281                assert_eq!(payload.len(), payload_bytes);
2282                assert!(payload.iter().all(|byte| *byte == b'x'));
2283                assert_eq!((*time, *diff), (0, 1));
2284                keys.push(*key);
2285            }
2286        }
2287        assert_eq!(keys, (0..u64::cast_from(expected)).collect::<Vec<_>>());
2288    }
2289
2290    #[mz_ore::test]
2291    fn chunker_enforces_byte_bound() {
2292        let mut chunker = ChunkChunker::default();
2293        let mut input: Column<WideUpdate> = wide_column((0..4000).rev(), 3000).into();
2294        chunker.push_into(&mut input);
2295        let mut chunks = VecDeque::new();
2296        if let Some(chunk) = chunker.extract() {
2297            chunks.push_back(std::mem::take(chunk));
2298        }
2299        while let Some(chunk) = chunker.finish() {
2300            chunks.push_back(std::mem::take(chunk));
2301        }
2302        assert_wide_byte_bound(chunks, 4000, 3000);
2303    }
2304
2305    #[mz_ore::test]
2306    fn merge_settle_enforces_byte_bound() {
2307        let mut left = VecDeque::from([WideChunk::from_body(wide_column(
2308            (0..2000).step_by(2),
2309            2100,
2310        ))]);
2311        let mut right = VecDeque::from([WideChunk::from_body(wide_column(
2312            (1..2000).step_by(2),
2313            2100,
2314        ))]);
2315        let mut merged = VecDeque::new();
2316        while !left.is_empty() && !right.is_empty() {
2317            WideChunk::merge(&mut left, &mut right, &mut merged);
2318        }
2319        merged.append(&mut left);
2320        merged.append(&mut right);
2321        let mut settled = VecDeque::new();
2322        WideChunk::settle(&mut merged, true, &mut settled);
2323        assert_wide_byte_bound(settled, 2000, 2100);
2324    }
2325
2326    #[mz_ore::test]
2327    fn settle_enforces_byte_bound_after_coalescing() {
2328        let mut input = VecDeque::from([
2329            WideChunk::from_body(wide_column(0..400, 3000)),
2330            WideChunk::from_body(wide_column(400..800, 3000)),
2331        ]);
2332        let mut settled = VecDeque::new();
2333        WideChunk::settle(&mut input, true, &mut settled);
2334        assert_wide_byte_bound(settled, 800, 3000);
2335    }
2336
2337    #[mz_ore::test]
2338    fn settle_byte_bound_allows_indivisible_update() {
2339        let mut input = VecDeque::from([WideChunk::from_body(wide_column(0..1, 2 * COMMIT_BYTES))]);
2340        let mut settled = VecDeque::new();
2341        WideChunk::settle(&mut input, true, &mut settled);
2342        assert_wide_byte_bound(settled, 1, 2 * COMMIT_BYTES);
2343    }
2344
2345    /// One record just under the bound among many narrow ones, so a cut at
2346    /// the average width leaves the piece holding the wide record over the
2347    /// bound.
2348    #[mz_ore::test]
2349    fn push_bounded_splits_uneven_widths() {
2350        let mut column: ColumnBody<WideUpdate> = ColumnBody::default();
2351        column.push_into(&((0, "x".repeat(COMMIT_BYTES - 4096)), 0, 1));
2352        for key in 1..20_000u64 {
2353            column.push_into(&((key, "x".to_string()), 0, 1));
2354        }
2355        assert!(column.length_in_bytes() > COMMIT_BYTES);
2356        let mut out = VecDeque::new();
2357        WideChunk::push_bounded(column, 3, &mut out);
2358        let mut keys = Vec::new();
2359        for chunk in out {
2360            let ColumnChunk::Resident(_, depth) = &chunk else {
2361                panic!("pieces are resident");
2362            };
2363            assert_eq!(*depth, 3, "pieces keep the input's depth");
2364            let column = chunk.into_body();
2365            assert!(
2366                column.length_in_bytes() <= COMMIT_BYTES || column.len() == 1,
2367                "{} bytes in a {}-record piece",
2368                column.length_in_bytes(),
2369                column.len(),
2370            );
2371            let view = column.borrow();
2372            for index in 0..view.len() {
2373                let ((key, _), _, _) = view.get(index);
2374                keys.push(*key);
2375            }
2376        }
2377        assert_eq!(keys, (0..20_000u64).collect::<Vec<_>>());
2378    }
2379
2380    /// Each side is under the bound, so any oversized output would have to
2381    /// come from `merge_from` itself. One input interleaves record by record,
2382    /// the other has runs long enough to gallop.
2383    #[mz_ore::test]
2384    fn merge_output_within_byte_bound_before_settle() {
2385        let interleaved = (
2386            wide_column((0..1600).step_by(2), 2100),
2387            wide_column((1..1600).step_by(2), 2100),
2388        );
2389        let runs = (
2390            wide_column((0..500).chain(1000..1100), 2100),
2391            wide_column(500..1000, 2100),
2392        );
2393        for (left, right) in [interleaved, runs] {
2394            let expected = left.len() + right.len();
2395            let mut left = VecDeque::from([WideChunk::from_body(left)]);
2396            let mut right = VecDeque::from([WideChunk::from_body(right)]);
2397            let mut merged = VecDeque::new();
2398            while !left.is_empty() && !right.is_empty() {
2399                WideChunk::merge(&mut left, &mut right, &mut merged);
2400            }
2401            assert!(merged.len() > 1, "merge never cut its output");
2402            merged.append(&mut left);
2403            merged.append(&mut right);
2404            assert_wide_byte_bound(merged, expected, 2100);
2405        }
2406    }
2407
2408    #[mz_ore::test]
2409    fn settle_split_preserves_depth() {
2410        let column = Rc::new(wide_column(0..2000, 2100));
2411        let shared = Rc::clone(&column);
2412        let mut input = VecDeque::from([WideChunk::Resident(column, 3)]);
2413        let mut settled = VecDeque::new();
2414        WideChunk::settle(&mut input, true, &mut settled);
2415        assert!(settled.len() > 1, "oversized chunk was not split");
2416        for chunk in &settled {
2417            assert_eq!(chunk.depth(), 3, "split piece changed generation");
2418        }
2419        assert_eq!(shared.len(), 2000, "shared body was mutated");
2420        assert_wide_byte_bound(settled, 2000, 2100);
2421    }
2422
2423    /// The last piece of a split coalesces with a shallower neighbor and
2424    /// commits at the deeper depth.
2425    #[mz_ore::test]
2426    fn settle_split_then_coalesce_keeps_depth() {
2427        let mut input = VecDeque::from([
2428            WideChunk::Resident(Rc::new(wide_column(0..800, 3000)), 2),
2429            WideChunk::Resident(Rc::new(wide_column(800..810, 3000)), 0),
2430        ]);
2431        let mut settled = VecDeque::new();
2432        WideChunk::settle(&mut input, true, &mut settled);
2433        assert!(settled.len() > 1, "oversized chunk was not split");
2434        for chunk in &settled {
2435            assert_eq!(chunk.depth(), 2, "coalescing lost the deeper generation");
2436        }
2437        assert_wide_byte_bound(settled, 810, 3000);
2438    }
2439
2440    /// Chunks of just over half the bound: unless settle moves part of each
2441    /// chunk into the column it is accumulating, every chunk commits at about
2442    /// half a slot.
2443    #[mz_ore::test]
2444    fn settle_fills_accumulated_column_from_next_chunk() {
2445        let per_chunk = 370;
2446        let mut input: VecDeque<WideChunk> = (0..8u64)
2447            .map(|i| {
2448                let start = i * per_chunk;
2449                WideChunk::from_body(wide_column(start..start + per_chunk, 3000))
2450            })
2451            .collect();
2452        let chunk_bytes = input[0].clone().into_body().length_in_bytes();
2453        assert!(2 * chunk_bytes > COMMIT_BYTES && chunk_bytes < COMMIT_BYTES);
2454        let mut settled = VecDeque::new();
2455        WideChunk::settle(&mut input, true, &mut settled);
2456        let sizes: Vec<usize> = settled
2457            .iter()
2458            .map(|chunk| chunk.clone().into_body().length_in_bytes())
2459            .collect();
2460        let (last, rest) = sizes.split_last().expect("settled output");
2461        assert!(*last <= COMMIT_BYTES);
2462        for size in rest {
2463            assert!(
2464                *size >= COMMIT_BYTES - COMMIT_BYTES / 10 && *size <= COMMIT_BYTES,
2465                "committed {size} bytes, sizes {sizes:?}",
2466            );
2467        }
2468        assert_wide_byte_bound(settled, 8 * 370, 3000);
2469    }
2470
2471    /// Advancing a large input cuts the output into several chunks near the
2472    /// ship threshold, and their concatenation is the reference result.
2473    #[mz_ore::test]
2474    #[cfg_attr(miri, ignore)]
2475    fn advance_cuts_large_output() {
2476        let records: Vec<Tuple> = (0..300_000u64).map(|k| ((k, 0), 0, 1)).collect();
2477        let mut input = VecDeque::from([ColumnChunk::from_body(build_column(&records))]);
2478        let frontier = Antichain::from_elem(0u64);
2479        let mut out = VecDeque::new();
2480        TestChunk::advance(&mut input, frontier.borrow(), true, &mut out);
2481        assert!(input.is_empty());
2482        assert!(
2483            out.len() >= 2,
2484            "expected a cut output, got {} chunk(s)",
2485            out.len()
2486        );
2487        assert_eq!(collect_bounded(out, 2 * COMMIT_BYTES), records);
2488    }
2489
2490    /// An input that is entirely one `D` group is withheld whole as the
2491    /// carry unless `done`: none of it is provably complete.
2492    #[mz_ore::test]
2493    #[cfg_attr(miri, ignore)]
2494    fn advance_withholds_giant_group() {
2495        let records: Vec<Tuple> = (0..100u64).map(|t| ((7, 7), t, 1)).collect();
2496        let mut input: VecDeque<TestChunk> = VecDeque::new();
2497        for piece in records.chunks(30) {
2498            input.push_back(ColumnChunk::from_body(build_column(piece)));
2499        }
2500        let frontier = Antichain::from_elem(50u64);
2501        let mut out = VecDeque::new();
2502        TestChunk::advance(&mut input, frontier.borrow(), false, &mut out);
2503        assert!(out.is_empty(), "nothing may ship from a single open group");
2504        assert_eq!(input.len(), 1, "the whole input becomes one carry chunk");
2505        // Sealing the carry advances and consolidates it.
2506        TestChunk::advance(&mut input, frontier.borrow(), true, &mut out);
2507        assert!(input.is_empty());
2508        let advanced = records.iter().map(|&(d, t, r)| (d, t.max(50), r)).collect();
2509        assert_eq!(collect_chunks(out), consolidate(advanced));
2510    }
2511
2512    /// Chunks the extract frontier does not split pass through whole from
2513    /// their resident time bounds: spilled bodies land on their side still
2514    /// spilled, with no load or re-commit, and a kept chunk's minimal times
2515    /// feed the residual frontier.
2516    #[mz_ore::test]
2517    #[cfg_attr(miri, ignore)]
2518    fn extract_passes_frontier_disjoint_chunks_through() {
2519        set_spill_override(Some(test_pool()));
2520        let low: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), i % 4, 1)).collect();
2521        let high: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 6 + i % 4, 1)).collect();
2522        let spilled_chunk = |data: &[Tuple]| {
2523            let chunk = TestChunk::commit(build_column(&consolidate(data.to_vec())), 1);
2524            assert!(chunk.is_spilled());
2525            chunk
2526        };
2527
2528        // A frontier between the two chunks' time ranges: the low chunk
2529        // ships whole and the high chunk keeps whole, both still spilled
2530        // (no load, no re-commit), and the residual is the kept chunk's
2531        // minimal time.
2532        let mut input = VecDeque::from([spilled_chunk(&low), spilled_chunk(&high)]);
2533        let frontier = Antichain::from_elem(5u64);
2534        let mut residual = Antichain::new();
2535        let (mut keep, mut ship) = (VecDeque::new(), VecDeque::new());
2536        while !input.is_empty() {
2537            TestChunk::extract(
2538                &mut input,
2539                frontier.borrow(),
2540                &mut residual,
2541                &mut keep,
2542                &mut ship,
2543            );
2544        }
2545        assert_eq!(ship.len(), 1);
2546        assert!(ship[0].is_spilled(), "shipped whole: body untouched");
2547        assert_eq!(keep.len(), 1);
2548        assert!(keep[0].is_spilled(), "kept whole: body untouched");
2549        assert_eq!(residual, Antichain::from_elem(6));
2550        let shipped = ship.pop_front().unwrap().into_body();
2551        assert_eq!(collect_column(&shipped), consolidate(low));
2552        let kept = keep.pop_front().unwrap().into_body();
2553        assert_eq!(collect_column(&kept), consolidate(high));
2554        set_spill_override(None);
2555    }
2556
2557    /// Extracting a large chunk at an intermediate frontier cuts both sides
2558    /// into several chunks and partitions exactly by time.
2559    #[mz_ore::test]
2560    #[cfg_attr(miri, ignore)]
2561    fn extract_cuts_large_output() {
2562        let records: Vec<Tuple> = (0..300_000u64).map(|k| ((k, 0), k % 2, 1)).collect();
2563        let mut input = VecDeque::from([ColumnChunk::from_body(build_column(&records))]);
2564        let frontier = Antichain::from_elem(1u64);
2565        let mut residual = Antichain::new();
2566        let (mut keep, mut ship) = (VecDeque::new(), VecDeque::new());
2567        while !input.is_empty() {
2568            TestChunk::extract(
2569                &mut input,
2570                frontier.borrow(),
2571                &mut residual,
2572                &mut keep,
2573                &mut ship,
2574            );
2575        }
2576        assert!(
2577            keep.len() >= 2,
2578            "expected a cut keep side, got {} chunk(s)",
2579            keep.len()
2580        );
2581        assert!(
2582            ship.len() >= 2,
2583            "expected a cut ship side, got {} chunk(s)",
2584            ship.len()
2585        );
2586        let kept: Vec<Tuple> = records.iter().copied().filter(|r| r.1 >= 1).collect();
2587        let shipped: Vec<Tuple> = records.iter().copied().filter(|r| r.1 < 1).collect();
2588        assert_eq!(collect_bounded(keep, 2 * COMMIT_BYTES), kept);
2589        assert_eq!(collect_bounded(ship, 2 * COMMIT_BYTES), shipped);
2590        assert_eq!(residual, Antichain::from_elem(1));
2591    }
2592
2593    /// `locate` answers from resident metadata on spilled chunks and follows
2594    /// the probe-relative-to-span convention.
2595    #[mz_ore::test]
2596    fn locate_uses_resident_bounds() {
2597        let pool = test_pool();
2598        let data: Vec<Tuple> = vec![((2, 0), 0, 1), ((4, 0), 0, 1)];
2599        let chunk = force_spill(ColumnChunk::from_body(build_column(&data)), &pool);
2600
2601        let mut probe_col = <u64 as Columnar>::Container::default();
2602        for key in [1u64, 3, 5] {
2603            probe_col.push(key);
2604        }
2605        let probes = probe_col.borrow();
2606        assert_eq!(chunk.locate(probes, 0), std::cmp::Ordering::Less);
2607        assert_eq!(chunk.locate(probes, 1), std::cmp::Ordering::Equal);
2608        assert_eq!(chunk.locate(probes, 2), std::cmp::Ordering::Greater);
2609    }
2610
2611    /// A body large enough to spill round-trips through the pool with resident
2612    /// metadata intact, and the batcher produces spilled sealed output.
2613    #[mz_ore::test]
2614    #[cfg_attr(miri, ignore)] // too slow
2615    fn spill_round_trip() {
2616        set_spill_override(Some(test_pool()));
2617
2618        let data: Vec<Tuple> = (0..40_000u64)
2619            .map(|i| ((i / 4, i % 4), i % 8, 1i64))
2620            .collect();
2621        let data = consolidate(data);
2622
2623        let column = build_column(&data);
2624        let committed = TestChunk::commit(column, 0);
2625        assert!(committed.is_spilled(), "large body must spill");
2626        assert_eq!(committed.len(), data.len());
2627        assert_eq!(collect_column(&committed.clone().into_body()), data);
2628
2629        let mut batcher = TestChunkBatcher::new(None, 0);
2630        for piece in data.chunks(10_000) {
2631            batcher.push_chunk(ColumnChunk::from_body(build_column(piece)));
2632        }
2633        let (sealed, _frontier) = batcher.extract(Antichain::new().borrow());
2634        let sealed = sealed.unwrap_or_default();
2635        assert!(
2636            sealed.iter().any(ColumnChunk::is_spilled),
2637            "sealed output should contain spilled chunks",
2638        );
2639        assert_eq!(collect_chunks(sealed), data);
2640
2641        set_spill_override(None);
2642    }
2643
2644    /// Merging spilled chains loads bodies call-scoped and consolidates
2645    /// correctly, and an untouched survivor keeps its spilled body.
2646    #[mz_ore::test]
2647    #[cfg_attr(miri, ignore)] // too slow
2648    fn merge_spilled_chains() {
2649        set_spill_override(Some(test_pool()));
2650
2651        let a: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect();
2652        let b: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 2i64)).collect();
2653
2654        let mut in1 = VecDeque::from([TestChunk::commit(build_column(&a), 0)]);
2655        let mut in2 = VecDeque::from([TestChunk::commit(build_column(&b), 0)]);
2656        assert!(in1[0].is_spilled() && in2[0].is_spilled());
2657
2658        let mut out = VecDeque::new();
2659        while !in1.is_empty() && !in2.is_empty() {
2660            TestChunk::merge(&mut in1, &mut in2, &mut out);
2661        }
2662        for tail in in1.drain(..).chain(in2.drain(..)) {
2663            out.push_back(tail);
2664        }
2665
2666        let expected: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 3i64)).collect();
2667        assert_eq!(collect_chunks(out), expected);
2668
2669        set_spill_override(None);
2670    }
2671
2672    /// A merge whose fronts have disjoint key ranges pushes the untouched
2673    /// survivor back in its original (spilled) form rather than rewriting it.
2674    #[mz_ore::test]
2675    fn merge_untouched_survivor_stays_spilled() {
2676        let pool = test_pool();
2677        let low: Vec<Tuple> = (0..100u64).map(|i| ((i, 0), 0, 1i64)).collect();
2678        let high: Vec<Tuple> = (1000..1100u64).map(|i| ((i, 0), 0, 1i64)).collect();
2679
2680        let mut in1 = VecDeque::from([force_spill(
2681            ColumnChunk::from_body(build_column(&low)),
2682            &pool,
2683        )]);
2684        let mut in2 = VecDeque::from([force_spill(
2685            ColumnChunk::from_body(build_column(&high)),
2686            &pool,
2687        )]);
2688        let mut out = VecDeque::new();
2689        TestChunk::merge(&mut in1, &mut in2, &mut out);
2690
2691        // `low` is fully consumed. `high` was never touched and must come
2692        // back spilled.
2693        assert!(in1.is_empty());
2694        assert_eq!(in2.len(), 1);
2695        assert!(in2[0].is_spilled(), "untouched survivor must stay spilled");
2696        let mut all = collect_chunks(out);
2697        Extend::extend(&mut all, collect_chunks(in2.drain(..)));
2698        let mut expected = low;
2699        Extend::extend(&mut expected, high);
2700        assert_eq!(all, expected);
2701    }
2702
2703    /// Merge output is one generation past its deepest input, a survivor
2704    /// rewritten from its remainder keeps its own depth, and a chunk passed
2705    /// through the disjoint fast path ages by its survival.
2706    #[mz_ore::test]
2707    fn merge_derives_generational_depth() {
2708        let low: Vec<Tuple> = (0..100u64).map(|i| ((i, 0), 0, 1i64)).collect();
2709        let high: Vec<Tuple> = (50..150u64).map(|i| ((i, 0), 0, 1i64)).collect();
2710        let mut in1 = VecDeque::from([ColumnChunk::from_body(build_column(&low))]);
2711        let mut in2 = VecDeque::from([ColumnChunk::from_body(build_column(&high))]);
2712        assert_eq!(in1[0].depth(), 0, "fresh chunks start at depth 0");
2713        let mut out = VecDeque::new();
2714        TestChunk::merge(&mut in1, &mut in2, &mut out);
2715        assert!(!out.is_empty());
2716        for chunk in &out {
2717            assert_eq!(chunk.depth(), 1, "merge output is one past its inputs");
2718        }
2719        // The merge runs through the shared horizon, so `high` survives with
2720        // its unmerged remainder at its original depth.
2721        assert!(in1.is_empty());
2722        assert_eq!(in2.len(), 1);
2723        assert_eq!(in2[0].depth(), 0, "rewritten survivor keeps its depth");
2724
2725        // A disjoint merge moves the lower front to the output with its data
2726        // unchanged, one generation older for having outlived the merge.
2727        let mut in1 = VecDeque::from([ColumnChunk::Resident(Rc::new(build_column(&low)), 3)]);
2728        let far: Vec<Tuple> = (1000..1100u64).map(|i| ((i, 0), 0, 1i64)).collect();
2729        let mut in2 = VecDeque::from([ColumnChunk::from_body(build_column(&far))]);
2730        let mut out = VecDeque::new();
2731        TestChunk::merge(&mut in1, &mut in2, &mut out);
2732        assert_eq!(out.len(), 1);
2733        assert_eq!(out[0].depth(), 4, "pass-through ages a generation");
2734        assert_eq!(collect_chunks(out), low);
2735    }
2736
2737    /// Advance output and carry keep the deepest input depth, since
2738    /// compaction rewrites within a generation.
2739    #[mz_ore::test]
2740    fn advance_preserves_depth() {
2741        let data: Vec<Tuple> = (0..100u64).map(|i| ((i, 0), 1, 1i64)).collect();
2742        let mut input = VecDeque::from([
2743            ColumnChunk::Resident(Rc::new(build_column(&data[..50])), 2),
2744            ColumnChunk::Resident(Rc::new(build_column(&data[50..])), 1),
2745        ]);
2746        let frontier = Antichain::from_elem(5u64);
2747        let mut out = VecDeque::new();
2748        TestChunk::advance(&mut input, frontier.borrow(), false, &mut out);
2749        for chunk in out.iter().chain(input.iter()) {
2750            assert_eq!(chunk.depth(), 2);
2751        }
2752        TestChunk::advance(&mut input, frontier.borrow(), true, &mut out);
2753        assert!(input.is_empty());
2754        assert!(!out.is_empty());
2755        for chunk in &out {
2756            assert_eq!(chunk.depth(), 2);
2757        }
2758    }
2759
2760    /// Settle commits at the deepest depth among coalesced chunks, and a
2761    /// commit large enough to spill carries the depth into its spilled
2762    /// metadata (and thus into the pool hints).
2763    #[mz_ore::test]
2764    #[cfg_attr(miri, ignore)] // too slow
2765    fn settle_commits_at_accumulated_depth() {
2766        set_spill_override(Some(test_pool()));
2767        let big: Vec<Tuple> = (0..60_000u64).map(|i| ((i, 0), 0, 1i64)).collect();
2768        let mut input = VecDeque::from([
2769            ColumnChunk::Resident(Rc::new(build_column(&big)), 1),
2770            ColumnChunk::Resident(Rc::new(build_column(&[((0, 0), 0, 1)])), 0),
2771            ColumnChunk::Resident(Rc::new(build_column(&[((1, 0), 0, 1)])), 2),
2772        ]);
2773        let mut out = VecDeque::new();
2774        TestChunk::settle(&mut input, true, &mut out);
2775        assert!(input.is_empty());
2776        assert_eq!(out.len(), 2);
2777        assert!(out[0].is_spilled(), "large commit must spill");
2778        assert_eq!(out[0].depth(), 1, "sole commit keeps its depth");
2779        assert!(!out[1].is_spilled(), "small commit stays resident");
2780        assert_eq!(out[1].depth(), 2, "coalesced commit takes the max depth");
2781        set_spill_override(None);
2782    }
2783
2784    /// The settle carry commits at a monotone size threshold rather than the
2785    /// periodic ship window, so mid-window chunk sizes cannot make it grow
2786    /// past the target unbounded.
2787    #[mz_ore::test]
2788    #[cfg_attr(miri, ignore)] // too slow
2789    fn settle_carry_commits_at_target() {
2790        // Two inputs fit in one slot. A third must start another chunk.
2791        let chunk_rows = u64::cast_from(800_000usize / 32);
2792        let mut input: VecDeque<TestChunk> = (0..4u64)
2793            .map(|c| {
2794                let data: Vec<Tuple> = (0..chunk_rows)
2795                    .map(|i| ((c * chunk_rows + i, 0), 0, 1i64))
2796                    .collect();
2797                ColumnChunk::from_body(build_column(&data))
2798            })
2799            .collect();
2800        let mut out = VecDeque::new();
2801        TestChunk::settle(&mut input, true, &mut out);
2802        // Catches the fixture drifting above `at_commit_size`, where settle
2803        // commits each chunk as-is and the size cap below holds vacuously.
2804        assert!(out.len() < 4, "nothing coalesced");
2805        for chunk in &out {
2806            let col = chunk.clone().into_body();
2807            assert!(
2808                col.length_in_bytes() <= COMMIT_BYTES,
2809                "settled chunk of {} bytes exceeds the commit target",
2810                col.length_in_bytes(),
2811            );
2812        }
2813        assert_eq!(
2814            collect_chunks(out).len(),
2815            usize::try_from(4 * chunk_rows).unwrap(),
2816        );
2817    }
2818
2819    #[mz_ore::test]
2820    fn small_chunks_stay_resident() {
2821        set_spill_override(Some(test_pool()));
2822        let committed = TestChunk::commit(build_column(&[((1, 1), 0, 1)]), 0);
2823        assert!(!committed.is_spilled());
2824        set_spill_override(None);
2825    }
2826
2827    /// The smallest column whose serialized size reaches `SPILL_MIN_BYTES`.
2828    /// One record less sits under the spill floor.
2829    fn column_at_spill_floor() -> (ColumnBody<Tuple>, u64) {
2830        let mut col: ColumnBody<Tuple> = ColumnBody::default();
2831        let mut n = 0u64;
2832        while col.length_in_bytes() < SPILL_MIN_BYTES {
2833            col.push_into(((n, n), 0, 1));
2834            n += 1;
2835        }
2836        (col, n)
2837    }
2838
2839    /// Bodies straddling the spill floor: one record under stays resident,
2840    /// at the floor spills.
2841    #[mz_ore::test]
2842    fn spill_floor_boundary() {
2843        set_spill_override(Some(test_pool()));
2844        let (col, n) = column_at_spill_floor();
2845        let mut under: ColumnBody<Tuple> = ColumnBody::default();
2846        for m in 0..n - 1 {
2847            under.push_into(((m, m), 0, 1));
2848        }
2849        assert!(under.length_in_bytes() < SPILL_MIN_BYTES);
2850        assert!(!TestChunk::commit(under, 0).is_spilled());
2851        assert!(TestChunk::commit(col, 0).is_spilled());
2852        set_spill_override(None);
2853    }
2854
2855    /// The compression depth floor picks the codec, not whether a body
2856    /// spills: shallow generations store at identity, the floor and deeper
2857    /// at lz4, and every depth spills and round-trips.
2858    #[mz_ore::test]
2859    fn spill_codec_depth_floor() {
2860        set_spill_override(Some(test_pool()));
2861        set_compress_min_depth_override(Some(2));
2862        // Codec identity via Debug: ZST statics and dyn vtables make
2863        // pointer comparison unreliable. The flag must agree with the codec,
2864        // since it is what decides whether a body wants migrating.
2865        let codec_name = |depth: u8| {
2866            let (codec, compressed) = codec_for_depth(depth);
2867            let name = format!("{:?}", codec);
2868            assert_eq!(compressed, name == "Lz4Codec", "flag tracks the codec");
2869            name
2870        };
2871        assert_eq!(codec_name(0), "IdentityCodec");
2872        assert_eq!(codec_name(1), "IdentityCodec");
2873        assert_eq!(codec_name(2), "Lz4Codec");
2874        assert_eq!(codec_name(u8::MAX), "Lz4Codec");
2875
2876        let data: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect();
2877        let data = consolidate(data);
2878        let column = build_column(&data);
2879        for depth in [0u8, 1, 2, 3] {
2880            let chunk = TestChunk::commit(column.clone(), depth);
2881            assert!(chunk.is_spilled(), "depth {depth} must spill");
2882            assert_eq!(collect_column(&chunk.into_body()), data);
2883        }
2884        set_spill_override(None);
2885        set_compress_min_depth_override(None);
2886
2887        // The default floor stores only fresh (depth 0) bodies at identity.
2888        set_compress_min_depth_override(Some(DEFAULT_COMPRESS_MIN_DEPTH));
2889        assert_eq!(codec_name(0), "IdentityCodec");
2890        assert_eq!(codec_name(1), "Lz4Codec");
2891        set_compress_min_depth_override(None);
2892    }
2893
2894    /// A chunk a merge carries forward untouched ages a generation, and a
2895    /// spilled body crossing the compression floor by doing so is re-spilled
2896    /// under the compressing codec. Key-disjoint input takes that path on
2897    /// every merge, so without the crossing its backlog would stay
2898    /// identity-coded for as long as it lived.
2899    #[mz_ore::test]
2900    #[cfg_attr(miri, ignore)] // too slow
2901    fn merge_survivor_crosses_compression_floor() {
2902        set_spill_override(Some(test_pool()));
2903        set_compress_min_depth_override(Some(1));
2904
2905        let low = consolidate((0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2906        let far = consolidate((100_000..120_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2907        let fresh_far = || VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
2908
2909        // Fresh spilled chunks sit below the floor, so both store identity
2910        // coded, and their data ranges are disjoint.
2911        let mut in1 = VecDeque::from([TestChunk::commit(build_column(&low), 0)]);
2912        let mut in2 = fresh_far();
2913        assert!(in1[0].is_spilled() && in2[0].is_spilled());
2914        assert!(
2915            !body_compressed(&in1[0]),
2916            "a fresh body below the floor is identity coded"
2917        );
2918
2919        let mut out = VecDeque::new();
2920        TestChunk::merge(&mut in1, &mut in2, &mut out);
2921        assert_eq!(out.len(), 1);
2922        let survived = out.pop_front().expect("the lower front passes through");
2923        assert_eq!(survived.depth(), 1, "survival ages across the floor");
2924        assert!(
2925            survived.is_spilled(),
2926            "the crossing re-spills, it does not evict"
2927        );
2928        assert!(
2929            body_compressed(&survived),
2930            "the survivor is re-spilled under the compressing codec"
2931        );
2932
2933        // Past the floor the next survival is a metadata-only bump: the body
2934        // is already compressed and stays where it is.
2935        let mut in1 = VecDeque::from([survived]);
2936        let mut in2 = fresh_far();
2937        let mut out = VecDeque::new();
2938        TestChunk::merge(&mut in1, &mut in2, &mut out);
2939        assert_eq!(out.len(), 1);
2940        assert_eq!(out[0].depth(), 2, "an aged survivor keeps aging");
2941        assert!(out[0].is_spilled());
2942        assert_eq!(
2943            collect_chunks(out),
2944            low,
2945            "the body reads back intact across both survivals"
2946        );
2947
2948        set_spill_override(None);
2949        set_compress_min_depth_override(None);
2950    }
2951
2952    /// Aging does not depend on holding the only reference to a body. The
2953    /// trace's compaction merger feeds `merge` clones of a source batch's
2954    /// chunks and keeps the batch alive throughout, so a shared body must
2955    /// still age. It must not be re-spilled: the other holder goes on
2956    /// storing the original whatever this reference does, and the merger
2957    /// rewrites its clone immediately.
2958    #[mz_ore::test]
2959    #[cfg_attr(miri, ignore)] // too slow
2960    fn merge_survivor_ages_while_shared() {
2961        set_spill_override(Some(test_pool()));
2962        set_compress_min_depth_override(Some(1));
2963
2964        let low = consolidate((0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2965        let far = consolidate((100_000..120_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2966
2967        // The source batch's chunk, held for the whole merge as the spine
2968        // holds it.
2969        let source = TestChunk::commit(build_column(&low), 0);
2970        let ColumnChunk::Spilled(source_body, 0) = &source else {
2971            panic!("a fresh commit above the spill floor is spilled at depth 0");
2972        };
2973        let source_body = Rc::clone(source_body);
2974
2975        let mut in1 = VecDeque::from([source.clone()]);
2976        let mut in2 = VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
2977        let mut out = VecDeque::new();
2978        TestChunk::merge(&mut in1, &mut in2, &mut out);
2979
2980        assert_eq!(out.len(), 1);
2981        assert_eq!(out[0].depth(), 1, "a shared body ages all the same");
2982        let ColumnChunk::Spilled(survived_body, _) = &out[0] else {
2983            panic!("the survivor stays spilled");
2984        };
2985        assert!(
2986            Rc::ptr_eq(&source_body, survived_body),
2987            "a shared body is aged in place, not re-spilled"
2988        );
2989        assert_eq!(source.depth(), 0, "the other holder is left as it was");
2990
2991        // Past the floor, where no re-spill is in question, a shared body
2992        // goes on aging rather than pinning at the crossing depth.
2993        let mut in1 = VecDeque::from([out.pop_front().expect("survivor observed above")]);
2994        let mut in2 = VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
2995        let mut out = VecDeque::new();
2996        TestChunk::merge(&mut in1, &mut in2, &mut out);
2997        assert_eq!(out.len(), 1);
2998        assert_eq!(out[0].depth(), 2, "aging past the floor is not pinned");
2999        assert_eq!(collect_chunks(out), low);
3000
3001        set_spill_override(None);
3002        set_compress_min_depth_override(None);
3003    }
3004
3005    /// A migration that cannot happen when a body first qualifies is retried
3006    /// at the next survival, never consumed. Each case leaves an
3007    /// identity-coded body at or past the floor, which would be stranded
3008    /// uncompressed for the rest of its life if the test were a depth
3009    /// transition rather than the body's stored codec.
3010    #[mz_ore::test]
3011    #[cfg_attr(miri, ignore)] // too slow
3012    fn survive_merge_retries_missed_migrations() {
3013        let low = consolidate((0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
3014        let far = consolidate((100_000..120_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
3015
3016        // Age a chunk one generation through a disjoint merge, which passes
3017        // the lower front through `survive_merge`.
3018        let survive = |chunk: TestChunk| {
3019            let mut in1 = VecDeque::from([chunk]);
3020            let mut in2 = VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
3021            let mut out = VecDeque::new();
3022            TestChunk::merge(&mut in1, &mut in2, &mut out);
3023            out.pop_front().expect("the lower front passes through")
3024        };
3025
3026        // No pool installed when the body qualifies: spilling can be toggled
3027        // off at runtime while existing handles stay valid.
3028        set_spill_override(Some(test_pool()));
3029        set_compress_min_depth_override(Some(1));
3030        let chunk = TestChunk::commit(build_column(&low), 0);
3031        assert!(!body_compressed(&chunk));
3032        set_spill_override(None);
3033        let chunk = survive(chunk);
3034        assert_eq!(chunk.depth(), 1, "aging does not need a pool");
3035        assert!(!body_compressed(&chunk), "no pool, no migration");
3036        set_spill_override(Some(test_pool()));
3037        let chunk = survive(chunk);
3038        assert!(
3039            body_compressed(&chunk),
3040            "the migration retries once a pool is back"
3041        );
3042
3043        // Shared when the body qualifies: the compaction merger holds the
3044        // source batch while merging clones of its chunks.
3045        let chunk = TestChunk::commit(build_column(&low), 0);
3046        let held = chunk.clone();
3047        let chunk = survive(chunk);
3048        assert!(!body_compressed(&chunk), "shared, so not migrated");
3049        drop(held);
3050        let chunk = survive(chunk);
3051        assert!(
3052            body_compressed(&chunk),
3053            "the migration retries once the body is unshared"
3054        );
3055
3056        // The floor lowered long after the body spilled, which is what an
3057        // operator reaches for under pool pressure. Nothing here is a
3058        // transition: the body is already several generations past the new
3059        // floor when it moves.
3060        set_compress_min_depth_override(Some(8));
3061        let chunk = TestChunk::commit(build_column(&low), 3);
3062        assert!(!body_compressed(&chunk));
3063        set_compress_min_depth_override(Some(1));
3064        let chunk = survive(chunk);
3065        assert_eq!(chunk.depth(), 4);
3066        assert!(
3067            body_compressed(&chunk),
3068            "lowering the floor migrates bodies already past it"
3069        );
3070
3071        set_spill_override(None);
3072        set_compress_min_depth_override(None);
3073    }
3074
3075    /// One test, because the gates are process-global.
3076    #[mz_ore::test]
3077    #[cfg_attr(miri, ignore)]
3078    fn spill_gates_are_independent() {
3079        let installed =
3080            crate::pool_config::apply_pool_config(crate::pool_config::PoolPagerConfig {
3081                budget_bytes: 32 << 20,
3082                spill_threads: 1,
3083                eager_backing: false,
3084                rss_target_bytes: 16 << 20,
3085            });
3086        assert!(installed, "pool reservation failed");
3087        // A body at the spill floor, so the gates alone decide.
3088        let (col, _) = column_at_spill_floor();
3089        let compute = |col: &ColumnBody<Tuple>| TestChunk::commit(col.clone(), 0).is_spilled();
3090        let storage = |col: &ColumnBody<Tuple>| {
3091            ColumnChunk::<(u64, u64), u64, i64, StorageSpill>::commit(col.clone(), 0).is_spilled()
3092        };
3093        let sink = |col: &ColumnBody<Tuple>| try_spill_ref(col, 0).is_some();
3094
3095        assert!(!compute(&col), "all gates off");
3096        assert!(!storage(&col), "all gates off");
3097        assert!(!sink(&col), "all gates off");
3098
3099        set_storage_spill_enabled(true);
3100        assert!(storage(&col), "the storage gate spills storage chunks");
3101        assert!(
3102            !compute(&col),
3103            "the storage gate must not spill compute chunks"
3104        );
3105        assert!(!sink(&col), "the storage gate must not spill sink bodies");
3106        set_storage_spill_enabled(false);
3107
3108        set_compute_spill_enabled(true);
3109        assert!(compute(&col), "the compute gate spills compute chunks");
3110        assert!(
3111            !storage(&col),
3112            "the compute gate must not spill storage chunks"
3113        );
3114        assert!(!sink(&col), "the compute gate must not spill sink bodies");
3115        set_compute_spill_enabled(false);
3116
3117        set_sink_spill_enabled(true);
3118        assert!(sink(&col), "the sink gate spills sink bodies");
3119        assert!(
3120            !compute(&col),
3121            "the sink gate must not spill compute chunks"
3122        );
3123        assert!(
3124            !storage(&col),
3125            "the sink gate must not spill storage chunks"
3126        );
3127        set_sink_spill_enabled(false);
3128        set_compress_min_depth_override(None);
3129    }
3130
3131    /// Re-spilling an already-serialized body exercises the `ColumnBody::Words`
3132    /// branch of the pooled body's serialization and round-trips
3133    /// byte-identically. A zero budget sends every insert straight to an
3134    /// extent, so each body is serialized and reads back as words.
3135    #[mz_ore::test]
3136    #[cfg_attr(miri, ignore)] // reserves a full-size pool
3137    fn spill_align_round_trip() {
3138        let pool = Pool::new().expect("pool reservation");
3139        pool.set_budget(0);
3140        let data: Vec<Tuple> = (0..64u64).map(|k| ((k, k), 0, 1)).collect();
3141        let spilled = force_spill(ColumnChunk::from_body(build_column(&data)), &pool);
3142        let column = spilled.into_body();
3143        let ColumnBody::Words(words) = &column else {
3144            panic!("an evicted body reads back as ColumnBody::Words");
3145        };
3146        let words = words.clone();
3147        let respilled = force_spill(ColumnChunk::from_body(column), &pool);
3148        let reread = respilled.into_body();
3149        let ColumnBody::Words(words2) = &reread else {
3150            panic!("an evicted body reads back as ColumnBody::Words");
3151        };
3152        assert_eq!(&words, words2, "byte-identical round trip");
3153        assert_eq!(collect_column(&reread), data);
3154    }
3155
3156    /// Reads every chunk's time bounds while spill threads evict it. The
3157    /// spill threads fill the bounds under the chunk's state lock, so a
3158    /// reader that takes the locks in the other order hangs this test.
3159    #[mz_ore::test]
3160    #[cfg_attr(miri, ignore)] // reserves a full-size pool
3161    fn time_bounds_race_spill_thread_eviction() {
3162        let (done_tx, done_rx) = std::sync::mpsc::channel();
3163        std::thread::spawn(move || {
3164            let pool = Pool::new().expect("pool reservation");
3165            pool.set_budget(0);
3166            pool.set_spill_threads(4);
3167            let data: Vec<Tuple> = (0..2048u64).map(|k| ((k, k), k % 7, 1)).collect();
3168            for _ in 0..2000 {
3169                let chunk = force_spill(ColumnChunk::from_body(build_column(&data)), &pool);
3170                let (lower, upper) = chunk.chunk_time_bounds();
3171                assert_eq!(lower.elements(), &[0]);
3172                assert_eq!(&upper[..], &[6]);
3173            }
3174            pool.set_spill_threads(0);
3175            done_tx.send(()).expect("receiver alive");
3176        });
3177        done_rx
3178            .recv_timeout(std::time::Duration::from_secs(120))
3179            .expect("time bounds reads deadlocked against spill-thread eviction");
3180    }
3181
3182    #[mz_ore::test]
3183    fn merge_depth_saturates() {
3184        let a = ColumnChunk::Resident(
3185            Rc::new(build_column(&[((1, 0), 0, 1), ((3, 0), 0, 1)])),
3186            u8::MAX,
3187        );
3188        let b = ColumnChunk::Resident(
3189            Rc::new(build_column(&[((2, 0), 0, 1), ((4, 0), 0, 1)])),
3190            u8::MAX,
3191        );
3192        let mut in1 = VecDeque::from([a]);
3193        let mut in2 = VecDeque::from([b]);
3194        let mut out = VecDeque::new();
3195        TestChunk::merge(&mut in1, &mut in2, &mut out);
3196        for chunk in out.iter().chain(in1.iter()).chain(in2.iter()) {
3197            assert_eq!(chunk.depth(), u8::MAX, "depth saturates");
3198        }
3199    }
3200
3201    #[mz_ore::test]
3202    fn into_body_copies_shared_resident() {
3203        let data: Vec<Tuple> = vec![((1, 1), 0, 1), ((2, 2), 0, 1)];
3204        let a = TestChunk::from_body(build_column(&data));
3205        let b = a.clone();
3206        assert_eq!(collect_column(&a.into_body()), data);
3207        assert_eq!(collect_column(&b.into_body()), data);
3208    }
3209}