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