1use std::borrow::Cow;
41#[cfg(test)]
42use std::cell::Cell;
43use std::cell::RefCell;
44use std::collections::VecDeque;
45use std::marker::PhantomData;
46use std::rc::Rc;
47use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
48use std::sync::{Arc, OnceLock};
49
50use columnar::bytes::indexed;
51use columnar::{Borrow, BorrowedOf, Columnar, Container as _, Index, Len, Push as _};
52use differential_dataflow::difference::Semigroup;
53use differential_dataflow::lattice::Lattice;
54use differential_dataflow::trace::chunk::Chunk;
55use mz_ore::cast::CastFrom;
56use mz_ore::pool::{ChunkHandle, ChunkHints, ExtentCodec, IDENTITY_CODEC, Pool, TypedBody};
57use smallvec::SmallVec;
58use timely::Accountable;
59use timely::PartialOrder;
60use timely::container::{ContainerBuilder, PushInto};
61use timely::progress::Timestamp;
62use timely::progress::frontier::{Antichain, AntichainRef};
63
64use crate::columnar::batcher::{ColumnChunker, gallop};
65use crate::columnar::body::{ColumnBody, borrow_words};
66use crate::columnar::unload::UnloadChunk;
67use crate::columnar::{Column, at_serialized_capacity};
68
69pub mod metrics;
70
71static COMPUTE_SPILL_ENABLED: AtomicBool = AtomicBool::new(false);
73
74static STORAGE_SPILL_ENABLED: AtomicBool = AtomicBool::new(false);
76
77static SINK_SPILL_ENABLED: AtomicBool = AtomicBool::new(false);
80
81thread_local! {
82 static SPILL_OVERRIDE: RefCell<Option<Pool>> = const { RefCell::new(None) };
86
87 #[cfg(test)]
91 static COMPRESS_MIN_DEPTH_OVERRIDE: Cell<Option<u8>> = const { Cell::new(None) };
92
93 static READ_SCRATCH: RefCell<Vec<u64>> = const { RefCell::new(Vec::new()) };
95}
96
97pub trait SpillGate: 'static {
106 fn enabled() -> bool;
108}
109
110pub enum ComputeSpill {}
112
113impl SpillGate for ComputeSpill {
114 fn enabled() -> bool {
115 COMPUTE_SPILL_ENABLED.load(Ordering::Relaxed)
116 }
117}
118
119pub enum StorageSpill {}
122
123impl SpillGate for StorageSpill {
124 fn enabled() -> bool {
125 STORAGE_SPILL_ENABLED.load(Ordering::Relaxed)
126 }
127}
128
129pub fn set_compute_spill_enabled(enabled: bool) {
132 COMPUTE_SPILL_ENABLED.store(enabled, Ordering::Relaxed);
133}
134
135pub fn set_storage_spill_enabled(enabled: bool) {
138 STORAGE_SPILL_ENABLED.store(enabled, Ordering::Relaxed);
139}
140
141pub fn set_sink_spill_enabled(enabled: bool) {
146 SINK_SPILL_ENABLED.store(enabled, Ordering::Relaxed);
147}
148
149pub fn set_spill_override(pool: Option<Pool>) {
153 SPILL_OVERRIDE.with(|cell| *cell.borrow_mut() = pool);
154}
155
156static COMPRESS_MIN_DEPTH: AtomicU8 = AtomicU8::new(DEFAULT_COMPRESS_MIN_DEPTH);
159
160pub fn set_compress_min_depth(depth: u8) {
179 COMPRESS_MIN_DEPTH.store(depth, Ordering::Relaxed);
180}
181
182#[cfg(test)]
186pub fn set_compress_min_depth_override(depth: Option<u8>) {
187 COMPRESS_MIN_DEPTH_OVERRIDE.with(|cell| cell.set(depth));
188}
189
190fn compress_min_depth() -> u8 {
192 #[cfg(test)]
193 if let Some(depth) = COMPRESS_MIN_DEPTH_OVERRIDE.with(|cell| cell.get()) {
194 return depth;
195 }
196 COMPRESS_MIN_DEPTH.load(Ordering::Relaxed)
197}
198
199fn codec_for_depth(depth: u8) -> (&'static dyn ExtentCodec, bool) {
204 if depth < compress_min_depth() {
205 (&IDENTITY_CODEC, false)
206 } else {
207 (&LZ4_CODEC, true)
208 }
209}
210
211fn spill_target(enabled: bool, len_bytes: usize) -> Option<Pool> {
217 resolve_pool(enabled).filter(|_| len_bytes >= SPILL_MIN_BYTES)
218}
219
220fn resolve_pool(enabled: bool) -> Option<Pool> {
222 if let Some(pool) = SPILL_OVERRIDE.with(|cell| cell.borrow().clone()) {
223 return Some(pool);
224 }
225 if enabled {
226 crate::pool_config::active_pool()
227 } else {
228 None
229 }
230}
231
232const SCRATCH_RETAIN_WORDS: usize = 1 << 18;
236
237fn with_scratch<Out>(f: impl FnOnce(&mut Vec<u64>) -> Out) -> Out {
239 READ_SCRATCH.with(|cell| {
240 let mut scratch = cell.take();
241 scratch.clear();
242 let out = f(&mut scratch);
243 if scratch.capacity() > SCRATCH_RETAIN_WORDS {
244 scratch.clear();
245 scratch.shrink_to_fit();
246 }
247 cell.replace(scratch);
248 out
249 })
250}
251
252const COMMIT_BYTES: usize = 2 << 20;
255
256const SPILL_MIN_BYTES: usize = 64 << 10;
265
266const DEFAULT_COMPRESS_MIN_DEPTH: u8 = 1;
275
276fn cut_records(len: usize, bytes: usize, space: usize) -> usize {
282 let space = space.saturating_sub(COMMIT_BYTES / 20);
283 (len.saturating_mul(space) / bytes.max(1)).min(len)
284}
285
286fn at_commit_size<C: Columnar>(body: &ColumnBody<C>) -> bool {
288 body.length_in_bytes() >= COMMIT_BYTES - COMMIT_BYTES / 10
289}
290
291#[inline(always)]
295fn rr<'b, 'a: 'b, C: Columnar>(item: columnar::Ref<'a, C>) -> columnar::Ref<'b, C> {
296 columnar::ContainerOf::<C>::reborrow_ref(item)
297}
298
299pub struct SpilledBody<D: Columnar, T, G> {
305 records: usize,
307 fences: D::Container,
311 time_bounds: Arc<OnceLock<TimeBounds<T>>>,
317 compressed: bool,
324 len_bytes: usize,
329 handle: ChunkHandle,
331 gate: PhantomData<G>,
334}
335
336pub enum ColumnChunk<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> {
355 Resident(Rc<ColumnBody<(D, T, R)>>, u8),
357 Spilled(Rc<SpilledBody<D, T, G>>, u8),
359}
360
361impl<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> Clone for ColumnChunk<D, T, R, G> {
362 fn clone(&self) -> Self {
363 match self {
364 ColumnChunk::Resident(col, depth) => ColumnChunk::Resident(Rc::clone(col), *depth),
365 ColumnChunk::Spilled(body, depth) => ColumnChunk::Spilled(Rc::clone(body), *depth),
366 }
367 }
368}
369
370impl<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> Default for ColumnChunk<D, T, R, G> {
371 fn default() -> Self {
372 ColumnChunk::Resident(Rc::new(ColumnBody::default()), 0)
373 }
374}
375
376impl<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> Accountable for ColumnChunk<D, T, R, G> {
377 fn record_count(&self) -> i64 {
378 i64::try_from(self.records()).expect("record count fits i64")
379 }
380}
381
382impl<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> ColumnChunk<D, T, R, G> {
383 pub fn from_body(body: ColumnBody<(D, T, R)>) -> Self {
386 mz_ore::soft_assert_no_log!(!body.is_empty(), "chunks must be non-empty");
387 ColumnChunk::Resident(Rc::new(body), 0)
388 }
389
390 fn push_bounded(body: ColumnBody<(D, T, R)>, depth: u8, out: &mut VecDeque<Self>) {
394 let len = body.len();
395 let bytes = body.length_in_bytes();
396 if len <= 1 || bytes <= COMMIT_BYTES {
397 if len > 0 {
398 out.push_back(Self::Resident(Rc::new(body), depth));
399 }
400 return;
401 }
402 Self::push_cuts(
403 &body,
404 cut_records(len, bytes, COMMIT_BYTES).max(1),
405 depth,
406 out,
407 );
408 }
409
410 fn push_cuts(
417 body: &ColumnBody<(D, T, R)>,
418 records: usize,
419 depth: u8,
420 out: &mut VecDeque<Self>,
421 ) {
422 let view = body.borrow();
423 Self::push_range(view, 0..view.len(), records, depth, out);
424 }
425
426 fn push_range(
431 view: BorrowedOf<'_, (D, T, R)>,
432 range: std::ops::Range<usize>,
433 records: usize,
434 depth: u8,
435 out: &mut VecDeque<Self>,
436 ) {
437 let mut start = range.start;
438 while start < range.end {
439 let end = range.end.min(start + records);
440 let mut part = <(D, T, R) as Columnar>::Container::default();
441 part.extend_from_self(view, start..end);
442 let part = ColumnBody::Typed(part);
443 if end - start > 1 && part.length_in_bytes() > COMMIT_BYTES {
444 drop(part);
445 Self::push_range(view, start..end, (end - start).div_ceil(2), depth, out);
446 } else {
447 out.push_back(Self::Resident(Rc::new(part), depth));
448 }
449 start = end;
450 }
451 }
452
453 pub fn into_body(self) -> ColumnBody<(D, T, R)> {
456 match self {
457 ColumnChunk::Resident(col, _) => {
458 Rc::try_unwrap(col).unwrap_or_else(|shared| shared.duplicate())
459 }
460 ColumnChunk::Spilled(body, _) => match Rc::try_unwrap(body) {
461 Ok(spilled) => match spilled.handle.take_typed() {
462 Ok(typed) => {
463 typed
464 .into_any()
465 .downcast::<PooledBody<D, T, R>>()
466 .expect("chunk bodies are pooled as `PooledBody`")
467 .body
468 }
469 Err(handle) => read_words(&handle),
470 },
471 Err(shared) => shared.handle.with_typed(|typed| match typed {
474 Some(typed) => pooled::<D, T, R>(typed).body.duplicate(),
475 None => read_words(&shared.handle),
476 }),
477 },
478 }
479 }
480
481 pub fn with_body<F, X>(&self, f: F) -> X
484 where
485 F: FnOnce(&ColumnBody<(D, T, R)>) -> X,
486 {
487 match self {
488 ColumnChunk::Resident(col, _) => f(col),
489 ColumnChunk::Spilled(body, _) => {
490 let mut f = Some(f);
491 let borrowed = body.handle.with_typed(|typed| {
492 typed.map(|typed| f.take().expect("unused")(&pooled::<D, T, R>(typed).body))
493 });
494 match borrowed {
495 Some(out) => out,
496 None => f.take().expect("unused")(&read_words(&body.handle)),
497 }
498 }
499 }
500 }
501
502 pub fn is_spilled(&self) -> bool {
504 matches!(self, ColumnChunk::Spilled(_, _))
505 }
506
507 fn records(&self) -> usize {
509 match self {
510 ColumnChunk::Resident(col, _) => col.len(),
511 ColumnChunk::Spilled(body, _) => body.records,
512 }
513 }
514
515 fn depth(&self) -> u8 {
517 match self {
518 ColumnChunk::Resident(_, depth) | ColumnChunk::Spilled(_, depth) => *depth,
519 }
520 }
521
522 fn data_span(&self) -> (columnar::Ref<'_, D>, columnar::Ref<'_, D>) {
524 match self {
525 ColumnChunk::Resident(col, _) => {
526 let data = col.borrow().0;
527 (data.get(0), data.get(data.len() - 1))
528 }
529 ColumnChunk::Spilled(body, _) => {
530 let fences = body.fences.borrow();
531 (fences.get(0), fences.get(1))
532 }
533 }
534 }
535
536 fn commit(body: ColumnBody<(D, T, R)>, depth: u8) -> Self
540 where
541 D: PoolColumnar,
542 T: Timestamp + PoolColumnar,
543 R: PoolColumnar,
544 {
545 let len_bytes = body.length_in_bytes();
546 metrics::record(metrics::Stage::Commit, body.len(), len_bytes);
547 mz_ore::soft_assert_no_log!(!body.is_empty(), "chunks must be non-empty");
548 match spill_target(G::enabled(), len_bytes) {
549 Some(pool) => Self::spill_body(body, &pool, depth),
550 None => ColumnChunk::Resident(Rc::new(body), depth),
551 }
552 }
553
554 fn spill_body(body: ColumnBody<(D, T, R)>, pool: &Pool, depth: u8) -> Self
562 where
563 D: PoolColumnar,
564 T: Timestamp + PoolColumnar,
565 R: PoolColumnar,
566 {
567 let (codec, compressed) = codec_for_depth(depth);
568 let len_bytes = body.length_in_bytes();
569 let view = body.borrow();
570 let records = view.len();
571 let mut fences = D::Container::default();
572 fences.push(view.0.get(0));
573 fences.push(view.0.get(records - 1));
574 let time_bounds = Arc::new(OnceLock::new());
575 let pooled = PooledBody {
576 body,
577 bounds: Arc::clone(&time_bounds),
578 compute: |body| {
579 let (lower, upper) = Self::time_bounds(body);
580 (lower, upper.into())
581 },
582 };
583 let handle = pool.insert_typed(Box::new(pooled), ChunkHints { depth }, codec);
584 ColumnChunk::Spilled(
585 Rc::new(SpilledBody {
586 records,
587 fences,
588 time_bounds,
589 compressed,
590 len_bytes,
591 handle,
592 gate: PhantomData,
593 }),
594 depth,
595 )
596 }
597
598 fn survive_merge(self) -> Self
614 where
615 D: PoolColumnar,
616 T: Timestamp + PoolColumnar,
617 R: PoolColumnar,
618 {
619 let depth = self.depth().saturating_add(1);
620 match self {
621 ColumnChunk::Resident(col, _) => ColumnChunk::Resident(col, depth),
622 ColumnChunk::Spilled(body, was) => {
623 let migrate = !body.compressed && depth >= compress_min_depth();
624 if !migrate || Rc::strong_count(&body) > 1 {
625 return ColumnChunk::Spilled(body, depth);
626 }
627 match resolve_pool(G::enabled()) {
628 Some(pool) => {
629 let body = ColumnChunk::Spilled(body, was).into_body();
630 Self::spill_body(body, &pool, depth)
631 }
632 None => ColumnChunk::Spilled(body, depth),
633 }
634 }
635 }
636 }
637
638 fn chunk_time_bounds(&self) -> (Cow<'_, Antichain<T>>, Cow<'_, [T]>)
643 where
644 T: Timestamp,
645 {
646 match self {
647 ColumnChunk::Resident(col, _) => {
648 let (lower, upper) = Self::time_bounds(col);
649 (Cow::Owned(lower), Cow::Owned(upper))
650 }
651 ColumnChunk::Spilled(body, _) => {
652 if body.time_bounds.get().is_none() {
657 let filled = body.handle.with_typed(|typed| {
658 typed
659 .map(|typed| pooled::<D, T, R>(typed).bounds())
660 .is_some()
661 });
662 if !filled {
665 let (lower, upper) = Self::time_bounds(&read_words(&body.handle));
666 let _ = body.time_bounds.set((lower, upper.into()));
667 }
668 }
669 let (lower, upper) = body.time_bounds.get().expect("time bounds filled above");
670 (Cow::Borrowed(lower), Cow::Borrowed(&upper[..]))
671 }
672 }
673 }
674
675 fn time_bounds(body: &ColumnBody<(D, T, R)>) -> (Antichain<T>, Vec<T>)
680 where
681 T: Timestamp,
682 {
683 let (_, times, _) = body.borrow();
684 let mut lower = Antichain::new();
685 let mut upper: Vec<T> = Vec::new();
686 let mut time = T::minimum();
690 for i in 0..times.len() {
691 time.copy_from(rr::<T>(times.get(i)));
692 if !upper.iter().any(|u| PartialOrder::less_equal(&time, u)) {
693 upper.retain(|u| !PartialOrder::less_equal(u, &time));
694 upper.push(time.clone());
695 }
696 lower.insert_ref(&time);
697 }
698 (lower, upper)
699 }
700}
701
702#[derive(Debug)]
707pub struct Lz4Codec;
708
709pub static LZ4_CODEC: Lz4Codec = Lz4Codec;
712
713impl ExtentCodec for Lz4Codec {
714 fn encode(&self, body: &[u8], out: &mut Vec<u8>) {
715 let max_out = lz4_flex::block::get_maximum_output_size(body.len());
716 out.resize(4 + max_out, 0);
717 let len = u32::try_from(body.len()).expect("chunk bodies are bounded by the size classes");
718 out[..4].copy_from_slice(&len.to_le_bytes());
719 let compressed = lz4_flex::block::compress_into(body, &mut out[4..])
720 .expect("output sized to the maximum");
721 out.truncate(4 + compressed);
722 }
723
724 fn decode(&self, stored: &[u8], body: &mut [u8]) {
725 let prefix: [u8; 4] = stored[..4].try_into().expect("prefix length");
726 let len = usize::try_from(u32::from_le_bytes(prefix)).expect("length fits usize");
727 assert_eq!(
728 len,
729 body.len(),
730 "destination must match the encoded body length"
731 );
732 let written = lz4_flex::block::decompress_into(&stored[4..], body)
733 .expect("stored bytes hold a valid lz4 block");
734 assert_eq!(written, body.len(), "decoded length mismatch");
735 }
736}
737
738type TimeBounds<T> = (Antichain<T>, SmallVec<[T; 1]>);
741
742pub trait PoolColumnar: Columnar<Container: Send> + Send + Sync {}
746
747impl<C: Columnar<Container: Send> + Send + Sync> PoolColumnar for C {}
748
749struct PooledBody<D: Columnar, T: Columnar, R: Columnar> {
757 body: ColumnBody<(D, T, R)>,
758 bounds: Arc<OnceLock<TimeBounds<T>>>,
759 compute: fn(&ColumnBody<(D, T, R)>) -> TimeBounds<T>,
760}
761
762impl<D: Columnar, T: Columnar, R: Columnar> PooledBody<D, T, R> {
763 fn bounds(&self) -> &TimeBounds<T> {
764 self.bounds.get_or_init(|| (self.compute)(&self.body))
765 }
766}
767
768impl<D: Columnar, T: Columnar, R: Columnar> std::fmt::Debug for PooledBody<D, T, R> {
769 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
770 f.debug_struct("PooledBody")
771 .field("len_bytes", &self.body.length_in_bytes())
772 .finish()
773 }
774}
775
776impl<D: PoolColumnar + 'static, T: PoolColumnar + 'static, R: PoolColumnar + 'static> TypedBody
777 for PooledBody<D, T, R>
778{
779 fn words(&self) -> usize {
780 self.body.length_in_bytes() / 8
781 }
782
783 fn serialize_into(&self, dst: &mut [u64]) {
784 self.bounds();
785 let len_bytes = self.body.length_in_bytes();
786 let bytes: &mut [u8] = bytemuck::cast_slice_mut(dst);
787 let mut cursor = std::io::Cursor::new(bytes);
788 self.body
789 .write_into(&mut cursor)
790 .expect("the destination is sized to the serialized body");
791 assert_eq!(
792 usize::try_from(cursor.position()).expect("usize position"),
793 len_bytes,
794 "serialized body must fill the destination exactly",
795 );
796 }
797
798 fn as_any(&self) -> &dyn std::any::Any {
799 self
800 }
801
802 fn into_any(self: Box<Self>) -> Box<dyn std::any::Any> {
803 self
804 }
805}
806
807fn read_words<C: Columnar>(handle: &ChunkHandle) -> ColumnBody<C> {
809 let mut words = Vec::new();
810 handle.read_into(&mut words);
811 ColumnBody::Words(words)
812}
813
814fn pooled<D: Columnar + 'static, T: Columnar + 'static, R: Columnar + 'static>(
816 typed: &dyn TypedBody,
817) -> &PooledBody<D, T, R> {
818 typed
819 .as_any()
820 .downcast_ref::<PooledBody<D, T, R>>()
821 .expect("chunk bodies are pooled as `PooledBody`")
822}
823
824fn spill_serialized<C: Columnar>(
828 body: &ColumnBody<C>,
829 pool: &Pool,
830 len_bytes: usize,
831 hints: ChunkHints,
832 codec: &'static dyn ExtentCodec,
833) -> ChunkHandle {
834 mz_ore::soft_assert_eq_no_log!(len_bytes % 8, 0);
835 pool.insert_with(len_bytes / 8, hints, codec, |dst| {
836 let bytes: &mut [u8] = bytemuck::cast_slice_mut(dst);
837 let mut cursor = std::io::Cursor::new(bytes);
838 body.write_into(&mut cursor)
839 .expect("the slot is sized to the serialized body");
840 assert_eq!(
841 usize::try_from(cursor.position()).expect("usize position"),
842 len_bytes,
843 "serialized body must fill the chunk exactly",
844 );
845 })
846}
847
848pub fn try_spill_ref<C: Columnar>(body: &ColumnBody<C>, depth: u8) -> Option<ChunkHandle> {
861 let len_bytes = body.length_in_bytes();
862 let pool = spill_target(SINK_SPILL_ENABLED.load(Ordering::Relaxed), len_bytes)?;
863 let (codec, _) = codec_for_depth(depth);
864 Some(spill_serialized(
865 body,
866 &pool,
867 len_bytes,
868 ChunkHints { depth },
869 codec,
870 ))
871}
872
873impl<D, T, R, G> Chunk for ColumnChunk<D, T, R, G>
874where
875 G: SpillGate,
876 D: PoolColumnar,
877 for<'a> columnar::Ref<'a, D>: Copy + Ord,
878 T: PoolColumnar + Default + Timestamp + Lattice + Ord,
879 for<'a> columnar::Ref<'a, T>: Copy + Ord,
880 R: PoolColumnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
881{
882 type Time = T;
883
884 const TARGET: usize = 65536;
890
891 fn len(&self) -> usize {
892 self.records()
893 }
894
895 fn merge(in1: &mut VecDeque<Self>, in2: &mut VecDeque<Self>, out: &mut VecDeque<Self>) {
902 let (a_first, a_last) = in1
907 .front()
908 .expect("caller guarantees non-empty input")
909 .data_span();
910 let (b_first, b_last) = in2
911 .front()
912 .expect("caller guarantees non-empty input")
913 .data_span();
914 let a_low = rr::<D>(a_last) < rr::<D>(b_first);
915 let b_low = rr::<D>(b_last) < rr::<D>(a_first);
916 if a_low {
917 let chunk = in1.pop_front().expect("front observed above");
918 out.push_back(chunk.survive_merge());
919 return;
920 }
921 if b_low {
922 let chunk = in2.pop_front().expect("front observed above");
923 out.push_back(chunk.survive_merge());
924 return;
925 }
926
927 let a = in1.pop_front().expect("caller guarantees non-empty input");
928 let b = in2.pop_front().expect("caller guarantees non-empty input");
929 let depths = [a.depth(), b.depth()];
932 let out_depth = depths[0].max(depths[1]).saturating_add(1);
933 let keep = |chunk: &Self| match chunk {
938 ColumnChunk::Spilled(body, _) if !body.handle.is_typed() => Some(Rc::clone(body)),
939 _ => None,
940 };
941 let respill = [a.is_spilled(), b.is_spilled()];
942 let mut spill_a = keep(&a);
943 let mut spill_b = keep(&b);
944 let mut cols = [a.into_body(), b.into_body()];
945 let mut positions = [0usize, 0usize];
946 loop {
947 let mut result: ColumnBody<(D, T, R)> = ColumnBody::default();
948 let yielded = result.merge_from(&mut cols, &mut positions);
949 if !result.is_empty() {
950 out.push_back(ColumnChunk::Resident(Rc::new(result), out_depth));
951 }
952 if !yielded {
953 break;
954 }
955 }
956 metrics::record(metrics::Stage::Merge, positions[0] + positions[1], 0);
959 let [col_a, col_b] = &mut cols;
960 for (col, pos, depth, spilled, respill, queue) in [
964 (
965 col_a,
966 positions[0],
967 depths[0],
968 &mut spill_a,
969 respill[0],
970 in1,
971 ),
972 (
973 col_b,
974 positions[1],
975 depths[1],
976 &mut spill_b,
977 respill[1],
978 in2,
979 ),
980 ] {
981 let len = col.len();
982 if pos == 0 && len > 0 {
983 let chunk = match spilled.take() {
986 Some(body) => ColumnChunk::Spilled(body, depth).survive_merge(),
987 None if respill => Self::commit(std::mem::take(col), depth.saturating_add(1)),
989 None => {
990 ColumnChunk::Resident(Rc::new(std::mem::take(col)), depth).survive_merge()
991 }
992 };
993 queue.push_front(chunk);
994 } else if pos < len {
995 let view = col.borrow();
996 let mut rest = <(D, T, R) as Columnar>::Container::default();
997 rest.extend_from_self(view, pos..len);
998 queue.push_front(ColumnChunk::Resident(
999 Rc::new(ColumnBody::Typed(rest)),
1000 depth,
1001 ));
1002 }
1003 }
1004 }
1005
1006 fn extract(
1010 input: &mut VecDeque<Self>,
1011 frontier: AntichainRef<T>,
1012 residual: &mut timely::progress::Antichain<T>,
1013 keep: &mut VecDeque<Self>,
1014 ship: &mut VecDeque<Self>,
1015 ) {
1016 let Some(chunk) = input.pop_front() else {
1017 return;
1018 };
1019 let (time_lower, time_upper) = chunk.chunk_time_bounds();
1025 if time_upper.iter().all(|t| !frontier.less_equal(t)) {
1026 ship.push_back(chunk);
1027 return;
1028 }
1029 if time_lower.elements().iter().all(|m| frontier.less_equal(m)) {
1030 for m in time_lower.elements() {
1033 residual.insert_ref(m);
1034 }
1035 keep.push_back(chunk);
1036 return;
1037 }
1038 let depth = chunk.depth();
1041 let mut col = chunk.into_body();
1042 let len = col.len();
1043 let mut pos = 0;
1044 let mut keep_col: ColumnBody<(D, T, R)> = ColumnBody::default();
1045 let mut ship_col: ColumnBody<(D, T, R)> = ColumnBody::default();
1046 let cut = |col: &mut ColumnBody<(D, T, R)>, queue: &mut VecDeque<Self>, force: bool| {
1052 if !col.is_empty() && (force || at_serialized_capacity(&col.borrow())) {
1053 queue.push_back(ColumnChunk::Resident(Rc::new(std::mem::take(col)), depth));
1054 }
1055 };
1056 while pos < len {
1057 col.extract(&mut pos, frontier, residual, &mut keep_col, &mut ship_col);
1058 if pos < len {
1059 cut(&mut keep_col, keep, false);
1060 cut(&mut ship_col, ship, false);
1061 }
1062 }
1063 cut(&mut keep_col, keep, true);
1064 cut(&mut ship_col, ship, true);
1065 }
1066
1067 fn advance(
1077 input: &mut VecDeque<Self>,
1078 frontier: AntichainRef<T>,
1079 done: bool,
1080 out: &mut VecDeque<Self>,
1081 ) {
1082 let Some(front) = input.pop_front() else {
1083 return;
1084 };
1085 let mut depth = front.depth();
1088 let mut base = front.into_body();
1092 {
1093 let base_c = base.typed_mut();
1094 for chunk in input.drain(..) {
1095 depth = depth.max(chunk.depth());
1096 let col = chunk.into_body();
1097 let view = col.borrow();
1098 base_c.extend_from_self(view, 0..view.len());
1099 }
1100 }
1101 let view = base.borrow();
1102 let total = view.len();
1103 if total == 0 {
1104 return;
1105 }
1106 let data = view.0;
1107
1108 if !done && data.get(0) == data.get(total - 1) {
1111 input.push_front(ColumnChunk::Resident(Rc::new(base), depth));
1112 return;
1113 }
1114
1115 let end = if done {
1118 total
1119 } else {
1120 let last = data.get(total - 1);
1121 let mut end = total - 1;
1122 while end > 0 && data.get(end - 1) == last {
1123 end -= 1;
1124 }
1125 end
1126 };
1127 metrics::record(metrics::Stage::Advance, end, 0);
1131
1132 let mut result = <(D, T, R) as Columnar>::Container::default();
1133 let mut scratch: Vec<(T, R)> = Vec::new();
1135 let mut index = 0;
1136 const CUT_CHECK_RECORDS: usize = 1024;
1145 let mut records_since_check = 0usize;
1146 while index < end {
1153 let group_d = data.get(index);
1154 scratch.clear();
1155 while index < end && data.get(index) == group_d {
1156 let (_, t, r) = view.get(index);
1157 let mut owned_t = T::into_owned(t);
1158 owned_t.advance_by(frontier);
1159 scratch.push((owned_t, R::into_owned(r)));
1160 index += 1;
1161 }
1162 scratch.sort_by(|a, b| a.0.cmp(&b.0));
1163 let mut run = scratch.drain(..).peekable();
1164 while let Some((t, mut r)) = run.next() {
1165 while run.peek().is_some_and(|(t2, _)| *t2 == t) {
1166 let (_, r2) = run.next().expect("peeked");
1167 r.plus_equals(&r2);
1168 }
1169 if !r.is_zero() {
1170 result.0.push(group_d);
1171 result.1.push(&t);
1172 result.2.push(&r);
1173 records_since_check += 1;
1174 if records_since_check >= CUT_CHECK_RECORDS {
1175 records_since_check = 0;
1176 if u64::cast_from(indexed::length_in_words(&result.borrow()))
1177 >= u64::cast_from(COMMIT_BYTES / 8)
1178 {
1179 out.push_back(ColumnChunk::Resident(
1180 Rc::new(ColumnBody::Typed(std::mem::take(&mut result))),
1181 depth,
1182 ));
1183 }
1184 }
1185 }
1186 }
1187 }
1188 if !result.is_empty() {
1189 out.push_back(ColumnChunk::Resident(
1190 Rc::new(ColumnBody::Typed(result)),
1191 depth,
1192 ));
1193 }
1194
1195 if end < total {
1197 let mut carry = <(D, T, R) as Columnar>::Container::default();
1198 carry.extend_from_self(view, end..total);
1199 input.push_front(ColumnChunk::Resident(
1200 Rc::new(ColumnBody::Typed(carry)),
1201 depth,
1202 ));
1203 }
1204 }
1205
1206 fn settle(input: &mut VecDeque<Self>, done: bool, out: &mut VecDeque<Self>) {
1212 Self::settle_graded(input, done, out, true)
1213 }
1214}
1215
1216impl<D, T, R, G> ColumnChunk<D, T, R, G>
1217where
1218 G: SpillGate,
1219 D: PoolColumnar,
1220 T: PoolColumnar + Timestamp,
1221 R: PoolColumnar,
1222 for<'a> columnar::Ref<'a, D>: Ord,
1223 for<'a> columnar::Ref<'a, T>: Ord,
1224{
1225 fn settle_graded(
1232 input: &mut VecDeque<Self>,
1233 done: bool,
1234 out: &mut VecDeque<Self>,
1235 commit: bool,
1236 ) {
1237 let mut acc: Option<(ColumnBody<(D, T, R)>, u8)> = None;
1241 while let Some(chunk) = input.pop_front() {
1242 let (rc, depth) = match chunk {
1243 spilled @ ColumnChunk::Spilled(_, _) => {
1244 if let Some((body, depth)) = acc.take() {
1245 out.push_back(Self::grade(body, depth, commit));
1246 }
1247 out.push_back(spilled);
1248 continue;
1249 }
1250 ColumnChunk::Resident(rc, depth) => (rc, depth),
1251 };
1252 if rc.length_in_bytes() > COMMIT_BYTES && rc.len() > 1 {
1253 let records = cut_records(rc.len(), rc.length_in_bytes(), COMMIT_BYTES).max(1);
1256 let mut pieces = VecDeque::new();
1257 Self::push_cuts(&rc, records, depth, &mut pieces);
1258 for piece in pieces.into_iter().rev() {
1259 input.push_front(piece);
1260 }
1261 continue;
1262 }
1263 let full = at_commit_size(&rc);
1264 let fits = acc.as_ref().is_some_and(|(body, _)| {
1267 body.length_in_bytes().saturating_add(rc.length_in_bytes()) <= COMMIT_BYTES
1268 });
1269 if !full
1270 && fits
1271 && let Some((mut body, acc_depth)) = acc.take()
1272 {
1273 let view = rc.borrow();
1274 body.typed_mut().extend_from_self(view, 0..view.len());
1275 let acc_depth = acc_depth.max(depth);
1276 if at_commit_size(&body) {
1277 out.push_back(Self::grade(body, acc_depth, commit));
1278 } else {
1279 acc = Some((body, acc_depth));
1280 }
1281 continue;
1282 }
1283 if let Some((mut body, acc_depth)) = acc.take() {
1289 let len = rc.len();
1290 let prefix = if full {
1291 0
1292 } else {
1293 let space = COMMIT_BYTES.saturating_sub(body.length_in_bytes());
1294 cut_records(len, rc.length_in_bytes(), space)
1295 };
1296 if prefix == 0 || prefix >= len {
1297 out.push_back(Self::grade(body, acc_depth, commit));
1298 } else {
1299 let view = rc.borrow();
1300 body.typed_mut().extend_from_self(view, 0..prefix);
1301 let mut rest = <(D, T, R) as Columnar>::Container::default();
1302 rest.extend_from_self(view, prefix..len);
1303 input.push_front(ColumnChunk::Resident(
1304 Rc::new(ColumnBody::Typed(rest)),
1305 depth,
1306 ));
1307 let acc_depth = acc_depth.max(depth);
1308 if body.length_in_bytes() > COMMIT_BYTES {
1309 let mut pieces = VecDeque::new();
1311 Self::push_bounded(body, acc_depth, &mut pieces);
1312 for piece in pieces {
1313 let ColumnChunk::Resident(piece, piece_depth) = piece else {
1314 unreachable!("push_bounded emits resident pieces");
1315 };
1316 let piece =
1317 Rc::try_unwrap(piece).unwrap_or_else(|piece| piece.duplicate());
1318 out.push_back(Self::grade(piece, piece_depth, commit));
1319 }
1320 } else {
1321 out.push_back(Self::grade(body, acc_depth, commit));
1322 }
1323 continue;
1324 }
1325 }
1326 let body = Rc::try_unwrap(rc).unwrap_or_else(|rc| rc.duplicate());
1327 if full {
1328 out.push_back(Self::grade(body, depth, commit));
1329 } else {
1330 acc = Some((body, depth));
1331 }
1332 }
1333 if let Some((body, depth)) = acc {
1334 if done {
1335 out.push_back(Self::grade(body, depth, commit));
1336 } else {
1337 input.push_front(ColumnChunk::Resident(Rc::new(body), depth));
1338 }
1339 }
1340 }
1341
1342 fn grade(body: ColumnBody<(D, T, R)>, depth: u8, commit: bool) -> Self {
1345 if commit {
1346 Self::commit(body, depth)
1347 } else {
1348 ColumnChunk::Resident(Rc::new(body), depth)
1349 }
1350 }
1351}
1352
1353fn extract_view_into<'v, 'p, K, V, T, R>(
1358 view: BorrowedOf<'v, ((K, V), T, R)>,
1359 probes: BorrowedOf<'p, K>,
1360 probe_index: &mut usize,
1361 staging: &mut <((K, V), T, R) as Columnar>::Container,
1362) where
1363 K: Columnar,
1364 V: Columnar,
1365 T: Columnar,
1366 R: Columnar,
1367 for<'b> columnar::Ref<'b, K>: Copy + Ord,
1368{
1369 let keys = view.0.0;
1370 let len = keys.len();
1371 let last = keys.get(len - 1);
1372 let count = probes.len();
1373 let mut pos = 0;
1374 while *probe_index < count {
1375 let probe = probes.get(*probe_index);
1376 mz_ore::soft_assert_no_log!(
1377 *probe_index == 0 || rr::<K>(probes.get(*probe_index - 1)) < rr::<K>(probe),
1378 "probe keys must be sorted and deduplicated"
1379 );
1380 if rr::<K>(probe) > rr::<K>(last) {
1381 return;
1382 }
1383 gallop(len, &mut pos, |i| rr::<K>(keys.get(i)) < rr::<K>(probe));
1384 let start = pos;
1385 while pos < len && rr::<K>(keys.get(pos)) == rr::<K>(probe) {
1386 pos += 1;
1387 }
1388 staging.extend_from_self(view, start..pos);
1389 if rr::<K>(probe) == rr::<K>(last) {
1390 return;
1391 }
1392 *probe_index += 1;
1393 }
1394}
1395
1396impl<K, V, T, R, G> UnloadChunk for ColumnChunk<(K, V), T, R, G>
1397where
1398 G: SpillGate,
1399 K: PoolColumnar,
1400 for<'a> columnar::Ref<'a, K>: Copy + Ord,
1401 V: PoolColumnar,
1402 for<'a> columnar::Ref<'a, V>: Copy + Ord,
1403 T: PoolColumnar + Default + Timestamp + Lattice + Ord,
1404 for<'a> columnar::Ref<'a, T>: Copy + Ord,
1405 R: PoolColumnar + Default + Semigroup + for<'a> Semigroup<columnar::Ref<'a, R>>,
1406{
1407 type Staging = <((K, V), T, R) as Columnar>::Container;
1410
1411 type Probes<'a> = BorrowedOf<'a, K>;
1414
1415 fn probe_count(probes: Self::Probes<'_>) -> usize {
1416 probes.len()
1417 }
1418
1419 fn locate(&self, probes: Self::Probes<'_>, probe_index: usize) -> std::cmp::Ordering {
1420 let probe = probes.get(probe_index);
1421 let (first, last) = self.data_span();
1424 let (first, last) = (first.0, last.0);
1425 if rr::<K>(probe) < rr::<K>(first) {
1426 std::cmp::Ordering::Less
1427 } else if rr::<K>(probe) > rr::<K>(last) {
1428 std::cmp::Ordering::Greater
1429 } else {
1430 std::cmp::Ordering::Equal
1431 }
1432 }
1433
1434 fn extract_into(
1435 &self,
1436 probes: Self::Probes<'_>,
1437 probe_index: &mut usize,
1438 staging: &mut Self::Staging,
1439 ) {
1440 match self {
1441 ColumnChunk::Resident(col, _) => {
1442 extract_view_into::<K, V, T, R>(col.borrow(), probes, probe_index, staging);
1443 }
1444 ColumnChunk::Spilled(body, _) => body.handle.with_typed(|typed| match typed {
1445 Some(typed) => {
1446 let view = pooled::<(K, V), T, R>(typed).body.borrow();
1447 extract_view_into::<K, V, T, R>(view, probes, probe_index, staging);
1448 }
1449 None => with_scratch(|scratch| {
1450 body.handle.read_into(scratch);
1456 let view = borrow_words::<((K, V), T, R)>(scratch);
1457 extract_view_into::<K, V, T, R>(view, probes, probe_index, staging);
1458 }),
1459 }),
1460 }
1461 }
1462
1463 fn fetch_into(&self, staging: &mut Self::Staging) {
1464 match self {
1465 ColumnChunk::Resident(col, _) => {
1466 let view = col.borrow();
1467 staging.extend_from_self(view, 0..view.len());
1468 }
1469 ColumnChunk::Spilled(body, _) => body.handle.with_typed(|typed| match typed {
1470 Some(typed) => {
1471 let view = pooled::<(K, V), T, R>(typed).body.borrow();
1472 staging.extend_from_self(view, 0..view.len());
1473 }
1474 None => with_scratch(|scratch| {
1475 body.handle.read_into(scratch);
1476 let view = borrow_words::<((K, V), T, R)>(scratch);
1477 staging.extend_from_self(view, 0..view.len());
1478 }),
1479 }),
1480 }
1481 }
1482}
1483
1484pub struct UnchunkBuilder<Bu, D: Columnar, T: Columnar, R: Columnar, G: SpillGate> {
1494 inner: Bu,
1495 _marker: std::marker::PhantomData<(D, T, R, G)>,
1496}
1497
1498impl<Bu: Default, D: Columnar, T: Columnar, R: Columnar, G: SpillGate> Default
1499 for UnchunkBuilder<Bu, D, T, R, G>
1500{
1501 fn default() -> Self {
1502 Self {
1503 inner: Bu::default(),
1504 _marker: std::marker::PhantomData,
1505 }
1506 }
1507}
1508
1509impl<Bu, D, T, R, G> differential_dataflow::trace::Builder for UnchunkBuilder<Bu, D, T, R, G>
1510where
1511 Bu: differential_dataflow::trace::Builder<Input = ColumnBody<(D, T, R)>> + ChainState,
1512 D: Columnar + 'static,
1513 T: Columnar + 'static,
1514 R: Columnar + 'static,
1515 G: SpillGate,
1516{
1517 type Input = ColumnChunk<D, T, R, G>;
1518 type Time = Bu::Time;
1519 type Output = Bu::Output;
1520
1521 fn push(&mut self, chunk: &mut Self::Input) {
1522 let mut body = std::mem::take(chunk).into_body();
1523 self.inner.push(&mut body);
1524 }
1525
1526 fn done(self) -> Option<Self::Output> {
1527 self.inner.done()
1528 }
1529}
1530
1531impl<Bu, D, T, R, G>
1532 differential_dataflow::trace::implementations::merge_batcher::Sealer<ColumnChunk<D, T, R, G>>
1533 for UnchunkBuilder<Bu, D, T, R, G>
1534where
1535 Bu: differential_dataflow::trace::Builder<Input = ColumnBody<(D, T, R)>> + ChainState,
1536 D: Columnar + 'static,
1537 T: Columnar + 'static,
1538 R: Columnar + 'static,
1539 G: SpillGate,
1540{
1541 type Output = Bu::Output;
1542
1543 fn seal(chain: &mut Vec<ColumnChunk<D, T, R, G>>) -> Option<Self::Output> {
1544 use differential_dataflow::trace::Builder;
1545
1546 let mut state = Bu::State::default();
1547 for chunk in chain.iter() {
1548 if !chunk.is_spilled() || Bu::wants_bodies() {
1553 chunk.with_body(|body| Bu::observe(&mut state, body));
1554 } else {
1555 Bu::observe_records(&mut state, chunk.records());
1556 }
1557 }
1558 let mut builder = Self {
1559 inner: Bu::from_state(state),
1560 _marker: std::marker::PhantomData,
1561 };
1562 for chunk in chain.iter_mut() {
1565 builder.push(chunk);
1566 }
1567 chain.clear();
1568 builder.done()
1569 }
1570}
1571
1572pub trait ChainState: differential_dataflow::trace::Builder {
1589 type State: Default;
1591
1592 fn wants_bodies() -> bool;
1596
1597 fn observe(state: &mut Self::State, input: &Self::Input);
1599
1600 fn observe_records(state: &mut Self::State, records: usize);
1603
1604 fn from_state(state: Self::State) -> Self;
1606}
1607
1608pub struct ChunkChunker<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> {
1611 inner: ColumnChunker<(D, T, R)>,
1612 ready: VecDeque<ColumnChunk<D, T, R, G>>,
1613 staged: ColumnChunk<D, T, R, G>,
1614}
1615
1616impl<D, T, R, G> Default for ChunkChunker<D, T, R, G>
1617where
1618 D: Columnar,
1619 T: Columnar,
1620 R: Columnar,
1621 G: SpillGate,
1622 ColumnChunker<(D, T, R)>: Default,
1623{
1624 fn default() -> Self {
1625 Self {
1626 inner: Default::default(),
1627 ready: VecDeque::new(),
1628 staged: Default::default(),
1629 }
1630 }
1631}
1632
1633impl<'a, D, T, R, G> PushInto<&'a mut Column<(D, T, R)>> for ChunkChunker<D, T, R, G>
1634where
1635 D: Columnar,
1636 T: Columnar,
1637 R: Columnar,
1638 G: SpillGate,
1639 ColumnChunker<(D, T, R)>: PushInto<&'a mut Column<(D, T, R)>>,
1640{
1641 fn push_into(&mut self, item: &'a mut Column<(D, T, R)>) {
1642 self.inner.push_into(item);
1643 }
1644}
1645
1646impl<D, T, R, G> ContainerBuilder for ChunkChunker<D, T, R, G>
1647where
1648 D: Columnar + 'static,
1649 T: Columnar + 'static,
1650 R: Columnar + 'static,
1651 G: SpillGate,
1652 ColumnChunker<(D, T, R)>: ContainerBuilder<Container = ColumnBody<(D, T, R)>>,
1653{
1654 type Container = ColumnChunk<D, T, R, G>;
1655
1656 fn extract(&mut self) -> Option<&mut Self::Container> {
1657 if self.ready.is_empty() {
1658 let body = self.inner.extract()?;
1659 ColumnChunk::push_bounded(std::mem::take(body), 0, &mut self.ready);
1660 }
1661 self.staged = self.ready.pop_front()?;
1662 Some(&mut self.staged)
1663 }
1664
1665 fn finish(&mut self) -> Option<&mut Self::Container> {
1666 if self.ready.is_empty() {
1667 let body = self.inner.finish()?;
1668 ColumnChunk::push_bounded(std::mem::take(body), 0, &mut self.ready);
1669 }
1670 self.staged = self.ready.pop_front()?;
1671 Some(&mut self.staged)
1672 }
1673}
1674
1675pub type AccountedChunkBatcher<D, T, R, G, S> =
1683 differential_dataflow::trace::implementations::merge_batcher::MergeBatcher<
1684 ChunkChunker<D, T, R, G>,
1685 AccountedChunkMerger<D, T, R, G>,
1686 S,
1687 >;
1688
1689pub type AccountedChunkChainBatcher<D, T, R, G> = crate::operator::ConsolidatingBatcher<
1692 ChunkChunker<D, T, R, G>,
1693 AccountedChunkMerger<D, T, R, G>,
1694>;
1695
1696pub struct AccountedChunkMerger<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> {
1709 inner: differential_dataflow::trace::chunk::ChunkMerger<ColumnChunk<D, T, R, G>>,
1710}
1711
1712impl<D: Columnar, T: Columnar, R: Columnar, G: SpillGate> Default
1713 for AccountedChunkMerger<D, T, R, G>
1714{
1715 fn default() -> Self {
1716 Self {
1717 inner: Default::default(),
1718 }
1719 }
1720}
1721
1722impl<D, T, R, G> differential_dataflow::trace::implementations::merge_batcher::Merger
1723 for AccountedChunkMerger<D, T, R, G>
1724where
1725 D: PoolColumnar + 'static,
1726 T: PoolColumnar + Timestamp + 'static,
1727 R: PoolColumnar + 'static,
1728 G: SpillGate,
1729 for<'a> columnar::Ref<'a, D>: Ord,
1730 for<'a> columnar::Ref<'a, T>: Ord,
1731 ColumnChunk<D, T, R, G>: Chunk,
1732 <ColumnChunk<D, T, R, G> as Chunk>::Time: Clone + PartialOrder + 'static,
1733{
1734 type Chunk = ColumnChunk<D, T, R, G>;
1735
1736 fn merge(
1737 &mut self,
1738 list1: Vec<Self::Chunk>,
1739 list2: Vec<Self::Chunk>,
1740 output: &mut Vec<Self::Chunk>,
1741 stash: &mut Vec<Self::Chunk>,
1742 ) {
1743 self.inner.merge(list1, list2, output, stash)
1744 }
1745
1746 fn len(chunk: &Self::Chunk) -> usize {
1747 Chunk::len(chunk)
1748 }
1749
1750 fn allocation(chunk: &Self::Chunk) -> (usize, usize, usize) {
1751 let bytes = match chunk {
1752 ColumnChunk::Resident(col, _) => col.length_in_bytes(),
1753 ColumnChunk::Spilled(body, _) => body.len_bytes,
1754 };
1755 (bytes, bytes, 1)
1756 }
1757}
1758
1759impl<D, T, R, G> differential_dataflow::trace::implementations::merge_batcher::Extractor
1760 for AccountedChunkMerger<D, T, R, G>
1761where
1762 D: PoolColumnar + 'static,
1763 T: PoolColumnar + Timestamp + 'static,
1764 R: PoolColumnar + 'static,
1765 G: SpillGate,
1766 for<'a> columnar::Ref<'a, D>: Ord,
1767 for<'a> columnar::Ref<'a, T>: Ord,
1768 ColumnChunk<D, T, R, G>: Chunk,
1769 <ColumnChunk<D, T, R, G> as Chunk>::Time: Clone + PartialOrder + 'static,
1770{
1771 type Time = <ColumnChunk<D, T, R, G> as Chunk>::Time;
1772
1773 fn extract(
1774 &mut self,
1775 merged: Vec<Self::Chunk>,
1776 upper: AntichainRef<Self::Time>,
1777 frontier: &mut Antichain<Self::Time>,
1778 readied: &mut Vec<Self::Chunk>,
1779 kept: &mut Vec<Self::Chunk>,
1780 _stash: &mut Vec<Self::Chunk>,
1781 ) {
1782 let mut input: VecDeque<Self::Chunk> = merged.into();
1790 let (mut keep, mut shipped) = (VecDeque::new(), VecDeque::new());
1791 let (mut kept_q, mut shipped_q) = (VecDeque::new(), VecDeque::new());
1792 while !input.is_empty() {
1793 Chunk::extract(&mut input, upper, frontier, &mut keep, &mut shipped);
1794 Chunk::settle(&mut keep, false, &mut kept_q);
1795 ColumnChunk::settle_graded(&mut shipped, false, &mut shipped_q, false);
1796 }
1797 Chunk::settle(&mut keep, true, &mut kept_q);
1798 ColumnChunk::settle_graded(&mut shipped, true, &mut shipped_q, false);
1799 Extend::extend(kept, kept_q);
1800 Extend::extend(readied, shipped_q);
1801 }
1802}
1803
1804#[cfg(test)]
1805mod tests {
1806 use differential_dataflow::batcher::Extract;
1816 use differential_dataflow::trace::chunk::ChunkBatch;
1817 use mz_ore::pool::Pool;
1818 use proptest::prelude::*;
1819 use timely::container::PushInto;
1820 use timely::progress::Antichain;
1821
1822 use crate::columnar::unload::UnloadBatch;
1823
1824 use super::*;
1825
1826 type Tuple = ((u64, u64), u64, i64);
1827 type TestChunk = ColumnChunk<(u64, u64), u64, i64, ComputeSpill>;
1828 type TestChunkBatcher = crate::operator::ConsolidatingBatcher<
1830 ChunkChunker<(u64, u64), u64, i64, ComputeSpill>,
1831 differential_dataflow::trace::chunk::ChunkMerger<TestChunk>,
1832 >;
1833
1834 #[mz_ore::test]
1839 fn lz4_codec_matches_the_previous_extent_framing() {
1840 let body: Vec<u8> = (0..100_000u32).flat_map(|i| i.to_le_bytes()).collect();
1841 let mut stored = Vec::new();
1842 LZ4_CODEC.encode(&body, &mut stored);
1843 assert_eq!(stored, lz4_flex::block::compress_prepend_size(&body));
1844 let mut round = vec![0u8; body.len()];
1845 LZ4_CODEC.decode(&stored, &mut round);
1846 assert_eq!(round, body);
1847 }
1848
1849 #[mz_ore::test]
1850 #[should_panic(expected = "destination must match")]
1851 fn lz4_codec_decode_length_mismatch_panics() {
1852 let mut stored = Vec::new();
1853 LZ4_CODEC.encode(&[7u8; 64], &mut stored);
1854 let mut short = vec![0u8; 32];
1855 LZ4_CODEC.decode(&stored, &mut short);
1856 }
1857
1858 fn consolidate(mut v: Vec<Tuple>) -> Vec<Tuple> {
1861 v.sort();
1862 let mut out: Vec<Tuple> = Vec::new();
1863 for (d, t, r) in v {
1864 if let Some(last) = out.last_mut() {
1865 if last.0 == d && last.1 == t {
1866 last.2 += r;
1867 continue;
1868 }
1869 }
1870 out.push((d, t, r));
1871 }
1872 out.retain(|x| x.2 != 0);
1873 out
1874 }
1875
1876 fn arb_consolidated() -> impl Strategy<Value = Vec<Tuple>> {
1877 prop::collection::vec(((0u64..5, 0u64..5), 0u64..4, -3i64..=3i64), 0..40)
1878 .prop_map(consolidate)
1879 }
1880
1881 fn build_column(v: &[Tuple]) -> ColumnBody<Tuple> {
1882 let mut col: ColumnBody<Tuple> = Default::default();
1883 for tup in v {
1884 col.push_into(*tup);
1885 }
1886 col
1887 }
1888
1889 fn collect_column(col: &ColumnBody<Tuple>) -> Vec<Tuple> {
1890 col.borrow()
1891 .into_index_iter()
1892 .map(|((k, v), t, r)| {
1893 (
1894 (u64::into_owned(k), u64::into_owned(v)),
1895 u64::into_owned(t),
1896 i64::into_owned(r),
1897 )
1898 })
1899 .collect()
1900 }
1901
1902 fn collect_chunks(chunks: impl IntoIterator<Item = TestChunk>) -> Vec<Tuple> {
1903 chunks
1904 .into_iter()
1905 .flat_map(|chunk| collect_column(&chunk.into_body()))
1906 .collect()
1907 }
1908
1909 fn collect_staging(staging: &<Tuple as Columnar>::Container) -> Vec<Tuple> {
1910 staging
1911 .borrow()
1912 .into_index_iter()
1913 .map(|((k, v), t, r)| {
1914 (
1915 (u64::into_owned(k), u64::into_owned(v)),
1916 u64::into_owned(t),
1917 i64::into_owned(r),
1918 )
1919 })
1920 .collect()
1921 }
1922
1923 fn chunked(data: &[Tuple], cuts: &[usize]) -> VecDeque<TestChunk> {
1925 let mut chunks = VecDeque::new();
1926 let mut start = 0;
1927 for cut in cuts {
1928 let end = (start + 1 + cut % 7).min(data.len());
1929 if end > start {
1930 chunks.push_back(ColumnChunk::from_body(build_column(&data[start..end])));
1931 start = end;
1932 }
1933 }
1934 if start < data.len() {
1935 chunks.push_back(ColumnChunk::from_body(build_column(&data[start..])));
1936 }
1937 chunks
1938 }
1939
1940 fn chunked_spilled(data: &[Tuple], cuts: &[usize], pool: &Pool) -> VecDeque<TestChunk> {
1943 chunked(data, cuts)
1944 .into_iter()
1945 .map(|chunk| force_spill(chunk, pool))
1946 .collect()
1947 }
1948
1949 fn body_compressed(chunk: &TestChunk) -> bool {
1953 match chunk {
1954 ColumnChunk::Spilled(body, _) => body.compressed,
1955 ColumnChunk::Resident(_, _) => panic!("chunk must be spilled"),
1956 }
1957 }
1958
1959 fn force_spill(chunk: TestChunk, pool: &Pool) -> TestChunk {
1962 let depth = chunk.depth();
1963 TestChunk::spill_body(chunk.into_body(), pool, depth)
1964 }
1965
1966 #[mz_ore::test]
1969 #[cfg_attr(miri, ignore)]
1970 fn allocation_reports_a_body_wherever_it_lives() {
1971 use differential_dataflow::trace::implementations::merge_batcher::Merger;
1972
1973 let data: Vec<Tuple> = (0..64u64).map(|i| ((i, i), 0, 1)).collect();
1974 let resident = TestChunk::from_body(build_column(&data));
1975 let (size, capacity, allocations) =
1976 AccountedChunkMerger::<(u64, u64), u64, i64, ComputeSpill>::allocation(&resident);
1977 assert!(size > 0, "a resident body reports its bytes");
1978 assert_eq!((capacity, allocations), (size, 1));
1979
1980 let spilled = force_spill(resident, &test_pool());
1981 assert_eq!(
1982 AccountedChunkMerger::<(u64, u64), u64, i64, ComputeSpill>::allocation(&spilled),
1983 (size, capacity, allocations),
1984 "spilling a body does not change what its owner holds"
1985 );
1986 }
1987
1988 #[mz_ore::test]
1997 #[cfg_attr(miri, ignore)]
1998 fn seal_readies_resident_chunks() {
1999 let data: Vec<Tuple> = (0..20_000u64).map(|i| ((i, i), 0, 1)).collect();
2003
2004 set_spill_override(Some(test_pool()));
2005 let mut batcher =
2006 AccountedChunkChainBatcher::<(u64, u64), u64, i64, ComputeSpill>::new(None, 0);
2007 batcher.push_chunk(TestChunk::from_body(build_column(&data)));
2008 let (chain, _frontier) = batcher.extract(Antichain::new().borrow());
2009 let chain = chain.unwrap_or_default();
2010 set_spill_override(None);
2011
2012 assert!(!chain.is_empty(), "the seal ships what was pushed");
2013 assert!(
2014 chain.iter().all(|chunk| !chunk.is_spilled()),
2015 "a chunk that arrived resident is readied resident"
2016 );
2017 }
2018
2019 fn test_pool() -> Pool {
2023 static POOL: std::sync::OnceLock<Pool> = std::sync::OnceLock::new();
2024 POOL.get_or_init(|| Pool::new().expect("pool creation"))
2025 .clone()
2026 }
2027
2028 proptest! {
2029 #[mz_ore::test]
2032 #[cfg_attr(miri, ignore)]
2033 fn batcher_round_trip(
2034 inputs in prop::collection::vec(arb_consolidated(), 1..6),
2035 cuts in prop::collection::vec(0usize..7, 0..8),
2036 ) {
2037 let mut batcher = TestChunkBatcher::new(None, 0);
2038 let mut union = Vec::new();
2039 for input in &inputs {
2040 Extend::extend(&mut union, input.iter().copied());
2041 for chunk in chunked(input, &cuts) {
2042 batcher.push_chunk(chunk);
2043 }
2044 }
2045 let (sealed, _frontier) = batcher.extract(Antichain::new().borrow());
2047 let sealed = sealed.unwrap_or_default();
2048 prop_assert_eq!(collect_chunks(sealed), consolidate(union));
2049 }
2050
2051 #[mz_ore::test]
2054 #[cfg_attr(miri, ignore)]
2055 fn batcher_round_trip_spilled(
2056 inputs in prop::collection::vec(arb_consolidated(), 1..4),
2057 cuts in prop::collection::vec(0usize..7, 0..6),
2058 ) {
2059 let pool = test_pool();
2060 let mut batcher = TestChunkBatcher::new(None, 0);
2061 let mut union = Vec::new();
2062 for input in &inputs {
2063 Extend::extend(&mut union, input.iter().copied());
2064 for chunk in chunked_spilled(input, &cuts, &pool) {
2065 batcher.push_chunk(chunk);
2066 }
2067 }
2068 let (sealed, _frontier) = batcher.extract(Antichain::new().borrow());
2069 let sealed = sealed.unwrap_or_default();
2070 prop_assert_eq!(collect_chunks(sealed), consolidate(union));
2071 }
2072
2073 #[mz_ore::test]
2076 #[cfg_attr(miri, ignore)]
2077 fn seal_partitions_by_time(
2078 input in arb_consolidated(),
2079 cuts in prop::collection::vec(0usize..7, 0..8),
2080 upper in 0u64..5,
2081 ) {
2082 let mut batcher = TestChunkBatcher::new(None, 0);
2083 for chunk in chunked(&input, &cuts) {
2084 batcher.push_chunk(chunk);
2085 }
2086 let (shipped, frontier) = batcher.extract(Antichain::from_elem(upper).borrow());
2087 let shipped = shipped.unwrap_or_default();
2088 let frontier = frontier.to_owned();
2089 let expected_shipped: Vec<Tuple> =
2090 input.iter().copied().filter(|(_, t, _)| *t < upper).collect();
2091 prop_assert_eq!(collect_chunks(shipped), consolidate(expected_shipped));
2092
2093 let kept_min = input.iter().filter(|(_, t, _)| *t >= upper).map(|(_, t, _)| *t).min();
2094 prop_assert_eq!(frontier.elements().first().copied(), kept_min);
2095
2096 let (rest, _frontier) = batcher.extract(Antichain::new().borrow());
2097 let rest = rest.unwrap_or_default();
2098 let expected_rest: Vec<Tuple> =
2099 input.iter().copied().filter(|(_, t, _)| *t >= upper).collect();
2100 prop_assert_eq!(collect_chunks(rest), consolidate(expected_rest));
2101 }
2102
2103 #[mz_ore::test]
2107 #[cfg_attr(miri, ignore)]
2108 fn seal_partitions_by_time_spilled(
2109 input in arb_consolidated(),
2110 cuts in prop::collection::vec(0usize..7, 0..8),
2111 upper in 0u64..5,
2112 ) {
2113 let pool = test_pool();
2114 let mut batcher = TestChunkBatcher::new(None, 0);
2115 for chunk in chunked_spilled(&input, &cuts, &pool) {
2116 batcher.push_chunk(chunk);
2117 }
2118 let (shipped, frontier) = batcher.extract(Antichain::from_elem(upper).borrow());
2119 let shipped = shipped.unwrap_or_default();
2120 let frontier = frontier.to_owned();
2121 let expected_shipped: Vec<Tuple> =
2122 input.iter().copied().filter(|(_, t, _)| *t < upper).collect();
2123 prop_assert_eq!(collect_chunks(shipped), consolidate(expected_shipped));
2124
2125 let kept_min = input.iter().filter(|(_, t, _)| *t >= upper).map(|(_, t, _)| *t).min();
2126 prop_assert_eq!(frontier.elements().first().copied(), kept_min);
2127
2128 let (rest, _frontier) = batcher.extract(Antichain::new().borrow());
2129 let rest = rest.unwrap_or_default();
2130 let expected_rest: Vec<Tuple> =
2131 input.iter().copied().filter(|(_, t, _)| *t >= upper).collect();
2132 prop_assert_eq!(collect_chunks(rest), consolidate(expected_rest));
2133 }
2134
2135 #[mz_ore::test]
2138 #[cfg_attr(miri, ignore)]
2139 fn advance_matches_reference(
2140 input in arb_consolidated(),
2141 cuts in prop::collection::vec(0usize..7, 0..8),
2142 frontier_elem in 0u64..5,
2143 ) {
2144 let frontier = Antichain::from_elem(frontier_elem);
2145 let mut chunks = chunked(&input, &cuts);
2146 let mut out = VecDeque::new();
2147 TestChunk::advance(&mut chunks, frontier.borrow(), false, &mut out);
2148 TestChunk::advance(&mut chunks, frontier.borrow(), true, &mut out);
2149 prop_assert!(chunks.is_empty());
2150
2151 let expected = consolidate(
2152 input
2153 .iter()
2154 .map(|&(d, mut t, r)| {
2155 t.advance_by(frontier.borrow());
2156 (d, t, r)
2157 })
2158 .collect(),
2159 );
2160 prop_assert_eq!(collect_chunks(out), expected);
2161 }
2162
2163 #[mz_ore::test]
2166 #[cfg_attr(miri, ignore)]
2167 fn settle_preserves_and_packs(
2168 input in arb_consolidated(),
2169 cuts in prop::collection::vec(0usize..7, 1..8),
2170 ) {
2171 let mut chunks = chunked(&input, &cuts);
2172 let mut out = VecDeque::new();
2173 TestChunk::settle(&mut chunks, true, &mut out);
2174 prop_assert!(chunks.is_empty());
2175 prop_assert!(out.len() <= 1);
2178 prop_assert_eq!(collect_chunks(out), input);
2179 }
2180
2181 #[mz_ore::test]
2185 #[cfg_attr(miri, ignore)]
2186 fn unload_extract_matches_filter(
2187 input in arb_consolidated(),
2188 cuts in prop::collection::vec(0usize..7, 0..8),
2189 probe_keys in prop::collection::btree_set(0u64..6, 0..6),
2190 spill in any::<bool>(),
2191 ) {
2192 prop_assume!(!input.is_empty());
2193 let pool = test_pool();
2194 let chunks: Vec<TestChunk> = if spill {
2195 chunked_spilled(&input, &cuts, &pool).into()
2196 } else {
2197 chunked(&input, &cuts).into()
2198 };
2199 let batch = ChunkBatch::new(chunks);
2200
2201 let mut probe_col = <u64 as Columnar>::Container::default();
2202 for key in &probe_keys {
2203 probe_col.push(*key);
2204 }
2205 let mut staging = <Tuple as Columnar>::Container::default();
2206 batch.extract_into(probe_col.borrow(), &mut staging);
2207
2208 let expected: Vec<Tuple> = input
2209 .iter()
2210 .copied()
2211 .filter(|((k, _), _, _)| probe_keys.contains(k))
2212 .collect();
2213 prop_assert_eq!(collect_staging(&staging), expected);
2214
2215 let mut staging = <Tuple as Columnar>::Container::default();
2218 batch.fetch_into(&mut staging);
2219 prop_assert_eq!(collect_staging(&staging), input);
2220 }
2221 }
2222
2223 #[mz_ore::test]
2226 fn locate_spans_keys() {
2227 let chunk = TestChunk::from_body(build_column(&[
2228 ((2, 0), 0, 1),
2229 ((4, 0), 0, 1),
2230 ((6, 0), 0, 1),
2231 ]));
2232 let mut probe_col = <u64 as Columnar>::Container::default();
2233 for key in [0u64, 2, 3, 6, 9] {
2234 probe_col.push(key);
2235 }
2236 let probes = probe_col.borrow();
2237 use std::cmp::Ordering::*;
2238 let expected = [Less, Equal, Equal, Equal, Greater];
2239 for (index, expected) in expected.iter().enumerate() {
2240 assert_eq!(chunk.locate(probes, index), *expected, "probe {index}");
2241 }
2242 }
2243
2244 fn collect_bounded(chunks: impl IntoIterator<Item = TestChunk>, bound: usize) -> Vec<Tuple> {
2247 let mut collected = Vec::new();
2248 for chunk in chunks {
2249 let col = chunk.into_body();
2250 let bytes = col.length_in_bytes();
2251 assert!(bytes <= bound, "chunk of {bytes} bytes exceeds {bound}");
2252 Extend::extend(&mut collected, collect_column(&col));
2253 }
2254 collected
2255 }
2256
2257 type WideUpdate = ((u64, String), u64, i64);
2258 type WideChunk = ColumnChunk<(u64, String), u64, i64, ComputeSpill>;
2259
2260 fn wide_column(keys: impl Iterator<Item = u64>, bytes: usize) -> ColumnBody<WideUpdate> {
2261 let mut column = ColumnBody::default();
2262 for key in keys {
2263 column.push_into(&((key, "x".repeat(bytes)), 0, 1));
2264 }
2265 column
2266 }
2267
2268 fn assert_wide_byte_bound(chunks: VecDeque<WideChunk>, expected: usize, payload_bytes: usize) {
2269 let mut keys = Vec::new();
2270 for chunk in chunks {
2271 let column = chunk.into_body();
2272 assert!(
2273 column.length_in_bytes() <= COMMIT_BYTES || column.len() == 1,
2274 "{} bytes in a {}-record chunk",
2275 column.length_in_bytes(),
2276 column.len(),
2277 );
2278 let view = column.borrow();
2279 for index in 0..view.len() {
2280 let ((key, payload), time, diff) = view.get(index);
2281 assert_eq!(payload.len(), payload_bytes);
2282 assert!(payload.iter().all(|byte| *byte == b'x'));
2283 assert_eq!((*time, *diff), (0, 1));
2284 keys.push(*key);
2285 }
2286 }
2287 assert_eq!(keys, (0..u64::cast_from(expected)).collect::<Vec<_>>());
2288 }
2289
2290 #[mz_ore::test]
2291 fn chunker_enforces_byte_bound() {
2292 let mut chunker = ChunkChunker::default();
2293 let mut input: Column<WideUpdate> = wide_column((0..4000).rev(), 3000).into();
2294 chunker.push_into(&mut input);
2295 let mut chunks = VecDeque::new();
2296 if let Some(chunk) = chunker.extract() {
2297 chunks.push_back(std::mem::take(chunk));
2298 }
2299 while let Some(chunk) = chunker.finish() {
2300 chunks.push_back(std::mem::take(chunk));
2301 }
2302 assert_wide_byte_bound(chunks, 4000, 3000);
2303 }
2304
2305 #[mz_ore::test]
2306 fn merge_settle_enforces_byte_bound() {
2307 let mut left = VecDeque::from([WideChunk::from_body(wide_column(
2308 (0..2000).step_by(2),
2309 2100,
2310 ))]);
2311 let mut right = VecDeque::from([WideChunk::from_body(wide_column(
2312 (1..2000).step_by(2),
2313 2100,
2314 ))]);
2315 let mut merged = VecDeque::new();
2316 while !left.is_empty() && !right.is_empty() {
2317 WideChunk::merge(&mut left, &mut right, &mut merged);
2318 }
2319 merged.append(&mut left);
2320 merged.append(&mut right);
2321 let mut settled = VecDeque::new();
2322 WideChunk::settle(&mut merged, true, &mut settled);
2323 assert_wide_byte_bound(settled, 2000, 2100);
2324 }
2325
2326 #[mz_ore::test]
2327 fn settle_enforces_byte_bound_after_coalescing() {
2328 let mut input = VecDeque::from([
2329 WideChunk::from_body(wide_column(0..400, 3000)),
2330 WideChunk::from_body(wide_column(400..800, 3000)),
2331 ]);
2332 let mut settled = VecDeque::new();
2333 WideChunk::settle(&mut input, true, &mut settled);
2334 assert_wide_byte_bound(settled, 800, 3000);
2335 }
2336
2337 #[mz_ore::test]
2338 fn settle_byte_bound_allows_indivisible_update() {
2339 let mut input = VecDeque::from([WideChunk::from_body(wide_column(0..1, 2 * COMMIT_BYTES))]);
2340 let mut settled = VecDeque::new();
2341 WideChunk::settle(&mut input, true, &mut settled);
2342 assert_wide_byte_bound(settled, 1, 2 * COMMIT_BYTES);
2343 }
2344
2345 #[mz_ore::test]
2349 fn push_bounded_splits_uneven_widths() {
2350 let mut column: ColumnBody<WideUpdate> = ColumnBody::default();
2351 column.push_into(&((0, "x".repeat(COMMIT_BYTES - 4096)), 0, 1));
2352 for key in 1..20_000u64 {
2353 column.push_into(&((key, "x".to_string()), 0, 1));
2354 }
2355 assert!(column.length_in_bytes() > COMMIT_BYTES);
2356 let mut out = VecDeque::new();
2357 WideChunk::push_bounded(column, 3, &mut out);
2358 let mut keys = Vec::new();
2359 for chunk in out {
2360 let ColumnChunk::Resident(_, depth) = &chunk else {
2361 panic!("pieces are resident");
2362 };
2363 assert_eq!(*depth, 3, "pieces keep the input's depth");
2364 let column = chunk.into_body();
2365 assert!(
2366 column.length_in_bytes() <= COMMIT_BYTES || column.len() == 1,
2367 "{} bytes in a {}-record piece",
2368 column.length_in_bytes(),
2369 column.len(),
2370 );
2371 let view = column.borrow();
2372 for index in 0..view.len() {
2373 let ((key, _), _, _) = view.get(index);
2374 keys.push(*key);
2375 }
2376 }
2377 assert_eq!(keys, (0..20_000u64).collect::<Vec<_>>());
2378 }
2379
2380 #[mz_ore::test]
2384 fn merge_output_within_byte_bound_before_settle() {
2385 let interleaved = (
2386 wide_column((0..1600).step_by(2), 2100),
2387 wide_column((1..1600).step_by(2), 2100),
2388 );
2389 let runs = (
2390 wide_column((0..500).chain(1000..1100), 2100),
2391 wide_column(500..1000, 2100),
2392 );
2393 for (left, right) in [interleaved, runs] {
2394 let expected = left.len() + right.len();
2395 let mut left = VecDeque::from([WideChunk::from_body(left)]);
2396 let mut right = VecDeque::from([WideChunk::from_body(right)]);
2397 let mut merged = VecDeque::new();
2398 while !left.is_empty() && !right.is_empty() {
2399 WideChunk::merge(&mut left, &mut right, &mut merged);
2400 }
2401 assert!(merged.len() > 1, "merge never cut its output");
2402 merged.append(&mut left);
2403 merged.append(&mut right);
2404 assert_wide_byte_bound(merged, expected, 2100);
2405 }
2406 }
2407
2408 #[mz_ore::test]
2409 fn settle_split_preserves_depth() {
2410 let column = Rc::new(wide_column(0..2000, 2100));
2411 let shared = Rc::clone(&column);
2412 let mut input = VecDeque::from([WideChunk::Resident(column, 3)]);
2413 let mut settled = VecDeque::new();
2414 WideChunk::settle(&mut input, true, &mut settled);
2415 assert!(settled.len() > 1, "oversized chunk was not split");
2416 for chunk in &settled {
2417 assert_eq!(chunk.depth(), 3, "split piece changed generation");
2418 }
2419 assert_eq!(shared.len(), 2000, "shared body was mutated");
2420 assert_wide_byte_bound(settled, 2000, 2100);
2421 }
2422
2423 #[mz_ore::test]
2426 fn settle_split_then_coalesce_keeps_depth() {
2427 let mut input = VecDeque::from([
2428 WideChunk::Resident(Rc::new(wide_column(0..800, 3000)), 2),
2429 WideChunk::Resident(Rc::new(wide_column(800..810, 3000)), 0),
2430 ]);
2431 let mut settled = VecDeque::new();
2432 WideChunk::settle(&mut input, true, &mut settled);
2433 assert!(settled.len() > 1, "oversized chunk was not split");
2434 for chunk in &settled {
2435 assert_eq!(chunk.depth(), 2, "coalescing lost the deeper generation");
2436 }
2437 assert_wide_byte_bound(settled, 810, 3000);
2438 }
2439
2440 #[mz_ore::test]
2444 fn settle_fills_accumulated_column_from_next_chunk() {
2445 let per_chunk = 370;
2446 let mut input: VecDeque<WideChunk> = (0..8u64)
2447 .map(|i| {
2448 let start = i * per_chunk;
2449 WideChunk::from_body(wide_column(start..start + per_chunk, 3000))
2450 })
2451 .collect();
2452 let chunk_bytes = input[0].clone().into_body().length_in_bytes();
2453 assert!(2 * chunk_bytes > COMMIT_BYTES && chunk_bytes < COMMIT_BYTES);
2454 let mut settled = VecDeque::new();
2455 WideChunk::settle(&mut input, true, &mut settled);
2456 let sizes: Vec<usize> = settled
2457 .iter()
2458 .map(|chunk| chunk.clone().into_body().length_in_bytes())
2459 .collect();
2460 let (last, rest) = sizes.split_last().expect("settled output");
2461 assert!(*last <= COMMIT_BYTES);
2462 for size in rest {
2463 assert!(
2464 *size >= COMMIT_BYTES - COMMIT_BYTES / 10 && *size <= COMMIT_BYTES,
2465 "committed {size} bytes, sizes {sizes:?}",
2466 );
2467 }
2468 assert_wide_byte_bound(settled, 8 * 370, 3000);
2469 }
2470
2471 #[mz_ore::test]
2474 #[cfg_attr(miri, ignore)]
2475 fn advance_cuts_large_output() {
2476 let records: Vec<Tuple> = (0..300_000u64).map(|k| ((k, 0), 0, 1)).collect();
2477 let mut input = VecDeque::from([ColumnChunk::from_body(build_column(&records))]);
2478 let frontier = Antichain::from_elem(0u64);
2479 let mut out = VecDeque::new();
2480 TestChunk::advance(&mut input, frontier.borrow(), true, &mut out);
2481 assert!(input.is_empty());
2482 assert!(
2483 out.len() >= 2,
2484 "expected a cut output, got {} chunk(s)",
2485 out.len()
2486 );
2487 assert_eq!(collect_bounded(out, 2 * COMMIT_BYTES), records);
2488 }
2489
2490 #[mz_ore::test]
2493 #[cfg_attr(miri, ignore)]
2494 fn advance_withholds_giant_group() {
2495 let records: Vec<Tuple> = (0..100u64).map(|t| ((7, 7), t, 1)).collect();
2496 let mut input: VecDeque<TestChunk> = VecDeque::new();
2497 for piece in records.chunks(30) {
2498 input.push_back(ColumnChunk::from_body(build_column(piece)));
2499 }
2500 let frontier = Antichain::from_elem(50u64);
2501 let mut out = VecDeque::new();
2502 TestChunk::advance(&mut input, frontier.borrow(), false, &mut out);
2503 assert!(out.is_empty(), "nothing may ship from a single open group");
2504 assert_eq!(input.len(), 1, "the whole input becomes one carry chunk");
2505 TestChunk::advance(&mut input, frontier.borrow(), true, &mut out);
2507 assert!(input.is_empty());
2508 let advanced = records.iter().map(|&(d, t, r)| (d, t.max(50), r)).collect();
2509 assert_eq!(collect_chunks(out), consolidate(advanced));
2510 }
2511
2512 #[mz_ore::test]
2517 #[cfg_attr(miri, ignore)]
2518 fn extract_passes_frontier_disjoint_chunks_through() {
2519 set_spill_override(Some(test_pool()));
2520 let low: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), i % 4, 1)).collect();
2521 let high: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 6 + i % 4, 1)).collect();
2522 let spilled_chunk = |data: &[Tuple]| {
2523 let chunk = TestChunk::commit(build_column(&consolidate(data.to_vec())), 1);
2524 assert!(chunk.is_spilled());
2525 chunk
2526 };
2527
2528 let mut input = VecDeque::from([spilled_chunk(&low), spilled_chunk(&high)]);
2533 let frontier = Antichain::from_elem(5u64);
2534 let mut residual = Antichain::new();
2535 let (mut keep, mut ship) = (VecDeque::new(), VecDeque::new());
2536 while !input.is_empty() {
2537 TestChunk::extract(
2538 &mut input,
2539 frontier.borrow(),
2540 &mut residual,
2541 &mut keep,
2542 &mut ship,
2543 );
2544 }
2545 assert_eq!(ship.len(), 1);
2546 assert!(ship[0].is_spilled(), "shipped whole: body untouched");
2547 assert_eq!(keep.len(), 1);
2548 assert!(keep[0].is_spilled(), "kept whole: body untouched");
2549 assert_eq!(residual, Antichain::from_elem(6));
2550 let shipped = ship.pop_front().unwrap().into_body();
2551 assert_eq!(collect_column(&shipped), consolidate(low));
2552 let kept = keep.pop_front().unwrap().into_body();
2553 assert_eq!(collect_column(&kept), consolidate(high));
2554 set_spill_override(None);
2555 }
2556
2557 #[mz_ore::test]
2560 #[cfg_attr(miri, ignore)]
2561 fn extract_cuts_large_output() {
2562 let records: Vec<Tuple> = (0..300_000u64).map(|k| ((k, 0), k % 2, 1)).collect();
2563 let mut input = VecDeque::from([ColumnChunk::from_body(build_column(&records))]);
2564 let frontier = Antichain::from_elem(1u64);
2565 let mut residual = Antichain::new();
2566 let (mut keep, mut ship) = (VecDeque::new(), VecDeque::new());
2567 while !input.is_empty() {
2568 TestChunk::extract(
2569 &mut input,
2570 frontier.borrow(),
2571 &mut residual,
2572 &mut keep,
2573 &mut ship,
2574 );
2575 }
2576 assert!(
2577 keep.len() >= 2,
2578 "expected a cut keep side, got {} chunk(s)",
2579 keep.len()
2580 );
2581 assert!(
2582 ship.len() >= 2,
2583 "expected a cut ship side, got {} chunk(s)",
2584 ship.len()
2585 );
2586 let kept: Vec<Tuple> = records.iter().copied().filter(|r| r.1 >= 1).collect();
2587 let shipped: Vec<Tuple> = records.iter().copied().filter(|r| r.1 < 1).collect();
2588 assert_eq!(collect_bounded(keep, 2 * COMMIT_BYTES), kept);
2589 assert_eq!(collect_bounded(ship, 2 * COMMIT_BYTES), shipped);
2590 assert_eq!(residual, Antichain::from_elem(1));
2591 }
2592
2593 #[mz_ore::test]
2596 fn locate_uses_resident_bounds() {
2597 let pool = test_pool();
2598 let data: Vec<Tuple> = vec![((2, 0), 0, 1), ((4, 0), 0, 1)];
2599 let chunk = force_spill(ColumnChunk::from_body(build_column(&data)), &pool);
2600
2601 let mut probe_col = <u64 as Columnar>::Container::default();
2602 for key in [1u64, 3, 5] {
2603 probe_col.push(key);
2604 }
2605 let probes = probe_col.borrow();
2606 assert_eq!(chunk.locate(probes, 0), std::cmp::Ordering::Less);
2607 assert_eq!(chunk.locate(probes, 1), std::cmp::Ordering::Equal);
2608 assert_eq!(chunk.locate(probes, 2), std::cmp::Ordering::Greater);
2609 }
2610
2611 #[mz_ore::test]
2614 #[cfg_attr(miri, ignore)] fn spill_round_trip() {
2616 set_spill_override(Some(test_pool()));
2617
2618 let data: Vec<Tuple> = (0..40_000u64)
2619 .map(|i| ((i / 4, i % 4), i % 8, 1i64))
2620 .collect();
2621 let data = consolidate(data);
2622
2623 let column = build_column(&data);
2624 let committed = TestChunk::commit(column, 0);
2625 assert!(committed.is_spilled(), "large body must spill");
2626 assert_eq!(committed.len(), data.len());
2627 assert_eq!(collect_column(&committed.clone().into_body()), data);
2628
2629 let mut batcher = TestChunkBatcher::new(None, 0);
2630 for piece in data.chunks(10_000) {
2631 batcher.push_chunk(ColumnChunk::from_body(build_column(piece)));
2632 }
2633 let (sealed, _frontier) = batcher.extract(Antichain::new().borrow());
2634 let sealed = sealed.unwrap_or_default();
2635 assert!(
2636 sealed.iter().any(ColumnChunk::is_spilled),
2637 "sealed output should contain spilled chunks",
2638 );
2639 assert_eq!(collect_chunks(sealed), data);
2640
2641 set_spill_override(None);
2642 }
2643
2644 #[mz_ore::test]
2647 #[cfg_attr(miri, ignore)] fn merge_spilled_chains() {
2649 set_spill_override(Some(test_pool()));
2650
2651 let a: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect();
2652 let b: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 2i64)).collect();
2653
2654 let mut in1 = VecDeque::from([TestChunk::commit(build_column(&a), 0)]);
2655 let mut in2 = VecDeque::from([TestChunk::commit(build_column(&b), 0)]);
2656 assert!(in1[0].is_spilled() && in2[0].is_spilled());
2657
2658 let mut out = VecDeque::new();
2659 while !in1.is_empty() && !in2.is_empty() {
2660 TestChunk::merge(&mut in1, &mut in2, &mut out);
2661 }
2662 for tail in in1.drain(..).chain(in2.drain(..)) {
2663 out.push_back(tail);
2664 }
2665
2666 let expected: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 3i64)).collect();
2667 assert_eq!(collect_chunks(out), expected);
2668
2669 set_spill_override(None);
2670 }
2671
2672 #[mz_ore::test]
2675 fn merge_untouched_survivor_stays_spilled() {
2676 let pool = test_pool();
2677 let low: Vec<Tuple> = (0..100u64).map(|i| ((i, 0), 0, 1i64)).collect();
2678 let high: Vec<Tuple> = (1000..1100u64).map(|i| ((i, 0), 0, 1i64)).collect();
2679
2680 let mut in1 = VecDeque::from([force_spill(
2681 ColumnChunk::from_body(build_column(&low)),
2682 &pool,
2683 )]);
2684 let mut in2 = VecDeque::from([force_spill(
2685 ColumnChunk::from_body(build_column(&high)),
2686 &pool,
2687 )]);
2688 let mut out = VecDeque::new();
2689 TestChunk::merge(&mut in1, &mut in2, &mut out);
2690
2691 assert!(in1.is_empty());
2694 assert_eq!(in2.len(), 1);
2695 assert!(in2[0].is_spilled(), "untouched survivor must stay spilled");
2696 let mut all = collect_chunks(out);
2697 Extend::extend(&mut all, collect_chunks(in2.drain(..)));
2698 let mut expected = low;
2699 Extend::extend(&mut expected, high);
2700 assert_eq!(all, expected);
2701 }
2702
2703 #[mz_ore::test]
2707 fn merge_derives_generational_depth() {
2708 let low: Vec<Tuple> = (0..100u64).map(|i| ((i, 0), 0, 1i64)).collect();
2709 let high: Vec<Tuple> = (50..150u64).map(|i| ((i, 0), 0, 1i64)).collect();
2710 let mut in1 = VecDeque::from([ColumnChunk::from_body(build_column(&low))]);
2711 let mut in2 = VecDeque::from([ColumnChunk::from_body(build_column(&high))]);
2712 assert_eq!(in1[0].depth(), 0, "fresh chunks start at depth 0");
2713 let mut out = VecDeque::new();
2714 TestChunk::merge(&mut in1, &mut in2, &mut out);
2715 assert!(!out.is_empty());
2716 for chunk in &out {
2717 assert_eq!(chunk.depth(), 1, "merge output is one past its inputs");
2718 }
2719 assert!(in1.is_empty());
2722 assert_eq!(in2.len(), 1);
2723 assert_eq!(in2[0].depth(), 0, "rewritten survivor keeps its depth");
2724
2725 let mut in1 = VecDeque::from([ColumnChunk::Resident(Rc::new(build_column(&low)), 3)]);
2728 let far: Vec<Tuple> = (1000..1100u64).map(|i| ((i, 0), 0, 1i64)).collect();
2729 let mut in2 = VecDeque::from([ColumnChunk::from_body(build_column(&far))]);
2730 let mut out = VecDeque::new();
2731 TestChunk::merge(&mut in1, &mut in2, &mut out);
2732 assert_eq!(out.len(), 1);
2733 assert_eq!(out[0].depth(), 4, "pass-through ages a generation");
2734 assert_eq!(collect_chunks(out), low);
2735 }
2736
2737 #[mz_ore::test]
2740 fn advance_preserves_depth() {
2741 let data: Vec<Tuple> = (0..100u64).map(|i| ((i, 0), 1, 1i64)).collect();
2742 let mut input = VecDeque::from([
2743 ColumnChunk::Resident(Rc::new(build_column(&data[..50])), 2),
2744 ColumnChunk::Resident(Rc::new(build_column(&data[50..])), 1),
2745 ]);
2746 let frontier = Antichain::from_elem(5u64);
2747 let mut out = VecDeque::new();
2748 TestChunk::advance(&mut input, frontier.borrow(), false, &mut out);
2749 for chunk in out.iter().chain(input.iter()) {
2750 assert_eq!(chunk.depth(), 2);
2751 }
2752 TestChunk::advance(&mut input, frontier.borrow(), true, &mut out);
2753 assert!(input.is_empty());
2754 assert!(!out.is_empty());
2755 for chunk in &out {
2756 assert_eq!(chunk.depth(), 2);
2757 }
2758 }
2759
2760 #[mz_ore::test]
2764 #[cfg_attr(miri, ignore)] fn settle_commits_at_accumulated_depth() {
2766 set_spill_override(Some(test_pool()));
2767 let big: Vec<Tuple> = (0..60_000u64).map(|i| ((i, 0), 0, 1i64)).collect();
2768 let mut input = VecDeque::from([
2769 ColumnChunk::Resident(Rc::new(build_column(&big)), 1),
2770 ColumnChunk::Resident(Rc::new(build_column(&[((0, 0), 0, 1)])), 0),
2771 ColumnChunk::Resident(Rc::new(build_column(&[((1, 0), 0, 1)])), 2),
2772 ]);
2773 let mut out = VecDeque::new();
2774 TestChunk::settle(&mut input, true, &mut out);
2775 assert!(input.is_empty());
2776 assert_eq!(out.len(), 2);
2777 assert!(out[0].is_spilled(), "large commit must spill");
2778 assert_eq!(out[0].depth(), 1, "sole commit keeps its depth");
2779 assert!(!out[1].is_spilled(), "small commit stays resident");
2780 assert_eq!(out[1].depth(), 2, "coalesced commit takes the max depth");
2781 set_spill_override(None);
2782 }
2783
2784 #[mz_ore::test]
2788 #[cfg_attr(miri, ignore)] fn settle_carry_commits_at_target() {
2790 let chunk_rows = u64::cast_from(800_000usize / 32);
2792 let mut input: VecDeque<TestChunk> = (0..4u64)
2793 .map(|c| {
2794 let data: Vec<Tuple> = (0..chunk_rows)
2795 .map(|i| ((c * chunk_rows + i, 0), 0, 1i64))
2796 .collect();
2797 ColumnChunk::from_body(build_column(&data))
2798 })
2799 .collect();
2800 let mut out = VecDeque::new();
2801 TestChunk::settle(&mut input, true, &mut out);
2802 assert!(out.len() < 4, "nothing coalesced");
2805 for chunk in &out {
2806 let col = chunk.clone().into_body();
2807 assert!(
2808 col.length_in_bytes() <= COMMIT_BYTES,
2809 "settled chunk of {} bytes exceeds the commit target",
2810 col.length_in_bytes(),
2811 );
2812 }
2813 assert_eq!(
2814 collect_chunks(out).len(),
2815 usize::try_from(4 * chunk_rows).unwrap(),
2816 );
2817 }
2818
2819 #[mz_ore::test]
2820 fn small_chunks_stay_resident() {
2821 set_spill_override(Some(test_pool()));
2822 let committed = TestChunk::commit(build_column(&[((1, 1), 0, 1)]), 0);
2823 assert!(!committed.is_spilled());
2824 set_spill_override(None);
2825 }
2826
2827 fn column_at_spill_floor() -> (ColumnBody<Tuple>, u64) {
2830 let mut col: ColumnBody<Tuple> = ColumnBody::default();
2831 let mut n = 0u64;
2832 while col.length_in_bytes() < SPILL_MIN_BYTES {
2833 col.push_into(((n, n), 0, 1));
2834 n += 1;
2835 }
2836 (col, n)
2837 }
2838
2839 #[mz_ore::test]
2842 fn spill_floor_boundary() {
2843 set_spill_override(Some(test_pool()));
2844 let (col, n) = column_at_spill_floor();
2845 let mut under: ColumnBody<Tuple> = ColumnBody::default();
2846 for m in 0..n - 1 {
2847 under.push_into(((m, m), 0, 1));
2848 }
2849 assert!(under.length_in_bytes() < SPILL_MIN_BYTES);
2850 assert!(!TestChunk::commit(under, 0).is_spilled());
2851 assert!(TestChunk::commit(col, 0).is_spilled());
2852 set_spill_override(None);
2853 }
2854
2855 #[mz_ore::test]
2859 fn spill_codec_depth_floor() {
2860 set_spill_override(Some(test_pool()));
2861 set_compress_min_depth_override(Some(2));
2862 let codec_name = |depth: u8| {
2866 let (codec, compressed) = codec_for_depth(depth);
2867 let name = format!("{:?}", codec);
2868 assert_eq!(compressed, name == "Lz4Codec", "flag tracks the codec");
2869 name
2870 };
2871 assert_eq!(codec_name(0), "IdentityCodec");
2872 assert_eq!(codec_name(1), "IdentityCodec");
2873 assert_eq!(codec_name(2), "Lz4Codec");
2874 assert_eq!(codec_name(u8::MAX), "Lz4Codec");
2875
2876 let data: Vec<Tuple> = (0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect();
2877 let data = consolidate(data);
2878 let column = build_column(&data);
2879 for depth in [0u8, 1, 2, 3] {
2880 let chunk = TestChunk::commit(column.clone(), depth);
2881 assert!(chunk.is_spilled(), "depth {depth} must spill");
2882 assert_eq!(collect_column(&chunk.into_body()), data);
2883 }
2884 set_spill_override(None);
2885 set_compress_min_depth_override(None);
2886
2887 set_compress_min_depth_override(Some(DEFAULT_COMPRESS_MIN_DEPTH));
2889 assert_eq!(codec_name(0), "IdentityCodec");
2890 assert_eq!(codec_name(1), "Lz4Codec");
2891 set_compress_min_depth_override(None);
2892 }
2893
2894 #[mz_ore::test]
2900 #[cfg_attr(miri, ignore)] fn merge_survivor_crosses_compression_floor() {
2902 set_spill_override(Some(test_pool()));
2903 set_compress_min_depth_override(Some(1));
2904
2905 let low = consolidate((0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2906 let far = consolidate((100_000..120_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2907 let fresh_far = || VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
2908
2909 let mut in1 = VecDeque::from([TestChunk::commit(build_column(&low), 0)]);
2912 let mut in2 = fresh_far();
2913 assert!(in1[0].is_spilled() && in2[0].is_spilled());
2914 assert!(
2915 !body_compressed(&in1[0]),
2916 "a fresh body below the floor is identity coded"
2917 );
2918
2919 let mut out = VecDeque::new();
2920 TestChunk::merge(&mut in1, &mut in2, &mut out);
2921 assert_eq!(out.len(), 1);
2922 let survived = out.pop_front().expect("the lower front passes through");
2923 assert_eq!(survived.depth(), 1, "survival ages across the floor");
2924 assert!(
2925 survived.is_spilled(),
2926 "the crossing re-spills, it does not evict"
2927 );
2928 assert!(
2929 body_compressed(&survived),
2930 "the survivor is re-spilled under the compressing codec"
2931 );
2932
2933 let mut in1 = VecDeque::from([survived]);
2936 let mut in2 = fresh_far();
2937 let mut out = VecDeque::new();
2938 TestChunk::merge(&mut in1, &mut in2, &mut out);
2939 assert_eq!(out.len(), 1);
2940 assert_eq!(out[0].depth(), 2, "an aged survivor keeps aging");
2941 assert!(out[0].is_spilled());
2942 assert_eq!(
2943 collect_chunks(out),
2944 low,
2945 "the body reads back intact across both survivals"
2946 );
2947
2948 set_spill_override(None);
2949 set_compress_min_depth_override(None);
2950 }
2951
2952 #[mz_ore::test]
2959 #[cfg_attr(miri, ignore)] fn merge_survivor_ages_while_shared() {
2961 set_spill_override(Some(test_pool()));
2962 set_compress_min_depth_override(Some(1));
2963
2964 let low = consolidate((0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2965 let far = consolidate((100_000..120_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
2966
2967 let source = TestChunk::commit(build_column(&low), 0);
2970 let ColumnChunk::Spilled(source_body, 0) = &source else {
2971 panic!("a fresh commit above the spill floor is spilled at depth 0");
2972 };
2973 let source_body = Rc::clone(source_body);
2974
2975 let mut in1 = VecDeque::from([source.clone()]);
2976 let mut in2 = VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
2977 let mut out = VecDeque::new();
2978 TestChunk::merge(&mut in1, &mut in2, &mut out);
2979
2980 assert_eq!(out.len(), 1);
2981 assert_eq!(out[0].depth(), 1, "a shared body ages all the same");
2982 let ColumnChunk::Spilled(survived_body, _) = &out[0] else {
2983 panic!("the survivor stays spilled");
2984 };
2985 assert!(
2986 Rc::ptr_eq(&source_body, survived_body),
2987 "a shared body is aged in place, not re-spilled"
2988 );
2989 assert_eq!(source.depth(), 0, "the other holder is left as it was");
2990
2991 let mut in1 = VecDeque::from([out.pop_front().expect("survivor observed above")]);
2994 let mut in2 = VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
2995 let mut out = VecDeque::new();
2996 TestChunk::merge(&mut in1, &mut in2, &mut out);
2997 assert_eq!(out.len(), 1);
2998 assert_eq!(out[0].depth(), 2, "aging past the floor is not pinned");
2999 assert_eq!(collect_chunks(out), low);
3000
3001 set_spill_override(None);
3002 set_compress_min_depth_override(None);
3003 }
3004
3005 #[mz_ore::test]
3011 #[cfg_attr(miri, ignore)] fn survive_merge_retries_missed_migrations() {
3013 let low = consolidate((0..20_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
3014 let far = consolidate((100_000..120_000u64).map(|i| ((i, 0), 0, 1i64)).collect());
3015
3016 let survive = |chunk: TestChunk| {
3019 let mut in1 = VecDeque::from([chunk]);
3020 let mut in2 = VecDeque::from([TestChunk::commit(build_column(&far), 0)]);
3021 let mut out = VecDeque::new();
3022 TestChunk::merge(&mut in1, &mut in2, &mut out);
3023 out.pop_front().expect("the lower front passes through")
3024 };
3025
3026 set_spill_override(Some(test_pool()));
3029 set_compress_min_depth_override(Some(1));
3030 let chunk = TestChunk::commit(build_column(&low), 0);
3031 assert!(!body_compressed(&chunk));
3032 set_spill_override(None);
3033 let chunk = survive(chunk);
3034 assert_eq!(chunk.depth(), 1, "aging does not need a pool");
3035 assert!(!body_compressed(&chunk), "no pool, no migration");
3036 set_spill_override(Some(test_pool()));
3037 let chunk = survive(chunk);
3038 assert!(
3039 body_compressed(&chunk),
3040 "the migration retries once a pool is back"
3041 );
3042
3043 let chunk = TestChunk::commit(build_column(&low), 0);
3046 let held = chunk.clone();
3047 let chunk = survive(chunk);
3048 assert!(!body_compressed(&chunk), "shared, so not migrated");
3049 drop(held);
3050 let chunk = survive(chunk);
3051 assert!(
3052 body_compressed(&chunk),
3053 "the migration retries once the body is unshared"
3054 );
3055
3056 set_compress_min_depth_override(Some(8));
3061 let chunk = TestChunk::commit(build_column(&low), 3);
3062 assert!(!body_compressed(&chunk));
3063 set_compress_min_depth_override(Some(1));
3064 let chunk = survive(chunk);
3065 assert_eq!(chunk.depth(), 4);
3066 assert!(
3067 body_compressed(&chunk),
3068 "lowering the floor migrates bodies already past it"
3069 );
3070
3071 set_spill_override(None);
3072 set_compress_min_depth_override(None);
3073 }
3074
3075 #[mz_ore::test]
3077 #[cfg_attr(miri, ignore)]
3078 fn spill_gates_are_independent() {
3079 let installed =
3080 crate::pool_config::apply_pool_config(crate::pool_config::PoolPagerConfig {
3081 budget_bytes: 32 << 20,
3082 spill_threads: 1,
3083 eager_backing: false,
3084 rss_target_bytes: 16 << 20,
3085 });
3086 assert!(installed, "pool reservation failed");
3087 let (col, _) = column_at_spill_floor();
3089 let compute = |col: &ColumnBody<Tuple>| TestChunk::commit(col.clone(), 0).is_spilled();
3090 let storage = |col: &ColumnBody<Tuple>| {
3091 ColumnChunk::<(u64, u64), u64, i64, StorageSpill>::commit(col.clone(), 0).is_spilled()
3092 };
3093 let sink = |col: &ColumnBody<Tuple>| try_spill_ref(col, 0).is_some();
3094
3095 assert!(!compute(&col), "all gates off");
3096 assert!(!storage(&col), "all gates off");
3097 assert!(!sink(&col), "all gates off");
3098
3099 set_storage_spill_enabled(true);
3100 assert!(storage(&col), "the storage gate spills storage chunks");
3101 assert!(
3102 !compute(&col),
3103 "the storage gate must not spill compute chunks"
3104 );
3105 assert!(!sink(&col), "the storage gate must not spill sink bodies");
3106 set_storage_spill_enabled(false);
3107
3108 set_compute_spill_enabled(true);
3109 assert!(compute(&col), "the compute gate spills compute chunks");
3110 assert!(
3111 !storage(&col),
3112 "the compute gate must not spill storage chunks"
3113 );
3114 assert!(!sink(&col), "the compute gate must not spill sink bodies");
3115 set_compute_spill_enabled(false);
3116
3117 set_sink_spill_enabled(true);
3118 assert!(sink(&col), "the sink gate spills sink bodies");
3119 assert!(
3120 !compute(&col),
3121 "the sink gate must not spill compute chunks"
3122 );
3123 assert!(
3124 !storage(&col),
3125 "the sink gate must not spill storage chunks"
3126 );
3127 set_sink_spill_enabled(false);
3128 set_compress_min_depth_override(None);
3129 }
3130
3131 #[mz_ore::test]
3136 #[cfg_attr(miri, ignore)] fn spill_align_round_trip() {
3138 let pool = Pool::new().expect("pool reservation");
3139 pool.set_budget(0);
3140 let data: Vec<Tuple> = (0..64u64).map(|k| ((k, k), 0, 1)).collect();
3141 let spilled = force_spill(ColumnChunk::from_body(build_column(&data)), &pool);
3142 let column = spilled.into_body();
3143 let ColumnBody::Words(words) = &column else {
3144 panic!("an evicted body reads back as ColumnBody::Words");
3145 };
3146 let words = words.clone();
3147 let respilled = force_spill(ColumnChunk::from_body(column), &pool);
3148 let reread = respilled.into_body();
3149 let ColumnBody::Words(words2) = &reread else {
3150 panic!("an evicted body reads back as ColumnBody::Words");
3151 };
3152 assert_eq!(&words, words2, "byte-identical round trip");
3153 assert_eq!(collect_column(&reread), data);
3154 }
3155
3156 #[mz_ore::test]
3160 #[cfg_attr(miri, ignore)] fn time_bounds_race_spill_thread_eviction() {
3162 let (done_tx, done_rx) = std::sync::mpsc::channel();
3163 std::thread::spawn(move || {
3164 let pool = Pool::new().expect("pool reservation");
3165 pool.set_budget(0);
3166 pool.set_spill_threads(4);
3167 let data: Vec<Tuple> = (0..2048u64).map(|k| ((k, k), k % 7, 1)).collect();
3168 for _ in 0..2000 {
3169 let chunk = force_spill(ColumnChunk::from_body(build_column(&data)), &pool);
3170 let (lower, upper) = chunk.chunk_time_bounds();
3171 assert_eq!(lower.elements(), &[0]);
3172 assert_eq!(&upper[..], &[6]);
3173 }
3174 pool.set_spill_threads(0);
3175 done_tx.send(()).expect("receiver alive");
3176 });
3177 done_rx
3178 .recv_timeout(std::time::Duration::from_secs(120))
3179 .expect("time bounds reads deadlocked against spill-thread eviction");
3180 }
3181
3182 #[mz_ore::test]
3183 fn merge_depth_saturates() {
3184 let a = ColumnChunk::Resident(
3185 Rc::new(build_column(&[((1, 0), 0, 1), ((3, 0), 0, 1)])),
3186 u8::MAX,
3187 );
3188 let b = ColumnChunk::Resident(
3189 Rc::new(build_column(&[((2, 0), 0, 1), ((4, 0), 0, 1)])),
3190 u8::MAX,
3191 );
3192 let mut in1 = VecDeque::from([a]);
3193 let mut in2 = VecDeque::from([b]);
3194 let mut out = VecDeque::new();
3195 TestChunk::merge(&mut in1, &mut in2, &mut out);
3196 for chunk in out.iter().chain(in1.iter()).chain(in2.iter()) {
3197 assert_eq!(chunk.depth(), u8::MAX, "depth saturates");
3198 }
3199 }
3200
3201 #[mz_ore::test]
3202 fn into_body_copies_shared_resident() {
3203 let data: Vec<Tuple> = vec![((1, 1), 0, 1), ((2, 2), 0, 1)];
3204 let a = TestChunk::from_body(build_column(&data));
3205 let b = a.clone();
3206 assert_eq!(collect_column(&a.into_body()), data);
3207 assert_eq!(collect_column(&b.into_body()), data);
3208 }
3209}