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