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