timely-util: spill young chunk generations uncompressed

A chunk at generational depth d is rewritten with frequency proportional
to 2^-d under geometric merging, so compressing a shallow generation buys
pool bytes back for only a short stay at a guaranteed near-term codec
round-trip: fresh (depth 0) chunks are consumed by their first merge with
certainty. Spill bodies below a configurable compression depth floor
(default 1) under a new identity ExtentCodec in mz_ore::pool instead:
still inserted into the pool, so they stay budget-accounted and
swap-backed like every extent, but encode and decode reduce to copies.

Keeping young generations out of the pool entirely was rejected: exempted
bytes are invisible to the pool budget, and coalesced carries at 2 MiB
each across many operators could accumulate unbudgeted resident state
exactly when the system is busiest.

The floor is wired through the replica-scoped
column_chunk_compress_min_depth dyncfg alongside the other chunk knobs.

Author: DAlperin

misc/python/materialize/mzcompose/__init__.py

@@ -592,6 +592,7 @@ def get_default_system_parameters(
     "linear_join_yielding",
     "enable_column_paged_batcher",
     "enable_column_paged_batcher_spill",
+    "column_chunk_compress_min_depth",
     "column_paged_batcher_budget_fraction",
     "column_paged_batcher_lz4",
     "column_paged_batcher_swap_pageout",

misc/python/materialize/parallel_workload/action.py

@@ -3052,6 +3052,11 @@ def __init__(
             "0.02",
         ]
         self.flags_with_values["enable_upsert_paged_spill"] = BOOLEAN_FLAG_VALUES
+        self.flags_with_values["column_chunk_compress_min_depth"] = [
+            "0",  # compress every spilled body
+            "1",  # the default: fresh chunks store uncompressed
+            "4",  # exempt several young generations
+        ]
         # 0 forces the estimated-size path for every table, the default forces
         # the exact COUNT(*) path for workload-sized tables.
         self.flags_with_values["mysql_source_snapshot_exact_count_max_rows"] = [

src/compute-types/src/dyncfgs.rs

@@ -65,6 +65,24 @@ pub const ENABLE_COLUMN_PAGED_BATCHER_SPILL: Config<bool> = Config::new(
 )
 .scoped(ParameterScope::Replica);
 
+/// The youngest chunk generation whose spilled bodies are compressed.
+///
+/// A chunk at generational depth `d` is rewritten with frequency
+/// proportional to `2^-d` under geometric merging, so compressing shallow
+/// generations buys pool bytes back for only a short stay at a guaranteed
+/// near-term codec round-trip. Generations below the floor spill under the
+/// identity codec: fully budgeted and swap-backed, with encode and decode
+/// reduced to copies. The default exempts only fresh (depth 0) chunks,
+/// which their first merge consumes with certainty. `0` compresses every
+/// spilled body.
+pub const COLUMN_CHUNK_COMPRESS_MIN_DEPTH: Config<u32> = Config::new(
+    "column_chunk_compress_min_depth",
+    1,
+    "The youngest chunk generation whose spilled bodies are lz4-compressed in the buffer \
+     pool; younger generations store uncompressed. 0 compresses every spilled body.",
+)
+.scoped(ParameterScope::Replica);
+
 /// Resident-bytes budget fraction for chunk spilling. Two consumers read
 /// it: the column pager's tiered policy multiplies it against the
 /// announced memory limit, and the buffer pool (`mz_ore::pool`)
@@ -601,4 +619,5 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet {
         .add(&COLUMN_PAGED_BATCHER_SPILL_WORKER_COUNT)
         .add(&COLUMN_PAGED_BATCHER_EAGER_BACKING)
         .add(&COLUMN_PAGED_BATCHER_POOL_RSS_TARGET_FRACTION)
+        .add(&COLUMN_CHUNK_COMPRESS_MIN_DEPTH)
 }

src/compute/src/compute_state.rs

@@ -428,6 +428,13 @@ impl ComputeState {
                     warn!("chunk spill: buffer pool unavailable; chunks stay resident");
                 }
             }
+
+            // The generational depth floor below which spilled bodies store
+            // uncompressed. Subsystem-independent, so applied here alongside
+            // the rest of the process-wide chunk configuration.
+            let compress_min_depth =
+                u8::try_from(COLUMN_CHUNK_COMPRESS_MIN_DEPTH.get(config)).unwrap_or(u8::MAX);
+            mz_timely_util::columnar::chunk::set_compress_min_depth(compress_min_depth);
         }
 
         // Remember the maintenance interval locally to avoid reading it from the config set on

src/ore/src/pool.rs

@@ -126,6 +126,28 @@ pub trait ExtentCodec: std::fmt::Debug + Send + Sync {
     fn decode(&self, stored: &[u8], body: &mut [u8]);
 }
 
