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