Skip to main content

mz_row_spine/
lib.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//! Types and traits in support of containers for row-encoded byte slices.
11//!
12//! This includes the vanilla `bytes_container` that holds byte slices in contiguous
13//! allocations, as well as a `dictionary` encoding wrapper that is able to rewrite
14//! the byte slices to use spare tags in each column to reference common values.
15
16pub use self::arc_batch::{ArcBatch, ArcBuilder};
17pub use self::dictionary::DatumContainer;
18pub use self::dictionary::DatumSeq;
19pub use self::dictionary::builders::RowRowColPagedState;
20pub use self::offset_opt::OffsetOptimized;
21pub use self::spines::{
22    ArcOrdKeyBatcher, ArcOrdKeyBuilder, ArcOrdKeySpine, ArcOrdValBatcher, ArcOrdValBuilder,
23    ArcOrdValSpine, FundedValRowSpine, RowBatcher, RowBuilder, RowRowBatcher, RowRowBuilder,
24    RowRowColPagedBuilder, RowRowSpine, RowSpine, RowValBatcher, RowValBuilder, RowValSpine,
25    ValRowBatcher, ValRowBuilder, ValRowColPagedBuilder, ValRowSpine,
26};
27
28mod arc_batch;
29
30use differential_dataflow::trace::implementations::OffsetList;
31
32/// Enable per-column dictionary compression in row containers.
33pub static DICTIONARY_COMPRESSION: std::sync::atomic::AtomicBool =
34    std::sync::atomic::AtomicBool::new(false);
35
36/// Spines specialized to contain `Row` types in keys and values.
37mod spines {
38    use columnation::Columnation;
39    use differential_dataflow::trace::implementations::BatchContainer;
40    use differential_dataflow::trace::implementations::Layout;
41    use differential_dataflow::trace::implementations::Update;
42    use differential_dataflow::trace::implementations::Vector;
43    use differential_dataflow::trace::implementations::merge_batcher::{
44        MergeBatcher, chunker::ContainerChunker, vec::VecMerger,
45    };
46    use differential_dataflow::trace::implementations::ord_neu::{
47        OrdKeyBatch, OrdKeyBuilder, OrdValBatch, OrdValBuilder,
48    };
49    use differential_dataflow::trace::implementations::spine_fueled::Spine;
50    use mz_repr::Row;
51    use mz_timely_util::columnation::{ColInternalMerger, ColumnationChunker, ColumnationStack};
52
53    use crate::arc_batch::{ArcBatch, ArcBuilder};
54    use crate::{DatumContainer, OffsetOptimized};
55
56    /// Batcher matching `mz_compute::typedefs::KeyValBatcher`, redeclared
57    /// locally so this crate does not need to depend on `mz_compute`.
58    type KeyValBatcher<K, V, T, D, Se> =
59        MergeBatcher<ColumnationChunker<((K, V), T, D)>, ColInternalMerger<(K, V), T, D>, Se>;
60    type KeyBatcher<K, T, D, Se> = KeyValBatcher<K, (), T, D, Se>;
61
62    pub type RowRowSpine<T, R> = Spine<ArcBatch<OrdValBatch<RowRowLayout<((Row, Row), T, R)>>>>;
63    pub type RowRowBatcher<T, R> = KeyValBatcher<Row, Row, T, R, RowRowBuilder<T, R>>;
64    pub type RowRowBuilder<T, R> = ArcBuilder<crate::dictionary::builders::RowRowBuilder<T, R>>;
65
66    /// `RowRowBuilder` variant that consumes [`ColumnBody`] chunks. Pairs with
67    /// any batcher whose chains are bodies, spillable ([`AccountedChunkBatcher`]
68    /// behind an [`UnchunkBuilder`]) or resident ([`Col2ValColBatcher`]) alike,
69    /// so the `Paged` in the name records where it started rather than a
70    /// restriction. Installs a dictionary codec at seal time, gathering
71    /// statistics from the sealed chain, so columnar arrangements compress on
72    /// the same footing as the columnation-fed [`RowRowBuilder`].
73    ///
74    /// [`AccountedChunkBatcher`]: mz_timely_util::columnar::chunk::AccountedChunkBatcher
75    /// [`Col2ValColBatcher`]: mz_timely_util::columnar::Col2ValColBatcher
76    /// [`ColumnBody`]: mz_timely_util::columnar::body::ColumnBody
77    /// [`UnchunkBuilder`]: mz_timely_util::columnar::chunk::UnchunkBuilder
78    pub type RowRowColPagedBuilder<T, R> =
79        ArcBuilder<crate::dictionary::builders::RowRowColPagedBuilder<T, R>>;
80
81    pub type RowValSpine<V, T, R> = Spine<ArcBatch<OrdValBatch<RowValLayout<((Row, V), T, R)>>>>;
82    pub type RowValBatcher<V, T, R> = KeyValBatcher<Row, V, T, R, RowValBuilder<V, T, R>>;
83    pub type RowValBuilder<V, T, R> =
84        ArcBuilder<crate::dictionary::builders::RowValBuilder<V, T, R>>;
85
86    /// Key-only `Row` spine. `DC` is the diff container, see `RowLayout`.
87    pub type RowSpine<T, R, DC = ColumnationStack<R>> =
88        Spine<ArcBatch<OrdKeyBatch<RowLayout<((Row, ()), T, R), DC>>>>;
89    pub type RowBatcher<T, R, DC = ColumnationStack<R>> =
90        KeyBatcher<Row, T, R, RowBuilder<T, R, DC>>;
91    pub type RowBuilder<T, R, DC = ColumnationStack<R>> =
92        ArcBuilder<crate::dictionary::builders::RowBuilder<T, R, DC>>;
93
94    pub type ValRowSpine<K, T, R> = Spine<ArcBatch<OrdValBatch<ValRowLayout<((K, Row), T, R)>>>>;
95    /// A `ValRowSpine` whose optional exertion is funded by its input.
96    pub type FundedValRowSpine<K, T, R> =
97        mz_timely_util::funded_spine::Spine<ArcBatch<OrdValBatch<ValRowLayout<((K, Row), T, R)>>>>;
98    pub type ValRowBatcher<K, T, R> = KeyValBatcher<K, Row, T, R, ValRowBuilder<K, T, R>>;
99    pub type ValRowBuilder<K, T, R> =
100        ArcBuilder<crate::dictionary::builders::ValRowBuilder<K, T, R>>;
101
102    /// `ValRowBuilder` variant that consumes [`Column`] chunks. Pairs with
103    /// `Col2ValPagedBatcher<K, Row, T, R, Chu, Se>` for the spillable arrange path where
104    /// keys are arbitrary `Columnar` values (e.g. `UpsertKey`) and values are
105    /// packed `Row` bytes. Installs a dictionary codec on the value container at
106    /// seal time, gathering statistics from the sealed `Column` chain; keys are
107    /// not `Row`-shaped and so are left uncompressed.
108    ///
109    /// [`Column`]: mz_timely_util::columnar::Column
110    pub type ValRowColPagedBuilder<K, T, R> =
111        ArcBuilder<crate::dictionary::builders::ValRowColPagedBuilder<K, T, R>>;
112
113    /// A generic `Arc`-backed key/value spine, for callers outside `mz_compute` that need an
114    /// arrangement over non-`Row`-specialized types. The `Arc` handle rides on the local
115    /// [`ArcBatch`] newtype, so no differential-side `Arc` batch impls are required.
116    pub type ArcOrdValSpine<K, V, T, R> = Spine<ArcBatch<OrdValBatch<Vector<((K, V), T, R)>>>>;
117    /// Generic `Arc`-backed key-only spine. See [`ArcOrdValSpine`].
118    pub type ArcOrdKeySpine<K, T, R> = Spine<ArcBatch<OrdKeyBatch<Vector<((K, ()), T, R)>>>>;
119    /// Builder pairing with [`ArcOrdValSpine`].
120    pub type ArcOrdValBuilder<K, V, T, R> =
121        ArcBuilder<OrdValBuilder<Vector<((K, V), T, R)>, Vec<((K, V), T, R)>>>;
122    /// Builder pairing with [`ArcOrdKeySpine`].
123    pub type ArcOrdKeyBuilder<K, T, R> =
124        ArcBuilder<OrdKeyBuilder<Vector<((K, ()), T, R)>, Vec<((K, ()), T, R)>>>;
125    /// Batcher pairing with [`ArcOrdValSpine`], over `Vec` input.
126    pub type ArcOrdValBatcher<K, V, T, R> = MergeBatcher<
127        ContainerChunker<Vec<((K, V), T, R)>>,
128        VecMerger<(K, V), T, R>,
129        ArcOrdValBuilder<K, V, T, R>,
130    >;
131    /// Batcher pairing with [`ArcOrdKeySpine`], over `Vec` input.
132    pub type ArcOrdKeyBatcher<K, T, R> = MergeBatcher<
133        ContainerChunker<Vec<((K, ()), T, R)>>,
134        VecMerger<(K, ()), T, R>,
135        ArcOrdKeyBuilder<K, T, R>,
136    >;
137
138    /// A layout based on timely stacks
139    pub struct RowRowLayout<U>
140    where
141        U: Update<Key = Row, Val = Row>,
142    {
143        phantom: std::marker::PhantomData<U>,
144    }
145    pub struct RowValLayout<U>
146    where
147        U: Update<Key = Row>,
148    {
149        phantom: std::marker::PhantomData<U>,
150    }
151    /// Layout for key-only `Row` updates. `DC` is the diff `BatchContainer`, a
152    /// columnation stack by default.
153    pub struct RowLayout<U, DC = ColumnationStack<<U as Update>::Diff>>
154    where
155        U: Update<Key = Row, Val = ()>,
156    {
157        phantom: std::marker::PhantomData<(U, DC)>,
158    }
159    /// Mirror of [`RowValLayout`] with the roles swapped: arbitrary `Columnation`
160    /// keys with `Row` values stored as packed bytes in a [`DatumContainer`].
161    pub struct ValRowLayout<U>
162    where
163        U: Update<Val = Row>,
164    {
165        phantom: std::marker::PhantomData<U>,
166    }
167
168    impl<U> Layout for RowRowLayout<U>
169    where
170        U: Update<Key = Row, Val = Row>,
171        U::Time: Columnation,
172        U::Diff: Columnation,
173    {
174        type KeyContainer = DatumContainer;
175        type ValContainer = DatumContainer;
176        type TimeContainer = ColumnationStack<U::Time>;
177        type DiffContainer = ColumnationStack<U::Diff>;
178        type OffsetContainer = OffsetOptimized;
179    }
180    impl<U> Layout for RowValLayout<U>
181    where
182        U: Update<Key = Row>,
183        U::Val: Columnation,
184        U::Time: Columnation,
185        U::Diff: Columnation,
186    {
187        type KeyContainer = DatumContainer;
188        type ValContainer = ColumnationStack<U::Val>;
189        type TimeContainer = ColumnationStack<U::Time>;
190        type DiffContainer = ColumnationStack<U::Diff>;
191        type OffsetContainer = OffsetOptimized;
192    }
193    impl<U, DC> Layout for RowLayout<U, DC>
194    where
195        U: Update<Key = Row, Val = ()>,
196        U::Time: Columnation,
197        DC: BatchContainer<Owned = U::Diff>,
198    {
199        type KeyContainer = DatumContainer;
200        type ValContainer = ColumnationStack<()>;
201        type TimeContainer = ColumnationStack<U::Time>;
202        type DiffContainer = DC;
203        type OffsetContainer = OffsetOptimized;
204    }
205    impl<U> Layout for ValRowLayout<U>
206    where
207        U: Update<Val = Row>,
208        U::Key: Columnation,
209        U::Time: Columnation,
210        U::Diff: Columnation,
211    {
212        type KeyContainer = ColumnationStack<U::Key>;
213        type ValContainer = DatumContainer;
214        type TimeContainer = ColumnationStack<U::Time>;
215        type DiffContainer = ColumnationStack<U::Diff>;
216        type OffsetContainer = OffsetOptimized;
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use crate::DatumContainer;
223    use crate::spines::{RowLayout, RowRowLayout, RowValLayout};
224    use differential_dataflow::trace::implementations::BatchContainer;
225    use differential_dataflow::trace::implementations::ord_neu::{OrdKeyBatch, OrdValBatch};
226    use mz_repr::adt::date::Date;
227    use mz_repr::adt::interval::Interval;
228    use mz_repr::{Datum, Diff, Row, SqlScalarType, Timestamp};
229    use mz_timely_util::columnation::ColumnationStack;
230
231    fn assert_send_sync<T: Send + Sync>() {}
232
233    /// The batch types backing our spines must stay `Send + Sync`, so that batches
234    /// can be shared across threads (for example behind an `Arc`) to serve reads
235    /// from outside the worker that maintains the trace. This holds because the
236    /// backing containers bottom out in `Vec`s, lgalloc regions, and `CompactBytes`,
237    /// all of which are thread-safe.
238    #[mz_ore::test]
239    fn batches_are_send_sync() {
240        assert_send_sync::<OrdValBatch<RowRowLayout<((Row, Row), Timestamp, Diff)>>>();
241        assert_send_sync::<OrdValBatch<RowValLayout<((Row, Row), Timestamp, Diff)>>>();
242        assert_send_sync::<OrdKeyBatch<RowLayout<((Row, ()), Timestamp, Diff)>>>();
243        assert_send_sync::<ColumnationStack<((Row, Row), Timestamp, Diff)>>();
244    }
245
246    /// A chain of chunks seals with the chain's dictionary codecs installed,
247    /// so its batch compresses like one sealed from columns directly. A seal
248    /// that skips the install stores raw bytes and, per
249    /// `DatumContainer::promote_stats_to_codec`, drags every merge it joins
250    /// down with it.
251    #[mz_ore::test]
252    #[cfg_attr(miri, ignore)] // integer-to-pointer casts in row decoding are unsupported under miri
253    fn chunked_seal_installs_codecs() {
254        use std::sync::atomic::Ordering;
255
256        use differential_dataflow::trace::implementations::merge_batcher::Sealer;
257        use differential_dataflow::trace::implementations::ord_neu::OrdValBatch;
258        use mz_timely_util::columnar::body::ColumnBody;
259        use mz_timely_util::columnar::chunk::{ColumnChunk, ComputeSpill, UnchunkBuilder};
260        use timely::container::PushInto;
261        use timely::progress::Timestamp as _;
262
263        use crate::ArcBatch;
264
265        // Gate the dictionary path on. Safe for other tests: the flag only
266        // controls whether codecs are built, never decode results.
267        crate::DICTIONARY_COMPRESSION.store(true, Ordering::Relaxed);
268
269        // Low-cardinality rows, well under `STATS_THRESHOLD` (64Ki pushes), so
270        // a codec can only come from the seal-time install.
271        let time = Timestamp::minimum();
272        let updates: Vec<((Row, Row), Timestamp, i64)> = (0..2_000i64)
273            .map(|i| {
274                (
275                    (
276                        Row::pack_slice(&[Datum::Int64(i)]),
277                        Row::pack_slice(&[Datum::String("a repeated string value")]),
278                    ),
279                    time,
280                    1i64,
281                )
282            })
283            .collect();
284
285        // Cut the sorted run the way a merge batcher's chain is cut.
286        let columns = || -> Vec<ColumnBody<((Row, Row), Timestamp, i64)>> {
287            updates
288                .chunks(250)
289                .map(|part| {
290                    let mut column: ColumnBody<((Row, Row), Timestamp, i64)> = Default::default();
291                    for update in part {
292                        column.push_into(update);
293                    }
294                    column
295                })
296                .collect()
297        };
298        type Paged = crate::RowRowColPagedBuilder<Timestamp, i64>;
299        type Chunked = UnchunkBuilder<Paged, (Row, Row), Timestamp, i64, ComputeSpill>;
300
301        let mut chain = columns();
302        let from_columns = <Paged as Sealer<_>>::seal(&mut chain).expect("chain holds updates");
303
304        let mut chain: Vec<ColumnChunk<(Row, Row), Timestamp, i64, ComputeSpill>> =
305            columns().into_iter().map(ColumnChunk::from_body).collect();
306        let from_chunks = <Chunked as Sealer<_>>::seal(&mut chain).expect("chain holds updates");
307
308        assert!(
309            from_chunks.0.storage.keys.has_codec() && from_chunks.0.storage.vals.vals.has_codec(),
310            "the sealed batch carries the chain's codecs"
311        );
312
313        // Same updates and same cuts, so the chunked seal has to reach the
314        // encoded size the column seal does.
315        let heap = |c: &DatumContainer| {
316            let mut size = 0;
317            c.heap_size(|_, cap| size += cap);
318            size
319        };
320        let encoded = |b: &ArcBatch<OrdValBatch<RowRowLayout<((Row, Row), Timestamp, i64)>>>| {
321            heap(&b.0.storage.keys) + heap(&b.0.storage.vals.vals)
322        };
323
324        assert!(
325            encoded(&from_chunks) <= encoded(&from_columns),
326            "chunked seal should compress like the column seal: chunks={} columns={}",
327            encoded(&from_chunks),
328            encoded(&from_columns),
329        );
330    }
331
332    #[mz_ore::test]
333    #[cfg_attr(miri, ignore)] // unsupported operation: integer-to-pointer casts and `ptr::with_exposed_provenance` are not supported
334    fn test_round_trip() {
335        fn round_trip(datums: Vec<Datum>) {
336            let row = Row::pack(datums.clone());
337
338            let mut container = DatumContainer::with_capacity(row.byte_len());
339            container.push_own(&row);
340
341            // When run under miri this catches undefined bytes written to data
342            // eg by calling push_copy! on a type which contains undefined padding values
343            println!("{:?}", container.index(0).iter.data);
344
345            let datums2 = container.index(0).collect::<Vec<_>>();
346            assert_eq!(datums, datums2);
347        }
348
349        round_trip(vec![]);
350        round_trip(
351            SqlScalarType::enumerate()
352                .iter()
353                .flat_map(|r#type| r#type.interesting_datums())
354                .collect(),
355        );
356        round_trip(vec![
357            Datum::Null,
358            Datum::Null,
359            Datum::False,
360            Datum::True,
361            Datum::Int16(-21),
362            Datum::Int32(-42),
363            Datum::Int64(-2_147_483_648 - 42),
364            Datum::UInt8(0),
365            Datum::UInt8(1),
366            Datum::UInt16(0),
367            Datum::UInt16(1),
368            Datum::UInt16(1 << 8),
369            Datum::UInt32(0),
370            Datum::UInt32(1),
371            Datum::UInt32(1 << 8),
372            Datum::UInt32(1 << 16),
373            Datum::UInt32(1 << 24),
374            Datum::UInt64(0),
375            Datum::UInt64(1),
376            Datum::UInt64(1 << 8),
377            Datum::UInt64(1 << 16),
378            Datum::UInt64(1 << 24),
379            Datum::UInt64(1 << 32),
380            Datum::UInt64(1 << 40),
381            Datum::UInt64(1 << 48),
382            Datum::UInt64(1 << 56),
383            Datum::Date(Date::from_pg_epoch(365 * 45 + 21).unwrap()),
384            Datum::Interval(Interval {
385                months: 312,
386                ..Default::default()
387            }),
388            Datum::Interval(Interval::new(0, 0, 1_012_312)),
389            Datum::Bytes(&[]),
390            Datum::Bytes(&[0, 2, 1, 255]),
391            Datum::String(""),
392            Datum::String("العَرَبِيَّة"),
393        ]);
394    }
395
396    /// Exercises the *compressed* encode→decode paths, which the dyncfg-gated
397    /// `test_round_trip` never reaches (it installs no codec). We drive the codec
398    /// directly: observe a sample, build a codec via both `new_from([c1, c2])`
399    /// (the merge path) and `new_safe` (the safe-tag path), then round-trip every
400    /// row through it. We additionally assert the dictionary actually engaged, so
401    /// the test keeps covering the compressed branch rather than silently
402    /// degrading to raw fall-through.
403    #[mz_ore::test]
404    #[cfg_attr(miri, ignore)] // integer-to-pointer casts in row decoding are unsupported under miri
405    fn test_codec_round_trip() {
406        use crate::row_codec::ColumnsCodec;
407
408        // Rows with a small set of repeated, multi-byte string values, so the
409        // dictionary installs entries (MisraGries keeps values with len > 1 and
410        // count > 1). Mixing in an integer column exercises the raw fall-through
411        // (and thus the new soundness `debug_assert`) alongside dictionary hits.
412        let values = ["apple", "banana", "cherry"];
413        let rows: Vec<Row> = (0..3_000)
414            .map(|i| {
415                Row::pack_slice(&[
416                    Datum::String(values[i % values.len()]),
417                    Datum::Int64(i64::try_from(i).unwrap()),
418                    Datum::String(values[(i / 7) % values.len()]),
419                ])
420            })
421            .collect();
422
423        // Accumulate statistics in two independent observers, so the merge in
424        // `new_from([&stats1, &stats2])` is actually exercised.
425        let mut stats1 = ColumnsCodec::default();
426        let mut stats2 = ColumnsCodec::default();
427        let mut scratch = Vec::new();
428        for (i, row) in rows.iter().enumerate() {
429            scratch.clear();
430            let stats = if i % 2 == 0 { &mut stats1 } else { &mut stats2 };
431            stats.encode(ColumnsCodec::borrow_row(row), &mut scratch);
432        }
433
434        let merged = ColumnsCodec::new_from([&stats1, &stats2]);
435        let safe = stats1.new_safe();
436        for mut codec in [merged, safe] {
437            let mut compressed_any = false;
438            for row in &rows {
439                let mut buf = Vec::new();
440                codec.encode(ColumnsCodec::borrow_row(row), &mut buf);
441
442                let decoded = codec.decode(&buf).collect::<Vec<_>>();
443                let expected = ColumnsCodec::borrow_row(row).collect::<Vec<_>>();
444                assert_eq!(decoded, expected, "round-trip mismatch for {row:?}");
445
446                compressed_any |= buf.len() < row.data().len();
447            }
448            assert!(
449                compressed_any,
450                "dictionary never engaged; test no longer covers the compressed path",
451            );
452        }
453    }
454
455    /// Regression test for a dictionary-codec soundness bug in the safe-install
456    /// path (`new_safe`), reachable with the paged batcher enabled.
457    ///
458    /// A from-scratch container stores its pre-install rows *raw* while gathering
459    /// statistics, then installs a *safe* codec via `new_safe`. `new_safe` used to
460    /// discard the first-byte bitmap gathered over those raw rows. That bitmap is
461    /// soundness-critical: a later `new_from` merge consults it to decide which
462    /// one-byte tags are free to hand out as dictionary keys. With the bitmap
463    /// dropped, the merge could assign a dictionary tag equal to a raw datum's
464    /// first byte, after which `decode` resolves that literal datum to the
465    /// dictionary entry — returning the wrong value.
466    ///
467    /// We drive the lifecycle directly: observe short strings (first byte
468    /// `StringTiny`) into the pre-install statistics, install a safe codec, then
469    /// feed it many distinct *long* strings (first byte `StringShort`)
470    /// post-install so the merge has heavy hitters to compress. Merging via
471    /// `new_from` and re-encoding the short strings then exercises the raw
472    /// fall-through whose first byte the merge must not have claimed as a tag.
473    /// Before the fix the `StringTiny` tag was handed out and the round-trip
474    /// produced a long string (and tripped `encode`'s soundness `debug_assert`).
475    #[mz_ore::test]
476    #[cfg_attr(miri, ignore)] // integer-to-pointer casts in row decoding are unsupported under miri
477    fn test_safe_codec_merge_bitmap_carryover() {
478        use crate::row_codec::ColumnsCodec;
479
480        // Short strings: length < 256, so they encode with the `StringTiny` tag.
481        // Unique, so MisraGries never makes them dictionary entries; they always
482        // fall through raw, exposing their first byte.
483        let short_rows: Vec<Row> = (0..256)
484            .map(|i| Row::pack_slice(&[Datum::String(&format!("s{i}"))]))
485            .collect();
486        // Long strings: length >= 256, so they encode with the `StringShort` tag —
487        // a *different* first byte than the short strings. Distinct values, each
488        // repeated, so the post-install codec accrues many heavy hitters and the
489        // merge assigns dictionary tags across the low byte range, reaching the
490        // short strings' `StringTiny` tag unless the bitmap reserves it.
491        let long_values: Vec<String> = (0..64).map(|i| format!("{i:0>300}")).collect();
492
493        // Pre-install statistics observe only the short strings' first bytes.
494        let mut stats = ColumnsCodec::default();
495        let mut scratch = Vec::new();
496        for row in &short_rows {
497            scratch.clear();
498            stats.encode(ColumnsCodec::borrow_row(row), &mut scratch);
499        }
500
501        // Install a safe codec, then feed it the long strings post-install so it
502        // accrues heavy hitters (and observes only the `StringShort` first byte).
503        let mut safe = stats.new_safe();
504        for _ in 0..8 {
505            for v in &long_values {
506                let row = Row::pack_slice(&[Datum::String(v)]);
507                scratch.clear();
508                safe.encode(ColumnsCodec::borrow_row(&row), &mut scratch);
509            }
510        }
511
512        // Merge, then round-trip the short strings. With the bitmap carried over,
513        // no dictionary tag collides with the short strings' first byte; without
514        // it, one does.
515        let mut merged = ColumnsCodec::new_from([&safe]);
516        for row in &short_rows {
517            let mut buf = Vec::new();
518            merged.encode(ColumnsCodec::borrow_row(row), &mut buf);
519            let decoded = merged.decode(&buf).collect::<Vec<_>>();
520            let expected = ColumnsCodec::borrow_row(row).collect::<Vec<_>>();
521            assert_eq!(decoded, expected, "round-trip mismatch for {row:?}");
522        }
523    }
524
525    /// Confirms the structural assumption underpinning `SAFE_TAG_BASE`: every
526    /// datum the row format produces encodes with a first byte strictly less
527    /// than `SAFE_TAG_BASE`. If `mz_repr` ever introduces a tag that crosses
528    /// the boundary, `DictionaryCodec::new_safe` would assign a dictionary tag
529    /// that collides with a literal datum first-byte, breaking decoding.
530    #[mz_ore::test]
531    fn test_safe_tag_base() {
532        use crate::row_codec::SAFE_TAG_BASE;
533        let check = |datum: Datum| {
534            let row = Row::pack_slice(&[datum]);
535            let data = row.data();
536            assert!(!data.is_empty(), "empty encoding for {datum:?}");
537            assert!(
538                data[0] < SAFE_TAG_BASE,
539                "datum {datum:?} encodes with first byte {} >= SAFE_TAG_BASE ({}); \
540                 a new row tag has crossed the safe boundary",
541                data[0],
542                SAFE_TAG_BASE,
543            );
544        };
545        for ty in SqlScalarType::enumerate().iter() {
546            for datum in ty.interesting_datums() {
547                check(datum);
548            }
549        }
550    }
551
552    /// A batch built via the builder's `push`/`done` path (as the `reduce` operator
553    /// does) that stays under `STATS_THRESHOLD` never installs a codec at build time.
554    /// `done` now promotes the gathered statistics into the codec slot, so the batch
555    /// carries a codec + heavy-hitter summary and does not poison a later merge.
556    ///
557    /// This drives that container lifecycle directly: gather raw (well under the
558    /// threshold), promote at "done", then merge two such containers the way a spine
559    /// compaction does. With promotion the merge takes the `new_from` path and
560    /// compresses; without it both inputs are codec-less and the merge stays raw.
561    /// Every merged row must still round-trip.
562    #[mz_ore::test]
563    #[cfg_attr(miri, ignore)] // integer-to-pointer casts in row decoding are unsupported under miri
564    fn push_done_promotion_avoids_merge_poison() {
565        use std::sync::atomic::Ordering;
566        use timely::container::PushInto;
567
568        // Gate the dictionary path on. Safe for other tests: the flag only controls
569        // whether `DatumContainer` gathers stats; it never changes decode results.
570        crate::DICTIONARY_COMPRESSION.store(true, Ordering::Relaxed);
571
572        // Low-cardinality rows, well under `STATS_THRESHOLD` (64Ki): a repeated
573        // multi-byte string the dictionary compresses, plus an integer column that
574        // exercises raw fall-through.
575        let rows: Vec<Row> = (0..2_000i64)
576            .map(|i| {
577                Row::pack_slice(&[
578                    Datum::Int64(i % 8),
579                    Datum::String("a repeated string value"),
580                ])
581            })
582            .collect();
583
584        // Build a container the way the push/done path does: gather raw without ever
585        // crossing `STATS_THRESHOLD`, optionally promoting at "done".
586        let build = |promote: bool| {
587            let mut c = DatumContainer::with_capacity(rows.len());
588            for row in &rows {
589                c.push_into(row);
590            }
591            if promote {
592                c.promote_stats_to_codec();
593            }
594            c
595        };
596
597        // Merge two containers as a spine compaction does: allocate via
598        // `merge_capacity`, then copy every row through.
599        let merge = |a: &DatumContainer, b: &DatumContainer| {
600            let mut m = DatumContainer::merge_capacity(a, b);
601            for i in 0..a.len() {
602                m.push_into(a.index(i));
603            }
604            for i in 0..b.len() {
605                m.push_into(b.index(i));
606            }
607            m
608        };
609
610        let heap = |c: &DatumContainer| {
611            let mut size = 0;
612            c.heap_size(|_, cap| size += cap);
613            size
614        };
615
616        // Codec-less inputs (no promotion): the merge cannot `new_from` and stays raw.
617        let poisoned = merge(&build(false), &build(false));
618        // Promoted inputs carry a codec + summary: the merge `new_from`s and compresses.
619        let compressed = merge(&build(true), &build(true));
620
621        // Round-trip: every merged row decodes back to the corresponding input row
622        // (the merge here concatenates a's rows then b's rows, no consolidation).
623        assert_eq!(compressed.len(), rows.len() * 2);
624        for i in 0..compressed.len() {
625            let got = compressed.index(i).collect::<Vec<_>>();
626            let want = rows[i % rows.len()].iter().collect::<Vec<_>>();
627            assert_eq!(got, want, "merged row {i} round-trips");
628        }
629
630        // The promoted merge must actually compress relative to the poisoned one,
631        // confirming promotion carried a usable summary into `new_from`.
632        assert!(
633            heap(&compressed) < heap(&poisoned),
634            "promotion should let the merge compress: compressed={} poisoned={}",
635            heap(&compressed),
636            heap(&poisoned),
637        );
638    }
639}
640
641/// A `[u8]`-specialized container.
642mod bytes_container {
643
644    use differential_dataflow::trace::implementations::BatchContainer;
645    use timely::container::PushInto;
646
647    use mz_ore::region::Region;
648
649    /// A slice container with four bytes overhead per slice.
650    pub struct BytesContainer {
651        /// Total length of `batches`, maintained because recomputation is expensive.
652        length: usize,
653        batches: Vec<BytesBatch>,
654    }
655
656    impl BytesContainer {
657        /// Visit contained allocations to determine their size and capacity.
658        #[inline]
659        pub fn heap_size(&self, mut callback: impl FnMut(usize, usize)) {
660            // Calculate heap size for local, stash, and stash entries
661            callback(
662                self.batches.len() * std::mem::size_of::<BytesBatch>(),
663                self.batches.capacity() * std::mem::size_of::<BytesBatch>(),
664            );
665            for batch in self.batches.iter() {
666                batch.offsets.heap_size(&mut callback);
667                callback(batch.storage.len(), batch.storage.capacity());
668            }
669        }
670    }
671
672    impl BatchContainer for BytesContainer {
673        type Owned = Vec<u8>;
674        type ReadItem<'a> = &'a [u8];
675
676        #[inline]
677        fn into_owned<'a>(item: Self::ReadItem<'a>) -> Self::Owned {
678            item.to_vec()
679        }
680
681        #[inline]
682        fn clone_onto<'a>(item: Self::ReadItem<'a>, other: &mut Self::Owned) {
683            other.clear();
684            other.extend_from_slice(item);
685        }
686
687        #[inline(always)]
688        fn push_ref(&mut self, item: Self::ReadItem<'_>) {
689            self.push_into(item);
690        }
691
692        #[inline(always)]
693        fn push_own(&mut self, item: &Self::Owned) {
694            self.push_into(item.as_slice())
695        }
696
697        fn clear(&mut self) {
698            self.batches.clear();
699            self.batches.push(BytesBatch::with_capacities(0, 0));
700            self.length = 0;
701        }
702
703        fn with_capacity(size: usize) -> Self {
704            Self {
705                length: 0,
706                batches: vec![BytesBatch::with_capacities(size, size)],
707            }
708        }
709
710        fn merge_capacity(cont1: &Self, cont2: &Self) -> Self {
711            let mut item_cap = 1;
712            let mut byte_cap = 0;
713            for batch in cont1.batches.iter() {
714                item_cap += batch.offsets.len() - 1;
715                byte_cap += batch.storage.len();
716            }
717            for batch in cont2.batches.iter() {
718                item_cap += batch.offsets.len() - 1;
719                byte_cap += batch.storage.len();
720            }
721            Self {
722                length: 0,
723                batches: vec![BytesBatch::with_capacities(item_cap, byte_cap)],
724            }
725        }
726
727        #[inline(always)]
728        fn reborrow<'b, 'a: 'b>(item: Self::ReadItem<'a>) -> Self::ReadItem<'b> {
729            item
730        }
731
732        #[inline]
733        fn index(&self, mut index: usize) -> Self::ReadItem<'_> {
734            for batch in self.batches.iter() {
735                if index < batch.len() {
736                    return batch.index(index);
737                }
738                index -= batch.len();
739            }
740            panic!("Index out of bounds");
741        }
742
743        #[inline(always)]
744        fn len(&self) -> usize {
745            self.length
746        }
747    }
748
749    impl PushInto<&[u8]> for BytesContainer {
750        #[inline]
751        fn push_into(&mut self, item: &[u8]) {
752            self.length += 1;
753            if let Some(batch) = self.batches.last_mut() {
754                let success = batch.try_push(item);
755                if !success {
756                    // double the lengths from `batch`.
757                    let item_cap = 2 * batch.offsets.len();
758                    let byte_cap = std::cmp::max(2 * batch.storage.capacity(), item.len());
759                    let mut new_batch = BytesBatch::with_capacities(item_cap, byte_cap);
760                    assert!(new_batch.try_push(item));
761                    self.batches.push(new_batch);
762                }
763            }
764        }
765    }
766
767    /// A batch of slice storage.
768    ///
769    /// The backing storage for this batch will not be resized.
770    pub struct BytesBatch {
771        offsets: crate::OffsetOptimized,
772        storage: Region<u8>,
773        len: usize,
774    }
775
776    impl BytesBatch {
777        /// Either accepts the slice and returns true,
778        /// or does not and returns false.
779        fn try_push(&mut self, slice: &[u8]) -> bool {
780            if self.storage.len() + slice.len() <= self.storage.capacity() {
781                self.storage.extend_from_slice(slice);
782                self.offsets.push_into(self.storage.len());
783                self.len += 1;
784                true
785            } else {
786                false
787            }
788        }
789        #[inline]
790        fn index(&self, index: usize) -> &[u8] {
791            let lower = self.offsets.index(index);
792            let upper = self.offsets.index(index + 1);
793            &self.storage[lower..upper]
794        }
795        #[inline(always)]
796        fn len(&self) -> usize {
797            mz_ore::soft_assert_eq_no_log!(self.len, self.offsets.len() - 1);
798            self.len
799        }
800
801        fn with_capacities(item_cap: usize, byte_cap: usize) -> Self {
802            // TODO: be wary of `byte_cap` greater than 2^32.
803            let mut offsets = crate::OffsetOptimized::with_capacity(item_cap + 1);
804            offsets.push_into(0);
805            Self {
806                offsets,
807                storage: Region::new_auto(byte_cap.next_power_of_two()),
808                len: 0,
809            }
810        }
811    }
812}
813
814mod offset_opt {
815    use differential_dataflow::trace::implementations::BatchContainer;
816    use differential_dataflow::trace::implementations::OffsetList;
817    use timely::container::PushInto;
818
819    enum OffsetStride {
820        Empty,
821        Zero,
822        Striding(usize, usize),
823        Saturated(usize, usize, usize),
824    }
825
826    impl OffsetStride {
827        /// Accepts or rejects a newly pushed element.
828        #[inline]
829        fn push(&mut self, item: usize) -> bool {
830            match self {
831                OffsetStride::Empty => {
832                    if item == 0 {
833                        *self = OffsetStride::Zero;
834                        true
835                    } else {
836                        false
837                    }
838                }
839                OffsetStride::Zero => {
840                    *self = OffsetStride::Striding(item, 2);
841                    true
842                }
843                OffsetStride::Striding(stride, count) => {
844                    if item == *stride * *count {
845                        *count += 1;
846                        true
847                    } else if item == *stride * (*count - 1) {
848                        *self = OffsetStride::Saturated(*stride, *count, 1);
849                        true
850                    } else {
851                        false
852                    }
853                }
854                OffsetStride::Saturated(stride, count, reps) => {
855                    if item == *stride * (*count - 1) {
856                        *reps += 1;
857                        true
858                    } else {
859                        false
860                    }
861                }
862            }
863        }
864
865        #[inline]
866        fn index(&self, index: usize) -> usize {
867            match self {
868                OffsetStride::Empty => {
869                    panic!("Empty OffsetStride")
870                }
871                OffsetStride::Zero => 0,
872                OffsetStride::Striding(stride, _steps) => *stride * index,
873                OffsetStride::Saturated(stride, steps, _reps) => {
874                    if index < *steps {
875                        *stride * index
876                    } else {
877                        *stride * (*steps - 1)
878                    }
879                }
880            }
881        }
882
883        #[inline]
884        fn len(&self) -> usize {
885            match self {
886                OffsetStride::Empty => 0,
887                OffsetStride::Zero => 1,
888                OffsetStride::Striding(_stride, steps) => *steps,
889                OffsetStride::Saturated(_stride, steps, reps) => *steps + *reps,
890            }
891        }
892    }
893
894    pub struct OffsetOptimized {
895        strided: OffsetStride,
896        spilled: OffsetList,
897    }
898
899    impl BatchContainer for OffsetOptimized {
900        type Owned = usize;
901        type ReadItem<'a> = usize;
902
903        #[inline]
904        fn into_owned<'a>(item: Self::ReadItem<'a>) -> Self::Owned {
905            item
906        }
907
908        #[inline]
909        fn push_ref(&mut self, item: Self::ReadItem<'_>) {
910            self.push_into(item)
911        }
912
913        #[inline]
914        fn push_own(&mut self, item: &Self::Owned) {
915            self.push_into(*item)
916        }
917
918        fn clear(&mut self) {
919            self.strided = OffsetStride::Empty;
920            self.spilled.clear();
921        }
922
923        fn with_capacity(_size: usize) -> Self {
924            Self {
925                strided: OffsetStride::Empty,
926                spilled: OffsetList::with_capacity(0),
927            }
928        }
929
930        fn merge_capacity(_cont1: &Self, _cont2: &Self) -> Self {
931            Self {
932                strided: OffsetStride::Empty,
933                spilled: OffsetList::with_capacity(0),
934            }
935        }
936
937        #[inline]
938        fn reborrow<'b, 'a: 'b>(item: Self::ReadItem<'a>) -> Self::ReadItem<'b> {
939            item
940        }
941
942        #[inline]
943        fn index(&self, index: usize) -> Self::ReadItem<'_> {
944            if index < self.strided.len() {
945                self.strided.index(index)
946            } else {
947                self.spilled.index(index - self.strided.len())
948            }
949        }
950
951        #[inline]
952        fn len(&self) -> usize {
953            self.strided.len() + self.spilled.len()
954        }
955    }
956
957    impl PushInto<usize> for OffsetOptimized {
958        #[inline]
959        fn push_into(&mut self, item: usize) {
960            if !self.spilled.is_empty() {
961                self.spilled.push(item);
962            } else {
963                let inserted = self.strided.push(item);
964                if !inserted {
965                    self.spilled.push(item);
966                }
967            }
968        }
969    }
970
971    impl OffsetOptimized {
972        pub fn heap_size(&self, callback: impl FnMut(usize, usize)) {
973            crate::offset_list_size(&self.spilled, callback);
974        }
975    }
976}
977
978/// Helper to compute the size of an [`OffsetList`] in memory.
979#[inline]
980pub(crate) fn offset_list_size(data: &OffsetList, mut callback: impl FnMut(usize, usize)) {
981    // Private `vec_size` because we should only use it where data isn't region-allocated.
982    // `T: Copy` makes sure the implementation is correct even if types change!
983    #[inline(always)]
984    fn vec_size<T: Copy>(data: &Vec<T>, mut callback: impl FnMut(usize, usize)) {
985        let size_of_t = std::mem::size_of::<T>();
986        callback(data.len() * size_of_t, data.capacity() * size_of_t);
987    }
988
989    vec_size(&data.smol, &mut callback);
990    vec_size(&data.chonk, callback);
991}
992
993/// A `Row`-specialized container using dictionary compression.
994///
995/// The approach is to establish for each column lists of common values, and to use "unoccupied"
996/// tags in the row encoding (e.g. where we would indicate types) to replace these common values.
997/// This substitution is opt-in, in that we don't need to do it, and in particular do not do it
998/// while we are collecting preliminary information about common values, and then start to use it
999/// once we believe we have enough information. Once we have started to use the substitutions we
1000/// cannot change the meaning of a reserved byte pattern, for the container we are populating.
1001///
1002/// Each from-scratch container observes `STATS_THRESHOLD` records before establishing a mapping
1003/// from spare tags to common values. Containers that are formed from merging other containers
1004/// use those input containers' common values to populate a codec and use it immediately.
1005///
1006/// The dictionary behavior is controlled by the `DICTIONARY_COMPRESSION` flag, which if disabled
1007/// prevents the construction of codecs, which when absent simply cause the wrapper to behave as
1008/// a no-op that fails to use any spare tags for common values. The flag is set once, when a
1009/// replica is created (from compute's `InstanceConfig::arrangement_dictionary_compression`, itself
1010/// captured from the `enable_arrangement_dictionary_compression_alpha` dyncfg at that moment), and is
1011/// not changed for the life of the process; flipping the dyncfg only affects replicas created
1012/// afterwards. Even with the flag fixed, a single replica can hold a mix of compressed and
1013/// uncompressed containers — e.g. containers that never observed enough records to install a
1014/// codec, or that were merged from uncompressed inputs.
1015mod dictionary {
1016
1017    use differential_dataflow::trace::implementations::BatchContainer;
1018
1019    use mz_repr::{Row, RowRef};
1020
1021    use super::row_codec::{ColumnsCodec, ColumnsIter};
1022
1023    /// Wrapper types that exist to support the creation of dictionary codecs.
1024    ///
1025    /// These types interpose at the seal() call, to traverse the data that is being sealed and
1026    /// then construct codecs that are used to encode the row-shaped keys and values. There are
1027    /// several variants, corresponding to the RowRow, RowVal, and Row-only spine types.
1028    pub mod builders {
1029        use columnar::{Columnar, Index};
1030        use columnation::Columnation;
1031        use differential_dataflow::difference::Semigroup;
1032        use differential_dataflow::lattice::Lattice;
1033        use differential_dataflow::trace::Builder;
1034        use differential_dataflow::trace::implementations::BatchContainer;
1035        use differential_dataflow::trace::implementations::merge_batcher::Sealer;
1036        use differential_dataflow::trace::implementations::ord_neu::{
1037            BuilderInput, OrdValBatch, OrdValBuilder,
1038        };
1039        use differential_dataflow::trace::implementations::ord_neu::{OrdKeyBatch, OrdKeyBuilder};
1040        use mz_timely_util::columnar::Column;
1041        use mz_timely_util::columnar::body::ColumnBody;
1042        use mz_timely_util::columnar::chunk::ChainState;
1043        use mz_timely_util::columnation::ColumnationStack as TimelyStack;
1044        use timely::progress::Timestamp;
1045
1046        use mz_repr::{Row, RowRef};
1047
1048        use super::super::row_codec::ColumnsCodec;
1049        use super::{DatumContainer, DatumSeq};
1050        use crate::DICTIONARY_COMPRESSION;
1051        use crate::spines::{RowLayout, RowRowLayout, RowValLayout, ValRowLayout};
1052
1053        /// Gather encoding statistics across `rows` and produce a codec from them.
1054        ///
1055        /// Accepts anything that borrows as a [`RowRef`], so it serves both the
1056        /// columnation-fed builders (which yield `&Row`) and the paged builders
1057        /// (which yield `&RowRef` straight out of a [`Column`] chunk).
1058        ///
1059        /// Returns `None` when dictionary compression is disabled.
1060        fn build_codec<'a, B>(rows: impl IntoIterator<Item = &'a B>) -> Option<ColumnsCodec>
1061        where
1062            B: std::borrow::Borrow<RowRef> + ?Sized + 'a,
1063        {
1064            if !DICTIONARY_COMPRESSION.load(std::sync::atomic::Ordering::Relaxed) {
1065                return None;
1066            }
1067            let mut stats = ColumnsCodec::default();
1068            observe_rows(&mut stats, rows);
1069            Some(ColumnsCodec::new_from([&stats]))
1070        }
1071
1072        /// Fold `rows` into `stats`, the accumulator a codec is built from.
1073        fn observe_rows<'a, B>(stats: &mut ColumnsCodec, rows: impl IntoIterator<Item = &'a B>)
1074        where
1075            B: std::borrow::Borrow<RowRef> + ?Sized + 'a,
1076        {
1077            for row in rows {
1078                let row = row.borrow();
1079                if !row.is_empty() {
1080                    // Gather stats only; the encoded output would be thrown away here, so
1081                    // `observe` skips the per-value lookup and the throwaway-buffer memcpy
1082                    // that `encode` would do (see `ColumnsCodec::observe`).
1083                    stats.observe(DatumSeq::borrow_as(row).bytes_iter());
1084                }
1085            }
1086        }
1087
1088        pub struct RowRowBuilder<T, R>
1089        where
1090            T: Lattice + Timestamp + Columnation,
1091            R: Ord + Semigroup + Columnation + 'static,
1092        {
1093            inner: OrdValBuilder<RowRowLayout<((Row, Row), T, R)>, TimelyStack<((Row, Row), T, R)>>,
1094        }
1095
1096        impl<T, R> Builder for RowRowBuilder<T, R>
1097        where
1098            T: Lattice + Timestamp + Columnation,
1099            R: Ord + Semigroup + Columnation + 'static,
1100        {
1101            type Input = TimelyStack<((Row, Row), T, R)>;
1102            type Time = T;
1103            type Output = OrdValBatch<RowRowLayout<((Row, Row), T, R)>>;
1104
1105            fn push(&mut self, chunk: &mut Self::Input) {
1106                self.inner.push(chunk)
1107            }
1108            fn done(self) -> Option<Self::Output> {
1109                // The push/done build path (e.g. the `reduce` operator, which builds
1110                // batches with `Builder::new()` + `push` + `done` rather than `seal`)
1111                // never runs `seal`'s codec install. Install a codec here from the
1112                // statistics gathered during `push`, mirroring `seal` — but without
1113                // building a dictionary or re-encoding the rows; see
1114                // `DatumContainer::promote_stats_to_codec` for why a codec-less batch
1115                // must be avoided even though its rows stay raw.
1116                let mut inner = self.inner;
1117                inner.result.keys.promote_stats_to_codec();
1118                inner.result.vals.vals.promote_stats_to_codec();
1119                inner.done()
1120            }
1121        }
1122
1123        impl<T, R> Default for RowRowBuilder<T, R>
1124        where
1125            T: Lattice + Timestamp + Columnation,
1126            R: Ord + Semigroup + Columnation + 'static,
1127        {
1128            fn default() -> Self {
1129                Self {
1130                    inner: Default::default(),
1131                }
1132            }
1133        }
1134
1135        impl<T, R> Sealer<TimelyStack<((Row, Row), T, R)>> for RowRowBuilder<T, R>
1136        where
1137            T: Lattice + Timestamp + Columnation,
1138            R: Ord + Semigroup + Columnation + 'static,
1139        {
1140            type Output = OrdValBatch<RowRowLayout<((Row, Row), T, R)>>;
1141
1142            fn seal(chain: &mut Vec<TimelyStack<((Row, Row), T, R)>>) -> Option<Self::Output> {
1143                let key_codec = build_codec(
1144                    chain
1145                        .iter()
1146                        .flat_map(|link| link.iter().map(|((k, _), _, _)| k)),
1147                );
1148                let val_codec = build_codec(
1149                    chain
1150                        .iter()
1151                        .flat_map(|link| link.iter().map(|((_, v), _, _)| v)),
1152                );
1153
1154                let (keys, vals, upds) =
1155                    <_ as BuilderInput<DatumContainer, DatumContainer>>::key_val_upd_counts(
1156                        &chain[..],
1157                    );
1158                let mut builder = Self {
1159                    inner: OrdValBuilder::with_capacity(keys, vals, upds),
1160                };
1161                // The seal path installs a codec directly, so the per-container stats
1162                // gatherer (which `with_capacity` may have allocated) is dead weight and
1163                // would contradict the `stats: None once codec installed` invariant.
1164                builder.inner.result.keys.codec = key_codec;
1165                builder.inner.result.keys.stats = None;
1166                builder.inner.result.vals.vals.codec = val_codec;
1167                builder.inner.result.vals.vals.stats = None;
1168
1169                for mut chunk in chain.drain(..) {
1170                    builder.push(&mut chunk);
1171                }
1172
1173                builder.done()
1174            }
1175        }
1176
1177        pub struct RowValBuilder<V, T, R>
1178        where
1179            V: Ord + Clone + Columnation + 'static,
1180            T: Lattice + Timestamp + Columnation,
1181            R: Ord + Semigroup + Columnation + 'static,
1182        {
1183            inner: OrdValBuilder<RowValLayout<((Row, V), T, R)>, TimelyStack<((Row, V), T, R)>>,
1184        }
1185
1186        impl<V, T, R> Builder for RowValBuilder<V, T, R>
1187        where
1188            V: Ord + Clone + Columnation,
1189            T: Lattice + Timestamp + Columnation,
1190            R: Ord + Semigroup + Columnation + 'static,
1191        {
1192            type Input = TimelyStack<((Row, V), T, R)>;
1193            type Time = T;
1194            type Output = OrdValBatch<RowValLayout<((Row, V), T, R)>>;
1195
1196            fn push(&mut self, chunk: &mut Self::Input) {
1197                self.inner.push(chunk)
1198            }
1199            fn done(self) -> Option<Self::Output> {
1200                // See `RowRowBuilder::done`: install a codec on the `Row`-shaped key
1201                // container for the push/done (e.g. `reduce`) path that skips `seal`.
1202                let mut inner = self.inner;
1203                inner.result.keys.promote_stats_to_codec();
1204                inner.done()
1205            }
1206        }
1207
1208        impl<V, T, R> Default for RowValBuilder<V, T, R>
1209        where
1210            V: Ord + Clone + Columnation,
1211            T: Lattice + Timestamp + Columnation,
1212            R: Ord + Semigroup + Columnation + 'static,
1213        {
1214            fn default() -> Self {
1215                Self {
1216                    inner: Default::default(),
1217                }
1218            }
1219        }
1220
1221        impl<V, T, R> Sealer<TimelyStack<((Row, V), T, R)>> for RowValBuilder<V, T, R>
1222        where
1223            V: Ord + Clone + Columnation,
1224            T: Lattice + Timestamp + Columnation,
1225            R: Ord + Semigroup + Columnation + 'static,
1226        {
1227            type Output = OrdValBatch<RowValLayout<((Row, V), T, R)>>;
1228
1229            fn seal(chain: &mut Vec<TimelyStack<((Row, V), T, R)>>) -> Option<Self::Output> {
1230                let key_codec = build_codec(
1231                    chain
1232                        .iter()
1233                        .flat_map(|link| link.iter().map(|((k, _), _, _)| k)),
1234                );
1235
1236                let (keys, vals, upds) =
1237                    <_ as BuilderInput<DatumContainer, TimelyStack<V>>>::key_val_upd_counts(
1238                        &chain[..],
1239                    );
1240                let mut builder = Self {
1241                    inner: OrdValBuilder::with_capacity(keys, vals, upds),
1242                };
1243                // See `RowRowBuilder::seal`: drop the now-redundant stats gatherer.
1244                builder.inner.result.keys.codec = key_codec;
1245                builder.inner.result.keys.stats = None;
1246
1247                for mut chunk in chain.drain(..) {
1248                    builder.push(&mut chunk);
1249                }
1250
1251                builder.done()
1252            }
1253        }
1254
1255        pub struct RowBuilder<T, R, DC = TimelyStack<R>>
1256        where
1257            T: Lattice + Timestamp + Columnation,
1258            R: Ord + Semigroup + Columnation + 'static,
1259            DC: BatchContainer<Owned = R>,
1260        {
1261            inner: OrdKeyBuilder<RowLayout<((Row, ()), T, R), DC>, TimelyStack<((Row, ()), T, R)>>,
1262        }
1263
1264        impl<T, R, DC> Builder for RowBuilder<T, R, DC>
1265        where
1266            T: Lattice + Timestamp + Columnation,
1267            R: Ord + Semigroup + Columnation + 'static,
1268            DC: BatchContainer<Owned = R>,
1269        {
1270            type Input = TimelyStack<((Row, ()), T, R)>;
1271            type Time = T;
1272            type Output = OrdKeyBatch<RowLayout<((Row, ()), T, R), DC>>;
1273
1274            fn push(&mut self, chunk: &mut Self::Input) {
1275                self.inner.push(chunk)
1276            }
1277            fn done(self) -> Option<Self::Output> {
1278                // See `RowRowBuilder::done`: install a codec on the `Row`-shaped key
1279                // container for the push/done (e.g. `reduce`) path that skips `seal`.
1280                let mut inner = self.inner;
1281                inner.result.keys.promote_stats_to_codec();
1282                inner.done()
1283            }
1284        }
1285
1286        impl<T, R, DC> Default for RowBuilder<T, R, DC>
1287        where
1288            T: Lattice + Timestamp + Columnation,
1289            R: Ord + Semigroup + Columnation + 'static,
1290            DC: BatchContainer<Owned = R>,
1291        {
1292            fn default() -> Self {
1293                Self {
1294                    inner: Default::default(),
1295                }
1296            }
1297        }
1298
1299        impl<T, R, DC> Sealer<TimelyStack<((Row, ()), T, R)>> for RowBuilder<T, R, DC>
1300        where
1301            T: Lattice + Timestamp + Columnation,
1302            R: Ord + Semigroup + Columnation + 'static,
1303            DC: BatchContainer<Owned = R>,
1304        {
1305            type Output = OrdKeyBatch<RowLayout<((Row, ()), T, R), DC>>;
1306
1307            fn seal(chain: &mut Vec<TimelyStack<((Row, ()), T, R)>>) -> Option<Self::Output> {
1308                let key_codec = build_codec(
1309                    chain
1310                        .iter()
1311                        .flat_map(|link| link.iter().map(|((k, _), _, _)| k)),
1312                );
1313
1314                let (keys, vals, upds) =
1315                    <_ as BuilderInput<DatumContainer, TimelyStack<()>>>::key_val_upd_counts(
1316                        &chain[..],
1317                    );
1318                let mut builder = Self {
1319                    inner: OrdKeyBuilder::with_capacity(keys, vals, upds),
1320                };
1321                // See `RowRowBuilder::seal`: drop the now-redundant stats gatherer.
1322                builder.inner.result.keys.codec = key_codec;
1323                builder.inner.result.keys.stats = None;
1324
1325                for mut chunk in chain.drain(..) {
1326                    builder.push(&mut chunk);
1327                }
1328
1329                builder.done()
1330            }
1331        }
1332
1333        /// Mirror of [`RowValBuilder`] with the roles swapped: arbitrary keys and
1334        /// `Row` *values*, so the dictionary codec is built for and installed on the
1335        /// value container.
1336        pub struct ValRowBuilder<K, T, R>
1337        where
1338            K: Ord + Clone + Columnation + 'static,
1339            T: Lattice + Timestamp + Columnation,
1340            R: Ord + Semigroup + Columnation + 'static,
1341        {
1342            inner: OrdValBuilder<ValRowLayout<((K, Row), T, R)>, TimelyStack<((K, Row), T, R)>>,
1343        }
1344
1345        impl<K, T, R> Builder for ValRowBuilder<K, T, R>
1346        where
1347            K: Ord + Clone + Columnation,
1348            T: Lattice + Timestamp + Columnation,
1349            R: Ord + Semigroup + Columnation + 'static,
1350        {
1351            type Input = TimelyStack<((K, Row), T, R)>;
1352            type Time = T;
1353            type Output = OrdValBatch<ValRowLayout<((K, Row), T, R)>>;
1354
1355            fn push(&mut self, chunk: &mut Self::Input) {
1356                self.inner.push(chunk)
1357            }
1358            fn done(self) -> Option<Self::Output> {
1359                // See `RowRowBuilder::done`: install a codec on the `Row`-shaped value
1360                // container for the push/done (e.g. `reduce`) path that skips `seal`.
1361                let mut inner = self.inner;
1362                inner.result.vals.vals.promote_stats_to_codec();
1363                inner.done()
1364            }
1365        }
1366
1367        impl<K, T, R> Default for ValRowBuilder<K, T, R>
1368        where
1369            K: Ord + Clone + Columnation,
1370            T: Lattice + Timestamp + Columnation,
1371            R: Ord + Semigroup + Columnation + 'static,
1372        {
1373            fn default() -> Self {
1374                Self {
1375                    inner: Default::default(),
1376                }
1377            }
1378        }
1379
1380        impl<K, T, R> Sealer<TimelyStack<((K, Row), T, R)>> for ValRowBuilder<K, T, R>
1381        where
1382            K: Ord + Clone + Columnation,
1383            T: Lattice + Timestamp + Columnation,
1384            R: Ord + Semigroup + Columnation + 'static,
1385        {
1386            type Output = OrdValBatch<ValRowLayout<((K, Row), T, R)>>;
1387
1388            fn seal(chain: &mut Vec<TimelyStack<((K, Row), T, R)>>) -> Option<Self::Output> {
1389                let val_codec = build_codec(
1390                    chain
1391                        .iter()
1392                        .flat_map(|link| link.iter().map(|((_, v), _, _)| v)),
1393                );
1394
1395                let (keys, vals, upds) =
1396                    <_ as BuilderInput<TimelyStack<K>, DatumContainer>>::key_val_upd_counts(
1397                        &chain[..],
1398                    );
1399                let mut builder = Self {
1400                    inner: OrdValBuilder::with_capacity(keys, vals, upds),
1401                };
1402                // See `RowRowBuilder::seal`: drop the now-redundant stats gatherer.
1403                builder.inner.result.vals.vals.codec = val_codec;
1404                builder.inner.result.vals.vals.stats = None;
1405
1406                for mut chunk in chain.drain(..) {
1407                    builder.push(&mut chunk);
1408                }
1409
1410                builder.done()
1411            }
1412        }
1413
1414        /// Counterpart of [`RowRowBuilder`] that consumes [`ColumnBody`] chunks
1415        /// instead of columnation stacks, whether or not the batcher that
1416        /// produced them spills. Mirrors `RowRowBuilder::seal`:
1417        /// it gathers key and value statistics from the sealed chain and
1418        /// installs codecs directly, then drops the per-container stats gatherer.
1419        pub struct RowRowColPagedBuilder<T, R>
1420        where
1421            T: Lattice + Timestamp + Columnation + Columnar,
1422            R: Ord + Semigroup + Columnation + Columnar + Clone + 'static,
1423        {
1424            inner: OrdValBuilder<RowRowLayout<((Row, Row), T, R)>, ColumnBody<((Row, Row), T, R)>>,
1425        }
1426
1427        impl<T, R> Builder for RowRowColPagedBuilder<T, R>
1428        where
1429            T: Lattice + Timestamp + Columnation + Columnar,
1430            R: Ord + Semigroup + Columnation + Columnar + Clone + 'static,
1431        {
1432            type Input = ColumnBody<((Row, Row), T, R)>;
1433            type Time = T;
1434            type Output = OrdValBatch<RowRowLayout<((Row, Row), T, R)>>;
1435
1436            fn push(&mut self, chunk: &mut Self::Input) {
1437                self.inner.push(chunk)
1438            }
1439            fn done(self) -> Option<Self::Output> {
1440                // See `RowRowBuilder::done`: a builder driven by `push` and `done`
1441                // alone gets its codec from the statistics gathered during the
1442                // pushes, since a codec-less batch poisons the merges it joins.
1443                // Installing a codec here is a no-op once one is installed.
1444                let mut inner = self.inner;
1445                inner.result.keys.promote_stats_to_codec();
1446                inner.result.vals.vals.promote_stats_to_codec();
1447                inner.done()
1448            }
1449        }
1450
1451        impl<T, R> Default for RowRowColPagedBuilder<T, R>
1452        where
1453            T: Lattice + Timestamp + Columnation + Columnar,
1454            R: Ord + Semigroup + Columnation + Columnar + Clone + 'static,
1455        {
1456            fn default() -> Self {
1457                Self {
1458                    inner: Default::default(),
1459                }
1460            }
1461        }
1462
1463        impl<T, R> Sealer<ColumnBody<((Row, Row), T, R)>> for RowRowColPagedBuilder<T, R>
1464        where
1465            T: Lattice + Timestamp + Columnation + Columnar,
1466            R: Ord + Semigroup + Columnation + Columnar + Clone + 'static,
1467        {
1468            type Output = OrdValBatch<RowRowLayout<((Row, Row), T, R)>>;
1469
1470            fn seal(chain: &mut Vec<ColumnBody<((Row, Row), T, R)>>) -> Option<Self::Output> {
1471                // `into_index_iter` yields the value column's `Row`s as `&RowRef`,
1472                // which `build_codec` consumes directly.
1473                let key_codec = build_codec(
1474                    chain
1475                        .iter()
1476                        .flat_map(|c| c.borrow().into_index_iter().map(|((k, _), _, _)| k)),
1477                );
1478                let val_codec = build_codec(
1479                    chain
1480                        .iter()
1481                        .flat_map(|c| c.borrow().into_index_iter().map(|((_, v), _, _)| v)),
1482                );
1483
1484                let (keys, vals, upds) =
1485                    <_ as BuilderInput<DatumContainer, DatumContainer>>::key_val_upd_counts(
1486                        &chain[..],
1487                    );
1488                let mut builder = Self {
1489                    inner: OrdValBuilder::with_capacity(keys, vals, upds),
1490                };
1491                // See `RowRowBuilder::seal`: install the codecs and drop the
1492                // now-redundant per-container stats gatherer.
1493                builder.inner.result.keys.codec = key_codec;
1494                builder.inner.result.keys.stats = None;
1495                builder.inner.result.vals.vals.codec = val_codec;
1496                builder.inner.result.vals.vals.stats = None;
1497
1498                for mut chunk in chain.drain(..) {
1499                    builder.push(&mut chunk);
1500                }
1501
1502                builder.done()
1503            }
1504        }
1505
1506        /// The chain-wide state of [`RowRowColPagedBuilder`]: the statistics its
1507        /// key and value codecs are built from, and the container sizes the
1508        /// chain implies.
1509        ///
1510        /// [`RowRowColPagedBuilder`]: crate::RowRowColPagedBuilder
1511        #[derive(Default)]
1512        pub struct RowRowColPagedState {
1513            keys: ColumnsCodec,
1514            vals: ColumnsCodec,
1515            /// Whether `keys` and `vals` hold statistics from any body.
1516            gathered: bool,
1517            counts: (usize, usize, usize),
1518        }
1519
1520        impl<T, R> ChainState for RowRowColPagedBuilder<T, R>
1521        where
1522            T: Lattice + Timestamp + Columnation + Columnar,
1523            R: Ord + Semigroup + Columnation + Columnar + Clone + 'static,
1524        {
1525            type State = RowRowColPagedState;
1526
1527            fn wants_bodies() -> bool {
1528                DICTIONARY_COMPRESSION.load(std::sync::atomic::Ordering::Relaxed)
1529            }
1530
1531            fn observe(state: &mut Self::State, input: &Self::Input) {
1532                if Self::wants_bodies() {
1533                    // `into_index_iter` yields the key and value columns' `Row`s
1534                    // as `&RowRef`, which `observe_rows` consumes directly.
1535                    observe_rows(
1536                        &mut state.keys,
1537                        input.borrow().into_index_iter().map(|((k, _), _, _)| k),
1538                    );
1539                    observe_rows(
1540                        &mut state.vals,
1541                        input.borrow().into_index_iter().map(|((_, v), _, _)| v),
1542                    );
1543                    state.gathered = true;
1544                }
1545
1546                let (keys, vals, upds) =
1547                    <_ as BuilderInput<DatumContainer, DatumContainer>>::key_val_upd_counts(
1548                        std::slice::from_ref(input),
1549                    );
1550                state.counts.0 += keys;
1551                state.counts.1 += vals;
1552                state.counts.2 += upds;
1553            }
1554
1555            fn observe_records(state: &mut Self::State, records: usize) {
1556                // Key and value counts need the bodies. Sizing the update
1557                // containers alone is what a record count buys.
1558                state.counts.2 += records;
1559            }
1560
1561            fn from_state(state: Self::State) -> Self {
1562                let mut builder = Self {
1563                    inner: OrdValBuilder::with_capacity(
1564                        state.counts.0,
1565                        state.counts.1,
1566                        state.counts.2,
1567                    ),
1568                };
1569                // Install from the statistics that were gathered, not from the
1570                // compression flag: the flag is process-global and a worker can
1571                // flip it mid-seal, and a codec built from no observations would
1572                // carry an empty summary and displace the gatherer that `done`
1573                // would otherwise promote.
1574                if state.gathered {
1575                    // See `RowRowBuilder::seal`: install the codecs the chain's
1576                    // statistics imply and drop the now-redundant per-container
1577                    // stats gatherer.
1578                    builder.inner.result.keys.codec = Some(ColumnsCodec::new_from([&state.keys]));
1579                    builder.inner.result.keys.stats = None;
1580                    builder.inner.result.vals.vals.codec =
1581                        Some(ColumnsCodec::new_from([&state.vals]));
1582                    builder.inner.result.vals.vals.stats = None;
1583                }
1584                builder
1585            }
1586        }
1587
1588        /// Paged counterpart of [`ValRowBuilder`] that consumes [`Column`]
1589        /// chunks. Keys are arbitrary `Columnar` values (not `Row`-shaped) and
1590        /// stay uncompressed; only the value container receives a codec.
1591        pub struct ValRowColPagedBuilder<K, T, R>
1592        where
1593            K: Ord + Clone + Columnation + Columnar + 'static,
1594            T: Lattice + Timestamp + Columnation + Columnar,
1595            R: Ord + Semigroup + Columnation + Columnar + Clone + 'static,
1596        {
1597            inner: OrdValBuilder<ValRowLayout<((K, Row), T, R)>, Column<((K, Row), T, R)>>,
1598        }
1599
1600        impl<K, T, R> Builder for ValRowColPagedBuilder<K, T, R>
1601        where
1602            K: Ord + Clone + Columnation + Columnar + 'static,
1603            T: Lattice + Timestamp + Columnation + Columnar,
1604            R: Ord + Semigroup + Columnation + Columnar + Clone + 'static,
1605            for<'a> columnar::Ref<'a, K>: Copy + Ord,
1606            for<'a, 'b> &'a K: PartialEq<columnar::Ref<'b, K>>,
1607            for<'a> TimelyStack<K>: timely::container::PushInto<columnar::Ref<'a, K>>,
1608        {
1609            type Input = Column<((K, Row), T, R)>;
1610            type Time = T;
1611            type Output = OrdValBatch<ValRowLayout<((K, Row), T, R)>>;
1612
1613            fn push(&mut self, chunk: &mut Self::Input) {
1614                self.inner.push(chunk)
1615            }
1616            fn done(self) -> Option<Self::Output> {
1617                self.inner.done()
1618            }
1619        }
1620
1621        impl<K, T, R> Default for ValRowColPagedBuilder<K, T, R>
1622        where
1623            K: Ord + Clone + Columnation + Columnar + 'static,
1624            T: Lattice + Timestamp + Columnation + Columnar,
1625            R: Ord + Semigroup + Columnation + Columnar + Clone + 'static,
1626            for<'a> columnar::Ref<'a, K>: Copy + Ord,
1627            for<'a, 'b> &'a K: PartialEq<columnar::Ref<'b, K>>,
1628            for<'a> TimelyStack<K>: timely::container::PushInto<columnar::Ref<'a, K>>,
1629        {
1630            fn default() -> Self {
1631                Self {
1632                    inner: Default::default(),
1633                }
1634            }
1635        }
1636
1637        impl<K, T, R> Sealer<Column<((K, Row), T, R)>> for ValRowColPagedBuilder<K, T, R>
1638        where
1639            K: Ord + Clone + Columnation + Columnar + 'static,
1640            T: Lattice + Timestamp + Columnation + Columnar,
1641            R: Ord + Semigroup + Columnation + Columnar + Clone + 'static,
1642            for<'a> columnar::Ref<'a, K>: Copy + Ord,
1643            for<'a, 'b> &'a K: PartialEq<columnar::Ref<'b, K>>,
1644            for<'a> TimelyStack<K>: timely::container::PushInto<columnar::Ref<'a, K>>,
1645        {
1646            type Output = OrdValBatch<ValRowLayout<((K, Row), T, R)>>;
1647
1648            fn seal(chain: &mut Vec<Column<((K, Row), T, R)>>) -> Option<Self::Output> {
1649                let val_codec = build_codec(
1650                    chain
1651                        .iter()
1652                        .flat_map(|c| c.borrow().into_index_iter().map(|((_, v), _, _)| v)),
1653                );
1654
1655                let (keys, vals, upds) =
1656                    <_ as BuilderInput<TimelyStack<K>, DatumContainer>>::key_val_upd_counts(
1657                        &chain[..],
1658                    );
1659                let mut builder = Self {
1660                    inner: OrdValBuilder::with_capacity(keys, vals, upds),
1661                };
1662                // See `RowRowBuilder::seal`: drop the now-redundant stats gatherer.
1663                builder.inner.result.vals.vals.codec = val_codec;
1664                builder.inner.result.vals.vals.stats = None;
1665
1666                for mut chunk in chain.drain(..) {
1667                    builder.push(&mut chunk);
1668                }
1669
1670                builder.done()
1671            }
1672        }
1673    }
1674
1675    pub struct DatumContainer {
1676        /// Encoder/decoder used to translate between row bytes and the stored bytes.
1677        /// `None` until enough pushes have been observed (or if compression is disabled).
1678        codec: Option<ColumnsCodec>,
1679        /// The stored, possibly-encoded, row bytes.
1680        inner: super::bytes_container::BytesContainer,
1681        /// Staging buffer for ingested `Row` types.
1682        staging: Vec<u8>,
1683        /// Statistics gatherer, used to build a safe codec after enough pushes.
1684        /// `None` once the codec has been installed or if compression is disabled.
1685        stats: Option<ColumnsCodec>,
1686    }
1687
1688    impl BatchContainer for DatumContainer {
1689        type Owned = Row;
1690        type ReadItem<'a> = DatumSeq<'a>;
1691
1692        fn with_capacity(size: usize) -> Self {
1693            let stats = if crate::DICTIONARY_COMPRESSION.load(std::sync::atomic::Ordering::Relaxed)
1694            {
1695                Some(Default::default())
1696            } else {
1697                None
1698            };
1699
1700            Self {
1701                codec: None,
1702                inner: BatchContainer::with_capacity(size),
1703                staging: Vec::new(),
1704                stats,
1705            }
1706        }
1707        fn merge_capacity(cont1: &Self, cont2: &Self) -> Self {
1708            // We only build a merged codec when *both* inputs carry one. A codec is
1709            // sound only for the data whose tag usage it observed, so we cannot reuse
1710            // one side's codec to decode the other side's rows. When exactly one side
1711            // is compressed we conservatively produce an uncompressed container rather
1712            // than risk a tag collision; the merged container re-gathers stats and may
1713            // install a fresh codec later via the `STATS_THRESHOLD` path.
1714            let codec = match (&cont1.codec, &cont2.codec) {
1715                (Some(c1), Some(c2)) => Some(ColumnsCodec::new_from([c1, c2])),
1716                _ => None,
1717            };
1718
1719            Self {
1720                codec,
1721                inner: BatchContainer::merge_capacity(&cont1.inner, &cont2.inner),
1722                staging: Vec::new(),
1723                stats: None,
1724            }
1725        }
1726        #[inline]
1727        fn index(&self, index: usize) -> Self::ReadItem<'_> {
1728            let data = self.inner.index(index);
1729            let iter = if let Some(codec) = &self.codec {
1730                codec.decode(data)
1731            } else {
1732                // Safety: without a codec we only push rows or datumseqs into `self.inner`.
1733                // Each retrieved byte slice should be row-encoded data, as long as we have
1734                // not unset the codec in the interim.
1735                unsafe { ColumnsIter::without_codec(data) }
1736            };
1737            DatumSeq { iter }
1738        }
1739        #[inline(always)]
1740        fn len(&self) -> usize {
1741            self.inner.len()
1742        }
1743
1744        #[inline(always)]
1745        fn reborrow<'b, 'a: 'b>(item: Self::ReadItem<'a>) -> Self::ReadItem<'b> {
1746            item
1747        }
1748
1749        #[inline(always)]
1750        fn into_owned<'a>(item: Self::ReadItem<'a>) -> Self::Owned {
1751            // Fast path: unencoded data is already row-formatted bytes.
1752            if item.iter.index.is_none() {
1753                // SAFETY: `iter.data` is raw row-encoded bytes when there is no codec.
1754                return unsafe { Row::from_bytes_unchecked(item.iter.data) };
1755            }
1756            Row::pack(item)
1757        }
1758
1759        #[inline(always)]
1760        fn clone_onto<'a>(item: Self::ReadItem<'a>, other: &mut Self::Owned) {
1761            // Fast path: unencoded data is already row-formatted bytes.
1762            if item.iter.index.is_none() {
1763                let mut packer = other.packer();
1764                // SAFETY: `iter.data` is raw row-encoded bytes when there is no codec.
1765                unsafe { packer.extend_by_slice_unchecked(item.iter.data) };
1766                return;
1767            }
1768            other.packer().extend(item);
1769        }
1770
1771        #[inline(always)]
1772        fn push_ref(&mut self, item: Self::ReadItem<'_>) {
1773            // Fast path: both sides unencoded — push raw bytes directly.
1774            if self.codec.is_none() && self.stats.is_none() && item.iter.index.is_none() {
1775                self.inner.push_ref(item.iter.data);
1776                return;
1777            }
1778            self.push_into(item);
1779        }
1780
1781        #[inline(always)]
1782        fn push_own(&mut self, item: &Self::Owned) {
1783            // Fast path: container is unencoded — push raw row bytes directly.
1784            if self.codec.is_none() && self.stats.is_none() {
1785                self.inner.push_ref(item.data());
1786                return;
1787            }
1788            self.push_into(item);
1789        }
1790
1791        #[inline(always)]
1792        fn clear(&mut self) {
1793            self.inner.clear();
1794            self.staging.clear();
1795            // Reset to the same state as a fresh `with_capacity`: drop any installed
1796            // codec and restore stats gathering (if compression is enabled). Keeping a
1797            // now-empty codec would leave `codec.is_some()`, which permanently routes
1798            // pushes down the encode path with an empty dictionary and prevents the
1799            // `STATS_THRESHOLD` install logic from ever re-engaging compression.
1800            self.codec = None;
1801            self.stats = if crate::DICTIONARY_COMPRESSION.load(std::sync::atomic::Ordering::Relaxed)
1802            {
1803                Some(Default::default())
1804            } else {
1805                None
1806            };
1807        }
1808    }
1809
1810    impl DatumContainer {
1811        /// Visit contained allocations to determine their size and capacity.
1812        #[inline]
1813        pub fn heap_size(&self, mut callback: impl FnMut(usize, usize)) {
1814            self.inner.heap_size(&mut callback);
1815            // The staging buffer and the (possibly absent) codec and stats gatherer all
1816            // hold heap allocations that the bare `inner` accounting misses.
1817            callback(self.staging.len(), self.staging.capacity());
1818            if let Some(codec) = &self.codec {
1819                codec.heap_size(&mut callback);
1820            }
1821            if let Some(stats) = &self.stats {
1822                stats.heap_size(&mut callback);
1823            }
1824        }
1825
1826        /// Promote a gathered-but-uninstalled statistics summary into the codec slot.
1827        ///
1828        /// A container filled via the builder's `push`/`done` path — as the `reduce`
1829        /// operator does, building batches with `Builder::new()` + `push` + `done`
1830        /// rather than `seal` — gathers statistics on every push but never reaches
1831        /// `seal`'s codec install, and only crosses the mid-formation
1832        /// `STATS_THRESHOLD` install if it grows past it. A smaller such container
1833        /// would otherwise be finalized with no codec at all, even with the flag on.
1834        ///
1835        /// That is a problem not because this batch needs compressing — its rows are
1836        /// already stored raw and we deliberately do *not* re-encode them here — but
1837        /// because a codec-less batch poisons future merges: [`Self::merge_capacity`]
1838        /// keys off the presence of a codec, so a codec-less input forces the merged
1839        /// container onto the uncompressed path. Moving the gathered statistics into
1840        /// the codec slot leaves the batch carrying a codec whose retained heavy-hitter
1841        /// summary a later merge can rebuild from via `ColumnsCodec::new_from`, while
1842        /// installing no dictionary: the empty `decode` map resolves every stored
1843        /// (raw) column through the literal-datum fall-through, so reads stay correct.
1844        ///
1845        /// We move the summary as-is rather than building a dictionary via `new_safe`
1846        /// / `new_from` (which reset the summary): unlike `seal` and the mid-formation
1847        /// install, `done` has no further rows to re-observe, so a reset summary would
1848        /// leave the eventual merge nothing to rebuild from.
1849        pub(crate) fn promote_stats_to_codec(&mut self) {
1850            if self.codec.is_none() {
1851                self.codec = self.stats.take();
1852            }
1853        }
1854
1855        /// Whether a codec is installed. A container without one stores raw
1856        /// bytes and blocks compression in every merge it takes part in.
1857        #[cfg(test)]
1858        pub(crate) fn has_codec(&self) -> bool {
1859            self.codec.is_some()
1860        }
1861    }
1862
1863    use timely::container::PushInto;
1864    impl PushInto<Row> for DatumContainer {
1865        #[inline(always)]
1866        fn push_into(&mut self, item: Row) {
1867            self.push_into(&item);
1868        }
1869    }
1870
1871    impl PushInto<&Row> for DatumContainer {
1872        #[inline(always)]
1873        fn push_into(&mut self, item: &Row) {
1874            self.push_into(DatumSeq::borrow_as(item));
1875        }
1876    }
1877
1878    impl PushInto<&RowRef> for DatumContainer {
1879        #[inline(always)]
1880        fn push_into(&mut self, item: &RowRef) {
1881            self.push_into(DatumSeq::borrow_as(item));
1882        }
1883    }
1884
1885    /// Number of pushes a from-scratch container observes before it turns its
1886    /// gathered stats into a safe codec.
1887    ///
1888    /// A safe codec has at most `256 - SAFE_TAG_BASE` (= 134) dictionary slots per
1889    /// column, so we only need to identify ~134 genuinely-popular values. The
1890    /// `MisraGries` summary retains up to `2 * k` (= 1024) distinct candidates
1891    /// between tidies and reduces to `k` (= 512), comfortably more than 134, so the
1892    /// threshold just needs to be large enough that heavy hitters accumulate counts
1893    /// well above 1 before we freeze the codec. 64Ki pushes gives that headroom while
1894    /// keeping the pre-codec (uncompressed) window short.
1895    const STATS_THRESHOLD: usize = 64 * 1024;
1896
1897    impl PushInto<DatumSeq<'_>> for DatumContainer {
1898        #[inline]
1899        fn push_into(&mut self, item: DatumSeq<'_>) {
1900            // Fast path: container and item are both unencoded.
1901            // This is the hot path when dictionary compression is disabled.
1902            if self.codec.is_none() && self.stats.is_none() && item.iter.index.is_none() {
1903                self.inner.push_ref(item.iter.data);
1904                return;
1905            }
1906
1907            // Check if we've gathered enough stats to install a safe codec.
1908            if self.codec.is_none() && self.stats.is_some() && self.inner.len() >= STATS_THRESHOLD {
1909                let stats = self.stats.take().unwrap();
1910                self.codec = Some(stats.new_safe());
1911            }
1912
1913            if let Some(codec) = &mut self.codec {
1914                // Encode using the installed codec.
1915                codec.encode(item.bytes_iter(), &mut self.staging);
1916            } else if let Some(stats) = &mut self.stats {
1917                // Stats-gathering phase: feed the statistics but store raw bytes.
1918                // `observe` updates the heavy-hitter/tag summaries without encoding, so
1919                // we copy each row exactly once (below) instead of also encoding it into
1920                // a buffer we would immediately discard.
1921                stats.observe(item.bytes_iter());
1922                for slice in item.bytes_iter() {
1923                    self.staging.extend_from_slice(slice);
1924                }
1925            } else {
1926                // No codec, no stats: raw copy.
1927                for slice in item.bytes_iter() {
1928                    self.staging.extend_from_slice(slice);
1929                }
1930            }
1931            self.inner.push_ref(&self.staging[..]);
1932            self.staging.clear();
1933        }
1934    }
1935
1936    use mz_repr::{Datum, read_datum};
1937
1938    /// A reference that can be resolved to a sequence of `Datum`s.
1939    ///
1940    /// This type must "compare" as if decoded to a `Row`, which means it needs to track
1941    /// various nuances of `Row::cmp`, which at the moment is first by length, and then by
1942    /// the raw binary slice backing the row. Neither of those are explicit in this struct.
1943    /// We will need to produce them in order to perform comparisons.
1944    #[derive(Debug)]
1945    pub struct DatumSeq<'a> {
1946        pub iter: ColumnsIter<'a>,
1947    }
1948
1949    impl<'a> DatumSeq<'a> {
1950        #[inline(always)]
1951        fn borrow_as(other: &'a RowRef) -> Self {
1952            Self {
1953                iter: ColumnsCodec::borrow_row(other),
1954            }
1955        }
1956
1957        /// Borrow a `Row` as a `DatumSeq` so that it can be used to seek into a
1958        /// trace whose key/value container is a [`DatumContainer`].
1959        #[inline]
1960        pub fn from_row(row: &'a Row) -> Self {
1961            Self::borrow_as(row)
1962        }
1963
1964        #[inline]
1965        pub fn to_row(&self) -> Row {
1966            // Fast path: unencoded data is already row-formatted bytes.
1967            if self.iter.index.is_none() {
1968                return unsafe { Row::from_bytes_unchecked(self.iter.data) };
1969            }
1970            Row::pack(*self)
1971        }
1972    }
1973
1974    impl<'a> Copy for DatumSeq<'a> {}
1975    impl<'a> Clone for DatumSeq<'a> {
1976        #[inline(always)]
1977        fn clone(&self) -> Self {
1978            *self
1979        }
1980    }
1981
1982    use std::cmp::Ordering;
1983    impl<'a, 'b> PartialEq<DatumSeq<'a>> for DatumSeq<'b> {
1984        #[inline(always)]
1985        fn eq(&self, other: &DatumSeq<'a>) -> bool {
1986            // Fast path: both sides are unencoded raw row bytes.
1987            if self.iter.index.is_none() && other.iter.index.is_none() {
1988                return self.iter.data == other.iter.data;
1989            }
1990            Iterator::eq(self.iter, other.iter)
1991        }
1992    }
1993    impl<'a> Eq for DatumSeq<'a> {}
1994    impl<'a, 'b> PartialOrd<DatumSeq<'a>> for DatumSeq<'b> {
1995        #[inline(always)]
1996        fn partial_cmp(&self, other: &DatumSeq<'a>) -> Option<Ordering> {
1997            // Fast path: both sides are unencoded raw row bytes.
1998            if self.iter.index.is_none() && other.iter.index.is_none() {
1999                let left = self.iter.data;
2000                let right = other.iter.data;
2001                return Some(match left.len().cmp(&right.len()) {
2002                    Ordering::Equal => left.cmp(right),
2003                    other => other,
2004                });
2005            }
2006            // Slow path: at least one side is dictionary-encoded.
2007            // Fused length + lexicographic comparison in a single pass per side.
2008            // Row ordering is: shorter < longer; equal lengths compared lexicographically.
2009            //
2010            // We compare byte-by-byte (via `flatten`) rather than slice-by-slice on
2011            // purpose: a dictionary tag expands to a multi-byte value on one side while
2012            // the other side may store those same bytes raw, so the per-column slice
2013            // boundaries do not line up between the two iterators. Decoding to a flat
2014            // byte stream is the only representation in which both sides are directly
2015            // comparable. This path is cold — it only runs when at least one operand is
2016            // dictionary-encoded; the common unencoded case is handled by the fast path
2017            // above with a single slice comparison.
2018            let mut left = self.iter.flatten();
2019            let mut right = other.iter.flatten();
2020            let mut first_diff = Ordering::Equal;
2021            loop {
2022                match (left.next(), right.next()) {
2023                    (Some(l), Some(r)) => {
2024                        if first_diff == Ordering::Equal {
2025                            first_diff = l.cmp(r);
2026                        }
2027                    }
2028                    // Left exhausted first: left is shorter, so Less.
2029                    (None, Some(_)) => return Some(Ordering::Less),
2030                    // Right exhausted first: right is shorter, so Greater.
2031                    (Some(_), None) => return Some(Ordering::Greater),
2032                    // Same length: use first lexicographic difference.
2033                    (None, None) => return Some(first_diff),
2034                }
2035            }
2036        }
2037    }
2038    impl<'a> Ord for DatumSeq<'a> {
2039        #[inline(always)]
2040        fn cmp(&self, other: &Self) -> Ordering {
2041            self.partial_cmp(other).unwrap()
2042        }
2043    }
2044
2045    impl<'a> PartialEq<&'a Row> for DatumSeq<'a> {
2046        #[inline(always)]
2047        fn eq(&self, other: &&'a Row) -> bool {
2048            self.eq(&Self::borrow_as(*other))
2049        }
2050    }
2051
2052    // Lifetimes decoupled (`'b` independent of `'a`): the arrange machinery
2053    // requires `for<'b> DatumSeq<'a>: PartialEq<&'b RowRef>`, i.e. a fixed
2054    // `DatumSeq` must compare against a `&RowRef` of any lifetime.
2055    impl<'a, 'b> PartialEq<&'b RowRef> for DatumSeq<'a> {
2056        #[inline(always)]
2057        fn eq(&self, other: &&'b RowRef) -> bool {
2058            self.eq(&DatumSeq::borrow_as(*other))
2059        }
2060    }
2061
2062    impl<'a> DatumSeq<'a> {
2063        #[inline(always)]
2064        pub fn bytes_iter(self) -> ColumnsIter<'a> {
2065            self.iter
2066        }
2067    }
2068
2069    impl<'a> Iterator for DatumSeq<'a> {
2070        type Item = Datum<'a>;
2071        #[inline(always)]
2072        fn next(&mut self) -> Option<Self::Item> {
2073            // Delegate to `ColumnsIter`, which handles both the codec and no-codec
2074            // cases. The no-codec scan hot path is served directly by `extend_datums`
2075            // (which decodes without going through this iterator), so the only callers
2076            // left here are the codec-encoded `extend_datums`/`to_row` paths and tests;
2077            // none warrant a dedicated no-codec fast path.
2078            self.iter
2079                .next()
2080                .map(|mut bytes| unsafe { read_datum(&mut bytes) })
2081        }
2082    }
2083
2084    use mz_repr::RowArena;
2085    use mz_repr::fixed_length::ExtendDatums;
2086    impl<'long> ExtendDatums for DatumSeq<'long> {
2087        #[inline]
2088        fn extend_datums<'a>(
2089            &'a self,
2090            _arena: &'a RowArena,
2091            target: &mut Vec<Datum<'a>>,
2092            max: Option<usize>,
2093        ) {
2094            // Branch on codec presence ONCE per row rather than once per datum.
2095            // With no codec (the common, feature-off case) push raw datums in a
2096            // tight loop, matching the pre-dictionary path; with a codec, fall
2097            // back to the per-column iterator. This keeps the codec check out of
2098            // the per-datum loop — the source of the feature-off scan overhead.
2099            if self.iter.index.is_none() {
2100                let mut data = self.iter.data;
2101                match max {
2102                    Some(max) => {
2103                        let mut n = 0;
2104                        while n < max && !data.is_empty() {
2105                            target.push(unsafe { read_datum(&mut data) });
2106                            n += 1;
2107                        }
2108                    }
2109                    None => {
2110                        while !data.is_empty() {
2111                            target.push(unsafe { read_datum(&mut data) });
2112                        }
2113                    }
2114                }
2115            } else {
2116                match max {
2117                    Some(max) => target.extend((*self).take(max)),
2118                    None => target.extend(*self),
2119                }
2120            }
2121        }
2122    }
2123}
2124
2125/// Traits abstracting the processes of encoding and decoding row-encoded byte sequences.
2126///
2127/// It is unsafe to use these types to encode byte sequences that are not row-encoded,
2128/// as they are parsed out of contiguous `[u8]` slices using `mz_repr::read_datum`.
2129mod row_codec {
2130
2131    pub use self::misra_gries::MisraGries;
2132    pub use columns::{ColumnsCodec, ColumnsIter};
2133    pub use dictionary::DictionaryCodec;
2134    #[cfg(test)]
2135    pub use dictionary::SAFE_TAG_BASE;
2136
2137    // Deterministic hasher state for the codecs' hash maps: a fixed-seed
2138    // `ahash::RandomState` shared with `mz_timely_util`'s consolidation hasher, so
2139    // the heavy-hitter summaries — and therefore which values each codec compresses
2140    // — are identical across runs and replicas, as the old `BTreeMap` backing was.
2141    use mz_timely_util::hash::fixed_state;
2142
2143    // The codecs encode and decode `[u8]` data specific to the `[Row]` encoding. They
2144    // soundly decode data they themselves encoded from valid `[Row]` data, but may be
2145    // unsound if asked to decode data that was not row-encoded, or was encoded with a
2146    // different codec. `ColumnsCodec` (a per-column wrapper around `DictionaryCodec`) is
2147    // the only codec the spine instantiates; the methods are inherent rather than behind
2148    // a `Codec` trait because nothing ever dispatches over codecs generically.
2149
2150    mod columns {
2151
2152        use mz_repr::{RowRef, read_datum};
2153
2154        use super::DictionaryCodec;
2155
2156        /// Independently encodes each column.
2157        #[derive(Default, Debug)]
2158        pub struct ColumnsCodec {
2159            columns: Vec<DictionaryCodec>,
2160        }
2161
2162        impl ColumnsCodec {
2163            /// Decode a row-encoded byte slice into per-column byte slices.
2164            pub(crate) fn decode<'a>(&'a self, bytes: &'a [u8]) -> ColumnsIter<'a> {
2165                ColumnsIter {
2166                    index: Some(self),
2167                    column: 0,
2168                    data: bytes,
2169                }
2170            }
2171            /// Encode a sequence of column byte slices, updating per-column statistics.
2172            pub(crate) fn encode<'a, I>(&mut self, iter: I, output: &mut Vec<u8>)
2173            where
2174                I: IntoIterator<Item = &'a [u8]>,
2175            {
2176                for (index, bytes) in iter.into_iter().enumerate() {
2177                    if self.columns.len() <= index {
2178                        self.columns.push(Default::default());
2179                    }
2180                    self.columns[index].encode(std::iter::once(bytes), output);
2181                }
2182            }
2183
2184            /// Construct a codec valid for the union of the supplied codecs' data.
2185            pub(crate) fn new_from<'a>(stats: impl IntoIterator<Item = &'a Self>) -> Self {
2186                // An empty `stats` iterator yields a zero-column codec, which encodes and
2187                // decodes nothing; callers merging no inputs get an inert (but sound) codec.
2188                let stats = stats.into_iter().collect::<Vec<_>>();
2189                let cols = stats.iter().map(|s| s.columns.len()).max().unwrap_or(0);
2190                let mut columns = Vec::with_capacity(cols);
2191                let default: DictionaryCodec = Default::default();
2192                for index in 0..cols {
2193                    columns.push(DictionaryCodec::new_from(
2194                        stats
2195                            .iter()
2196                            .map(|s| s.columns.get(index).unwrap_or(&default)),
2197                    ));
2198                }
2199                Self { columns }
2200            }
2201
2202            /// Reveal a row's bytes for fast-path comparison, with no codec to consult.
2203            #[inline(always)]
2204            pub(crate) fn borrow_row(row: &RowRef) -> ColumnsIter<'_> {
2205                ColumnsIter {
2206                    index: None,
2207                    column: 0,
2208                    data: row.data(),
2209                }
2210            }
2211        }
2212
2213        impl ColumnsCodec {
2214            /// Visit contained allocations to determine their size and capacity.
2215            pub(crate) fn heap_size(&self, callback: &mut impl FnMut(usize, usize)) {
2216                let elem = std::mem::size_of::<DictionaryCodec>();
2217                callback(self.columns.len() * elem, self.columns.capacity() * elem);
2218                for column in &self.columns {
2219                    column.heap_size(callback);
2220                }
2221            }
2222        }
2223
2224        impl ColumnsCodec {
2225            /// Record a row's column values in the statistics without encoding.
2226            ///
2227            /// Used during the stats-gathering phase, where we want the heavy-hitter
2228            /// and tag-usage information but store the row raw, so encoding into a
2229            /// throwaway buffer would be pure waste.
2230            #[inline]
2231            pub(crate) fn observe<'a, I>(&mut self, iter: I)
2232            where
2233                I: IntoIterator<Item = &'a [u8]>,
2234            {
2235                for (index, bytes) in iter.into_iter().enumerate() {
2236                    if self.columns.len() <= index {
2237                        self.columns.push(Default::default());
2238                    }
2239                    self.columns[index].observe(bytes);
2240                }
2241            }
2242        }
2243
2244        impl ColumnsCodec {
2245            /// Construct a codec using only structurally safe tags.
2246            ///
2247            /// Consumes `self`: this is only ever called on stats that have just been
2248            /// `take`n out of a container and are about to be discarded, so we move the
2249            /// per-column `MisraGries` summaries through rather than cloning them.
2250            pub(crate) fn new_safe(self) -> Self {
2251                let columns = self
2252                    .columns
2253                    .into_iter()
2254                    .map(DictionaryCodec::new_safe)
2255                    .collect();
2256                Self { columns }
2257            }
2258        }
2259
2260        #[derive(Debug, Copy, Clone)]
2261        pub struct ColumnsIter<'a> {
2262            // `None` when iterating an owned row directly, with no codec to consult.
2263            pub index: Option<&'a ColumnsCodec>,
2264            pub column: usize,
2265            pub data: &'a [u8],
2266        }
2267
2268        impl<'a> Iterator for ColumnsIter<'a> {
2269            type Item = &'a [u8];
2270            #[inline(always)]
2271            fn next(&mut self) -> Option<Self::Item> {
2272                if self.data.is_empty() {
2273                    None
2274                } else if let Some(bytes) = self
2275                    .index
2276                    .as_ref()
2277                    .and_then(|i| i.columns.get(self.column))
2278                    .and_then(|i| i.decode.get(self.data[0].into()))
2279                {
2280                    self.data = &self.data[1..];
2281                    self.column += 1;
2282                    Some(bytes)
2283                } else {
2284                    let mut data = self.data;
2285                    let data_len = data.len();
2286                    unsafe {
2287                        read_datum(&mut data);
2288                    }
2289                    let (prev, next) = self.data.split_at(data_len - data.len());
2290                    self.data = next;
2291                    self.column += 1;
2292                    Some(prev)
2293                }
2294            }
2295        }
2296
2297        impl<'a> ColumnsIter<'a> {
2298            /// Create a column iterator without a codec.
2299            ///
2300            /// This requires the data to be row-formatted, and it will be erroneous otherwise.
2301            #[inline(always)]
2302            pub unsafe fn without_codec(data: &'a [u8]) -> Self {
2303                Self {
2304                    index: None,
2305                    column: 0,
2306                    data,
2307                }
2308            }
2309        }
2310    }
2311
2312    /// A dictionary encoding codec for `[Row]` data.
2313    ///
2314    /// The dictionary harvests unused tags within each column and uses them to
2315    /// represent popular values within that column. There are two mechanisms it
2316    /// uses to accomplish this:
2317    ///
2318    /// 1. Statically free tags: `SAFE_TAG_BASE` is taken as an exclusive upper bound
2319    ///    on the tags that will be used by `[Row]`, and tags greater or equal to this
2320    ///    value are always safe to use.
2321    /// 2. Dynamically free tags: having seen an entire collection, we can use any
2322    ///    tag not otherwise used by the collection, as it would not be ambiguous.
2323    ///
2324    /// It goes without saying that if either of these approaches are incorrect,
2325    /// there are calamitous unsoundness implications.
2326    mod dictionary {
2327        // The `encode` map is a pure value->tag lookup table (never iterated for logic),
2328        // so `mz_ore::collections::HashMap`'s order-hiding would suffice — but it offers
2329        // no fixed-seed constructor, and we want the same deterministic hasher as the
2330        // summary above. `heap_size`'s `keys()` walk is an order-insensitive sum.
2331        #![allow(clippy::disallowed_types)]
2332
2333        use std::collections::HashMap;
2334
2335        use super::fixed_state;
2336        pub use super::{BytesMap, MisraGries};
2337
2338        /// First byte value that is structurally unused by the datum encoding.
2339        /// All byte values >= this are safe to use as dictionary tags without
2340        /// observing the data, since no datum's first byte can have this value.
2341        ///
2342        /// `mz_repr`'s `Row` `Tag` enum currently has 84 variants (discriminants
2343        /// 0..=83), so the truly tight bound is 84. We deliberately pick a larger,
2344        /// round-ish constant to leave headroom for new tags without having to also
2345        /// bump the safe set, and the `test_safe_tag_base` test pins the real
2346        /// invariant: every datum the row format produces must encode with a first
2347        /// byte strictly less than this value. If a future tag crosses the boundary
2348        /// that test fails loudly rather than silently corrupting decoding.
2349        pub const SAFE_TAG_BASE: u8 = 122;
2350
2351        /// Per-column dictionary codec. Encodes column byte slices, replacing popular
2352        /// values with spare tags; decoding is performed by `ColumnsIter` reading the
2353        /// `decode` map directly.
2354        #[derive(Default, Debug)]
2355        pub struct DictionaryCodec {
2356            // Looked up once per value on the encode path; mostly misses (only popular
2357            // values compress), so a hash map beats a `BTreeMap`'s byte-slice walk. The
2358            // map is only ever read via `get` — never iterated — so its hasher seed has
2359            // no observable effect; the populated maps are built with `fixed_state` in
2360            // `new_from`/`new_safe` for consistency, while the derived-`Default` (stats
2361            // accumulator) variant stays empty and is never consulted.
2362            encode: HashMap<Vec<u8>, u8, ahash::RandomState>,
2363            pub decode: BytesMap,
2364            stats: (MisraGries<Vec<u8>>, [u64; 4]),
2365        }
2366
2367        impl DictionaryCodec {
2368            /// Encode a sequence of byte slices.
2369            ///
2370            /// Encoding also records statistics about the structure of the input.
2371            ///
2372            /// Decoding has no symmetric method here: a column's bytes are decoded by
2373            /// `ColumnsIter`, which consults the `decode` map directly.
2374            pub(super) fn encode<'a, I>(&mut self, iter: I, output: &mut Vec<u8>)
2375            where
2376                I: IntoIterator<Item = &'a [u8]>,
2377            {
2378                for bytes in iter.into_iter() {
2379                    mz_ore::soft_assert_no_log!(
2380                        !bytes.is_empty(),
2381                        "row encoding never yields empty column slices",
2382                    );
2383                    // If we have an index referencing `bytes`, use the index key.
2384                    if let Some(b) = self.encode.get(bytes) {
2385                        output.push(*b);
2386                    } else {
2387                        // Raw fall-through. Soundness rests on `bytes[0]` never being a
2388                        // tag we hand out as a dictionary key: `new_from`/`new_safe` only
2389                        // assign dictionary tags from first-byte values that were never
2390                        // observed (or are `>= SAFE_TAG_BASE`, which no datum first-byte
2391                        // can equal). If a literal datum's first byte collided with a
2392                        // dictionary tag, `decode` would resolve it to the dictionary
2393                        // entry instead of reading the datum. This `debug_assert` makes
2394                        // the load-bearing "no later first-byte outside the observed
2395                        // union" invariant self-checking.
2396                        mz_ore::soft_assert_no_log!(
2397                            self.decode.get(bytes[0].into()).is_none(),
2398                            "raw datum first-byte {} collides with a dictionary tag; \
2399                             decode would be ambiguous",
2400                            bytes[0],
2401                        );
2402                        output.extend(bytes);
2403                    }
2404                    self.observe(bytes);
2405                }
2406            }
2407
2408            /// Construct a new encoder from supplied statistics.
2409            pub(super) fn new_from<'a>(stats: impl IntoIterator<Item = &'a Self>) -> Self {
2410                // Collect most popular bytes from combined containers.
2411                let mut mg = MisraGries::default();
2412                let mut tags: [u64; 4] = [0; 4];
2413                for stat in stats.into_iter() {
2414                    for (thing, count) in stat.stats.0.clone().done() {
2415                        mg.update(thing, count);
2416                    }
2417                    tags[0] |= stat.stats.1[0];
2418                    tags[1] |= stat.stats.1[1];
2419                    tags[2] |= stat.stats.1[2];
2420                    tags[3] |= stat.stats.1[3];
2421                }
2422                let mut mg = mg
2423                    .done()
2424                    .into_iter()
2425                    .filter(|(next_bytes, count)| next_bytes.len() > 1 && count > &1);
2426                // Establish encoding and decoding rules.
2427                let mut encode = HashMap::with_hasher(fixed_state());
2428                let mut decode = BytesMap::default();
2429                for tag in 0..=255 {
2430                    let tag_idx: usize = (tag % 4).into();
2431                    let shift = tag >> 2;
2432                    if (tags[tag_idx] >> shift) & 0x01 != 0 {
2433                        // Tag is used by a literal datum first-byte; reserve the slot.
2434                        decode.push(None);
2435                    } else if let Some((next_bytes, _count)) = mg.next() {
2436                        decode.push(Some(&next_bytes[..]));
2437                        encode.insert(next_bytes, tag);
2438                    } else {
2439                        // Unused tag, but the heavy-hitter supply is exhausted. We must
2440                        // still push a slot so that `decode`'s index stays aligned with
2441                        // the tag value: every iteration pushes exactly once, keeping the
2442                        // map length 256 and `decode.get(tag)` addressable by tag.
2443                        decode.push(None);
2444                    }
2445                }
2446
2447                Self {
2448                    encode,
2449                    decode,
2450                    stats: (MisraGries::default(), [0u64; 4]),
2451                }
2452            }
2453        }
2454
2455        impl DictionaryCodec {
2456            /// Visit contained allocations to determine their size and capacity.
2457            ///
2458            /// The `encode` table is approximated as one logical entry's worth of bytes
2459            /// per element for size and its reserved `capacity()` for capacity; the
2460            /// dominant terms (the owned key bytes and the `decode` map's byte arena)
2461            /// are accounted exactly.
2462            pub fn heap_size(&self, callback: &mut impl FnMut(usize, usize)) {
2463                let entry = std::mem::size_of::<(Vec<u8>, u8)>();
2464                callback(self.encode.len() * entry, self.encode.capacity() * entry);
2465                for key in self.encode.keys() {
2466                    callback(key.len(), key.capacity());
2467                }
2468                self.decode.heap_size(callback);
2469                self.stats.0.heap_size(callback);
2470            }
2471
2472            /// Record a single column value in this codec's statistics without
2473            /// producing any encoded output.
2474            ///
2475            /// Statistics come in two decoupled parts, with very different costs and
2476            /// purposes:
2477            ///
2478            /// 1. The tag bitmap (`stats.1`) records which first-byte values have been
2479            ///    observed. It is cheap (four `u64` ORs) and *soundness critical*:
2480            ///    `new_from`'s dynamic-tag path only hands out tags that this bitmap
2481            ///    reports as unused, so it must stay accurate for the entire life of the
2482            ///    codec, including on the hot encode path.
2483            /// 2. The MisraGries summary (`stats.0`) tracks heavy hitters and only
2484            ///    affects *which* values a future codec compresses, never correctness.
2485            ///    It is the expensive part (a `BTreeMap` insert per column per row). We
2486            ///    keep feeding it after install, on the hot encode path, on purpose: a
2487            ///    later merge rebuilds the merged codec from these summaries via
2488            ///    `new_from`. If we froze the summary at install time, then as the
2489            ///    collection evolves — records cancel under consolidation, the popular
2490            ///    set drifts — the codec could never reclaim slots for newly-popular
2491            ///    values and would eventually be left compressing values that no longer
2492            ///    occur, ceasing to compress the ones that do.
2493            #[inline]
2494            pub fn observe(&mut self, bytes: &[u8]) {
2495                mz_ore::soft_assert_no_log!(
2496                    !bytes.is_empty(),
2497                    "row encoding never yields empty column slices",
2498                );
2499                let tag = bytes[0];
2500                let tag_idx: usize = (tag % 4).into();
2501                self.stats.1[tag_idx] |= 1 << (tag >> 2);
2502                self.stats.0.insert_ref(bytes);
2503            }
2504
2505            /// Construct a codec using only structurally safe tags (>= SAFE_TAG_BASE).
2506            /// These tags never collide with datum first-bytes, so the codec can be
2507            /// installed without observing all data first.
2508            pub(super) fn new_safe(stats: Self) -> Self {
2509                // The container stores its pre-install rows raw, so the first-byte
2510                // bitmap (`stats.1`) gathered while observing them must carry over to
2511                // the installed codec. The bitmap is soundness-critical: a later
2512                // `new_from` merge consults it to decide which one-byte tags are free
2513                // to hand out as dictionary keys. If we dropped it here, the merge
2514                // could assign a dictionary tag equal to a pre-install datum's first
2515                // byte, after which `decode` would resolve that literal datum to the
2516                // dictionary entry. The MisraGries summary (`stats.0`), by contrast,
2517                // is consumed below to seed the dictionary and is reset, since the
2518                // installed codec re-accumulates it from rows it sees post-install.
2519                let (mg, observed_tags) = stats.stats;
2520                let mut mg = mg
2521                    .done()
2522                    .into_iter()
2523                    .filter(|(next_bytes, count)| next_bytes.len() > 1 && count > &1);
2524                let mut encode = HashMap::with_hasher(fixed_state());
2525                let mut decode = BytesMap::default();
2526                // Fill slots 0..SAFE_TAG_BASE with None (reserved for datum tags).
2527                for _ in 0..SAFE_TAG_BASE {
2528                    decode.push(None);
2529                }
2530                // Assign dictionary entries to safe tags.
2531                for tag in SAFE_TAG_BASE..=255 {
2532                    if let Some((next_bytes, _count)) = mg.next() {
2533                        decode.push(Some(&next_bytes[..]));
2534                        encode.insert(next_bytes, tag);
2535                    }
2536                }
2537                Self {
2538                    encode,
2539                    decode,
2540                    stats: (MisraGries::default(), observed_tags),
2541                }
2542            }
2543        }
2544    }
2545
2546    /// A map from `0 .. something` to `Option<&[u8]>`.
2547    ///
2548    /// Non-empty slices are pushed in order, and can be retrieved by index.
2549    /// Pushing an empty slice is equivalent to pushing `None`.
2550    #[derive(Debug)]
2551    pub struct BytesMap {
2552        offsets: Vec<usize>,
2553        bytes: Vec<u8>,
2554    }
2555    impl Default for BytesMap {
2556        #[inline(always)]
2557        fn default() -> Self {
2558            Self {
2559                offsets: vec![0],
2560                bytes: Vec::new(),
2561            }
2562        }
2563    }
2564    impl BytesMap {
2565        #[inline]
2566        fn push(&mut self, input: Option<&[u8]>) {
2567            if let Some(bytes) = input {
2568                self.bytes.extend(bytes);
2569            }
2570            self.offsets.push(self.bytes.len());
2571        }
2572        /// Visit contained allocations to determine their size and capacity.
2573        fn heap_size(&self, callback: &mut impl FnMut(usize, usize)) {
2574            let off = std::mem::size_of::<usize>();
2575            callback(self.offsets.len() * off, self.offsets.capacity() * off);
2576            callback(self.bytes.len(), self.bytes.capacity());
2577        }
2578        #[inline]
2579        fn get(&self, index: usize) -> Option<&[u8]> {
2580            if index < self.offsets.len() - 1 {
2581                let lower = self.offsets[index];
2582                let upper = self.offsets[index + 1];
2583                if lower < upper {
2584                    Some(&self.bytes[lower..upper])
2585                } else {
2586                    None
2587                }
2588            } else {
2589                None
2590            }
2591        }
2592    }
2593
2594    mod misra_gries {
2595        // The summary must iterate its entries (to extract heavy hitters in `done`, to
2596        // `tidy`, and to size itself), which `mz_ore::collections::HashMap` deliberately
2597        // forbids. We instead get determinism from the fixed-seed hasher (`fixed_state`)
2598        // plus the total-order sort in `done`; `tidy`/`heap_size` are order-insensitive.
2599        #![allow(clippy::disallowed_types)]
2600
2601        use std::collections::HashMap;
2602        use std::hash::Hash;
2603
2604        use super::fixed_state;
2605
2606        /// Maintains a summary of "heavy hitters" in a presented collection of items.
2607        ///
2608        /// Uses a hash map internally so that repeated observations of the same
2609        /// element only allocate once (on first sighting), and so the per-element
2610        /// `insert_ref` is an O(1) hash rather than an O(log n) walk of byte-slice
2611        /// comparisons. This is the hot path: one lookup per column per row, fed both
2612        /// while gathering stats and on the steady-state encode path. The hasher is
2613        /// fixed-seed (see [`fixed_state`]) so the summary — and thus which values a
2614        /// codec compresses — stays deterministic across runs and replicas.
2615        ///
2616        /// Tidy is performed when the number of *distinct* elements exceeds `2 * k`,
2617        /// reducing to at most `k` entries.
2618        #[derive(Clone, Debug)]
2619        pub struct MisraGries<T>
2620        where
2621            T: Ord + Hash,
2622        {
2623            inner: HashMap<T, usize, ahash::RandomState>,
2624            k: usize,
2625        }
2626
2627        impl<T> Default for MisraGries<T>
2628        where
2629            T: Ord + Hash,
2630        {
2631            #[inline(always)]
2632            fn default() -> Self {
2633                Self {
2634                    inner: HashMap::with_hasher(fixed_state()),
2635                    k: 512,
2636                }
2637            }
2638        }
2639
2640        impl<T> MisraGries<T>
2641        where
2642            T: Ord + Hash,
2643        {
2644            /// Inserts an additional element to the summary.
2645            #[inline(always)]
2646            pub fn insert(&mut self, element: T) {
2647                self.update(element, 1);
2648            }
2649            /// Inserts multiple copies of an element to the summary.
2650            #[inline]
2651            pub fn update(&mut self, element: T, count: usize) {
2652                *self.inner.entry(element).or_insert(0) += count;
2653                if self.inner.len() > 2 * self.k {
2654                    self.tidy();
2655                }
2656            }
2657
2658            /// Completes the summary, and extracts the items and their counts.
2659            pub fn done(self) -> Vec<(T, usize)> {
2660                let mut result: Vec<_> = self.inner.into_iter().collect();
2661                // Descending count, ties broken by key, so the values a codec selects
2662                // are deterministic regardless of hash-map iteration order.
2663                result.sort_by(|x, y| y.1.cmp(&x.1).then_with(|| x.0.cmp(&y.0)));
2664                result
2665            }
2666
2667            /// Reduces the summary down to at most `k` distinct items by
2668            /// subtracting the (k+1)-th largest count from all entries and
2669            /// discarding those that drop to zero or below.
2670            fn tidy(&mut self) {
2671                let mut counts: Vec<usize> = self.inner.values().copied().collect();
2672                counts.sort_unstable_by(|a, b| b.cmp(a));
2673                // The (k+1)-th largest count, or 0 if fewer than k+1 entries.
2674                let sub_weight = counts.get(self.k).copied().unwrap_or(0);
2675                if sub_weight > 0 {
2676                    self.inner.retain(|_, count| {
2677                        *count = count.saturating_sub(sub_weight);
2678                        *count > 0
2679                    });
2680                }
2681            }
2682        }
2683
2684        impl MisraGries<Vec<u8>> {
2685            /// Visit contained allocations to determine their size and capacity.
2686            ///
2687            /// The hash table is approximated as one logical entry per element for
2688            /// size and its reserved `capacity()` for capacity; the owned key bytes
2689            /// are accounted exactly.
2690            pub fn heap_size(&self, callback: &mut impl FnMut(usize, usize)) {
2691                let entry = std::mem::size_of::<(Vec<u8>, usize)>();
2692                callback(self.inner.len() * entry, self.inner.capacity() * entry);
2693                for key in self.inner.keys() {
2694                    callback(key.len(), key.capacity());
2695                }
2696            }
2697
2698            /// Insert a borrowed byte slice, only allocating if the key is new.
2699            #[inline]
2700            pub fn insert_ref(&mut self, element: &[u8]) {
2701                if let Some(count) = self.inner.get_mut(element) {
2702                    *count += 1;
2703                } else {
2704                    self.insert(element.to_owned());
2705                }
2706            }
2707        }
2708
2709        impl<T> std::ops::AddAssign for MisraGries<T>
2710        where
2711            T: Ord + Hash,
2712        {
2713            fn add_assign(&mut self, rhs: Self) {
2714                for (element, count) in rhs.done() {
2715                    self.update(element, count);
2716                }
2717            }
2718        }
2719    }
2720}