1use std::borrow::Cow;
42#[cfg(test)]
43use std::cell::Cell;
44use std::cell::RefCell;
45use std::collections::VecDeque;
46use std::rc::Rc;
47use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
48
49use columnar::bytes::indexed;
50use columnar::{Borrow, BorrowedOf, Columnar, Container as _, FromBytes, Index, Len, Push as _};
51use differential_dataflow::difference::Semigroup;
52use differential_dataflow::lattice::Lattice;
53use differential_dataflow::trace::chunk::Chunk;
54use mz_ore::cast::CastFrom;
55use mz_ore::pool::{ChunkHandle, ChunkHints, ExtentCodec, IDENTITY_CODEC, Pool};
56use smallvec::SmallVec;
57use timely::Accountable;
58use timely::PartialOrder;
59use timely::container::{ContainerBuilder, PushInto};
60use timely::dataflow::channels::ContainerBytes;
61use timely::progress::Timestamp;
62use timely::progress::frontier::{Antichain, AntichainRef};
63
64use crate::columnar::batcher::{ColumnChunker, gallop};
65use crate::columnar::unload::UnloadChunk;
66use crate::columnar::{Column, at_serialized_capacity};
67
68static COMPUTE_SPILL_ENABLED: AtomicBool = AtomicBool::new(false);
70
71static STORAGE_SPILL_ENABLED: AtomicBool = AtomicBool::new(false);
73
74thread_local! {
75 static SPILL_OVERRIDE: RefCell<Option<Pool>> = const { RefCell::new(None) };
79
80 #[cfg(test)]
84 static COMPRESS_MIN_DEPTH_OVERRIDE: Cell<Option<u8>> = const { Cell::new(None) };
85
86 static READ_SCRATCH: RefCell<Vec<u64>> = const { RefCell::new(Vec::new()) };
88}
89
90pub fn set_compute_spill_enabled(enabled: bool) {
104 COMPUTE_SPILL_ENABLED.store(enabled, Ordering::Relaxed);
105}
106
107pub fn set_storage_spill_enabled(enabled: bool) {
111 STORAGE_SPILL_ENABLED.store(enabled, Ordering::Relaxed);
112}
113
114pub fn set_spill_override(pool: Option<Pool>) {
118 SPILL_OVERRIDE.with(|cell| *cell.borrow_mut() = pool);
119}
120
121static COMPRESS_MIN_DEPTH: AtomicU8 = AtomicU8::new(DEFAULT_COMPRESS_MIN_DEPTH);
124
125pub fn set_compress_min_depth(depth: u8) {
144 COMPRESS_MIN_DEPTH.store(depth, Ordering::Relaxed);
145}
146
147#[cfg(test)]
151pub fn set_compress_min_depth_override(depth: Option<u8>) {
152 COMPRESS_MIN_DEPTH_OVERRIDE.with(|cell| cell.set(depth));
153}
154
155fn compress_min_depth() -> u8 {
157 #[cfg(test)]
158 if let Some(depth) = COMPRESS_MIN_DEPTH_OVERRIDE.with(|cell| cell.get()) {
159 return depth;
160 }
161 COMPRESS_MIN_DEPTH.load(Ordering::Relaxed)
162}
163
164fn codec_for_depth(depth: u8) -> (&'static dyn ExtentCodec, bool) {
169 if depth < compress_min_depth() {
170 (&IDENTITY_CODEC, false)
171 } else {
172 (&LZ4_CODEC, true)
173 }
174}
175
176fn spill_pool() -> Option<Pool> {
178 if let Some(pool) = SPILL_OVERRIDE.with(|cell| cell.borrow().clone()) {
179 return Some(pool);
180 }
181 let enabled = COMPUTE_SPILL_ENABLED.load(Ordering::Relaxed)
182 || STORAGE_SPILL_ENABLED.load(Ordering::Relaxed);
183 if enabled {
184 crate::pool_config::active_pool()
185 } else {
186 None
187 }
188}
189
190const SCRATCH_RETAIN_WORDS: usize = 1 << 18;
194
195fn with_scratch<Out>(f: impl FnOnce(&mut Vec<u64>) -> Out) -> Out {
197 READ_SCRATCH.with(|cell| {
198 let mut scratch = cell.take();
199 scratch.clear();
200 let out = f(&mut scratch);
201 if scratch.capacity() > SCRATCH_RETAIN_WORDS {
202 scratch.clear();
203 scratch.shrink_to_fit();
204 }
205 cell.replace(scratch);
206 out
207 })
208}
209
210const COMMIT_BYTES: usize = 2 << 20;
213
214const SPILL_MIN_BYTES: usize = 64 << 10;
223
224const DEFAULT_COMPRESS_MIN_DEPTH: u8 = 1;
233
234fn at_commit_size<C: Columnar>(column: &Column<C>) -> bool {
238 column.length_in_bytes() >= COMMIT_BYTES - COMMIT_BYTES / 10
239}
240
241fn borrow_words<C: Columnar>(words: &[u64]) -> BorrowedOf<'_, C> {
244 <BorrowedOf<'_, C>>::from_bytes(&mut indexed::decode(words))
245}
246
247#[inline(always)]
251fn rr<'b, 'a: 'b, C: Columnar>(item: columnar::Ref<'a, C>) -> columnar::Ref<'b, C> {
252 columnar::ContainerOf::<C>::reborrow_ref(item)
253}
254
255pub struct SpilledBody<D: Columnar, T> {
261 records: usize,
263 fences: D::Container,
267 time_lower: Antichain<T>,
271 time_upper: SmallVec<[T; 1]>,
276 compressed: bool,
283 handle: ChunkHandle,
285}
286
287pub enum ColumnChunk<D: Columnar, T: Columnar, R: Columnar> {
304 Resident(Rc<Column<(D, T, R)>>, u8),
306 Spilled(Rc<SpilledBody<D, T>>, u8),
308}
309
310impl<D: Columnar, T: Columnar, R: Columnar> Clone for ColumnChunk<D, T, R> {
311 fn clone(&self) -> Self {
312 match self {
313 ColumnChunk::Resident(col, depth) => ColumnChunk::Resident(Rc::clone(col), *depth),
314 ColumnChunk::Spilled(body, depth) => ColumnChunk::Spilled(Rc::clone(body), *depth),
315 }
316 }
317}
318
319impl<D: Columnar, T: Columnar, R: Columnar> Default for ColumnChunk<D, T, R> {
320 fn default() -> Self {
321 ColumnChunk::Resident(Rc::new(Column::default()), 0)
322 }
323}
324
325impl<D: Columnar, T: Columnar, R: Columnar> Accountable for ColumnChunk<D, T, R> {
326 fn record_count(&self) -> i64 {
327 i64::try_from(self.records()).expect("record count fits i64")
328 }
329}
330
331impl<D: Columnar, T: Columnar, R: Columnar> ColumnChunk<D, T, R> {
332 pub fn from_column(column: Column<(D, T, R)>) -> Self {
335 mz_ore::soft_assert_no_log!(!column.is_empty(), "chunks must be non-empty");
336 ColumnChunk::Resident(Rc::new(column), 0)
337 }
338
339 pub fn into_column(self) -> Column<(D, T, R)> {
342 match self {
343 ColumnChunk::Resident(col, _) => {
344 Rc::try_unwrap(col).unwrap_or_else(|shared| copy_column(&shared))
345 }
346 ColumnChunk::Spilled(body, _) => {
347 let mut words = Vec::new();
348 body.handle.read_into(&mut words);
349 Column::Align(words)
350 }
351 }
352 }
353
354 pub fn is_spilled(&self) -> bool {
356 matches!(self, ColumnChunk::Spilled(_, _))
357 }
358
359 fn records(&self) -> usize {
361 match self {
362 ColumnChunk::Resident(col, _) => col.borrow().len(),
363 ColumnChunk::Spilled(body, _) => body.records,
364 }
365 }
366
367 fn depth(&self) -> u8 {
369 match self {
370 ColumnChunk::Resident(_, depth) | ColumnChunk::Spilled(_, depth) => *depth,
371 }
372 }
373
374 fn data_span(&self) -> (columnar::Ref<'_, D>, columnar::Ref<'_, D>) {
376 match self {
377 ColumnChunk::Resident(col, _) => {
378 let data = col.borrow().0;
379 (data.get(0), data.get(data.len() - 1))
380 }
381 ColumnChunk::Spilled(body, _) => {
382 let fences = body.fences.borrow();
383 (fences.get(0), fences.get(1))
384 }
385 }
386 }
387
388 fn commit(column: Column<(D, T, R)>, depth: u8) -> Self
392 where
393 T: Timestamp,
394 {
395 mz_ore::soft_assert_no_log!(!column.is_empty(), "chunks must be non-empty");
396 if let Some(pool) = spill_pool() {
397 if column.length_in_bytes() >= SPILL_MIN_BYTES {
398 return Self::spill_body(column, &pool, depth);
399 }
400 }
401 ColumnChunk::Resident(Rc::new(column), depth)
402 }
403
404 fn spill_body(column: Column<(D, T, R)>, pool: &Pool, depth: u8) -> Self
412 where
413 T: Timestamp,
414 {
415 let (codec, compressed) = codec_for_depth(depth);
416 let len_bytes = column.length_in_bytes();
417 let (time_lower, time_upper) = Self::time_bounds(&column);
418 let view = column.borrow();
419 let records = view.len();
420 let mut fences = D::Container::default();
421 fences.push(view.0.get(0));
422 fences.push(view.0.get(records - 1));
423 let handle = spill_column(column, pool, len_bytes, ChunkHints { depth }, codec);
424 ColumnChunk::Spilled(
425 Rc::new(SpilledBody {
426 records,
427 fences,
428 time_lower,
429 time_upper: time_upper.into(),
430 compressed,
431 handle,
432 }),
433 depth,
434 )
435 }
436
437 fn survive_merge(self) -> Self
453 where
454 T: Timestamp,
455 {
456 let depth = self.depth().saturating_add(1);
457 match self {
458 ColumnChunk::Resident(col, _) => ColumnChunk::Resident(col, depth),
459 ColumnChunk::Spilled(body, was) => {
460 let migrate = !body.compressed && depth >= compress_min_depth();
461 if !migrate || Rc::strong_count(&body) > 1 {
462 return ColumnChunk::Spilled(body, depth);
463 }
464 match spill_pool() {
465 Some(pool) => {
466 let column = ColumnChunk::Spilled(body, was).into_column();
467 Self::spill_body(column, &pool, depth)
468 }
469 None => ColumnChunk::Spilled(body, depth),
470 }
471 }
472 }
473 }
474
475 fn chunk_time_bounds(&self) -> (Cow<'_, Antichain<T>>, Cow<'_, [T]>)
480 where
481 T: Timestamp,
482 {
483 match self {
484 ColumnChunk::Resident(col, _) => {
485 let (lower, upper) = Self::time_bounds(col);
486 (Cow::Owned(lower), Cow::Owned(upper))
487 }
488 ColumnChunk::Spilled(body, _) => (
489 Cow::Borrowed(&body.time_lower),
490 Cow::Borrowed(&body.time_upper[..]),
491 ),
492 }
493 }
494
495 fn time_bounds(column: &Column<(D, T, R)>) -> (Antichain<T>, Vec<T>)
500 where
501 T: Timestamp,
502 {
503 let (_, times, _) = column.borrow();
504 let mut lower = Antichain::new();
505 let mut upper: Vec<T> = Vec::new();
506 let mut time = T::minimum();
510 for i in 0..times.len() {
511 time.copy_from(rr::<T>(times.get(i)));
512 if !upper.iter().any(|u| PartialOrder::less_equal(&time, u)) {
513 upper.retain(|u| !PartialOrder::less_equal(u, &time));
514 upper.push(time.clone());
515 }
516 lower.insert_ref(&time);
517 }
518 (lower, upper)
519 }
520}
521
522fn copy_column<C: Columnar>(column: &Column<C>) -> Column<C> {
524 let view = column.borrow();
525 let mut fresh = C::Container::default();
526 fresh.extend_from_self(view, 0..view.len());
527 Column::Typed(fresh)
528}
529
530#[derive(Debug)]
535pub struct Lz4Codec;
536
537pub static LZ4_CODEC: Lz4Codec = Lz4Codec;
540
541impl ExtentCodec for Lz4Codec {
542 fn encode(&self, body: &[u8], out: &mut Vec<u8>) {
543 let max_out = lz4_flex::block::get_maximum_output_size(body.len());
544 out.resize(4 + max_out, 0);
545 let len = u32::try_from(body.len()).expect("chunk bodies are bounded by the size classes");
546 out[..4].copy_from_slice(&len.to_le_bytes());
547 let compressed = lz4_flex::block::compress_into(body, &mut out[4..])
548 .expect("output sized to the maximum");
549 out.truncate(4 + compressed);
550 }
551
552 fn decode(&self, stored: &[u8], body: &mut [u8]) {
553 let prefix: [u8; 4] = stored[..4].try_into().expect("prefix length");
554 let len = usize::try_from(u32::from_le_bytes(prefix)).expect("length fits usize");
555 assert_eq!(
556 len,
557 body.len(),
558 "destination must match the encoded body length"
559 );
560 let written = lz4_flex::block::decompress_into(&stored[4..], body)
561 .expect("stored bytes hold a valid lz4 block");
562 assert_eq!(written, body.len(), "decoded length mismatch");
563 }
564}
565
566fn spill_column<C: Columnar>(
571 column: Column<C>,
572 pool: &Pool,
573 len_bytes: usize,
574 hints: ChunkHints,
575 codec: &'static dyn ExtentCodec,
576) -> ChunkHandle {
577 mz_ore::soft_assert_eq_no_log!(len_bytes % 8, 0);
578 match column {
579 Column::Align(words) => {
580 pool.insert_with(words.len(), hints, codec, |dst| dst.copy_from_slice(&words))
581 }
582 other => pool.insert_with(len_bytes / 8, hints, codec, |dst| {
583 let bytes: &mut [u8] = bytemuck::cast_slice_mut(dst);
584 let mut cursor = std::io::Cursor::new(bytes);
585 other.into_bytes(&mut cursor);
586 assert_eq!(
587 usize::try_from(cursor.position()).expect("usize position"),
588 len_bytes,
589 "serialized body must fill the chunk exactly",
590 );
591 }),
592 }
593}
594
595fn to_typed<C: Columnar>(column: Column<C>) -> Column<C> {
599 match column {
600 typed @ Column::Typed(_) => typed,
601 other => copy_column(&other),
602 }
603}
604
605impl<D, T, R> Chunk for ColumnChunk<D, T, R>
606where
607 D: Columnar,
608 for<'a> columnar::Ref<'a, D>: Copy + Ord,
609 T: Columnar + Default + Timestamp + Lattice + Ord,
610 for<'a> columnar::Ref<'a, T>: Copy + Ord,
611 R: Columnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
612{
613 type Time = T;
614
615 const TARGET: usize = 65536;
621
622 fn len(&self) -> usize {
623 self.records()
624 }
625
626 fn merge(in1: &mut VecDeque<Self>, in2: &mut VecDeque<Self>, out: &mut VecDeque<Self>) {
633 let (a_first, a_last) = in1
638 .front()
639 .expect("caller guarantees non-empty input")
640 .data_span();
641 let (b_first, b_last) = in2
642 .front()
643 .expect("caller guarantees non-empty input")
644 .data_span();
645 let a_low = rr::<D>(a_last) < rr::<D>(b_first);
646 let b_low = rr::<D>(b_last) < rr::<D>(a_first);
647 if a_low {
648 let chunk = in1.pop_front().expect("front observed above");
649 out.push_back(chunk.survive_merge());
650 return;
651 }
652 if b_low {
653 let chunk = in2.pop_front().expect("front observed above");
654 out.push_back(chunk.survive_merge());
655 return;
656 }
657
658 let a = in1.pop_front().expect("caller guarantees non-empty input");
659 let b = in2.pop_front().expect("caller guarantees non-empty input");
660 let depths = [a.depth(), b.depth()];
663 let out_depth = depths[0].max(depths[1]).saturating_add(1);
664 let mut spill_a = match &a {
665 ColumnChunk::Spilled(body, _) => Some(Rc::clone(body)),
666 ColumnChunk::Resident(_, _) => None,
667 };
668 let mut spill_b = match &b {
669 ColumnChunk::Spilled(body, _) => Some(Rc::clone(body)),
670 ColumnChunk::Resident(_, _) => None,
671 };
672 let mut cols = [a.into_column(), b.into_column()];
673 let mut positions = [0usize, 0usize];
674 loop {
675 let mut result: Column<(D, T, R)> = Column::default();
676 let yielded = result.merge_from(&mut cols, &mut positions);
677 if !result.is_empty() {
678 out.push_back(ColumnChunk::Resident(Rc::new(result), out_depth));
679 }
680 if !yielded {
681 break;
682 }
683 }
684 let [col_a, col_b] = &mut cols;
685 for (col, pos, depth, spilled, queue) in [
689 (col_a, positions[0], depths[0], &mut spill_a, in1),
690 (col_b, positions[1], depths[1], &mut spill_b, in2),
691 ] {
692 let len = col.borrow().len();
693 if pos == 0 && len > 0 {
694 let chunk = match spilled.take() {
697 Some(body) => ColumnChunk::Spilled(body, depth),
698 None => ColumnChunk::Resident(Rc::new(std::mem::take(col)), depth),
699 };
700 queue.push_front(chunk.survive_merge());
701 } else if pos < len {
702 let view = col.borrow();
703 let mut rest = <(D, T, R) as Columnar>::Container::default();
704 rest.extend_from_self(view, pos..len);
705 queue.push_front(ColumnChunk::Resident(Rc::new(Column::Typed(rest)), depth));
706 }
707 }
708 }
709
710 fn extract(
714 input: &mut VecDeque<Self>,
715 frontier: AntichainRef<T>,
716 residual: &mut timely::progress::Antichain<T>,
717 keep: &mut VecDeque<Self>,
718 ship: &mut VecDeque<Self>,
719 ) {
720 let Some(chunk) = input.pop_front() else {
721 return;
722 };
723 let (time_lower, time_upper) = chunk.chunk_time_bounds();
729 if time_upper.iter().all(|t| !frontier.less_equal(t)) {
730 ship.push_back(chunk);
731 return;
732 }
733 if time_lower.elements().iter().all(|m| frontier.less_equal(m)) {
734 for m in time_lower.elements() {
737 residual.insert_ref(m);
738 }
739 keep.push_back(chunk);
740 return;
741 }
742 let depth = chunk.depth();
745 let mut col = chunk.into_column();
746 let len = col.borrow().len();
747 let mut pos = 0;
748 let mut keep_col: Column<(D, T, R)> = Column::default();
749 let mut ship_col: Column<(D, T, R)> = Column::default();
750 let cut = |col: &mut Column<(D, T, R)>, queue: &mut VecDeque<Self>, force: bool| {
756 if !col.is_empty() && (force || at_serialized_capacity(&col.borrow())) {
757 queue.push_back(ColumnChunk::Resident(Rc::new(std::mem::take(col)), depth));
758 }
759 };
760 while pos < len {
761 col.extract(&mut pos, frontier, residual, &mut keep_col, &mut ship_col);
762 if pos < len {
763 cut(&mut keep_col, keep, false);
764 cut(&mut ship_col, ship, false);
765 }
766 }
767 cut(&mut keep_col, keep, true);
768 cut(&mut ship_col, ship, true);
769 }
770
771 fn advance(
781 input: &mut VecDeque<Self>,
782 frontier: AntichainRef<T>,
783 done: bool,
784 out: &mut VecDeque<Self>,
785 ) {
786 let Some(front) = input.pop_front() else {
787 return;
788 };
789 let mut depth = front.depth();
792 let mut base = to_typed(front.into_column());
796 {
797 let Column::Typed(base_c) = &mut base else {
798 unreachable!("to_typed returns Typed");
799 };
800 for chunk in input.drain(..) {
801 depth = depth.max(chunk.depth());
802 let col = chunk.into_column();
803 let view = col.borrow();
804 base_c.extend_from_self(view, 0..view.len());
805 }
806 }
807 let view = base.borrow();
808 let total = view.len();
809 if total == 0 {
810 return;
811 }
812 let data = view.0;
813
814 if !done && data.get(0) == data.get(total - 1) {
817 input.push_front(ColumnChunk::Resident(Rc::new(base), depth));
818 return;
819 }
820
821 let end = if done {
824 total
825 } else {
826 let last = data.get(total - 1);
827 let mut end = total - 1;
828 while end > 0 && data.get(end - 1) == last {
829 end -= 1;
830 }
831 end
832 };
833
834 let mut result = <(D, T, R) as Columnar>::Container::default();
835 let mut scratch: Vec<(T, R)> = Vec::new();
837 let mut index = 0;
838 const CUT_CHECK_RECORDS: usize = 1024;
847 let mut records_since_check = 0usize;
848 while index < end {
855 let group_d = data.get(index);
856 scratch.clear();
857 while index < end && data.get(index) == group_d {
858 let (_, t, r) = view.get(index);
859 let mut owned_t = T::into_owned(t);
860 owned_t.advance_by(frontier);
861 scratch.push((owned_t, R::into_owned(r)));
862 index += 1;
863 }
864 scratch.sort_by(|a, b| a.0.cmp(&b.0));
865 let mut run = scratch.drain(..).peekable();
866 while let Some((t, mut r)) = run.next() {
867 while run.peek().is_some_and(|(t2, _)| *t2 == t) {
868 let (_, r2) = run.next().expect("peeked");
869 r.plus_equals(&r2);
870 }
871 if !r.is_zero() {
872 result.0.push(group_d);
873 result.1.push(&t);
874 result.2.push(&r);
875 records_since_check += 1;
876 if records_since_check >= CUT_CHECK_RECORDS {
877 records_since_check = 0;
878 if u64::cast_from(indexed::length_in_words(&result.borrow()))
879 >= u64::cast_from(COMMIT_BYTES / 8)
880 {
881 out.push_back(ColumnChunk::Resident(
882 Rc::new(Column::Typed(std::mem::take(&mut result))),
883 depth,
884 ));
885 }
886 }
887 }
888 }
889 }
890 if !result.is_empty() {
891 out.push_back(ColumnChunk::Resident(Rc::new(Column::Typed(result)), depth));
892 }
893
894 if end < total {
896 let mut carry = <(D, T, R) as Columnar>::Container::default();
897 carry.extend_from_self(view, end..total);
898 input.push_front(ColumnChunk::Resident(Rc::new(Column::Typed(carry)), depth));
899 }
900 }
901
902 fn settle(input: &mut VecDeque<Self>, done: bool, out: &mut VecDeque<Self>) {
908 let mut carry: Option<(Column<(D, T, R)>, u8)> = None;
911 while let Some(chunk) = input.pop_front() {
912 let (rc, depth) = match chunk {
913 spilled @ ColumnChunk::Spilled(_, _) => {
914 if let Some((col, depth)) = carry.take() {
915 out.push_back(ColumnChunk::commit(col, depth));
916 }
917 out.push_back(spilled);
918 continue;
919 }
920 ColumnChunk::Resident(rc, depth) => (rc, depth),
921 };
922 let full = at_commit_size(&rc);
923 if !full && let Some((mut acc, acc_depth)) = carry.take() {
926 let Column::Typed(acc_c) = &mut acc else {
927 unreachable!("carry is always Typed");
928 };
929 let view = rc.borrow();
930 acc_c.extend_from_self(view, 0..view.len());
931 let acc_depth = acc_depth.max(depth);
932 if at_commit_size(&acc) {
933 out.push_back(ColumnChunk::commit(acc, acc_depth));
934 } else {
935 carry = Some((acc, acc_depth));
936 }
937 continue;
938 }
939 if let Some((acc, acc_depth)) = carry.take() {
942 out.push_back(ColumnChunk::commit(acc, acc_depth));
943 }
944 let col = Rc::try_unwrap(rc).unwrap_or_else(|rc| copy_column(&rc));
945 if full {
946 out.push_back(ColumnChunk::commit(col, depth));
947 } else {
948 carry = Some((to_typed(col), depth));
949 }
950 }
951 if let Some((col, depth)) = carry {
952 if done {
953 out.push_back(ColumnChunk::commit(col, depth));
954 } else {
955 input.push_front(ColumnChunk::Resident(Rc::new(col), depth));
956 }
957 }
958 }
959}
960
961fn extract_view_into<'v, 'p, K, V, T, R>(
966 view: BorrowedOf<'v, ((K, V), T, R)>,
967 probes: BorrowedOf<'p, K>,
968 probe_index: &mut usize,
969 staging: &mut <((K, V), T, R) as Columnar>::Container,
970) where
971 K: Columnar,
972 V: Columnar,
973 T: Columnar,
974 R: Columnar,
975 for<'b> columnar::Ref<'b, K>: Copy + Ord,
976{
977 let keys = view.0.0;
978 let len = keys.len();
979 let last = keys.get(len - 1);
980 let count = probes.len();
981 let mut pos = 0;
982 while *probe_index < count {
983 let probe = probes.get(*probe_index);
984 mz_ore::soft_assert_no_log!(
985 *probe_index == 0 || rr::<K>(probes.get(*probe_index - 1)) < rr::<K>(probe),
986 "probe keys must be sorted and deduplicated"
987 );
988 if rr::<K>(probe) > rr::<K>(last) {
989 return;
990 }
991 gallop(len, &mut pos, |i| rr::<K>(keys.get(i)) < rr::<K>(probe));
992 let start = pos;
993 while pos < len && rr::<K>(keys.get(pos)) == rr::<K>(probe) {
994 pos += 1;
995 }
996 staging.extend_from_self(view, start..pos);
997 if rr::<K>(probe) == rr::<K>(last) {
998 return;
999 }
1000 *probe_index += 1;
1001 }
1002}
1003
1004impl<K, V, T, R> UnloadChunk for ColumnChunk<(K, V), T, R>
1005where
1006 K: Columnar,
1007 for<'a> columnar::Ref<'a, K>: Copy + Ord,
1008 V: Columnar,
1009 for<'a> columnar::Ref<'a, V>: Copy + Ord,
1010 T: Columnar + Default + Timestamp + Lattice + Ord,
1011 for<'a> columnar::Ref<'a, T>: Copy + Ord,
1012 R: Columnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
1013{
1014 type Staging = <((K, V), T, R) as Columnar>::Container;
1017
1018 type Probes<'a> = BorrowedOf<'a, K>;
1021
1022 fn probe_count(probes: Self::Probes<'_>) -> usize {
1023 probes.len()
1024 }
1025
1026 fn locate(&self, probes: Self::Probes<'_>, probe_index: usize) -> std::cmp::Ordering {
1027 let probe = probes.get(probe_index);
1028 let (first, last) = self.data_span();
1031 let (first, last) = (first.0, last.0);
1032 if rr::<K>(probe) < rr::<K>(first) {
1033 std::cmp::Ordering::Less
1034 } else if rr::<K>(probe) > rr::<K>(last) {
1035 std::cmp::Ordering::Greater
1036 } else {
1037 std::cmp::Ordering::Equal
1038 }
1039 }
1040
1041 fn extract_into(
1042 &self,
1043 probes: Self::Probes<'_>,
1044 probe_index: &mut usize,
1045 staging: &mut Self::Staging,
1046 ) {
1047 match self {
1048 ColumnChunk::Resident(col, _) => {
1049 extract_view_into::<K, V, T, R>(col.borrow(), probes, probe_index, staging);
1050 }
1051 ColumnChunk::Spilled(body, _) => with_scratch(|scratch| {
1052 body.handle.read_into(scratch);
1058 let view = borrow_words::<((K, V), T, R)>(scratch);
1059 extract_view_into::<K, V, T, R>(view, probes, probe_index, staging);
1060 }),
1061 }
1062 }
1063
1064 fn fetch_into(&self, staging: &mut Self::Staging) {
1065 match self {
1066 ColumnChunk::Resident(col, _) => {
1067 let view = col.borrow();
1068 staging.extend_from_self(view, 0..view.len());
1069 }
1070 ColumnChunk::Spilled(body, _) => with_scratch(|scratch| {
1071 body.handle.read_into(scratch);
1072 let view = borrow_words::<((K, V), T, R)>(scratch);
1073 staging.extend_from_self(view, 0..view.len());
1074 }),
1075 }
1076 }
1077}
1078
1079pub struct UnchunkBuilder<Bu, D: Columnar, T: Columnar, R: Columnar> {
1089 inner: Bu,
1090 _marker: std::marker::PhantomData<(D, T, R)>,
1091}
1092
1093impl<Bu, D, T, R> differential_dataflow::trace::Builder for UnchunkBuilder<Bu, D, T, R>
1094where
1095 Bu: differential_dataflow::trace::Builder<Input = Column<(D, T, R)>>,
1096 D: Columnar + 'static,
1097 T: Columnar + 'static,
1098 R: Columnar + 'static,
1099{
1100 type Input = ColumnChunk<D, T, R>;
1101 type Time = Bu::Time;
1102 type Output = Bu::Output;
1103
1104 fn with_capacity(keys: usize, vals: usize, upds: usize) -> Self {
1105 Self {
1106 inner: Bu::with_capacity(keys, vals, upds),
1107 _marker: std::marker::PhantomData,
1108 }
1109 }
1110
1111 fn push(&mut self, chunk: &mut Self::Input) {
1112 let mut column = std::mem::take(chunk).into_column();
1113 self.inner.push(&mut column);
1114 }
1115
1116 fn done(
1117 self,
1118 description: differential_dataflow::trace::Description<Self::Time>,
1119 ) -> Self::Output {
1120 self.inner.done(description)
1121 }
1122
1123 fn seal(
1124 chain: &mut Vec<Self::Input>,
1125 description: differential_dataflow::trace::Description<Self::Time>,
1126 ) -> Self::Output {
1127 let mut builder = Self::new();
1130 for chunk in chain.iter_mut() {
1131 builder.push(chunk);
1132 }
1133 chain.clear();
1134 builder.done(description)
1135 }
1136}
1137
1138pub struct ChunkChunker<D: Columnar, T: Columnar, R: Columnar> {
1141 inner: ColumnChunker<(D, T, R)>,
1142 staged: ColumnChunk<D, T, R>,
1143}
1144
1145impl<D, T, R> Default for ChunkChunker<D, T, R>
1146where
1147 D: Columnar,
1148 T: Columnar,
1149 R: Columnar,
1150 ColumnChunker<(D, T, R)>: Default,
1151{
1152 fn default() -> Self {
1153 Self {
1154 inner: Default::default(),
1155 staged: Default::default(),
1156 }
1157 }
1158}
1159
1160impl<'a, D, T, R> PushInto<&'a mut Column<(D, T, R)>> for ChunkChunker<D, T, R>
1161where
1162 D: Columnar,
1163 T: Columnar,
1164 R: Columnar,
1165 ColumnChunker<(D, T, R)>: PushInto<&'a mut Column<(D, T, R)>>,
1166{
1167 fn push_into(&mut self, item: &'a mut Column<(D, T, R)>) {
1168 self.inner.push_into(item);
1169 }
1170}
1171
1172impl<D, T, R> ContainerBuilder for ChunkChunker<D, T, R>
1173where
1174 D: Columnar + 'static,
1175 T: Columnar + 'static,
1176 R: Columnar + 'static,
1177 ColumnChunker<(D, T, R)>: ContainerBuilder<Container = Column<(D, T, R)>>,
1178{
1179 type Container = ColumnChunk<D, T, R>;
1180
1181 fn extract(&mut self) -> Option<&mut Self::Container> {
1182 let col = self.inner.extract()?;
1183 self.staged = ColumnChunk::from_column(std::mem::take(col));
1184 Some(&mut self.staged)
1185 }
1186
1187 fn finish(&mut self) -> Option<&mut Self::Container> {
1188 let col = self.inner.finish()?;
1189 self.staged = ColumnChunk::from_column(std::mem::take(col));
1190 Some(&mut self.staged)
1191 }
1192}
1193
1194#[cfg(test)]
1195mod tests {
1196 use differential_dataflow::trace::chunk::{ChunkBatch, ChunkBatcher};
1206 use differential_dataflow::trace::{Batcher, Description};
1207 use mz_ore::pool::Pool;
1208 use proptest::prelude::*;
1209 use timely::container::PushInto;
1210 use timely::progress::Antichain;
1211
1212 use crate::columnar::unload::UnloadBatch;
1213
1214 use super::*;
1215
1216 type Tuple = ((u64, u64), u64, i64);
1217 type TestChunk = ColumnChunk<(u64, u64), u64, i64>;
1218
1219 #[mz_ore::test]
1224 fn lz4_codec_matches_the_previous_extent_framing() {
1225 let body: Vec<u8> = (0..100_000u32).flat_map(|i| i.to_le_bytes()).collect();
1226 let mut stored = Vec::new();
1227 LZ4_CODEC.encode(&body, &mut stored);
1228 assert_eq!(stored, lz4_flex::block::compress_prepend_size(&body));
1229 let mut round = vec![0u8; body.len()];
1230 LZ4_CODEC.decode(&stored, &mut round);
1231 assert_eq!(round, body);
1232 }
1233
1234 #[mz_ore::test]
1235 #[should_panic(expected = "destination must match")]
1236 fn lz4_codec_decode_length_mismatch_panics() {
1237 let mut stored = Vec::new();
1238 LZ4_CODEC.encode(&[7u8; 64], &mut stored);
1239 let mut short = vec![0u8; 32];
1240 LZ4_CODEC.decode(&stored, &mut short);
1241 }
1242
1243 fn consolidate(mut v: Vec<Tuple>) -> Vec<Tuple> {
1246 v.sort();
1247 let mut out: Vec<Tuple> = Vec::new();
1248 for (d, t, r) in v {
1249 if let Some(last) = out.last_mut() {
1250 if last.0 == d && last.1 == t {
1251 last.2 += r;
1252 continue;
1253 }
1254 }
1255 out.push((d, t, r));
1256 }
1257 out.retain(|x| x.2 != 0);
1258 out
1259 }
1260
1261 fn arb_consolidated() -> impl Strategy<Value = Vec<Tuple>> {
1262 prop::collection::vec(((0u64..5, 0u64..5), 0u64..4, -3i64..=3i64), 0..40)
1263 .prop_map(consolidate)
1264 }
1265
1266 fn build_column(v: &[Tuple]) -> Column<Tuple> {
1267 let mut col: Column<Tuple> = Default::default();
1268 for tup in v {
1269 col.push_into(*tup);
1270 }
1271 col
1272 }
1273
1274 fn collect_column(col: &Column<Tuple>) -> Vec<Tuple> {
1275 col.borrow()
1276 .into_index_iter()
1277 .map(|((k, v), t, r)| {
1278 (
1279 (u64::into_owned(k), u64::into_owned(v)),
1280 u64::into_owned(t),
1281 i64::into_owned(r),
1282 )
1283 })
1284 .collect()
1285 }
1286
1287 fn collect_chunks(chunks: impl IntoIterator<Item = TestChunk>) -> Vec<Tuple> {
1288 chunks
1289 .into_iter()
1290 .flat_map(|chunk| collect_column(&chunk.into_column()))
1291 .collect()
1292 }
1293
1294 fn collect_staging(staging: &<Tuple as Columnar>::Container) -> Vec<Tuple> {
1295 staging
1296 .borrow()
1297 .into_index_iter()
1298 .map(|((k, v), t, r)| {
1299 (
1300 (u64::into_owned(k), u64::into_owned(v)),
1301 u64::into_owned(t),
1302 i64::into_owned(r),
1303 )
1304 })
1305 .collect()
1306 }
1307
1308 fn chunked(data: &[Tuple], cuts: &[usize]) -> VecDeque<TestChunk> {
1310 let mut chunks = VecDeque::new();
1311 let mut start = 0;
1312 for cut in cuts {
1313 let end = (start + 1 + cut % 7).min(data.len());
1314 if end > start {
1315 chunks.push_back(ColumnChunk::from_column(build_column(&data[start..end])));
1316 start = end;
1317 }
1318 }
1319 if start < data.len() {
1320 chunks.push_back(ColumnChunk::from_column(build_column(&data[start..])));
1321 }
1322 chunks
1323 }
1324
1325 fn chunked_spilled(data: &[Tuple], cuts: &[usize], pool: &Pool) -> VecDeque<TestChunk> {
1328 chunked(data, cuts)
1329 .into_iter()
1330 .map(|chunk| force_spill(chunk, pool))
1331 .collect()
1332 }
1333
1334 fn body_compressed(chunk: &TestChunk) -> bool {
1338 match chunk {
1339 ColumnChunk::Spilled(body, _) => body.compressed,
1340 ColumnChunk::Resident(_, _) => panic!("chunk must be spilled"),
1341 }
1342 }
1343
1344 fn force_spill(chunk: TestChunk, pool: &Pool) -> TestChunk {
1347 let depth = chunk.depth();
1348 TestChunk::spill_body(chunk.into_column(), pool, depth)
1349 }
1350
1351 fn test_pool() -> Pool {
1355 static POOL: std::sync::OnceLock<Pool> = std::sync::OnceLock::new();
1356 POOL.get_or_init(|| Pool::new().expect("pool creation"))
1357 .clone()
1358 }
1359
1360 proptest! {
1361 #[mz_ore::test]
1364 #[cfg_attr(miri, ignore)]
1365 fn batcher_round_trip(
1366 inputs in prop::collection::vec(arb_consolidated(), 1..6),
1367 cuts in prop::collection::vec(0usize..7, 0..8),
1368 ) {
1369 let mut batcher: ChunkBatcher<TestChunk> = Batcher::new(None, 0);
1370 let mut union = Vec::new();
1371 for input in &inputs {
1372 Extend::extend(&mut union, input.iter().copied());
1373 for chunk in chunked(input, &cuts) {
1374 batcher.push_into(chunk);
1375 }
1376 }
1377 let (sealed, _description) = batcher.seal(Antichain::new());
1379 prop_assert_eq!(collect_chunks(sealed), consolidate(union));
1380 }
1381
1382 #[mz_ore::test]
1385 #[cfg_attr(miri, ignore)]
1386 fn batcher_round_trip_spilled(
1387 inputs in prop::collection::vec(arb_consolidated(), 1..4),
1388 cuts in prop::collection::vec(0usize..7, 0..6),
1389 ) {
1390 let pool = test_pool();
1391 let mut batcher: ChunkBatcher<TestChunk> = Batcher::new(None, 0);
1392 let mut union = Vec::new();
1393 for input in &inputs {
1394 Extend::extend(&mut union, input.iter().copied());
1395 for chunk in chunked_spilled(input, &cuts, &pool) {
1396 batcher.push_into(chunk);
1397 }
1398 }
1399 let (sealed, _description) = batcher.seal(Antichain::new());
1400 prop_assert_eq!(collect_chunks(sealed), consolidate(union));
1401 }
1402
1403 #[mz_ore::test]
1406 #[cfg_attr(miri, ignore)]
1407 fn seal_partitions_by_time(
1408 input in arb_consolidated(),
1409 cuts in prop::collection::vec(0usize..7, 0..8),
1410 upper in 0u64..5,
1411 ) {
1412 let mut batcher: ChunkBatcher<TestChunk> = Batcher::new(None, 0);
1413 for chunk in chunked(&input, &cuts) {
1414 batcher.push_into(chunk);
1415 }
1416 let (shipped, _) = batcher.seal(Antichain::from_elem(upper));
1417 let expected_shipped: Vec<Tuple> =
1418 input.iter().copied().filter(|(_, t, _)| *t < upper).collect();
1419 prop_assert_eq!(collect_chunks(shipped), consolidate(expected_shipped));
1420
1421 let kept_min = input.iter().filter(|(_, t, _)| *t >= upper).map(|(_, t, _)| *t).min();
1422 let frontier = batcher.frontier().to_owned();
1423 prop_assert_eq!(frontier.elements().first().copied(), kept_min);
1424
1425 let (rest, _) = batcher.seal(Antichain::new());
1426 let expected_rest: Vec<Tuple> =
1427 input.iter().copied().filter(|(_, t, _)| *t >= upper).collect();
1428 prop_assert_eq!(collect_chunks(rest), consolidate(expected_rest));
1429 }
1430
1431 #[mz_ore::test]
1435 #[cfg_attr(miri, ignore)]
1436 fn seal_partitions_by_time_spilled(
1437 input in arb_consolidated(),
1438 cuts in prop::collection::vec(0usize..7, 0..8),
1439 upper in 0u64..5,
1440 ) {
1441 let pool = test_pool();
1442 let mut batcher: ChunkBatcher<TestChunk> = Batcher::new(None, 0);
1443 for chunk in chunked_spilled(&input, &cuts, &pool) {
1444 batcher.push_into(chunk);
1445 }
1446 let (shipped, _) = batcher.seal(Antichain::from_elem(upper));
1447 let expected_shipped: Vec<Tuple> =
1448 input.iter().copied().filter(|(_, t, _)| *t < upper).collect();
1449 prop_assert_eq!(collect_chunks(shipped), consolidate(expected_shipped));
1450
1451 let kept_min = input.iter().filter(|(_, t, _)| *t >= upper).map(|(_, t, _)| *t).min();
1452 let frontier = batcher.frontier().to_owned();
1453 prop_assert_eq!(frontier.elements().first().copied(), kept_min);
1454
1455 let (rest, _) = batcher.seal(Antichain::new());
1456 let expected_rest: Vec<Tuple> =
1457 input.iter().copied().filter(|(_, t, _)| *t >= upper).collect();
1458 prop_assert_eq!(collect_chunks(rest), consolidate(expected_rest));
1459 }
1460
1461 #[mz_ore::test]
1464 #[cfg_attr(miri, ignore)]
1465 fn advance_matches_reference(
1466 input in arb_consolidated(),
1467 cuts in prop::collection::vec(0usize..7, 0..8),
1468 frontier_elem in 0u64..5,
1469 ) {
1470 let frontier = Antichain::from_elem(frontier_elem);
1471 let mut chunks = chunked(&input, &cuts);
1472 let mut out = VecDeque::new();
1473 TestChunk::advance(&mut chunks, frontier.borrow(), false, &mut out);
1474 TestChunk::advance(&mut chunks, frontier.borrow(), true, &mut out);
1475 prop_assert!(chunks.is_empty());
1476
1477 let expected = consolidate(
1478 input
1479 .iter()
1480 .map(|&(d, mut t, r)| {
1481 t.advance_by(frontier.borrow());
1482 (d, t, r)
1483 })
1484 .collect(),
1485 );
1486 prop_assert_eq!(collect_chunks(out), expected);
1487 }
1488
1489 #[mz_ore::test]
1492 #[cfg_attr(miri, ignore)]
1493 fn settle_preserves_and_packs(
1494 input in arb_consolidated(),
1495 cuts in prop::collection::vec(0usize..7, 1..8),
1496 ) {
1497 let mut chunks = chunked(&input, &cuts);
1498 let mut out = VecDeque::new();
1499 TestChunk::settle(&mut chunks, true, &mut out);
1500 prop_assert!(chunks.is_empty());
1501 prop_assert!(out.len() <= 1);
1504 prop_assert_eq!(collect_chunks(out), input);
1505 }
1506
1507 #[mz_ore::test]
1511 #[cfg_attr(miri, ignore)]
1512 fn unload_extract_matches_filter(
1513 input in arb_consolidated(),
1514 cuts in prop::collection::vec(0usize..7, 0..8),
1515 probe_keys in prop::collection::btree_set(0u64..6, 0..6),
1516 spill in any::<bool>(),
1517 ) {
1518 prop_assume!(!input.is_empty());
1519 let pool = test_pool();
1520 let chunks: Vec<TestChunk> = if spill {
1521 chunked_spilled(&input, &cuts, &pool).into()
1522 } else {
1523 chunked(&input, &cuts).into()
1524 };
1525 let description = Description::new(
1526 Antichain::from_elem(0u64),
1527 Antichain::new(),
1528 Antichain::from_elem(0u64),
1529 );
1530 let batch = ChunkBatch::new(chunks, description);
1531
1532 let mut probe_col = <u64 as Columnar>::Container::default();
1533 for key in &probe_keys {
1534 probe_col.push(*key);
1535 }
1536 let mut staging = <Tuple as Columnar>::Container::default();
1537 batch.extract_into(probe_col.borrow(), &mut staging);
1538
1539 let expected: Vec<Tuple> = input
1540 .iter()
1541 .copied()
1542 .filter(|((k, _), _, _)| probe_keys.contains(k))
1543 .collect();
1544 prop_assert_eq!(collect_staging(&staging), expected);
1545
1546 let mut staging = <Tuple as Columnar>::Container::default();
1549 batch.fetch_into(&mut staging);
1550 prop_assert_eq!(collect_staging(&staging), input);
1551 }
1552 }
1553
1554 #[mz_ore::test]
1557 fn locate_spans_keys() {
1558 let chunk = ColumnChunk::from_column(build_column(&[
1559 ((2, 0), 0, 1),
1560 ((4, 0), 0, 1),
1561 ((6, 0), 0, 1),
1562 ]));
1563 let mut probe_col = <u64 as Columnar>::Container::default();
1564 for key in [0u64, 2, 3, 6, 9] {
1565 probe_col.push(key);
1566 }
1567 let probes = probe_col.borrow();
1568 use std::cmp::Ordering::*;
1569 let expected = [Less, Equal, Equal, Equal, Greater];
1570 for (index, expected) in expected.iter().enumerate() {
1571 assert_eq!(chunk.locate(probes, index), *expected, "probe {index}");
1572 }
1573 }
1574
1575 fn collect_bounded(chunks: impl IntoIterator<Item = TestChunk>, bound: usize) -> Vec<Tuple> {
1578 let mut collected = Vec::new();
1579 for chunk in chunks {
1580 let col = chunk.into_column();
1581 let bytes = col.length_in_bytes();
1582 assert!(bytes <= bound, "chunk of {bytes} bytes exceeds {bound}");
1583 Extend::extend(&mut collected, collect_column(&col));
1584 }
1585 collected
1586 }
1587
1588 #[mz_ore::test]
1591 #[cfg_attr(miri, ignore)]
1592 fn advance_cuts_large_output() {
1593 let records: Vec<Tuple> = (0..300_000u64).map(|k| ((k, 0), 0, 1)).collect();
1594 let mut input = VecDeque::from([ColumnChunk::from_column(build_column(&records))]);
1595 let frontier = Antichain::from_elem(0u64);
1596 let mut out = VecDeque::new();
1597 TestChunk::advance(&mut input, frontier.borrow(), true, &mut out);
1598 assert!(input.is_empty());
1599 assert!(
1600 out.len() >= 2,
1601 "expected a cut output, got {} chunk(s)",
1602 out.len()
1603 );
1604 assert_eq!(collect_bounded(out, 2 * COMMIT_BYTES), records);
1605 }
1606
1607 #[mz_ore::test]
1610 #[cfg_attr(miri, ignore)]
1611 fn advance_withholds_giant_group() {
1612 let records: Vec<Tuple> = (0..100u64).map(|t| ((7, 7), t, 1)).collect();
1613 let mut input: VecDeque<TestChunk> = VecDeque::new();
1614 for piece in records.chunks(30) {
1615 input.push_back(ColumnChunk::from_column(build_column(piece)));
1616 }
1617 let frontier = Antichain::from_elem(50u64);
1618 let mut out = VecDeque::new();
1619 TestChunk::advance(&mut input, frontier.borrow(), false, &mut out);
1620 assert!(out.is_empty(), "nothing may ship from a single open group");
1621 assert_eq!(input.len(), 1, "the whole input becomes one carry chunk");
1622 TestChunk::advance(&mut input, frontier.borrow(), true, &mut out);
1624 assert!(input.is_empty());
1625 let advanced = records.iter().map(|&(d, t, r)| (d, t.max(50), r)).collect();
1626 assert_eq!(collect_chunks(out), consolidate(advanced));
1627 }
1628
1629 #[mz_ore::test]
1634 #[cfg_attr(miri, ignore)]
1635 fn extract_passes_frontier_disjoint_chunks_through() {
1636 set_spill_override(Some(test_pool()));
1637 let low: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), i % 4, 1)).collect();
1638 let high: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 6 + i % 4, 1)).collect();
1639 let spilled_chunk = |data: &[Tuple]| {
1640 let chunk = TestChunk::commit(build_column(&consolidate(data.to_vec())), 1);
1641 assert!(chunk.is_spilled());
1642 chunk
1643 };
1644
1645 let mut input = VecDeque::from([spilled_chunk(&low), spilled_chunk(&high)]);
1650 let frontier = Antichain::from_elem(5u64);
1651 let mut residual = Antichain::new();
1652 let (mut keep, mut ship) = (VecDeque::new(), VecDeque::new());
1653 while !input.is_empty() {
1654 TestChunk::extract(
1655 &mut input,
1656 frontier.borrow(),
1657 &mut residual,
1658 &mut keep,
1659 &mut ship,
1660 );
1661 }
1662 assert_eq!(ship.len(), 1);
1663 assert!(ship[0].is_spilled(), "shipped whole: body untouched");
1664 assert_eq!(keep.len(), 1);
1665 assert!(keep[0].is_spilled(), "kept whole: body untouched");
1666 assert_eq!(residual, Antichain::from_elem(6));
1667 let shipped = ship.pop_front().unwrap().into_column();
1668 assert_eq!(collect_column(&shipped), consolidate(low));
1669 let kept = keep.pop_front().unwrap().into_column();
1670 assert_eq!(collect_column(&kept), consolidate(high));
1671 set_spill_override(None);
1672 }
1673
1674 #[mz_ore::test]
1677 #[cfg_attr(miri, ignore)]
1678 fn extract_cuts_large_output() {
1679 let records: Vec<Tuple> = (0..300_000u64).map(|k| ((k, 0), k % 2, 1)).collect();
1680 let mut input = VecDeque::from([ColumnChunk::from_column(build_column(&records))]);
1681 let frontier = Antichain::from_elem(1u64);
1682 let mut residual = Antichain::new();
1683 let (mut keep, mut ship) = (VecDeque::new(), VecDeque::new());
1684 while !input.is_empty() {
1685 TestChunk::extract(
1686 &mut input,
1687 frontier.borrow(),
1688 &mut residual,
1689 &mut keep,
1690 &mut ship,
1691 );
1692 }
1693 assert!(
1694 keep.len() >= 2,
1695 "expected a cut keep side, got {} chunk(s)",
1696 keep.len()
1697 );
1698 assert!(
1699 ship.len() >= 2,
1700 "expected a cut ship side, got {} chunk(s)",
1701 ship.len()
1702 );
1703 let kept: Vec<Tuple> = records.iter().copied().filter(|r| r.1 >= 1).collect();
1704 let shipped: Vec<Tuple> = records.iter().copied().filter(|r| r.1 < 1).collect();
1705 assert_eq!(collect_bounded(keep, 2 * COMMIT_BYTES), kept);
1706 assert_eq!(collect_bounded(ship, 2 * COMMIT_BYTES), shipped);
1707 assert_eq!(residual, Antichain::from_elem(1));
1708 }
1709
1710 #[mz_ore::test]
1713 fn locate_uses_resident_bounds() {
1714 let pool = test_pool();
1715 let data: Vec<Tuple> = vec![((2, 0), 0, 1), ((4, 0), 0, 1)];
1716 let chunk = force_spill(ColumnChunk::from_column(build_column(&data)), &pool);
1717
1718 let mut probe_col = <u64 as Columnar>::Container::default();
1719 for key in [1u64, 3, 5] {
1720 probe_col.push(key);
1721 }
1722 let probes = probe_col.borrow();
1723 assert_eq!(chunk.locate(probes, 0), std::cmp::Ordering::Less);
1724 assert_eq!(chunk.locate(probes, 1), std::cmp::Ordering::Equal);
1725 assert_eq!(chunk.locate(probes, 2), std::cmp::Ordering::Greater);
1726 }
1727
1728 #[mz_ore::test]
1731 #[cfg_attr(miri, ignore)] fn spill_round_trip() {
1733 set_spill_override(Some(test_pool()));
1734
1735 let data: Vec<Tuple> = (0..40_000u64)
1736 .map(|i| ((i / 4, i % 4), i % 8, 1i64))
1737 .collect();
1738 let data = consolidate(data);
1739
1740 let column = build_column(&data);
1741 let committed = TestChunk::commit(column, 0);
1742 assert!(committed.is_spilled(), "large body must spill");
1743 assert_eq!(committed.len(), data.len());
1744 assert_eq!(collect_column(&committed.clone().into_column()), data);
1745
1746 let mut batcher: ChunkBatcher<TestChunk> = Batcher::new(None, 0);
1747 for piece in data.chunks(10_000) {
1748 batcher.push_into(ColumnChunk::from_column(build_column(piece)));
1749 }
1750 let (sealed, _) = batcher.seal(Antichain::new());
1751 assert!(
1752 sealed.iter().any(ColumnChunk::is_spilled),
1753 "sealed output should contain spilled chunks",
1754 );
1755 assert_eq!(collect_chunks(sealed), data);
1756
1757 set_spill_override(None);
1758 }
1759
1760 #[mz_ore::test]
1763 #[cfg_attr(miri, ignore)] fn merge_spilled_chains() {
1765 set_spill_override(Some(test_pool()));
1766
1767 let a: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect();
1768 let b: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 2i64)).collect();
1769
1770 let mut in1 = VecDeque::from([TestChunk::commit(build_column(&a), 0)]);
1771 let mut in2 = VecDeque::from([TestChunk::commit(build_column(&b), 0)]);
1772 assert!(in1[0].is_spilled() && in2[0].is_spilled());
1773
1774 let mut out = VecDeque::new();
1775 while !in1.is_empty() && !in2.is_empty() {
1776 TestChunk::merge(&mut in1, &mut in2, &mut out);
1777 }
1778 for tail in in1.drain(..).chain(in2.drain(..)) {
1779 out.push_back(tail);
1780 }
1781
1782 let expected: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 3i64)).collect();
1783 assert_eq!(collect_chunks(out), expected);
1784
1785 set_spill_override(None);
1786 }
1787
1788 #[mz_ore::test]
1791 fn merge_untouched_survivor_stays_spilled() {
1792 let pool = test_pool();
1793 let low: Vec<Tuple> = (0..100u64).map(|i| ((i, 0), 0, 1i64)).collect();
1794 let high: Vec<Tuple> = (1000..1100u64).map(|i| ((i, 0), 0, 1i64)).collect();
1795
1796 let mut in1 = VecDeque::from([force_spill(
1797 ColumnChunk::from_column(build_column(&low)),
1798 &pool,
1799 )]);
1800 let mut in2 = VecDeque::from([force_spill(
1801 ColumnChunk::from_column(build_column(&high)),
1802 &pool,
1803 )]);
1804 let mut out = VecDeque::new();
1805 TestChunk::merge(&mut in1, &mut in2, &mut out);
1806
1807 assert!(in1.is_empty());
1810 assert_eq!(in2.len(), 1);
1811 assert!(in2[0].is_spilled(), "untouched survivor must stay spilled");
1812 let mut all = collect_chunks(out);
1813 Extend::extend(&mut all, collect_chunks(in2.drain(..)));
1814 let mut expected = low;
1815 Extend::extend(&mut expected, high);
1816 assert_eq!(all, expected);
1817 }
1818
1819 #[mz_ore::test]
1823 fn merge_derives_generational_depth() {
1824 let low: Vec<Tuple> = (0..100u64).map(|i| ((i, 0), 0, 1i64)).collect();
1825 let high: Vec<Tuple> = (50..150u64).map(|i| ((i, 0), 0, 1i64)).collect();
1826 let mut in1 = VecDeque::from([ColumnChunk::from_column(build_column(&low))]);
1827 let mut in2 = VecDeque::from([ColumnChunk::from_column(build_column(&high))]);
1828 assert_eq!(in1[0].depth(), 0, "fresh chunks start at depth 0");
1829 let mut out = VecDeque::new();
1830 TestChunk::merge(&mut in1, &mut in2, &mut out);
1831 assert!(!out.is_empty());
1832 for chunk in &out {
1833 assert_eq!(chunk.depth(), 1, "merge output is one past its inputs");
1834 }
1835 assert!(in1.is_empty());
1838 assert_eq!(in2.len(), 1);
1839 assert_eq!(in2[0].depth(), 0, "rewritten survivor keeps its depth");
1840
1841 let mut in1 = VecDeque::from([ColumnChunk::Resident(Rc::new(build_column(&low)), 3)]);
1844 let far: Vec<Tuple> = (1000..1100u64).map(|i| ((i, 0), 0, 1i64)).collect();
1845 let mut in2 = VecDeque::from([ColumnChunk::from_column(build_column(&far))]);
1846 let mut out = VecDeque::new();
1847 TestChunk::merge(&mut in1, &mut in2, &mut out);
1848 assert_eq!(out.len(), 1);
1849 assert_eq!(out[0].depth(), 4, "pass-through ages a generation");
1850 assert_eq!(collect_chunks(out), low);
1851 }
1852
1853 #[mz_ore::test]
1856 fn advance_preserves_depth() {
1857 let data: Vec<Tuple> = (0..100u64).map(|i| ((i, 0), 1, 1i64)).collect();
1858 let mut input = VecDeque::from([
1859 ColumnChunk::Resident(Rc::new(build_column(&data[..50])), 2),
1860 ColumnChunk::Resident(Rc::new(build_column(&data[50..])), 1),
1861 ]);
1862 let frontier = Antichain::from_elem(5u64);
1863 let mut out = VecDeque::new();
1864 TestChunk::advance(&mut input, frontier.borrow(), false, &mut out);
1865 for chunk in out.iter().chain(input.iter()) {
1866 assert_eq!(chunk.depth(), 2);
1867 }
1868 TestChunk::advance(&mut input, frontier.borrow(), true, &mut out);
1869 assert!(input.is_empty());
1870 assert!(!out.is_empty());
1871 for chunk in &out {
1872 assert_eq!(chunk.depth(), 2);
1873 }
1874 }
1875
1876 #[mz_ore::test]
1880 #[cfg_attr(miri, ignore)] fn settle_commits_at_accumulated_depth() {
1882 set_spill_override(Some(test_pool()));
1883 let big: Vec<Tuple> = (0..100_000u64).map(|i| ((i, 0), 0, 1i64)).collect();
1884 let mut input = VecDeque::from([
1885 ColumnChunk::Resident(Rc::new(build_column(&big)), 1),
1886 ColumnChunk::Resident(Rc::new(build_column(&[((0, 0), 0, 1)])), 0),
1887 ColumnChunk::Resident(Rc::new(build_column(&[((1, 0), 0, 1)])), 2),
1888 ]);
1889 let mut out = VecDeque::new();
1890 TestChunk::settle(&mut input, true, &mut out);
1891 assert!(input.is_empty());
1892 assert_eq!(out.len(), 2);
1893 assert!(out[0].is_spilled(), "large commit must spill");
1894 assert_eq!(out[0].depth(), 1, "sole commit keeps its depth");
1895 assert!(!out[1].is_spilled(), "small commit stays resident");
1896 assert_eq!(out[1].depth(), 2, "coalesced commit takes the max depth");
1897 set_spill_override(None);
1898 }
1899
1900 #[mz_ore::test]
1904 #[cfg_attr(miri, ignore)] fn settle_carry_commits_at_target() {
1906 let chunk_rows = u64::cast_from(1_500_000usize / 32);
1910 let mut input: VecDeque<TestChunk> = (0..4u64)
1911 .map(|c| {
1912 let data: Vec<Tuple> = (0..chunk_rows)
1913 .map(|i| ((c * chunk_rows + i, 0), 0, 1i64))
1914 .collect();
1915 ColumnChunk::from_column(build_column(&data))
1916 })
1917 .collect();
1918 let mut out = VecDeque::new();
1919 TestChunk::settle(&mut input, true, &mut out);
1920 assert!(out.len() < 4, "nothing coalesced");
1923 for chunk in &out {
1924 let col = chunk.clone().into_column();
1925 assert!(
1926 col.length_in_bytes() < 2 * COMMIT_BYTES,
1927 "settled chunk of {} bytes exceeds twice the commit target",
1928 col.length_in_bytes(),
1929 );
1930 }
1931 assert_eq!(
1932 collect_chunks(out).len(),
1933 usize::try_from(4 * chunk_rows).unwrap(),
1934 );
1935 }
1936
1937 #[mz_ore::test]
1938 fn small_chunks_stay_resident() {
1939 set_spill_override(Some(test_pool()));
1940 let committed = TestChunk::commit(build_column(&[((1, 1), 0, 1)]), 0);
1941 assert!(!committed.is_spilled());
1942 set_spill_override(None);
1943 }
1944
1945 fn column_at_spill_floor() -> (Column<Tuple>, u64) {
1948 let mut col: Column<Tuple> = Column::default();
1949 let mut n = 0u64;
1950 while col.length_in_bytes() < SPILL_MIN_BYTES {
1951 col.push_into(((n, n), 0, 1));
1952 n += 1;
1953 }
1954 (col, n)
1955 }
1956
1957 #[mz_ore::test]
1960 fn spill_floor_boundary() {
1961 set_spill_override(Some(test_pool()));
1962 let (col, n) = column_at_spill_floor();
1963 let mut under: Column<Tuple> = Column::default();
1964 for m in 0..n - 1 {
1965 under.push_into(((m, m), 0, 1));
1966 }
1967 assert!(under.length_in_bytes() < SPILL_MIN_BYTES);
1968 assert!(!TestChunk::commit(under, 0).is_spilled());
1969 assert!(TestChunk::commit(col, 0).is_spilled());
1970 set_spill_override(None);
1971 }
1972
1973 #[mz_ore::test]
1977 fn spill_codec_depth_floor() {
1978 set_spill_override(Some(test_pool()));
1979 set_compress_min_depth_override(Some(2));
1980 let codec_name = |depth: u8| {
1984 let (codec, compressed) = codec_for_depth(depth);
1985 let name = format!("{:?}", codec);
1986 assert_eq!(compressed, name == "Lz4Codec", "flag tracks the codec");
1987 name
1988 };
1989 assert_eq!(codec_name(0), "IdentityCodec");
1990 assert_eq!(codec_name(1), "IdentityCodec");
1991 assert_eq!(codec_name(2), "Lz4Codec");
1992 assert_eq!(codec_name(u8::MAX), "Lz4Codec");
1993
1994 let data: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect();
1995 let data = consolidate(data);
1996 let column = build_column(&data);
1997 for depth in [0u8, 1, 2, 3] {
1998 let chunk = TestChunk::commit(column.clone(), depth);
1999 assert!(chunk.is_spilled(), "depth {depth} must spill");
2000 assert_eq!(collect_column(&chunk.into_column()), data);
2001 }
2002 set_spill_override(None);
2003 set_compress_min_depth_override(None);
2004
2005 set_compress_min_depth_override(Some(DEFAULT_COMPRESS_MIN_DEPTH));
2007 assert_eq!(codec_name(0), "IdentityCodec");
2008 assert_eq!(codec_name(1), "Lz4Codec");
2009 set_compress_min_depth_override(None);
2010 }
2011
2012 #[mz_ore::test]
2018 #[cfg_attr(miri, ignore)] fn merge_survivor_crosses_compression_floor() {
2020 set_spill_override(Some(test_pool()));
2021 set_compress_min_depth_override(Some(1));
2022
2023 let low = consolidate((0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2024 let far = consolidate((100_000..120_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2025 let fresh_far = || VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
2026
2027 let mut in1 = VecDeque::from([TestChunk::commit(build_column(&low), 0)]);
2030 let mut in2 = fresh_far();
2031 assert!(in1[0].is_spilled() && in2[0].is_spilled());
2032 assert!(
2033 !body_compressed(&in1[0]),
2034 "a fresh body below the floor is identity coded"
2035 );
2036
2037 let mut out = VecDeque::new();
2038 TestChunk::merge(&mut in1, &mut in2, &mut out);
2039 assert_eq!(out.len(), 1);
2040 let survived = out.pop_front().expect("the lower front passes through");
2041 assert_eq!(survived.depth(), 1, "survival ages across the floor");
2042 assert!(
2043 survived.is_spilled(),
2044 "the crossing re-spills, it does not evict"
2045 );
2046 assert!(
2047 body_compressed(&survived),
2048 "the survivor is re-spilled under the compressing codec"
2049 );
2050
2051 let mut in1 = VecDeque::from([survived]);
2054 let mut in2 = fresh_far();
2055 let mut out = VecDeque::new();
2056 TestChunk::merge(&mut in1, &mut in2, &mut out);
2057 assert_eq!(out.len(), 1);
2058 assert_eq!(out[0].depth(), 2, "an aged survivor keeps aging");
2059 assert!(out[0].is_spilled());
2060 assert_eq!(
2061 collect_chunks(out),
2062 low,
2063 "the body reads back intact across both survivals"
2064 );
2065
2066 set_spill_override(None);
2067 set_compress_min_depth_override(None);
2068 }
2069
2070 #[mz_ore::test]
2077 #[cfg_attr(miri, ignore)] fn merge_survivor_ages_while_shared() {
2079 set_spill_override(Some(test_pool()));
2080 set_compress_min_depth_override(Some(1));
2081
2082 let low = consolidate((0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2083 let far = consolidate((100_000..120_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2084
2085 let source = TestChunk::commit(build_column(&low), 0);
2088 let ColumnChunk::Spilled(source_body, 0) = &source else {
2089 panic!("a fresh commit above the spill floor is spilled at depth 0");
2090 };
2091 let source_body = Rc::clone(source_body);
2092
2093 let mut in1 = VecDeque::from([source.clone()]);
2094 let mut in2 = VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
2095 let mut out = VecDeque::new();
2096 TestChunk::merge(&mut in1, &mut in2, &mut out);
2097
2098 assert_eq!(out.len(), 1);
2099 assert_eq!(out[0].depth(), 1, "a shared body ages all the same");
2100 let ColumnChunk::Spilled(survived_body, _) = &out[0] else {
2101 panic!("the survivor stays spilled");
2102 };
2103 assert!(
2104 Rc::ptr_eq(&source_body, survived_body),
2105 "a shared body is aged in place, not re-spilled"
2106 );
2107 assert_eq!(source.depth(), 0, "the other holder is left as it was");
2108
2109 let mut in1 = VecDeque::from([out.pop_front().expect("survivor observed above")]);
2112 let mut in2 = VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
2113 let mut out = VecDeque::new();
2114 TestChunk::merge(&mut in1, &mut in2, &mut out);
2115 assert_eq!(out.len(), 1);
2116 assert_eq!(out[0].depth(), 2, "aging past the floor is not pinned");
2117 assert_eq!(collect_chunks(out), low);
2118
2119 set_spill_override(None);
2120 set_compress_min_depth_override(None);
2121 }
2122
2123 #[mz_ore::test]
2129 #[cfg_attr(miri, ignore)] fn survive_merge_retries_missed_migrations() {
2131 let low = consolidate((0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2132 let far = consolidate((100_000..120_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2133
2134 let survive = |chunk: TestChunk| {
2137 let mut in1 = VecDeque::from([chunk]);
2138 let mut in2 = VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
2139 let mut out = VecDeque::new();
2140 TestChunk::merge(&mut in1, &mut in2, &mut out);
2141 out.pop_front().expect("the lower front passes through")
2142 };
2143
2144 set_spill_override(Some(test_pool()));
2147 set_compress_min_depth_override(Some(1));
2148 let chunk = TestChunk::commit(build_column(&low), 0);
2149 assert!(!body_compressed(&chunk));
2150 set_spill_override(None);
2151 let chunk = survive(chunk);
2152 assert_eq!(chunk.depth(), 1, "aging does not need a pool");
2153 assert!(!body_compressed(&chunk), "no pool, no migration");
2154 set_spill_override(Some(test_pool()));
2155 let chunk = survive(chunk);
2156 assert!(
2157 body_compressed(&chunk),
2158 "the migration retries once a pool is back"
2159 );
2160
2161 let chunk = TestChunk::commit(build_column(&low), 0);
2164 let held = chunk.clone();
2165 let chunk = survive(chunk);
2166 assert!(!body_compressed(&chunk), "shared, so not migrated");
2167 drop(held);
2168 let chunk = survive(chunk);
2169 assert!(
2170 body_compressed(&chunk),
2171 "the migration retries once the body is unshared"
2172 );
2173
2174 set_compress_min_depth_override(Some(8));
2179 let chunk = TestChunk::commit(build_column(&low), 3);
2180 assert!(!body_compressed(&chunk));
2181 set_compress_min_depth_override(Some(1));
2182 let chunk = survive(chunk);
2183 assert_eq!(chunk.depth(), 4);
2184 assert!(
2185 body_compressed(&chunk),
2186 "lowering the floor migrates bodies already past it"
2187 );
2188
2189 set_spill_override(None);
2190 set_compress_min_depth_override(None);
2191 }
2192
2193 #[mz_ore::test]
2197 #[cfg_attr(miri, ignore)]
2198 fn spill_gates_compose_as_or() {
2199 let installed =
2200 crate::pool_config::apply_pool_config(crate::pool_config::PoolPagerConfig {
2201 budget_bytes: 32 << 20,
2202 spill_threads: 1,
2203 eager_backing: false,
2204 rss_target_bytes: 16 << 20,
2205 });
2206 assert!(installed, "pool reservation failed");
2207 let (col, _) = column_at_spill_floor();
2209 let commit = |col: &Column<Tuple>| TestChunk::commit(col.clone(), 0).is_spilled();
2210
2211 assert!(!commit(&col), "both gates off");
2212 set_storage_spill_enabled(true);
2213 assert!(commit(&col), "the storage gate alone spills");
2214 set_compute_spill_enabled(false);
2215 assert!(
2216 commit(&col),
2217 "the compute setter must not clobber the storage gate"
2218 );
2219 set_compute_spill_enabled(true);
2220 set_storage_spill_enabled(false);
2221 assert!(commit(&col), "the compute gate alone spills");
2222 set_compute_spill_enabled(false);
2223 assert!(!commit(&col), "both gates off again");
2224 set_compress_min_depth_override(None);
2225 }
2226
2227 #[mz_ore::test]
2230 fn spill_align_round_trip() {
2231 let pool = test_pool();
2232 let data: Vec<Tuple> = (0..64u64).map(|k| ((k, k), 0, 1)).collect();
2233 let spilled = force_spill(ColumnChunk::from_column(build_column(&data)), &pool);
2234 let column = spilled.into_column();
2235 let Column::Align(words) = &column else {
2236 panic!("a spilled body reads back as Column::Align");
2237 };
2238 let words = words.clone();
2239 let respilled = force_spill(ColumnChunk::from_column(column), &pool);
2240 let reread = respilled.into_column();
2241 let Column::Align(words2) = &reread else {
2242 panic!("a spilled body reads back as Column::Align");
2243 };
2244 assert_eq!(&words, words2, "byte-identical round trip");
2245 assert_eq!(collect_column(&reread), data);
2246 }
2247
2248 #[mz_ore::test]
2249 fn merge_depth_saturates() {
2250 let a = ColumnChunk::Resident(
2251 Rc::new(build_column(&[((1, 0), 0, 1), ((3, 0), 0, 1)])),
2252 u8::MAX,
2253 );
2254 let b = ColumnChunk::Resident(
2255 Rc::new(build_column(&[((2, 0), 0, 1), ((4, 0), 0, 1)])),
2256 u8::MAX,
2257 );
2258 let mut in1 = VecDeque::from([a]);
2259 let mut in2 = VecDeque::from([b]);
2260 let mut out = VecDeque::new();
2261 TestChunk::merge(&mut in1, &mut in2, &mut out);
2262 for chunk in out.iter().chain(in1.iter()).chain(in2.iter()) {
2263 assert_eq!(chunk.depth(), u8::MAX, "depth saturates");
2264 }
2265 }
2266
2267 #[mz_ore::test]
2268 fn into_column_copies_shared_resident() {
2269 let data: Vec<Tuple> = vec![((1, 1), 0, 1), ((2, 2), 0, 1)];
2270 let a = ColumnChunk::from_column(build_column(&data));
2271 let b = a.clone();
2272 assert_eq!(collect_column(&a.into_column()), data);
2273 assert_eq!(collect_column(&b.into_column()), data);
2274 }
2275}