diff --git a/src/timely-util/proptest-regressions/columnar/chunk.txt b/src/timely-util/proptest-regressions/columnar/chunk.txt new file mode 100644 index 0000000000000..c6be3107fea94 --- /dev/null +++ b/src/timely-util/proptest-regressions/columnar/chunk.txt @@ -0,0 +1,8 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 79f9e0f3cf01ba590155d536a36d67864533c44d17ad3c38da58886159837681 # shrinks to input = [((0, 0), 0, 1)], cuts = [], probe_keys = {}, spill = false +cc 6a70b97870c5a545462247056ceac3dd0c69f0ff2fb855502d3ed8d87201c8ea # shrinks to inputs = [[((0, 3), 3, -2), ((0, 4), 0, 3), ((3, 3), 1, 2)]], cuts = [4, 0] diff --git a/src/timely-util/src/columnar.rs b/src/timely-util/src/columnar.rs index 03596b0c9c1fb..e66c66b07198e 100644 --- a/src/timely-util/src/columnar.rs +++ b/src/timely-util/src/columnar.rs @@ -20,6 +20,7 @@ pub mod batcher; pub mod builder; pub mod builder_input; +pub mod chunk; pub mod consolidate; pub mod merge_batcher; pub mod unload; @@ -175,26 +176,29 @@ where /// merger and chunks shipped from the builder are sized comparably. const SHIP_WORDS: usize = 1 << 18; -/// Returns true once the serialized size of `borrow` is within 10% of the next -/// `SHIP_WORDS` boundary. +/// Returns true once the serialized size of `borrow` reaches 10% under +/// `SHIP_WORDS`. /// -/// Same heuristic as `ColumnBuilder::push_into`; lifted out so the merger and -/// the `SizableContainer` impl agree on the ship signal. +/// Monotone in size, deliberately not a window below the boundary. A single +/// record wider than a window steps clear over it, and a ship signal that +/// un-fires past the boundary lets a chunk grow until it exceeds the buffer +/// pool's largest size class, past which a spilled body degrades to +/// permanently resident. The same heuristic as [`builder::ColumnBuilder`]'s +/// ship point, lifted out so the builder, the merger, and the +/// `SizableContainer` impl agree on the signal. #[inline] pub(crate) fn at_serialized_capacity<'a, A>(borrow: &A) -> bool where A: columnar::AsBytes<'a>, { - let words = indexed::length_in_words(borrow); - let round = (words + (SHIP_WORDS - 1)) & !(SHIP_WORDS - 1); - round - words < round / 10 + indexed::length_in_words(borrow) >= SHIP_WORDS - SHIP_WORDS / 10 } impl SizableContainer for Column { fn at_capacity(&self) -> bool { - // Match `ColumnBuilder`'s ship heuristic: serialized size within 10% - // of the next 2 MiB. Aligns chunk-size choices across the two paths - // and keeps recipients dealing with a single granularity. + // Match `ColumnBuilder`'s ship heuristic: serialized size at the + // 2 MiB ship threshold. Aligns chunk-size choices across the two + // paths and keeps recipients dealing with a single granularity. // // Serialized chunks (`Bytes` / `Align`) have no typed builder to push // into, so they're trivially "at capacity" — there's no further work @@ -365,4 +369,24 @@ mod tests { vec![&1, &2, &3] ); } + + /// The ship signal is monotone: once it fires it stays fired, even when + /// a single wide record steps far past the 2 MiB boundary in one push. + #[mz_ore::test] + fn ship_threshold_monotone() { + use columnar::Push; + let mut container = as Columnar>::Container::default(); + // Wider than 10% of any boundary a 25 MiB run can reach. + let wide: Vec = vec![0u64; 50_000]; + let mut fired = false; + for pushes in 1..=64 { + container.push(&wide); + let now = at_serialized_capacity(&container.borrow()); + if fired { + assert!(now, "ship signal un-fired at {pushes} records"); + } + fired = fired || now; + } + assert!(fired, "ship signal never fired"); + } } diff --git a/src/timely-util/src/columnar/batcher.rs b/src/timely-util/src/columnar/batcher.rs index 4d0d49ea0be56..2cb20107f1bbb 100644 --- a/src/timely-util/src/columnar/batcher.rs +++ b/src/timely-util/src/columnar/batcher.rs @@ -259,7 +259,7 @@ where /// Compared to a linear scan, this is `O(log K)` for a run of length `K` /// satisfying `cmp` — useful when one side of a sorted merge has long runs /// dominated by the other side. -fn gallop(upper: usize, lower: &mut usize, mut cmp: impl FnMut(usize) -> bool) { +pub(crate) fn gallop(upper: usize, lower: &mut usize, mut cmp: impl FnMut(usize) -> bool) { // If `cmp` is already false at `*lower`, the run is empty — nothing to do. if *lower < upper && cmp(*lower) { // Phase 1 (overshoot): advance by exponentially growing steps as long diff --git a/src/timely-util/src/columnar/builder.rs b/src/timely-util/src/columnar/builder.rs index 84102f7198fa3..55714d54cad5e 100644 --- a/src/timely-util/src/columnar/builder.rs +++ b/src/timely-util/src/columnar/builder.rs @@ -44,28 +44,25 @@ where #[inline] fn push_into(&mut self, item: T) { self.current.push(item); - // If there is less than 10% slop with 2MB backing allocations, mint a container. + // Mint a container once the serialized size reaches the ship threshold. use columnar::Borrow; - let words = indexed::length_in_words(&self.current.borrow()); - let round = (words + ((1 << 18) - 1)) & !((1 << 18) - 1); - if round - words < round / 10 { + if crate::columnar::at_serialized_capacity(&self.current.borrow()) { /// Move the contents from `current` to a `Vec` allocation built via /// `indexed::encode` (so no zero-init pre-pass), and push it to `pending`. #[cold] - fn outlined_align( - current: &mut C::Container, - words: usize, - pending: &mut VecDeque>, - ) where + fn outlined_align(current: &mut C::Container, pending: &mut VecDeque>) + where C: Columnar, { + use columnar::Borrow; + let words = indexed::length_in_words(¤t.borrow()); let mut alloc: Vec = Vec::with_capacity(words); indexed::encode(&mut alloc, ¤t.borrow()); pending.push_back(Column::Align(alloc)); current.clear(); } - outlined_align(&mut self.current, words, &mut self.pending); + outlined_align(&mut self.current, &mut self.pending); } } } diff --git a/src/timely-util/src/columnar/chunk.rs b/src/timely-util/src/columnar/chunk.rs new file mode 100644 index 0000000000000..002c0f012c072 --- /dev/null +++ b/src/timely-util/src/columnar/chunk.rs @@ -0,0 +1,1779 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +//! [`ColumnChunk`]: differential's [`Chunk`] over [`Column`]-shaped updates. +//! +//! A chunk is a sorted, consolidated run of `(D, T, R)` updates in the flat +//! columnar layout, in one of two homes: +//! +//! * **Resident**: an `Rc`-shared [`Column`] on the heap. Fresh input, merge +//! output, and small tails live here. +//! * **Spilled**: the serialized body in the process [`Pool`], with the record +//! count and the first and last data items resident. The pool owns residency +//! from there, with slots under a memory budget and compression and device +//! pageout under pressure, and a body that dies before pressure reaches it +//! is freed without I/O. +//! +//! Reads of a spilled body are copy-out and scoped to the call that needs +//! them: the body is read into caller-owned memory and no reference into pool +//! memory ever exists outside the pool. That contract is what lets the pool +//! evict with no reader accounting at all. +//! +//! Spilling happens in [`Chunk::settle`], the trait's designated commit point: +//! chunks moved to settled output are handed to the pool when spilling is +//! enabled (see [`set_compute_spill_enabled`] and [`set_storage_spill_enabled`]). +//! The spill destination resolves per commit from three pieces of mutable +//! state. A thread-local pool override, for tests and benches, wins outright. +//! Otherwise the compute and storage gates, composed as an OR, route commits +//! to the process pool installed by [`crate::pool_config`], and with no pool +//! installed chunks stay resident. A second thread-local holds the reusable +//! scratch that call-scoped reads of spilled bodies copy into. +//! Grading is by serialized bytes, the ship size +//! [`Column`] already targets, rather than by the record-count `TARGET`, +//! since record count does not bound bytes for variable-width data. +//! +//! Chunks whose data is a `(key, val)` pair additionally implement +//! [`UnloadChunk`], the bulk-read capability: sorted probe keys in, matching +//! updates appended to caller-owned staging, with `locate` answered from the +//! resident fence metadata so a probe set faults only the chunk bodies it +//! actually touches. + +use std::cell::RefCell; +use std::collections::VecDeque; +use std::rc::Rc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use columnar::bytes::indexed; +use columnar::{Borrow, BorrowedOf, Columnar, Container as _, FromBytes, Index, Len, Push as _}; +use differential_dataflow::difference::Semigroup; +use differential_dataflow::lattice::Lattice; +use differential_dataflow::trace::chunk::Chunk; +use mz_ore::cast::CastFrom; +use mz_ore::pool::{ChunkHandle, ChunkHints, ExtentCodec, Pool}; +use timely::Accountable; +use timely::container::{ContainerBuilder, PushInto}; +use timely::dataflow::channels::ContainerBytes; +use timely::progress::Timestamp; +use timely::progress::frontier::AntichainRef; + +use crate::columnar::batcher::{ColumnChunker, gallop}; +use crate::columnar::unload::UnloadChunk; +use crate::columnar::{Column, at_serialized_capacity}; + +/// Compute's leg of the process spill gate. See [`set_compute_spill_enabled`]. +static COMPUTE_SPILL_ENABLED: AtomicBool = AtomicBool::new(false); + +/// Storage's leg of the process spill gate. See [`set_storage_spill_enabled`]. +static STORAGE_SPILL_ENABLED: AtomicBool = AtomicBool::new(false); + +thread_local! { + /// A thread-scoped pool override, taking precedence over the global + /// enable flag and pool. Lets tests and benches spill through a private + /// pool without touching process-global state. + static SPILL_OVERRIDE: RefCell> = const { RefCell::new(None) }; + + /// Reusable staging for call-scoped reads of spilled bodies. + static READ_SCRATCH: RefCell> = const { RefCell::new(Vec::new()) }; +} + +/// Enable or disable chunk spilling on behalf of compute's arrangement +/// batchers. +/// +/// Chunks carry no subsystem identity, so the spill decision is process-wide: +/// committed chunks spill while *either* the compute or the storage gate is +/// set. Each subsystem's config application writes only its own gate, so the +/// two dyncfg flags compose as an OR instead of clobbering each other. +/// +/// Takes effect at the next `settle`. Already-spilled chunks are unaffected +/// either way. The pool is resolved per commit through +/// [`crate::pool_config::active_pool`], so chunks spill only once +/// `apply_pool_config` has installed and budgeted the pool. With no pool +/// installed chunks stay resident regardless of the gates. +pub fn set_compute_spill_enabled(enabled: bool) { + COMPUTE_SPILL_ENABLED.store(enabled, Ordering::Relaxed); +} + +/// Enable or disable chunk spilling on behalf of storage's upsert dataflows. +/// +/// See [`set_compute_spill_enabled`] for the shared-gate semantics. +pub fn set_storage_spill_enabled(enabled: bool) { + STORAGE_SPILL_ENABLED.store(enabled, Ordering::Relaxed); +} + +/// Set or unset the pool through which this thread's chunk spills are +/// routed, taking precedence over the gates and the process pool. `None` +/// restores the global resolution. +pub fn set_spill_override(pool: Option) { + SPILL_OVERRIDE.with(|cell| *cell.borrow_mut() = pool); +} + +/// The pool committed chunks spill to, if any. +fn spill_pool() -> Option { + if let Some(pool) = SPILL_OVERRIDE.with(|cell| cell.borrow().clone()) { + return Some(pool); + } + let enabled = COMPUTE_SPILL_ENABLED.load(Ordering::Relaxed) + || STORAGE_SPILL_ENABLED.load(Ordering::Relaxed); + if enabled { + crate::pool_config::active_pool() + } else { + None + } +} + +/// Scratch capacity retained across reads, in words. A read larger than this +/// releases the buffer afterward, so a thread's scratch does not ratchet to +/// the largest body it ever carried (heap no pool gauge can see). +const SCRATCH_RETAIN_WORDS: usize = 1 << 18; + +/// Run `f` with this thread's read scratch, cleared of any previous use. +fn with_scratch(f: impl FnOnce(&mut Vec) -> Out) -> Out { + READ_SCRATCH.with(|cell| { + let mut scratch = cell.take(); + scratch.clear(); + let out = f(&mut scratch); + if scratch.capacity() > SCRATCH_RETAIN_WORDS { + scratch.clear(); + scratch.shrink_to_fit(); + } + cell.replace(scratch); + out + }) +} + +/// The serialized-byte size committed chunks aim for, matching the ship size +/// of the columnar merge machinery. +const COMMIT_BYTES: usize = 2 << 20; + +/// Bodies smaller than this stay resident: the pool's smallest size class is +/// 64 KiB, so spilling below it trades no meaningful memory for slot waste. +/// +/// Sub-floor bodies are invisible to the pool's budget, which is safe only +/// while they are rare. `settle` coalesces toward `COMMIT_BYTES` before +/// committing, so in the harness only a final `done` tail commits below the +/// floor. A caller that commits many small chunks directly accumulates +/// unbudgeted heap, and no accounting here would catch it. +const SPILL_MIN_BYTES: usize = 64 << 10; + +/// Whether a column is big enough to commit on its own. A monotone +/// threshold, so settle's carry, which grows by whole chunks, cannot step +/// over it. +fn at_commit_size(column: &Column) -> bool { + column.length_in_bytes() >= COMMIT_BYTES - COMMIT_BYTES / 10 +} + +/// Reconstructs the borrowed columnar view from serialized words, the same +/// zero-copy decode [`Column::borrow`] performs on its `Align` variant. +fn borrow_words(words: &[u64]) -> BorrowedOf<'_, C> { + >::from_bytes(&mut indexed::decode(words)) +} + +/// Narrow a columnar ref to a shorter lifetime, so refs from different +/// borrows, such as a probe column and a chunk's own columns, can be compared +/// (the refs are lifetime-invariant). +fn rr<'b, 'a: 'b, C: Columnar>(item: columnar::Ref<'a, C>) -> columnar::Ref<'b, C> { + columnar::ContainerOf::::reborrow_ref(item) +} + +/// A spilled chunk body: the serialized column in the pool, plus the resident +/// metadata every [`Chunk`] must answer without fetching. That metadata is +/// the record count and the first and last data items (the fence entries +/// [`UnloadChunk::locate`] consults). +pub struct SpilledBody { + /// Number of updates in the body. + records: usize, + /// The first and last data items, as a two-element container. One + /// container rather than two singletons, so the leaf allocations are not + /// duplicated per fence. + fences: D::Container, + /// The chunk's generational depth, mirrored into the pool's + /// [`ChunkHints`] at spill time. + depth: u8, + /// The pool chunk holding the serialized column. + handle: ChunkHandle, +} + +/// A sorted, consolidated run of `(D, T, R)` updates, resident or spilled. +/// +/// Every chunk carries a generational depth, fixed at creation: fresh chunks +/// are depth 0, a merge output is one generation past its deepest input +/// (saturating at `u8::MAX`, where remerged long-lived chunks stay), and +/// rewrites within a generation (extract, advance, settle coalescing) +/// preserve depth. At spill time the depth becomes the pool's [`ChunkHints`], +/// so repeatedly merged (older, colder) data lands in deeper eviction bands. +pub enum ColumnChunk { + /// Body on the heap, shared via `Rc`, with its generational depth. + Resident(Rc>, u8), + /// Body in the pool. See [`SpilledBody`]. + Spilled(Rc>), +} + +impl Clone for ColumnChunk { + fn clone(&self) -> Self { + match self { + ColumnChunk::Resident(col, depth) => ColumnChunk::Resident(Rc::clone(col), *depth), + ColumnChunk::Spilled(body) => ColumnChunk::Spilled(Rc::clone(body)), + } + } +} + +impl Default for ColumnChunk { + fn default() -> Self { + ColumnChunk::Resident(Rc::new(Column::default()), 0) + } +} + +impl Accountable for ColumnChunk { + fn record_count(&self) -> i64 { + i64::try_from(self.records()).expect("record count fits i64") + } +} + +impl ColumnChunk { + /// Wrap a sorted, consolidated, non-empty column as a resident chunk of + /// the youngest generation. + pub fn from_column(column: Column<(D, T, R)>) -> Self { + debug_assert!(column.borrow().len() > 0, "chunks must be non-empty"); + ColumnChunk::Resident(Rc::new(column), 0) + } + + /// The body as an owned column. A spilled body is copied out of the pool + /// within this call. A shared resident body is copied. + pub fn into_column(self) -> Column<(D, T, R)> { + match self { + ColumnChunk::Resident(col, _) => { + Rc::try_unwrap(col).unwrap_or_else(|shared| copy_column(&shared)) + } + ColumnChunk::Spilled(body) => { + let mut words = Vec::new(); + body.handle.read_into(&mut words); + Column::Align(words) + } + } + } + + /// True when the body lives in the pool. + pub fn is_spilled(&self) -> bool { + matches!(self, ColumnChunk::Spilled(_)) + } + + /// The number of updates, from resident state only. + fn records(&self) -> usize { + match self { + ColumnChunk::Resident(col, _) => col.borrow().len(), + ColumnChunk::Spilled(body) => body.records, + } + } + + /// The generational depth, from resident state only. + fn depth(&self) -> u8 { + match self { + ColumnChunk::Resident(_, depth) => *depth, + ColumnChunk::Spilled(body) => body.depth, + } + } + + /// The first and last data items, from resident state only. + fn data_span(&self) -> (columnar::Ref<'_, D>, columnar::Ref<'_, D>) { + match self { + ColumnChunk::Resident(col, _) => { + let data = col.borrow().0; + (data.get(0), data.get(data.len() - 1)) + } + ColumnChunk::Spilled(body) => { + let fences = body.fences.borrow(); + (fences.get(0), fences.get(1)) + } + } + } + + /// Commit a non-empty column at the given generational depth: spill it to + /// the pool when spilling is on and the body is worth a slot, else keep it + /// resident. + fn commit(column: Column<(D, T, R)>, depth: u8) -> Self { + debug_assert!(column.borrow().len() > 0, "chunks must be non-empty"); + if let Some(pool) = spill_pool() { + if column.length_in_bytes() >= SPILL_MIN_BYTES { + return Self::spill_body(column, &pool, depth); + } + } + ColumnChunk::Resident(Rc::new(column), depth) + } + + /// Spill a non-empty column into `pool` unconditionally, capturing the + /// resident fence metadata. + fn spill_body(column: Column<(D, T, R)>, pool: &Pool, depth: u8) -> Self { + let len_bytes = column.length_in_bytes(); + let view = column.borrow(); + let records = view.len(); + let mut fences = D::Container::default(); + fences.push(view.0.get(0)); + fences.push(view.0.get(records - 1)); + let handle = spill_column(column, pool, len_bytes, ChunkHints { depth }); + ColumnChunk::Spilled(Rc::new(SpilledBody { + records, + fences, + depth, + handle, + })) + } +} + +/// Copy a column into a fresh `Typed` column via bulk per-leaf extension. +fn copy_column(column: &Column) -> Column { + let view = column.borrow(); + let mut fresh = C::Container::default(); + fresh.extend_from_self(view, 0..view.len()); + Column::Typed(fresh) +} + +/// The chunk-side [`ExtentCodec`]: a little-endian `u32` body-length prefix +/// followed by one lz4 block, the framing +/// `lz4_flex::block::compress_prepend_size` produces. Every chunk consumer +/// passes [`LZ4_CODEC`] at insert; the pool itself has no codec opinion. +#[derive(Debug)] +pub struct Lz4Codec; + +/// The [`Lz4Codec`] instance chunk consumers pass to +/// [`Pool::insert_with`]. +pub static LZ4_CODEC: Lz4Codec = Lz4Codec; + +impl ExtentCodec for Lz4Codec { + fn encode(&self, body: &[u8], out: &mut Vec) { + let max_out = lz4_flex::block::get_maximum_output_size(body.len()); + out.resize(4 + max_out, 0); + let len = u32::try_from(body.len()).expect("chunk bodies are bounded by the size classes"); + out[..4].copy_from_slice(&len.to_le_bytes()); + let compressed = lz4_flex::block::compress_into(body, &mut out[4..]) + .expect("output sized to the maximum"); + out.truncate(4 + compressed); + } + + fn decode(&self, stored: &[u8], body: &mut [u8]) { + let prefix: [u8; 4] = stored[..4].try_into().expect("prefix length"); + let len = usize::try_from(u32::from_le_bytes(prefix)).expect("length fits usize"); + assert_eq!( + len, + body.len(), + "destination must match the encoded body length" + ); + let written = lz4_flex::block::decompress_into(&stored[4..], body) + .expect("stored bytes hold a valid lz4 block"); + assert_eq!(written, body.len(), "decoded length mismatch"); + } +} + +/// Serialize a column into a pool slot. The `Align` variant is already the +/// serialized form and copies in directly. Other variants write their +/// [`ContainerBytes`] encoding through a cursor over the slot memory. Sizing +/// is exact, so a short or overlong write is a contract violation and panics. +fn spill_column( + column: Column, + pool: &Pool, + len_bytes: usize, + hints: ChunkHints, +) -> ChunkHandle { + debug_assert_eq!(len_bytes % 8, 0); + match column { + Column::Align(words) => pool.insert_with(words.len(), hints, &LZ4_CODEC, |dst| { + dst.copy_from_slice(&words) + }), + other => pool.insert_with(len_bytes / 8, hints, &LZ4_CODEC, |dst| { + let bytes: &mut [u8] = bytemuck::cast_slice_mut(dst); + let mut cursor = std::io::Cursor::new(bytes); + other.into_bytes(&mut cursor); + assert_eq!( + usize::try_from(cursor.position()).expect("usize position"), + len_bytes, + "serialized body must fill the chunk exactly", + ); + }), + } +} + +/// A column is `Typed`, or becomes one by copy. Merge and settle accumulate +/// into `Typed` targets. Serialized variants arrive from spill reads and +/// remote channels. +fn to_typed(column: Column) -> Column { + match column { + typed @ Column::Typed(_) => typed, + other => copy_column(&other), + } +} + +impl Chunk for ColumnChunk +where + D: Columnar, + for<'a> columnar::Ref<'a, D>: Copy + Ord, + T: Columnar + Default + Timestamp + Lattice + Ord, + for<'a> columnar::Ref<'a, T>: Copy + Ord, + R: Columnar + Default + Semigroup + for<'a> Semigroup>, +{ + type Time = T; + + /// A nominal record count for the harness's fuel and ladder accounting, + /// not a bound. Actual chunk sizing is by serialized bytes: `merge` and + /// `extract` cut output at the [`Column`] ship threshold, and `settle` + /// grades by `at_commit_size`, so a chunk of narrow records can hold more + /// records than this and nothing here consults it. + const TARGET: usize = 65536; + + fn len(&self) -> usize { + self.records() + } + + /// [`Column::merge_from`] does the work: gallop bulk-copies for disjoint + /// runs, semigroup consolidation on equal `(data, time)`, output cut at + /// the ship threshold. A survivor pushed back untouched keeps its + /// original form, in particular a spilled body is neither rebuilt nor + /// re-spilled. + /// + /// Fronts whose data ranges are disjoint never load at all: the resident + /// fence entries decide, and the lower front moves to the output verbatim. + fn merge(in1: &mut VecDeque, in2: &mut VecDeque, out: &mut VecDeque) { + // Disjoint fast path: when one front lies strictly below the other's + // first data item (equal boundary data could still interleave on + // time), the merged prefix through the shared horizon is exactly that + // front, unchanged. + let (a_first, a_last) = in1 + .front() + .expect("caller guarantees non-empty input") + .data_span(); + let (b_first, b_last) = in2 + .front() + .expect("caller guarantees non-empty input") + .data_span(); + let a_low = rr::(a_last) < rr::(b_first); + let b_low = rr::(b_last) < rr::(a_first); + if a_low { + out.push_back(in1.pop_front().expect("front observed above")); + return; + } + if b_low { + out.push_back(in2.pop_front().expect("front observed above")); + return; + } + + let a = in1.pop_front().expect("caller guarantees non-empty input"); + let b = in2.pop_front().expect("caller guarantees non-empty input"); + // Merged output is one generation past its deepest input. A survivor + // (untouched or rewritten from its remainder) keeps its own depth. + let depths = [a.depth(), b.depth()]; + let out_depth = depths[0].max(depths[1]).saturating_add(1); + let mut spill_a = match &a { + ColumnChunk::Spilled(body) => Some(Rc::clone(body)), + ColumnChunk::Resident(_, _) => None, + }; + let mut spill_b = match &b { + ColumnChunk::Spilled(body) => Some(Rc::clone(body)), + ColumnChunk::Resident(_, _) => None, + }; + let mut cols = [a.into_column(), b.into_column()]; + let mut positions = [0usize, 0usize]; + loop { + let mut result: Column<(D, T, R)> = Column::default(); + let yielded = result.merge_from(&mut cols, &mut positions); + if result.borrow().len() > 0 { + out.push_back(ColumnChunk::Resident(Rc::new(result), out_depth)); + } + if !yielded { + break; + } + } + let [col_a, col_b] = &mut cols; + // Per input side: the loaded column and the merge's consumed position + // within it, the side's pre-merge depth, its original spilled body + // when it had one, and the deque a survivor returns to. + for (col, pos, depth, spilled, queue) in [ + (col_a, positions[0], depths[0], &mut spill_a, in1), + (col_b, positions[1], depths[1], &mut spill_b, in2), + ] { + let len = col.borrow().len(); + if pos == 0 && len > 0 { + // Untouched survivor: restore it as it was, spilled bodies + // included (the loaded copy is dropped). + let chunk = match spilled.take() { + Some(body) => ColumnChunk::Spilled(body), + None => ColumnChunk::Resident(Rc::new(std::mem::take(col)), depth), + }; + queue.push_front(chunk); + } else if pos < len { + let view = col.borrow(); + let mut rest = <(D, T, R) as Columnar>::Container::default(); + rest.extend_from_self(view, pos..len); + queue.push_front(ColumnChunk::Resident(Rc::new(Column::Typed(rest)), depth)); + } + } + } + + /// Partition one front chunk by `frontier`, folding kept times into + /// `residual`. One chunk per call, so the harness settles both sides + /// between chunks. Output is cut at the ship threshold. + fn extract( + input: &mut VecDeque, + frontier: AntichainRef, + residual: &mut timely::progress::Antichain, + keep: &mut VecDeque, + ship: &mut VecDeque, + ) { + let Some(chunk) = input.pop_front() else { + return; + }; + // Partitioning rewrites within a generation, so both sides keep the + // input chunk's depth. + let depth = chunk.depth(); + let mut col = chunk.into_column(); + let len = col.borrow().len(); + let mut pos = 0; + let mut keep_col: Column<(D, T, R)> = Column::default(); + let mut ship_col: Column<(D, T, R)> = Column::default(); + // TODO: rewrite the underlying `Column::extract` as two passes, the + // time column first to find run boundaries, then bulk per-range + // copies of the remaining leaves. + // Move a side's accumulation to its queue, at the ship threshold + // mid-loop, or any non-empty remainder at the end. + let cut = |col: &mut Column<(D, T, R)>, queue: &mut VecDeque, force: bool| { + if col.borrow().len() > 0 && (force || at_serialized_capacity(&col.borrow())) { + queue.push_back(ColumnChunk::Resident(Rc::new(std::mem::take(col)), depth)); + } + }; + while pos < len { + col.extract(&mut pos, frontier, residual, &mut keep_col, &mut ship_col); + if pos < len { + cut(&mut keep_col, keep, false); + cut(&mut ship_col, ship, false); + } + } + cut(&mut keep_col, keep, true); + cut(&mut ship_col, ship, true); + } + + /// Advance times by `frontier` and consolidate, withholding the trailing + /// `D` group as the carry unless `done` (its updates may continue in input + /// this call has not seen). + /// + /// The input concatenates into the carry's container, so a group that + /// grows across many calls is appended to, not rebuilt. Each record is + /// copied once on arrival, keeping the run linear. Advancing is + /// lattice-monotone but not order-monotone, so each group's advanced + /// times are re-sorted before adjacent equal times fold. + fn advance( + input: &mut VecDeque, + frontier: AntichainRef, + done: bool, + out: &mut VecDeque, + ) { + let Some(front) = input.pop_front() else { + return; + }; + // Advancing rewrites within a generation, so output and carry keep + // the deepest input depth. Only merges increment. + let mut depth = front.depth(); + // Concatenate the input into one column, reusing the front chunk's + // storage when it is exclusively owned (the usual case: it is last + // call's carry). + let mut base = to_typed(front.into_column()); + { + let Column::Typed(base_c) = &mut base else { + unreachable!("to_typed returns Typed"); + }; + for chunk in input.drain(..) { + depth = depth.max(chunk.depth()); + let col = chunk.into_column(); + let view = col.borrow(); + base_c.extend_from_self(view, 0..view.len()); + } + } + let view = base.borrow(); + let total = view.len(); + if total == 0 { + return; + } + let data = view.0; + + // Giant-group early-out: if the whole input is one `D` group, nothing + // is provably complete. Unless `done`, push it all back as the carry. + if !done && data.get(0) == data.get(total - 1) { + input.push_front(ColumnChunk::Resident(Rc::new(base), depth)); + return; + } + + // The processing bound: everything, or everything before the trailing + // `D` group when it must be withheld. + let end = if done { + total + } else { + let last = data.get(total - 1); + let mut end = total - 1; + while end > 0 && data.get(end - 1) == last { + end -= 1; + } + end + }; + + let mut result = <(D, T, R) as Columnar>::Container::default(); + // Per-group scratch: advanced owned times with owned diffs. + let mut scratch: Vec<(T, R)> = Vec::new(); + let mut index = 0; + // Cut output at the commit size, checked amortized by emitted records + // (the size test walks the container's leaves, so probing it per + // record would be quadratic). Records, not groups: a single group may + // carry arbitrarily many advanced times, and a cut is legal anywhere + // in the sorted sequence, so bounding by records keeps the largest + // possible output chunk within one check period of the target. It + // must not outgrow the pool's largest size class, past which a body + // degrades to a permanently resident heap chunk. + const CUT_CHECK_RECORDS: usize = 1024; + let mut records_since_check = 0usize; + // TODO: the output leaves are addressed independently, so a group + // that folds nothing (no time collisions, no zeroed diffs) could bulk + // `extend_from_self` the D leaf over the whole group range and push + // only the advanced times and diffs per record, and a singleton group + // (the common case for mostly-unique D) could skip the scratch and + // sort round trip entirely. + while index < end { + let group_d = data.get(index); + scratch.clear(); + while index < end && data.get(index) == group_d { + let (_, t, r) = view.get(index); + let mut owned_t = T::into_owned(t); + owned_t.advance_by(frontier); + scratch.push((owned_t, R::into_owned(r))); + index += 1; + } + scratch.sort_by(|a, b| a.0.cmp(&b.0)); + let mut run = scratch.drain(..).peekable(); + while let Some((t, mut r)) = run.next() { + while run.peek().is_some_and(|(t2, _)| *t2 == t) { + let (_, r2) = run.next().expect("peeked"); + r.plus_equals(&r2); + } + if !r.is_zero() { + result.0.push(group_d); + result.1.push(&t); + result.2.push(&r); + records_since_check += 1; + if records_since_check >= CUT_CHECK_RECORDS { + records_since_check = 0; + if u64::cast_from(indexed::length_in_words(&result.borrow())) + >= u64::cast_from(COMMIT_BYTES / 8) + { + out.push_back(ColumnChunk::Resident( + Rc::new(Column::Typed(std::mem::take(&mut result))), + depth, + )); + } + } + } + } + } + if result.borrow().len() > 0 { + out.push_back(ColumnChunk::Resident(Rc::new(Column::Typed(result)), depth)); + } + + // Rebuild the withheld trailing group as the carry. + if end < total { + let mut carry = <(D, T, R) as Columnar>::Container::default(); + carry.extend_from_self(view, end..total); + input.push_front(ColumnChunk::Resident(Rc::new(Column::Typed(carry)), depth)); + } + } + + /// Grade by serialized bytes and commit: spilled chunks pass through + /// untouched, resident chunks at the commit size commit as they are, and + /// smaller neighbors coalesce until the accumulation reaches it. + /// Committing is the spill hook (see `ColumnChunk::commit`). A + /// sub-threshold tail is withheld as the carry unless `done`. + fn settle(input: &mut VecDeque, done: bool, out: &mut VecDeque) { + // Coalescing rewrites within a generation, so the carry commits at + // the deepest depth among its constituent chunks. + let mut carry: Option<(Column<(D, T, R)>, u8)> = None; + while let Some(chunk) = input.pop_front() { + let (rc, depth) = match chunk { + spilled @ ColumnChunk::Spilled(_) => { + if let Some((col, depth)) = carry.take() { + out.push_back(ColumnChunk::commit(col, depth)); + } + out.push_back(spilled); + continue; + } + ColumnChunk::Resident(rc, depth) => (rc, depth), + }; + let full = at_commit_size(&rc); + // A sub-threshold chunk coalesces into the open carry by borrow, + // never unwrapping a shared body. + if !full && let Some((mut acc, acc_depth)) = carry.take() { + let Column::Typed(acc_c) = &mut acc else { + unreachable!("carry is always Typed"); + }; + let view = rc.borrow(); + acc_c.extend_from_self(view, 0..view.len()); + let acc_depth = acc_depth.max(depth); + if at_commit_size(&acc) { + out.push_back(ColumnChunk::commit(acc, acc_depth)); + } else { + carry = Some((acc, acc_depth)); + } + continue; + } + // Otherwise any open carry flushes, and the chunk either commits + // whole or opens the next carry. + if let Some((acc, acc_depth)) = carry.take() { + out.push_back(ColumnChunk::commit(acc, acc_depth)); + } + let col = Rc::try_unwrap(rc).unwrap_or_else(|rc| copy_column(&rc)); + if full { + out.push_back(ColumnChunk::commit(col, depth)); + } else { + carry = Some((to_typed(col), depth)); + } + } + if let Some((col, depth)) = carry { + if done { + out.push_back(ColumnChunk::commit(col, depth)); + } else { + input.push_front(ColumnChunk::Resident(Rc::new(col), depth)); + } + } + } +} + +/// Append every update in `view` whose key matches a probe at or after +/// `*probe_index` into `staging`, per the [`UnloadChunk`] consume-index +/// protocol: probes strictly below the view's last key are consumed, a probe +/// equal to it is extracted but left for the next chunk. +fn extract_view_into<'v, 'p, K, V, T, R>( + view: BorrowedOf<'v, ((K, V), T, R)>, + probes: BorrowedOf<'p, K>, + probe_index: &mut usize, + staging: &mut <((K, V), T, R) as Columnar>::Container, +) where + K: Columnar, + V: Columnar, + T: Columnar, + R: Columnar, + for<'b> columnar::Ref<'b, K>: Copy + Ord, +{ + let keys = view.0.0; + let len = keys.len(); + let last = keys.get(len - 1); + let count = probes.len(); + let mut pos = 0; + while *probe_index < count { + let probe = probes.get(*probe_index); + debug_assert!( + *probe_index == 0 || rr::(probes.get(*probe_index - 1)) < rr::(probe), + "probe keys must be sorted and deduplicated" + ); + if rr::(probe) > rr::(last) { + return; + } + gallop(len, &mut pos, |i| rr::(keys.get(i)) < rr::(probe)); + let start = pos; + while pos < len && rr::(keys.get(pos)) == rr::(probe) { + pos += 1; + } + staging.extend_from_self(view, start..pos); + if rr::(probe) == rr::(last) { + return; + } + *probe_index += 1; + } +} + +impl UnloadChunk for ColumnChunk<(K, V), T, R> +where + K: Columnar, + for<'a> columnar::Ref<'a, K>: Copy + Ord, + V: Columnar, + for<'a> columnar::Ref<'a, V>: Copy + Ord, + T: Columnar + Default + Timestamp + Lattice + Ord, + for<'a> columnar::Ref<'a, T>: Copy + Ord, + R: Columnar + Default + Semigroup + for<'a> Semigroup>, +{ + /// The flat columnar accumulation. Appends are bulk column-range copies, + /// and a group straddling chunks stitches by plain concatenation. + type Staging = <((K, V), T, R) as Columnar>::Container; + + /// A borrowed key column, e.g. of a `Column` the consumer assembled + /// from its sorted, deduplicated probe keys. + type Probes<'a> = BorrowedOf<'a, K>; + + fn probe_count(probes: Self::Probes<'_>) -> usize { + probes.len() + } + + fn locate(&self, probes: Self::Probes<'_>, probe_index: usize) -> std::cmp::Ordering { + let probe = probes.get(probe_index); + // A data ref is a `(key ref, val ref)` tuple, so the key fences are a + // projection of the data fences. + let (first, last) = self.data_span(); + let (first, last) = (first.0, last.0); + if rr::(probe) < rr::(first) { + std::cmp::Ordering::Less + } else if rr::(probe) > rr::(last) { + std::cmp::Ordering::Greater + } else { + std::cmp::Ordering::Equal + } + } + + fn extract_into( + &self, + probes: Self::Probes<'_>, + probe_index: &mut usize, + staging: &mut Self::Staging, + ) { + match self { + ColumnChunk::Resident(col, _) => { + extract_view_into::(col.borrow(), probes, probe_index, staging); + } + ColumnChunk::Spilled(body) => with_scratch(|scratch| { + // NOTE: deliberately the non-admitting read. One probe set + // touching a chunk is weak evidence it will be touched again, + // and probing a spilled trace must not accrete it back into + // residency. The cost is a full decode per probe set against + // an evicted chunk. + body.handle.read_into(scratch); + let view = borrow_words::<((K, V), T, R)>(scratch); + extract_view_into::(view, probes, probe_index, staging); + }), + } + } + + fn fetch_into(&self, staging: &mut Self::Staging) { + match self { + ColumnChunk::Resident(col, _) => { + let view = col.borrow(); + staging.extend_from_self(view, 0..view.len()); + } + ColumnChunk::Spilled(body) => with_scratch(|scratch| { + body.handle.read_into(scratch); + let view = borrow_words::<((K, V), T, R)>(scratch); + staging.extend_from_self(view, 0..view.len()); + }), + } + } +} + +/// A batch builder over [`ColumnChunk`] input that delegates to a builder +/// over [`Column`] input, loading each chunk's body as it is pushed. +/// +/// This is the adapter that lets a [`ChunkBatcher`] feed the existing +/// column-input batch builders (and through them the existing spine layouts): +/// the batcher's chains carry pool-spillable chunks, and bodies are read back +/// copy-out only at the seal, one chunk at a time. +/// +/// [`ChunkBatcher`]: differential_dataflow::trace::chunk::ChunkBatcher +pub struct UnchunkBuilder { + inner: Bu, + _marker: std::marker::PhantomData<(D, T, R)>, +} + +impl differential_dataflow::trace::Builder for UnchunkBuilder +where + Bu: differential_dataflow::trace::Builder>, + D: Columnar + 'static, + T: Columnar + 'static, + R: Columnar + 'static, +{ + type Input = ColumnChunk; + type Time = Bu::Time; + type Output = Bu::Output; + + fn with_capacity(keys: usize, vals: usize, upds: usize) -> Self { + Self { + inner: Bu::with_capacity(keys, vals, upds), + _marker: std::marker::PhantomData, + } + } + + fn push(&mut self, chunk: &mut Self::Input) { + let mut column = std::mem::take(chunk).into_column(); + self.inner.push(&mut column); + } + + fn done( + self, + description: differential_dataflow::trace::Description, + ) -> Self::Output { + self.inner.done(description) + } + + fn seal( + chain: &mut Vec, + description: differential_dataflow::trace::Description, + ) -> Self::Output { + // One chunk at a time through `push`, so peak transient memory is a + // single loaded body rather than the whole chain at once. + let mut builder = Self::new(); + for chunk in chain.iter_mut() { + builder.push(chunk); + } + chain.clear(); + builder.done(description) + } +} + +/// A chunker for `arrange_core` over [`ColumnChunk`]s: sorts and consolidates +/// raw input columns through a [`ColumnChunker`] and wraps its output chunks. +pub struct ChunkChunker { + inner: ColumnChunker<(D, T, R)>, + staged: ColumnChunk, +} + +impl Default for ChunkChunker +where + D: Columnar, + T: Columnar, + R: Columnar, + ColumnChunker<(D, T, R)>: Default, +{ + fn default() -> Self { + Self { + inner: Default::default(), + staged: Default::default(), + } + } +} + +impl<'a, D, T, R> PushInto<&'a mut Column<(D, T, R)>> for ChunkChunker +where + D: Columnar, + T: Columnar, + R: Columnar, + ColumnChunker<(D, T, R)>: PushInto<&'a mut Column<(D, T, R)>>, +{ + fn push_into(&mut self, item: &'a mut Column<(D, T, R)>) { + self.inner.push_into(item); + } +} + +impl ContainerBuilder for ChunkChunker +where + D: Columnar + 'static, + T: Columnar + 'static, + R: Columnar + 'static, + ColumnChunker<(D, T, R)>: ContainerBuilder>, +{ + type Container = ColumnChunk; + + fn extract(&mut self) -> Option<&mut Self::Container> { + let col = self.inner.extract()?; + self.staged = ColumnChunk::from_column(std::mem::take(col)); + Some(&mut self.staged) + } + + fn finish(&mut self) -> Option<&mut Self::Container> { + let col = self.inner.finish()?; + self.staged = ColumnChunk::from_column(std::mem::take(col)); + Some(&mut self.staged) + } +} + +#[cfg(test)] +mod tests { + //! Property tests for the [`Chunk`] and [`UnloadChunk`] contracts on + //! [`ColumnChunk`]. + //! + //! Strategy: generate sorted+consolidated inputs (the chunk invariant), + //! drive the trait methods the way the differential harness does, and + //! compare against brute-force references on owned tuples. Test types are + //! `D = (u64, u64)`, `T = u64`, `R = i64` from small ranges so equal-key + //! collisions are common and consolidation actually runs. + + use differential_dataflow::trace::chunk::{ChunkBatch, ChunkBatcher}; + use differential_dataflow::trace::{Batcher, Description}; + use mz_ore::pool::Pool; + use proptest::prelude::*; + use timely::container::PushInto; + use timely::progress::Antichain; + + use crate::columnar::unload::UnloadBatch; + + use super::*; + + type Tuple = ((u64, u64), u64, i64); + type TestChunk = ColumnChunk<(u64, u64), u64, i64>; + + /// The delegated codec's stored form is byte-identical to the extent + /// store's previous hard-coded framing: a little-endian `u32` + /// body-length prefix followed by one lz4 block, which is exactly what + /// `compress_prepend_size` produces. + #[mz_ore::test] + fn lz4_codec_matches_the_previous_extent_framing() { + let body: Vec = (0..100_000u32).flat_map(|i| i.to_le_bytes()).collect(); + let mut stored = Vec::new(); + LZ4_CODEC.encode(&body, &mut stored); + assert_eq!(stored, lz4_flex::block::compress_prepend_size(&body)); + let mut round = vec![0u8; body.len()]; + LZ4_CODEC.decode(&stored, &mut round); + assert_eq!(round, body); + } + + #[mz_ore::test] + #[should_panic(expected = "destination must match")] + fn lz4_codec_decode_length_mismatch_panics() { + let mut stored = Vec::new(); + LZ4_CODEC.encode(&[7u8; 64], &mut stored); + let mut short = vec![0u8; 32]; + LZ4_CODEC.decode(&stored, &mut short); + } + + /// Reference consolidation: sort by `(data, time)`, sum diffs over equal + /// pairs, drop zeros. + fn consolidate(mut v: Vec) -> Vec { + v.sort(); + let mut out: Vec = Vec::new(); + for (d, t, r) in v { + if let Some(last) = out.last_mut() { + if last.0 == d && last.1 == t { + last.2 += r; + continue; + } + } + out.push((d, t, r)); + } + out.retain(|x| x.2 != 0); + out + } + + fn arb_consolidated() -> impl Strategy> { + prop::collection::vec(((0u64..5, 0u64..5), 0u64..4, -3i64..=3i64), 0..40) + .prop_map(consolidate) + } + + fn build_column(v: &[Tuple]) -> Column { + let mut col: Column = Default::default(); + for tup in v { + col.push_into(*tup); + } + col + } + + fn collect_column(col: &Column) -> Vec { + col.borrow() + .into_index_iter() + .map(|((k, v), t, r)| { + ( + (u64::into_owned(k), u64::into_owned(v)), + u64::into_owned(t), + i64::into_owned(r), + ) + }) + .collect() + } + + fn collect_chunks(chunks: impl IntoIterator) -> Vec { + chunks + .into_iter() + .flat_map(|chunk| collect_column(&chunk.into_column())) + .collect() + } + + fn collect_staging(staging: &::Container) -> Vec { + staging + .borrow() + .into_index_iter() + .map(|((k, v), t, r)| { + ( + (u64::into_owned(k), u64::into_owned(v)), + u64::into_owned(t), + i64::into_owned(r), + ) + }) + .collect() + } + + /// Cut consolidated data into non-empty chunks at the given points. + fn chunked(data: &[Tuple], cuts: &[usize]) -> VecDeque { + let mut chunks = VecDeque::new(); + let mut start = 0; + for cut in cuts { + let end = (start + 1 + cut % 7).min(data.len()); + if end > start { + chunks.push_back(ColumnChunk::from_column(build_column(&data[start..end]))); + start = end; + } + } + if start < data.len() { + chunks.push_back(ColumnChunk::from_column(build_column(&data[start..]))); + } + chunks + } + + /// The chunked cut, with every chunk force-spilled through a private pool + /// (bounds captured, bodies in the pool) regardless of size thresholds. + fn chunked_spilled(data: &[Tuple], cuts: &[usize], pool: &Pool) -> VecDeque { + chunked(data, cuts) + .into_iter() + .map(|chunk| force_spill(chunk, pool)) + .collect() + } + + /// Spill one chunk through `pool`, bypassing the size threshold and + /// keeping the chunk's depth. + fn force_spill(chunk: TestChunk, pool: &Pool) -> TestChunk { + let depth = chunk.depth(); + TestChunk::spill_body(chunk.into_column(), pool, depth) + } + + /// A single pool shared by every test in the module. A pool reserves a + /// large slab of address space, so one per test (let alone per proptest + /// case) exhausts the VM map under parallel test threads. + fn test_pool() -> Pool { + static POOL: std::sync::OnceLock = std::sync::OnceLock::new(); + POOL.get_or_init(|| Pool::new().expect("pool creation")) + .clone() + } + + proptest! { + /// A full batcher round trip: push chunked inputs, seal everything, + /// and compare with the reference consolidation of the union. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] + fn batcher_round_trip( + inputs in prop::collection::vec(arb_consolidated(), 1..6), + cuts in prop::collection::vec(0usize..7, 0..8), + ) { + let mut batcher: ChunkBatcher = Batcher::new(None, 0); + let mut union = Vec::new(); + for input in &inputs { + Extend::extend(&mut union, input.iter().copied()); + for chunk in chunked(input, &cuts) { + batcher.push_into(chunk); + } + } + // An empty upper ships everything. + let (sealed, _description) = batcher.seal(Antichain::new()); + prop_assert_eq!(collect_chunks(sealed), consolidate(union)); + } + + /// The same round trip over force-spilled inputs: merge and extract + /// read bodies back from the pool call-scoped. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] + fn batcher_round_trip_spilled( + inputs in prop::collection::vec(arb_consolidated(), 1..4), + cuts in prop::collection::vec(0usize..7, 0..6), + ) { + let pool = test_pool(); + let mut batcher: ChunkBatcher = Batcher::new(None, 0); + let mut union = Vec::new(); + for input in &inputs { + Extend::extend(&mut union, input.iter().copied()); + for chunk in chunked_spilled(input, &cuts, &pool) { + batcher.push_into(chunk); + } + } + let (sealed, _description) = batcher.seal(Antichain::new()); + prop_assert_eq!(collect_chunks(sealed), consolidate(union)); + } + + /// Sealing at an intermediate upper partitions by time and reports + /// the kept lower envelope as the frontier. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] + fn seal_partitions_by_time( + input in arb_consolidated(), + cuts in prop::collection::vec(0usize..7, 0..8), + upper in 0u64..5, + ) { + let mut batcher: ChunkBatcher = Batcher::new(None, 0); + for chunk in chunked(&input, &cuts) { + batcher.push_into(chunk); + } + let (shipped, _) = batcher.seal(Antichain::from_elem(upper)); + let expected_shipped: Vec = + input.iter().copied().filter(|(_, t, _)| *t < upper).collect(); + prop_assert_eq!(collect_chunks(shipped), consolidate(expected_shipped)); + + let kept_min = input.iter().filter(|(_, t, _)| *t >= upper).map(|(_, t, _)| *t).min(); + let frontier = batcher.frontier().to_owned(); + prop_assert_eq!(frontier.elements().first().copied(), kept_min); + + let (rest, _) = batcher.seal(Antichain::new()); + let expected_rest: Vec = + input.iter().copied().filter(|(_, t, _)| *t >= upper).collect(); + prop_assert_eq!(collect_chunks(rest), consolidate(expected_rest)); + } + + /// The intermediate-upper partition of `seal_partitions_by_time`, over + /// force-spilled inputs: bodies read back from the pool and split by + /// time in one seal. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] + fn seal_partitions_by_time_spilled( + input in arb_consolidated(), + cuts in prop::collection::vec(0usize..7, 0..8), + upper in 0u64..5, + ) { + let pool = test_pool(); + let mut batcher: ChunkBatcher = Batcher::new(None, 0); + for chunk in chunked_spilled(&input, &cuts, &pool) { + batcher.push_into(chunk); + } + let (shipped, _) = batcher.seal(Antichain::from_elem(upper)); + let expected_shipped: Vec = + input.iter().copied().filter(|(_, t, _)| *t < upper).collect(); + prop_assert_eq!(collect_chunks(shipped), consolidate(expected_shipped)); + + let kept_min = input.iter().filter(|(_, t, _)| *t >= upper).map(|(_, t, _)| *t).min(); + let frontier = batcher.frontier().to_owned(); + prop_assert_eq!(frontier.elements().first().copied(), kept_min); + + let (rest, _) = batcher.seal(Antichain::new()); + let expected_rest: Vec = + input.iter().copied().filter(|(_, t, _)| *t >= upper).collect(); + prop_assert_eq!(collect_chunks(rest), consolidate(expected_rest)); + } + + /// `advance` equals per-record time advancement plus reference + /// consolidation, including across a `done = false` carry. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] + fn advance_matches_reference( + input in arb_consolidated(), + cuts in prop::collection::vec(0usize..7, 0..8), + frontier_elem in 0u64..5, + ) { + let frontier = Antichain::from_elem(frontier_elem); + let mut chunks = chunked(&input, &cuts); + let mut out = VecDeque::new(); + TestChunk::advance(&mut chunks, frontier.borrow(), false, &mut out); + TestChunk::advance(&mut chunks, frontier.borrow(), true, &mut out); + prop_assert!(chunks.is_empty()); + + let expected = consolidate( + input + .iter() + .map(|&(d, mut t, r)| { + t.advance_by(frontier.borrow()); + (d, t, r) + }) + .collect(), + ); + prop_assert_eq!(collect_chunks(out), expected); + } + + /// `settle` preserves contents and order, moves everything on `done`, + /// and coalesces small neighbors. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] + fn settle_preserves_and_packs( + input in arb_consolidated(), + cuts in prop::collection::vec(0usize..7, 1..8), + ) { + let mut chunks = chunked(&input, &cuts); + let mut out = VecDeque::new(); + TestChunk::settle(&mut chunks, true, &mut out); + prop_assert!(chunks.is_empty()); + // Test chunks are far below the byte threshold, so maximal + // packing coalesces everything into a single chunk. + prop_assert!(out.len() <= 1); + prop_assert_eq!(collect_chunks(out), input); + } + + /// `ChunkBatch::extract_into` over sorted, deduplicated probe keys + /// equals the reference filter, resident and spilled alike, straddled + /// keys included. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] + fn unload_extract_matches_filter( + input in arb_consolidated(), + cuts in prop::collection::vec(0usize..7, 0..8), + probe_keys in prop::collection::btree_set(0u64..6, 0..6), + spill in any::(), + ) { + prop_assume!(!input.is_empty()); + let pool = test_pool(); + let chunks: Vec = if spill { + chunked_spilled(&input, &cuts, &pool).into() + } else { + chunked(&input, &cuts).into() + }; + let description = Description::new( + Antichain::from_elem(0u64), + Antichain::new(), + Antichain::from_elem(0u64), + ); + let batch = ChunkBatch::new(chunks, description); + + let mut probe_col = ::Container::default(); + for key in &probe_keys { + probe_col.push(*key); + } + let mut staging = ::Container::default(); + batch.extract_into(probe_col.borrow(), &mut staging); + + let expected: Vec = input + .iter() + .copied() + .filter(|((k, _), _, _)| probe_keys.contains(k)) + .collect(); + prop_assert_eq!(collect_staging(&staging), expected); + + // The scan path reproduces the batch exactly, resident and + // spilled alike. + let mut staging = ::Container::default(); + batch.fetch_into(&mut staging); + prop_assert_eq!(collect_staging(&staging), input); + } + } + + /// `locate` answers the three-way span comparison for every probe + /// placement: below, within, and past the chunk's keys. + #[mz_ore::test] + fn locate_spans_keys() { + let chunk = ColumnChunk::from_column(build_column(&[ + ((2, 0), 0, 1), + ((4, 0), 0, 1), + ((6, 0), 0, 1), + ])); + let mut probe_col = ::Container::default(); + for key in [0u64, 2, 3, 6, 9] { + probe_col.push(key); + } + let probes = probe_col.borrow(); + use std::cmp::Ordering::*; + let expected = [Less, Equal, Equal, Equal, Greater]; + for (index, expected) in expected.iter().enumerate() { + assert_eq!(chunk.locate(probes, index), *expected, "probe {index}"); + } + } + + /// Collect chunk contents while asserting each chunk's serialized size + /// stays within `bound` bytes. + fn collect_bounded(chunks: impl IntoIterator, bound: usize) -> Vec { + let mut collected = Vec::new(); + for chunk in chunks { + let col = chunk.into_column(); + let bytes = col.length_in_bytes(); + assert!(bytes <= bound, "chunk of {bytes} bytes exceeds {bound}"); + Extend::extend(&mut collected, collect_column(&col)); + } + collected + } + + /// Advancing a large input cuts the output into several chunks near the + /// ship threshold, and their concatenation is the reference result. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] + fn advance_cuts_large_output() { + let records: Vec = (0..300_000u64).map(|k| ((k, 0), 0, 1)).collect(); + let mut input = VecDeque::from([ColumnChunk::from_column(build_column(&records))]); + let frontier = Antichain::from_elem(0u64); + let mut out = VecDeque::new(); + TestChunk::advance(&mut input, frontier.borrow(), true, &mut out); + assert!(input.is_empty()); + assert!( + out.len() >= 2, + "expected a cut output, got {} chunk(s)", + out.len() + ); + assert_eq!(collect_bounded(out, 2 * COMMIT_BYTES), records); + } + + /// An input that is entirely one `D` group is withheld whole as the + /// carry unless `done`: none of it is provably complete. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] + fn advance_withholds_giant_group() { + let records: Vec = (0..100u64).map(|t| ((7, 7), t, 1)).collect(); + let mut input: VecDeque = VecDeque::new(); + for piece in records.chunks(30) { + input.push_back(ColumnChunk::from_column(build_column(piece))); + } + let frontier = Antichain::from_elem(50u64); + let mut out = VecDeque::new(); + TestChunk::advance(&mut input, frontier.borrow(), false, &mut out); + assert!(out.is_empty(), "nothing may ship from a single open group"); + assert_eq!(input.len(), 1, "the whole input becomes one carry chunk"); + // Sealing the carry advances and consolidates it. + TestChunk::advance(&mut input, frontier.borrow(), true, &mut out); + assert!(input.is_empty()); + let advanced = records.iter().map(|&(d, t, r)| (d, t.max(50), r)).collect(); + assert_eq!(collect_chunks(out), consolidate(advanced)); + } + + /// Extracting a large chunk at an intermediate frontier cuts both sides + /// into several chunks and partitions exactly by time. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] + fn extract_cuts_large_output() { + let records: Vec = (0..300_000u64).map(|k| ((k, 0), k % 2, 1)).collect(); + let mut input = VecDeque::from([ColumnChunk::from_column(build_column(&records))]); + let frontier = Antichain::from_elem(1u64); + let mut residual = Antichain::new(); + let (mut keep, mut ship) = (VecDeque::new(), VecDeque::new()); + while !input.is_empty() { + TestChunk::extract( + &mut input, + frontier.borrow(), + &mut residual, + &mut keep, + &mut ship, + ); + } + assert!( + keep.len() >= 2, + "expected a cut keep side, got {} chunk(s)", + keep.len() + ); + assert!( + ship.len() >= 2, + "expected a cut ship side, got {} chunk(s)", + ship.len() + ); + let kept: Vec = records.iter().copied().filter(|r| r.1 >= 1).collect(); + let shipped: Vec = records.iter().copied().filter(|r| r.1 < 1).collect(); + assert_eq!(collect_bounded(keep, 2 * COMMIT_BYTES), kept); + assert_eq!(collect_bounded(ship, 2 * COMMIT_BYTES), shipped); + assert_eq!(residual, Antichain::from_elem(1)); + } + + /// `locate` answers from resident metadata on spilled chunks and follows + /// the probe-relative-to-span convention. + #[mz_ore::test] + fn locate_uses_resident_bounds() { + let pool = test_pool(); + let data: Vec = vec![((2, 0), 0, 1), ((4, 0), 0, 1)]; + let chunk = force_spill(ColumnChunk::from_column(build_column(&data)), &pool); + + let mut probe_col = ::Container::default(); + for key in [1u64, 3, 5] { + probe_col.push(key); + } + let probes = probe_col.borrow(); + assert_eq!(chunk.locate(probes, 0), std::cmp::Ordering::Less); + assert_eq!(chunk.locate(probes, 1), std::cmp::Ordering::Equal); + assert_eq!(chunk.locate(probes, 2), std::cmp::Ordering::Greater); + } + + /// A body large enough to spill round-trips through the pool with resident + /// metadata intact, and the batcher produces spilled sealed output. + #[mz_ore::test] + fn spill_round_trip() { + set_spill_override(Some(test_pool())); + + let data: Vec = (0..40_000u64) + .map(|i| ((i / 4, i % 4), i % 8, 1i64)) + .collect(); + let data = consolidate(data); + + let column = build_column(&data); + let committed = TestChunk::commit(column, 0); + assert!(committed.is_spilled(), "large body must spill"); + assert_eq!(committed.len(), data.len()); + assert_eq!(collect_column(&committed.clone().into_column()), data); + + let mut batcher: ChunkBatcher = Batcher::new(None, 0); + for piece in data.chunks(10_000) { + batcher.push_into(ColumnChunk::from_column(build_column(piece))); + } + let (sealed, _) = batcher.seal(Antichain::new()); + assert!( + sealed.iter().any(ColumnChunk::is_spilled), + "sealed output should contain spilled chunks", + ); + assert_eq!(collect_chunks(sealed), data); + + set_spill_override(None); + } + + /// Merging spilled chains loads bodies call-scoped and consolidates + /// correctly, and an untouched survivor keeps its spilled body. + #[mz_ore::test] + fn merge_spilled_chains() { + set_spill_override(Some(test_pool())); + + let a: Vec = (0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect(); + let b: Vec = (0..20_000u64).map(|i| ((i, 0), 0, 2i64)).collect(); + + let mut in1 = VecDeque::from([TestChunk::commit(build_column(&a), 0)]); + let mut in2 = VecDeque::from([TestChunk::commit(build_column(&b), 0)]); + assert!(in1[0].is_spilled() && in2[0].is_spilled()); + + let mut out = VecDeque::new(); + while !in1.is_empty() && !in2.is_empty() { + TestChunk::merge(&mut in1, &mut in2, &mut out); + } + for tail in in1.drain(..).chain(in2.drain(..)) { + out.push_back(tail); + } + + let expected: Vec = (0..20_000u64).map(|i| ((i, 0), 0, 3i64)).collect(); + assert_eq!(collect_chunks(out), expected); + + set_spill_override(None); + } + + /// A merge whose fronts have disjoint key ranges pushes the untouched + /// survivor back in its original (spilled) form rather than rewriting it. + #[mz_ore::test] + fn merge_untouched_survivor_stays_spilled() { + let pool = test_pool(); + let low: Vec = (0..100u64).map(|i| ((i, 0), 0, 1i64)).collect(); + let high: Vec = (1000..1100u64).map(|i| ((i, 0), 0, 1i64)).collect(); + + let mut in1 = VecDeque::from([force_spill( + ColumnChunk::from_column(build_column(&low)), + &pool, + )]); + let mut in2 = VecDeque::from([force_spill( + ColumnChunk::from_column(build_column(&high)), + &pool, + )]); + let mut out = VecDeque::new(); + TestChunk::merge(&mut in1, &mut in2, &mut out); + + // `low` is fully consumed. `high` was never touched and must come + // back spilled. + assert!(in1.is_empty()); + assert_eq!(in2.len(), 1); + assert!(in2[0].is_spilled(), "untouched survivor must stay spilled"); + let mut all = collect_chunks(out); + Extend::extend(&mut all, collect_chunks(in2.drain(..))); + let mut expected = low; + Extend::extend(&mut expected, high); + assert_eq!(all, expected); + } + + /// Merge output is one generation past its deepest input, a survivor + /// rewritten from its remainder keeps its own depth, and a chunk passed + /// through the disjoint fast path keeps its depth. + #[mz_ore::test] + fn merge_derives_generational_depth() { + let low: Vec = (0..100u64).map(|i| ((i, 0), 0, 1i64)).collect(); + let high: Vec = (50..150u64).map(|i| ((i, 0), 0, 1i64)).collect(); + let mut in1 = VecDeque::from([ColumnChunk::from_column(build_column(&low))]); + let mut in2 = VecDeque::from([ColumnChunk::from_column(build_column(&high))]); + assert_eq!(in1[0].depth(), 0, "fresh chunks start at depth 0"); + let mut out = VecDeque::new(); + TestChunk::merge(&mut in1, &mut in2, &mut out); + assert!(!out.is_empty()); + for chunk in &out { + assert_eq!(chunk.depth(), 1, "merge output is one past its inputs"); + } + // The merge runs through the shared horizon, so `high` survives with + // its unmerged remainder at its original depth. + assert!(in1.is_empty()); + assert_eq!(in2.len(), 1); + assert_eq!(in2[0].depth(), 0, "rewritten survivor keeps its depth"); + + // A disjoint merge moves the lower front to the output unchanged. + let mut in1 = VecDeque::from([ColumnChunk::Resident(Rc::new(build_column(&low)), 3)]); + let far: Vec = (1000..1100u64).map(|i| ((i, 0), 0, 1i64)).collect(); + let mut in2 = VecDeque::from([ColumnChunk::from_column(build_column(&far))]); + let mut out = VecDeque::new(); + TestChunk::merge(&mut in1, &mut in2, &mut out); + assert_eq!(out.len(), 1); + assert_eq!(out[0].depth(), 3, "pass-through keeps its depth"); + } + + /// Advance output and carry keep the deepest input depth, since + /// compaction rewrites within a generation. + #[mz_ore::test] + fn advance_preserves_depth() { + let data: Vec = (0..100u64).map(|i| ((i, 0), 1, 1i64)).collect(); + let mut input = VecDeque::from([ + ColumnChunk::Resident(Rc::new(build_column(&data[..50])), 2), + ColumnChunk::Resident(Rc::new(build_column(&data[50..])), 1), + ]); + let frontier = Antichain::from_elem(5u64); + let mut out = VecDeque::new(); + TestChunk::advance(&mut input, frontier.borrow(), false, &mut out); + for chunk in out.iter().chain(input.iter()) { + assert_eq!(chunk.depth(), 2); + } + TestChunk::advance(&mut input, frontier.borrow(), true, &mut out); + assert!(input.is_empty()); + assert!(!out.is_empty()); + for chunk in &out { + assert_eq!(chunk.depth(), 2); + } + } + + /// Settle commits at the deepest depth among coalesced chunks, and a + /// commit large enough to spill carries the depth into its spilled + /// metadata (and thus into the pool hints). + #[mz_ore::test] + fn settle_commits_at_accumulated_depth() { + set_spill_override(Some(test_pool())); + let big: Vec = (0..100_000u64).map(|i| ((i, 0), 0, 1i64)).collect(); + let mut input = VecDeque::from([ + ColumnChunk::Resident(Rc::new(build_column(&big)), 1), + ColumnChunk::Resident(Rc::new(build_column(&[((0, 0), 0, 1)])), 0), + ColumnChunk::Resident(Rc::new(build_column(&[((1, 0), 0, 1)])), 2), + ]); + let mut out = VecDeque::new(); + TestChunk::settle(&mut input, true, &mut out); + assert!(input.is_empty()); + assert_eq!(out.len(), 2); + assert!(out[0].is_spilled(), "large commit must spill"); + assert_eq!(out[0].depth(), 1, "sole commit keeps its depth"); + assert!(!out[1].is_spilled(), "small commit stays resident"); + assert_eq!(out[1].depth(), 2, "coalesced commit takes the max depth"); + set_spill_override(None); + } + + /// The settle carry commits at a monotone size threshold rather than the + /// periodic ship window, so mid-window chunk sizes cannot make it grow + /// past the target unbounded. + #[mz_ore::test] + fn settle_carry_commits_at_target() { + // ~1.5 MiB per chunk: inside the dead zone of the periodic window + // check (see `at_commit_size`). + let chunk_rows = u64::cast_from(1_500_000usize / 24); + let mut input: VecDeque = (0..4u64) + .map(|c| { + let data: Vec = (0..chunk_rows) + .map(|i| ((c * chunk_rows + i, 0), 0, 1i64)) + .collect(); + ColumnChunk::from_column(build_column(&data)) + }) + .collect(); + let mut out = VecDeque::new(); + TestChunk::settle(&mut input, true, &mut out); + for chunk in &out { + let col = chunk.clone().into_column(); + assert!( + col.length_in_bytes() < 2 * COMMIT_BYTES, + "settled chunk of {} bytes exceeds twice the commit target", + col.length_in_bytes(), + ); + } + assert_eq!( + collect_chunks(out).len(), + usize::try_from(4 * chunk_rows).unwrap(), + ); + } + + /// A tiny chunk stays resident regardless of the spill gate. + #[mz_ore::test] + fn small_chunks_stay_resident() { + set_spill_override(Some(test_pool())); + let committed = TestChunk::commit(build_column(&[((1, 1), 0, 1)]), 0); + assert!(!committed.is_spilled()); + set_spill_override(None); + } + + /// The smallest column whose serialized size reaches `SPILL_MIN_BYTES`. + /// One record less sits under the spill floor. + fn column_at_spill_floor() -> (Column, u64) { + let mut col: Column = Column::default(); + let mut n = 0u64; + while col.length_in_bytes() < SPILL_MIN_BYTES { + col.push_into(((n, n), 0, 1)); + n += 1; + } + (col, n) + } + + /// Bodies straddling the spill floor: one record under stays resident, + /// at the floor spills. + #[mz_ore::test] + fn spill_floor_boundary() { + set_spill_override(Some(test_pool())); + let (col, n) = column_at_spill_floor(); + let mut under: Column = Column::default(); + for m in 0..n - 1 { + under.push_into(((m, m), 0, 1)); + } + assert!(under.length_in_bytes() < SPILL_MIN_BYTES); + assert!(!TestChunk::commit(under, 0).is_spilled()); + assert!(TestChunk::commit(col, 0).is_spilled()); + set_spill_override(None); + } + + /// The compute and storage spill gates compose as an OR: either gate + /// routes commits to the installed pool, and each setter writes only its + /// own gate. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] + fn spill_gates_compose_as_or() { + let installed = + crate::pool_config::apply_pool_config(crate::pool_config::PoolPagerConfig { + budget_bytes: 32 << 20, + spill_threads: 1, + eager_backing: false, + rss_target_bytes: 16 << 20, + }); + assert!(installed, "pool reservation failed"); + // A body at the spill floor, so the gates alone decide. + let (col, _) = column_at_spill_floor(); + let commit = |col: &Column| TestChunk::commit(col.clone(), 0).is_spilled(); + + assert!(!commit(&col), "both gates off"); + set_storage_spill_enabled(true); + assert!(commit(&col), "the storage gate alone spills"); + set_compute_spill_enabled(false); + assert!( + commit(&col), + "the compute setter must not clobber the storage gate" + ); + set_compute_spill_enabled(true); + set_storage_spill_enabled(false); + assert!(commit(&col), "the compute gate alone spills"); + set_compute_spill_enabled(false); + assert!(!commit(&col), "both gates off again"); + } + + /// Re-spilling an already-serialized body exercises the `Column::Align` + /// branch of `spill_column` and round-trips byte-identically. + #[mz_ore::test] + fn spill_align_round_trip() { + let pool = test_pool(); + let data: Vec = (0..64u64).map(|k| ((k, k), 0, 1)).collect(); + let spilled = force_spill(ColumnChunk::from_column(build_column(&data)), &pool); + let column = spilled.into_column(); + let Column::Align(words) = &column else { + panic!("a spilled body reads back as Column::Align"); + }; + let words = words.clone(); + let respilled = force_spill(ColumnChunk::from_column(column), &pool); + let reread = respilled.into_column(); + let Column::Align(words2) = &reread else { + panic!("a spilled body reads back as Column::Align"); + }; + assert_eq!(&words, words2, "byte-identical round trip"); + assert_eq!(collect_column(&reread), data); + } + + /// Merge depth saturates at `u8::MAX` instead of wrapping. + #[mz_ore::test] + fn merge_depth_saturates() { + let a = ColumnChunk::Resident( + Rc::new(build_column(&[((1, 0), 0, 1), ((3, 0), 0, 1)])), + u8::MAX, + ); + let b = ColumnChunk::Resident( + Rc::new(build_column(&[((2, 0), 0, 1), ((4, 0), 0, 1)])), + u8::MAX, + ); + let mut in1 = VecDeque::from([a]); + let mut in2 = VecDeque::from([b]); + let mut out = VecDeque::new(); + TestChunk::merge(&mut in1, &mut in2, &mut out); + for chunk in out.iter().chain(in1.iter()).chain(in2.iter()) { + assert_eq!(chunk.depth(), u8::MAX, "depth saturates"); + } + } + + /// `into_column` on a shared resident chunk copies instead of stealing + /// the shared body. + #[mz_ore::test] + fn into_column_copies_shared_resident() { + let data: Vec = vec![((1, 1), 0, 1), ((2, 2), 0, 1)]; + let a = ColumnChunk::from_column(build_column(&data)); + let b = a.clone(); + assert_eq!(collect_column(&a.into_column()), data); + assert_eq!(collect_column(&b.into_column()), data); + } +}