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