+/// The identity [`ExtentCodec`]: the stored form is the body. Encode and
+/// decode are copies, and range reads copy the range directly, so a chunk
+/// stored under this codec pays no compression work in either direction
+/// while remaining fully budgeted and swap-backed like any other extent.
+#[derive(Debug)]
+pub struct IdentityCodec;
+
+/// The [`IdentityCodec`] instance to pass to [`Pool::insert_with`].
+pub static IDENTITY_CODEC: IdentityCodec = IdentityCodec;
+
+impl ExtentCodec for IdentityCodec {
+    fn encode(&self, body: &[u8], out: &mut Vec<u8>) {
+        out.clear();
+        out.extend_from_slice(body);
+    }
+
+    fn decode(&self, stored: &[u8], body: &mut [u8]) {
+        assert_eq!(stored.len(), body.len(), "identity stored form is the body");
+        body.copy_from_slice(stored);
+    }
+}
+
 /// The largest stored form [`ExtentCodec::encode`] may produce for a
 /// `body_len`-byte body: an incompressible-input expansion matching lz4's
 /// worst case plus a four-byte length prefix. The extent store's size-class
@@ -3703,6 +3725,24 @@ mod tests {
         assert_eq!(pool.stats().resident_bytes, 0);
     }
 
+    /// The identity codec stores the body verbatim: eviction and reads,
+    /// whole and by range, reconstruct it unchanged.
+    #[mz_ore::test]
+    fn identity_codec_round_trips() {
+        let pool = test_pool(usize::MAX);
+        let want = payload(SMALL, 601);
+        let h = pool.insert_with(SMALL, ChunkHints::default(), &IDENTITY_CODEC, |dst| {
+            dst.copy_from_slice(&want);
+        });
+        assert_eq!(read(&h), want);
+        pool.evict(&h);
+        assert_eq!(read(&h), want, "round-trips through the extent");
+        pool.evict(&h);
+        let mut range = Vec::new();
+        h.read_range_into(8..24, &mut range);
+        assert_eq!(range, want[8..24], "range reads copy the range directly");
+    }
+
     #[mz_ore::test]
     fn insert_with_fills_in_place() {
         let pool = test_pool(usize::MAX);

src/timely-util/src/columnar/chunk.rs

@@ -44,18 +44,18 @@
 //! resident fence metadata so a probe set faults only the chunk bodies it
 //! actually touches.
 
-use std::cell::RefCell;
+use std::cell::{Cell, RefCell};
 use std::collections::VecDeque;
 use std::rc::Rc;
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU8, 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 mz_ore::pool::{ChunkHandle, ChunkHints, ExtentCodec, IDENTITY_CODEC, Pool};
 use timely::Accountable;
 use timely::container::{ContainerBuilder, PushInto};
 use timely::dataflow::channels::ContainerBytes;
@@ -78,6 +78,11 @@ thread_local! {
     /// pool without touching process-global state.
     static SPILL_OVERRIDE: RefCell<Option<Pool>> = const { RefCell::new(None) };
 
+    /// A thread-scoped depth-floor override, taking precedence over the
+    /// global value. Lets tests and benches pin the floor without racing
+    /// concurrently running tests on the process-global state.
+    static COMPRESS_MIN_DEPTH_OVERRIDE: Cell<Option<u8>> = const { Cell::new(None) };
+
     /// Reusable staging for call-scoped reads of spilled bodies.
     static READ_SCRATCH: RefCell<Vec<u64>> = const { RefCell::new(Vec::new()) };
 }
@@ -113,6 +118,51 @@ pub fn set_spill_override(pool: Option<Pool>) {
     SPILL_OVERRIDE.with(|cell| *cell.borrow_mut() = pool);
 }
 
+/// The youngest generational depth whose spilled bodies are compressed. See
+/// [`set_compress_min_depth`].
+static COMPRESS_MIN_DEPTH: AtomicU8 = AtomicU8::new(DEFAULT_COMPRESS_MIN_DEPTH);
+
+/// Set the youngest generational depth whose spilled bodies are compressed.
+///
+/// A chunk at depth `d` is rewritten (merged, extracted, advanced) with
+/// frequency proportional to `2^-d` under geometric merging, so compressing
+/// a shallow chunk buys a short stay in the pool at the cost of a guaranteed
+/// near-term codec round-trip: the body is encoded only to be read back and
+/// decoded by the next rewrite. Generations below the floor spill under the
+/// identity codec instead: still budgeted and swap-backed like every extent,
+/// but encode and decode are copies. The floor never exempts a body from the
+/// pool, so it cannot grow unbudgeted resident state.
+///
+/// `0` compresses every spilled body. Consulted at every commit, so changes
+/// apply to running dataflows.
+pub fn set_compress_min_depth(depth: u8) {
+    COMPRESS_MIN_DEPTH.store(depth, Ordering::Relaxed);
+}
+
+/// Set or unset a thread-scoped depth-floor override, taking precedence over
+/// [`set_compress_min_depth`]. For tests and benches, which run concurrently
+/// and must not race on the process-global floor.
+pub fn set_compress_min_depth_override(depth: Option<u8>) {
+    COMPRESS_MIN_DEPTH_OVERRIDE.with(|cell| cell.set(depth));
+}
+
+/// The depth floor in effect for this thread's commits.
+fn compress_min_depth() -> u8 {
+    COMPRESS_MIN_DEPTH_OVERRIDE
+        .with(|cell| cell.get())
+        .unwrap_or_else(|| COMPRESS_MIN_DEPTH.load(Ordering::Relaxed))
+}
+
+/// The codec a body at `depth` stores under: identity below the compression
+/// floor, lz4 at and past it.
+fn codec_for_depth(depth: u8) -> &'static dyn ExtentCodec {
+    if depth < compress_min_depth() {
+        &IDENTITY_CODEC
+    } else {
+        &LZ4_CODEC
+    }
+}
+
 /// The pool committed chunks spill to, if any.
 fn spill_pool() -> Option<Pool> {
     if let Some(pool) = SPILL_OVERRIDE.with(|cell| cell.borrow().clone()) {
@@ -161,6 +211,16 @@ const COMMIT_BYTES: usize = 2 << 20;
 /// unbudgeted heap, and no accounting here would catch it.
 const SPILL_MIN_BYTES: usize = 64 << 10;
 
+/// The default compression depth floor: fresh (depth 0) bodies spill
+/// uncompressed.
+///
+/// A fresh chunk is consumed by its first merge with certainty, so
+/// compressing it can never save pool bytes for longer than one merge
+/// cadence and always costs a full encode plus decode. Depth 1 and beyond
+/// have survived a merge and wait geometrically longer for the next, so
+/// their compression amortizes.
+const DEFAULT_COMPRESS_MIN_DEPTH: u8 = 1;
+
 /// 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.
@@ -309,14 +369,20 @@ impl<D: Columnar, T: Columnar, R: Columnar> ColumnChunk<D, T, R> {
 
     /// Spill a non-empty column into `pool` unconditionally, capturing the
     /// resident fence metadata.
+    ///
+    /// Generations below the compression depth floor store under the
+    /// identity codec: rewritten too soon for compression to amortize, they
+    /// stay budgeted and swap-backed while encode and decode reduce to
+    /// copies.
     fn spill_body(column: Column<(D, T, R)>, pool: &Pool, depth: u8) -> Self {
+        let codec = codec_for_depth(depth);
         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 });
+        let handle = spill_column(column, pool, len_bytes, ChunkHints { depth }, codec);
         ColumnChunk::Spilled(Rc::new(SpilledBody {
             records,
             fences,
@@ -379,13 +445,14 @@ fn spill_column<C: Columnar>(
     pool: &Pool,
     len_bytes: usize,
     hints: ChunkHints,
+    codec: &'static dyn ExtentCodec,
 ) -> ChunkHandle {
     mz_ore::soft_assert_eq_no_log!(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| {
+        Column::Align(words) => {
+            pool.insert_with(words.len(), hints, codec, |dst| dst.copy_from_slice(&words))
+        }
+        other => pool.insert_with(len_bytes / 8, hints, codec, |dst| {
             let bytes: &mut [u8] = bytemuck::cast_slice_mut(dst);
             let mut cursor = std::io::Cursor::new(bytes);
             other.into_bytes(&mut cursor);
@@ -1701,6 +1768,39 @@ mod tests {
         set_spill_override(None);
     }
 
+    /// The compression depth floor picks the codec, not whether a body
+    /// spills: shallow generations store at identity, the floor and deeper
+    /// at lz4, and every depth spills and round-trips.
+    #[mz_ore::test]
+    fn spill_codec_depth_floor() {
+        set_spill_override(Some(test_pool()));
+        set_compress_min_depth_override(Some(2));
+        // Codec identity via Debug: ZST statics and dyn vtables make
+        // pointer comparison unreliable.
+        let codec_name = |depth: u8| format!("{:?}", codec_for_depth(depth));
+        assert_eq!(codec_name(0), "IdentityCodec");
+        assert_eq!(codec_name(1), "IdentityCodec");
+        assert_eq!(codec_name(2), "Lz4Codec");
+        assert_eq!(codec_name(u8::MAX), "Lz4Codec");
+
+        let data: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect();
+        let data = consolidate(data);
+        let column = build_column(&data);
+        for depth in [0u8, 1, 2, 3] {
+            let chunk = TestChunk::commit(column.clone(), depth);
+            assert!(chunk.is_spilled(), "depth {depth} must spill");
+            assert_eq!(collect_column(&chunk.into_column()), data);
+        }
+        set_spill_override(None);
+        set_compress_min_depth_override(None);
+
+        // The default floor stores only fresh (depth 0) bodies at identity.
+        set_compress_min_depth_override(Some(DEFAULT_COMPRESS_MIN_DEPTH));
+        assert_eq!(codec_name(0), "IdentityCodec");
+        assert_eq!(codec_name(1), "Lz4Codec");
+        set_compress_min_depth_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.
@@ -1732,6 +1832,7 @@ mod tests {
         assert!(commit(&col), "the compute gate alone spills");
         set_compute_spill_enabled(false);
         assert!(!commit(&col), "both gates off again");
+        set_compress_min_depth_override(None);
     }
 
     /// Re-spilling an already-serialized body exercises the `Column::Align`