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