1use std::cell::OnceCell;
31use std::collections::VecDeque;
32use std::marker::PhantomData;
33use std::rc::Rc;
34
35use columnar::{Borrow, Columnar, Container, ContainerOf, Index, Len, Push};
36use timely::Accountable;
37use timely::container::{PushInto, SizableContainer};
38use timely::progress::Antichain;
39use timely::progress::frontier::AntichainRef;
40
41use crate::consolidation::Consolidate;
42use crate::lattice::Lattice;
43use crate::trace::Navigable;
44use crate::trace::cursor::Cursor;
45use crate::trace::implementations::{BatchContainer, Layout, WithLayout};
46
47use crate::columnar::layout::{ColumnarLayout, ColumnarUpdate, Coltainer};
48use crate::columnar::updates::{child_range, UpdatesBuilder, UpdatesTyped};
49use crate::columnar::trie_merger;
50
51use super::spill::{self, BytesSource};
52
53use crate::trace::chunk::Chunk;
54
55const TARGET: usize = 8192;
60
61pub struct ChunkMeta<U: ColumnarUpdate> {
65 fk: ContainerOf<U::Key>, fv: ContainerOf<U::Val>, ft: ContainerOf<U::Time>,
66 lk: ContainerOf<U::Key>, lv: ContainerOf<U::Val>, lt: ContainerOf<U::Time>,
67 len: usize,
68}
69
70pub struct PagedChunk<U: ColumnarUpdate> {
73 meta: ChunkMeta<U>,
74 source: Box<dyn BytesSource>,
75 cache: OnceCell<Rc<UpdatesTyped<U>>>,
76}
77
78fn singleton<C: Container + Default>(borrowed: <C as Borrow>::Borrowed<'_>, i: usize) -> C {
80 let mut out = C::default();
81 out.push(borrowed.get(i));
82 out
83}
84
85fn meta_of<U: ColumnarUpdate>(t: &UpdatesTyped<U>) -> ChunkMeta<U> {
87 let v = t.view();
88 let (nk, nv, nt) = (v.keys.values.len(), v.vals.values.len(), v.times.values.len());
89 ChunkMeta {
90 fk: singleton(v.keys.values, 0), lk: singleton(v.keys.values, nk - 1),
91 fv: singleton(v.vals.values, 0), lv: singleton(v.vals.values, nv - 1),
92 ft: singleton(v.times.values, 0), lt: singleton(v.times.values, nt - 1),
93 len: t.len(),
94 }
95}
96
97pub enum ColChunk<U: ColumnarUpdate> {
100 Resident(Rc<UpdatesTyped<U>>),
102 Paged(Rc<PagedChunk<U>>),
104}
105
106impl<U: ColumnarUpdate> Clone for ColChunk<U> {
107 fn clone(&self) -> Self {
108 match self {
109 ColChunk::Resident(rc) => ColChunk::Resident(Rc::clone(rc)),
110 ColChunk::Paged(p) => ColChunk::Paged(Rc::clone(p)),
111 }
112 }
113}
114impl<U: ColumnarUpdate> Default for ColChunk<U> {
115 fn default() -> Self { ColChunk::Resident(Rc::new(UpdatesTyped::default())) }
116}
117
118impl<U: ColumnarUpdate> ColChunk<U> {
119 pub fn from_trie(updates: UpdatesTyped<U>) -> Self { ColChunk::Resident(Rc::new(updates)) }
122
123 fn trie(&self) -> &UpdatesTyped<U> {
126 match self {
127 ColChunk::Resident(rc) => &**rc,
128 ColChunk::Paged(p) => &**p.cache.get_or_init(|| {
129 spill::note_fetched();
130 Rc::new(spill::decode::<U>(&*p.source))
131 }),
132 }
133 }
134
135 pub fn updates_mut(&mut self) -> &mut UpdatesTyped<U> {
139 if let ColChunk::Paged(_) = self {
140 *self = ColChunk::Resident(Rc::new(into_trie(std::mem::take(self))));
141 }
142 match self {
143 ColChunk::Resident(rc) => Rc::make_mut(rc),
144 ColChunk::Paged(_) => unreachable!(),
145 }
146 }
147}
148
149fn into_trie<U: ColumnarUpdate>(chunk: ColChunk<U>) -> UpdatesTyped<U> {
152 match chunk {
153 ColChunk::Resident(rc) => Rc::try_unwrap(rc).unwrap_or_else(|rc| (*rc).clone()),
154 ColChunk::Paged(p) => match p.cache.get() {
155 Some(rc) => (**rc).clone(),
156 None => { spill::note_fetched(); spill::decode::<U>(&*p.source) }
157 },
158 }
159}
160
161impl<U: ColumnarUpdate> Accountable for ColChunk<U> {
164 fn record_count(&self) -> i64 { Chunk::len(self) as i64 }
165}
166
167impl<U: ColumnarUpdate> SizableContainer for ColChunk<U> {
168 fn at_capacity(&self) -> bool { Chunk::len(self) >= TARGET }
170 fn ensure_capacity(&mut self, _stash: &mut Option<Self>) { }
172}
173
174impl<U: ColumnarUpdate> Consolidate for ColChunk<U> {
175 fn len(&self) -> usize { Chunk::len(self) }
176 fn clear(&mut self) { *self.updates_mut() = UpdatesTyped::default(); }
177 fn consolidate_into(&mut self, target: &mut Self) {
178 let taken = std::mem::take(self.updates_mut());
179 *target.updates_mut() = taken.consolidate();
180 }
181}
182
183impl<U: ColumnarUpdate> PushInto<((U::Key, U::Val), U::Time, U::Diff)> for ColChunk<U> {
184 fn push_into(&mut self, item: ((U::Key, U::Val), U::Time, U::Diff)) {
185 self.updates_mut().push_into(item);
186 }
187}
188
189pub struct ColChunkCursor<U: ColumnarUpdate> {
194 key_cursor: usize,
195 val_cursor: usize,
196 phantom: PhantomData<U>,
197}
198
199impl<U: ColumnarUpdate> WithLayout for ColChunk<U> {
200 type Layout = ColumnarLayout<U>;
201}
202impl<U: ColumnarUpdate> WithLayout for ColChunkCursor<U> {
203 type Layout = ColumnarLayout<U>;
204}
205
206impl<U: ColumnarUpdate> Cursor for ColChunkCursor<U> {
207 type Storage = ColChunk<U>;
208
209 type KeyContainer = <ColumnarLayout<U> as Layout>::KeyContainer;
210 type Key<'a> = <<ColumnarLayout<U> as Layout>::KeyContainer as BatchContainer>::ReadItem<'a>;
211 type ValContainer = <ColumnarLayout<U> as Layout>::ValContainer;
212 type Val<'a> = <<ColumnarLayout<U> as Layout>::ValContainer as BatchContainer>::ReadItem<'a>;
213 type ValOwn = <<ColumnarLayout<U> as Layout>::ValContainer as BatchContainer>::Owned;
214 type TimeContainer = <ColumnarLayout<U> as Layout>::TimeContainer;
215 type TimeGat<'a> = <<ColumnarLayout<U> as Layout>::TimeContainer as BatchContainer>::ReadItem<'a>;
216 type Time = <<ColumnarLayout<U> as Layout>::TimeContainer as BatchContainer>::Owned;
217 type DiffContainer = <ColumnarLayout<U> as Layout>::DiffContainer;
218 type DiffGat<'a> = <<ColumnarLayout<U> as Layout>::DiffContainer as BatchContainer>::ReadItem<'a>;
219 type Diff = <<ColumnarLayout<U> as Layout>::DiffContainer as BatchContainer>::Owned;
220
221 fn key_valid(&self, s: &Self::Storage) -> bool { self.key_cursor < s.trie().view().keys.values.len() }
222 fn val_valid(&self, s: &Self::Storage) -> bool {
223 let view = s.trie().view();
224 self.key_cursor < view.keys.values.len()
225 && self.val_cursor < child_range(view.vals.bounds, self.key_cursor).end
226 }
227 fn key<'a>(&self, s: &'a Self::Storage) -> Self::Key<'a> { s.trie().view().keys.values.get(self.key_cursor) }
228 fn val<'a>(&self, s: &'a Self::Storage) -> Self::Val<'a> { s.trie().view().vals.values.get(self.val_cursor) }
229 fn get_key<'a>(&self, s: &'a Self::Storage) -> Option<Self::Key<'a>> {
230 if self.key_valid(s) { Some(self.key(s)) } else { None }
231 }
232 fn get_val<'a>(&self, s: &'a Self::Storage) -> Option<Self::Val<'a>> {
233 if self.val_valid(s) { Some(self.val(s)) } else { None }
234 }
235 fn map_times<L: FnMut(Self::TimeGat<'_>, Self::DiffGat<'_>)>(&mut self, s: &Self::Storage, mut logic: L) {
236 if !self.val_valid(s) { return; }
237 let view = s.trie().view();
238 for t in child_range(view.times.bounds, self.val_cursor) {
239 let time = view.times.values.get(t);
240 for d in child_range(view.diffs.bounds, t) {
241 logic(time, view.diffs.values.get(d));
242 }
243 }
244 }
245 fn step_key(&mut self, s: &Self::Storage) {
246 self.key_cursor += 1;
247 if self.key_valid(s) { self.rewind_vals(s); }
248 else { self.key_cursor = s.trie().view().keys.values.len(); }
249 }
250 fn seek_key(&mut self, s: &Self::Storage, key: Self::Key<'_>) {
251 let view = s.trie().view();
252 let n = view.keys.values.len();
253 let mut lo = self.key_cursor;
254 trie_merger::gallop(view.keys.values, &mut lo, n, |x|
255 <Coltainer<U::Key> as BatchContainer>::reborrow(x).lt(&<Coltainer<U::Key> as BatchContainer>::reborrow(key)));
256 self.key_cursor = lo;
257 if self.key_valid(s) { self.rewind_vals(s); }
258 }
259 fn step_val(&mut self, s: &Self::Storage) {
260 self.val_cursor += 1;
261 if !self.val_valid(s) {
262 self.val_cursor = child_range(s.trie().view().vals.bounds, self.key_cursor).end;
263 }
264 }
265 fn seek_val(&mut self, s: &Self::Storage, val: Self::Val<'_>) {
266 if !self.key_valid(s) { return; }
267 let view = s.trie().view();
268 let upper = child_range(view.vals.bounds, self.key_cursor).end;
269 let mut lo = self.val_cursor;
270 trie_merger::gallop(view.vals.values, &mut lo, upper, |x|
271 <Coltainer<U::Val> as BatchContainer>::reborrow(x).lt(&<Coltainer<U::Val> as BatchContainer>::reborrow(val)));
272 self.val_cursor = lo;
273 }
274 fn rewind_keys(&mut self, s: &Self::Storage) { self.key_cursor = 0; self.rewind_vals(s); }
275 fn rewind_vals(&mut self, s: &Self::Storage) {
276 if self.key_valid(s) {
277 self.val_cursor = child_range(s.trie().view().vals.bounds, self.key_cursor).start;
278 }
279 }
280}
281
282fn emit<U: ColumnarUpdate>(updates: UpdatesTyped<U>, out: &mut VecDeque<ColChunk<U>>) {
284 if updates.len() > 0 { out.push_back(ColChunk::Resident(Rc::new(updates))); }
285}
286
287fn consolidated<U: ColumnarUpdate>(merged: UpdatesTyped<U>) -> UpdatesTyped<U> {
292 if merged.diffs.values.len() == merged.times.values.len() { merged } else { merged.filter_zero() }
293}
294
295impl<U: ColumnarUpdate> Navigable for ColChunk<U>
296where U::Time: 'static {
297 type Cursor = ColChunkCursor<U>;
298
299 fn cursor(&self) -> Self::Cursor {
300 ColChunkCursor { key_cursor: 0, val_cursor: 0, phantom: PhantomData }
301 }
302}
303
304impl<U: ColumnarUpdate> crate::trace::chunk::NavigableChunk for ColChunk<U>
305where U::Time: 'static {
306 fn bounds(&self) -> (
307 (<Self::Cursor as Cursor>::Key<'_>, <Self::Cursor as Cursor>::Val<'_>, <Self::Cursor as Cursor>::TimeGat<'_>),
308 (<Self::Cursor as Cursor>::Key<'_>, <Self::Cursor as Cursor>::Val<'_>, <Self::Cursor as Cursor>::TimeGat<'_>),
309 ) {
310 match self {
311 ColChunk::Resident(rc) => {
312 let view = rc.view();
313 let (nk, nv, nt) = (view.keys.values.len(), view.vals.values.len(), view.times.values.len());
314 ((view.keys.values.get(0), view.vals.values.get(0), view.times.values.get(0)),
316 (view.keys.values.get(nk - 1), view.vals.values.get(nv - 1), view.times.values.get(nt - 1)))
317 }
318 ColChunk::Paged(p) => {
319 let m = &p.meta;
321 ((m.fk.borrow().get(0), m.fv.borrow().get(0), m.ft.borrow().get(0)),
322 (m.lk.borrow().get(0), m.lv.borrow().get(0), m.lt.borrow().get(0)))
323 }
324 }
325 }
326}
327
328impl<U: ColumnarUpdate> Chunk for ColChunk<U>
329where U::Time: 'static {
330 type Time = <<ColumnarLayout<U> as Layout>::TimeContainer as BatchContainer>::Owned;
331
332 const TARGET: usize = TARGET;
333
334 fn len(&self) -> usize {
335 match self {
336 ColChunk::Resident(rc) => rc.len(),
337 ColChunk::Paged(p) => p.meta.len,
338 }
339 }
340
341 fn merge(in1: &mut VecDeque<Self>, in2: &mut VecDeque<Self>, out: &mut VecDeque<Self>) {
352 let mut cursor1 = Some(((0, 0, 0), into_trie(in1.pop_front().unwrap())));
353 let mut cursor2 = Some(((0, 0, 0), into_trie(in2.pop_front().unwrap())));
354 emit(consolidated(trie_merger::merge_pair(&mut cursor1, &mut cursor2)), out);
355 if let Some((cursor, batch)) = cursor1 {
357 let suffix = trie_merger::suffix_chunk(cursor, &batch);
358 if suffix.len() > 0 { in1.push_front(ColChunk::Resident(Rc::new(suffix))); }
359 }
360 if let Some((cursor, batch)) = cursor2 {
361 let suffix = trie_merger::suffix_chunk(cursor, &batch);
362 if suffix.len() > 0 { in2.push_front(ColChunk::Resident(Rc::new(suffix))); }
363 }
364 }
365
366 fn extract(
369 input: &mut VecDeque<Self>,
370 frontier: AntichainRef<U::Time>,
371 residual: &mut Antichain<U::Time>,
372 keep: &mut VecDeque<Self>,
373 ship: &mut VecDeque<Self>,
374 ) {
375 let Some(chunk) = input.pop_front() else { return };
376 trie_merger::extract(
377 std::iter::once(into_trie(chunk)),
378 frontier,
379 residual,
380 |c| emit(c, ship),
381 |c| emit(c, keep),
382 );
383 }
384
385 fn advance(
397 input: &mut VecDeque<Self>,
398 frontier: AntichainRef<U::Time>,
399 done: bool,
400 out: &mut VecDeque<Self>,
401 ) {
402 let mut carry: Option<UpdatesTyped<U>> = None;
405 while let Some(chunk) = input.pop_front() {
406 let combined = match carry.take() {
409 None => into_trie(chunk),
410 Some(c) => {
411 let mut builder = UpdatesBuilder::new_from(c);
412 builder.meld(&into_trie(chunk));
413 builder.done()
414 }
415 };
416 if combined.len() == 0 { continue; }
417
418 let tail = {
421 let v = combined.view();
422 let last_val = child_range(v.vals.bounds, v.keys.values.len() - 1).end - 1;
423 let times = child_range(v.times.bounds, last_val);
424 child_range(v.diffs.bounds, times.end - 1).end
425 - child_range(v.diffs.bounds, times.start).start
426 };
427 if tail == combined.len() {
428 carry = Some(combined);
430 continue;
431 }
432 let split = combined.len() - tail;
433 let (keep, rest) = trie_merger::split_at(combined, split);
434 carry = Some(rest);
435 emit(advance_trie(keep, frontier), out);
437 }
438 if let Some(c) = carry {
440 if done { emit(advance_trie(c, frontier), out); }
441 else { input.push_front(ColChunk::Resident(Rc::new(c))); }
442 }
443 }
444
445 fn settle(input: &mut VecDeque<Self>, done: bool, out: &mut VecDeque<Self>) {
451 crate::trace::chunk::pack(
452 input, done, out,
453 |acc, next| {
454 let mut build = UpdatesBuilder::new_from(into_trie(std::mem::take(acc)));
455 build.meld(&into_trie(next));
456 *acc = ColChunk::Resident(Rc::new(build.done()));
457 },
458 |chunk, n| {
459 let (first, rest) = trie_merger::split_at(into_trie(chunk), n);
460 (ColChunk::Resident(Rc::new(first)), ColChunk::Resident(Rc::new(rest)))
461 },
462 seal_chunk,
463 );
464 }
465}
466
467fn seal_chunk<U: ColumnarUpdate>(chunk: ColChunk<U>) -> ColChunk<U> {
471 let ColChunk::Resident(rc) = chunk else { return chunk };
472 if !spill::active() { return ColChunk::Resident(rc); }
473 let updates = Rc::try_unwrap(rc).unwrap_or_else(|rc| (*rc).clone());
474 let meta = meta_of(&updates);
475 match spill::try_page(updates) {
476 Ok(source) => ColChunk::Paged(Rc::new(PagedChunk { meta, source, cache: OnceCell::new() })),
477 Err(updates) => ColChunk::Resident(Rc::new(updates)),
478 }
479}
480
481fn advance_trie<U: ColumnarUpdate>(
488 keep: UpdatesTyped<U>,
489 frontier: AntichainRef<U::Time>,
490) -> UpdatesTyped<U> {
491 use crate::difference::{IsZero, Semigroup};
492 let view = keep.view();
493 let mut out = UpdatesTyped::<U>::default();
494 let mut run: Vec<(U::Time, U::Diff)> = Vec::new();
495 let mut time = U::Time::default();
496 let mut any_key = false;
497
498 for key_idx in 0..view.keys.values.len() {
499 let mut key_emitted = false;
500 for val_idx in child_range(view.vals.bounds, key_idx) {
501 run.clear();
503 for time_idx in child_range(view.times.bounds, val_idx) {
504 <U::Time as Columnar>::copy_from(&mut time, view.times.values.get(time_idx));
505 time.advance_by(frontier);
506 let mut diff = U::Diff::default();
507 for diff_idx in child_range(view.diffs.bounds, time_idx) {
508 diff.plus_equals(&<U::Diff as Columnar>::into_owned(view.diffs.values.get(diff_idx)));
509 }
510 run.push((time.clone(), diff));
511 }
512 run.sort_by(|a, b| a.0.cmp(&b.0));
514 let (mut w, mut r) = (0, 0);
515 while r < run.len() {
516 let t = run[r].0.clone();
517 let mut d = run[r].1.clone();
518 r += 1;
519 while r < run.len() && run[r].0 == t { d.plus_equals(&run[r].1); r += 1; }
520 if !d.is_zero() { run[w] = (t, d); w += 1; }
521 }
522 run.truncate(w);
523 if run.is_empty() { continue; }
524
525 if !key_emitted {
527 out.keys.values.push(view.keys.values.get(key_idx));
528 key_emitted = true;
529 }
530 out.vals.values.push(view.vals.values.get(val_idx));
531 for (t, d) in &run {
532 out.times.values.push(t);
533 out.diffs.values.push(d);
534 out.diffs.bounds.push(out.diffs.values.len() as u64);
535 }
536 out.times.bounds.push(out.times.values.len() as u64);
537 }
538 if key_emitted {
539 out.vals.bounds.push(out.vals.values.len() as u64);
540 any_key = true;
541 }
542 }
543 if any_key { out.keys.bounds.push(out.keys.values.len() as u64); }
544 out
545}
546
547#[cfg(test)]
548mod test {
549 use std::collections::VecDeque;
550 use columnar::Push;
551 use super::{ColChunk, Chunk};
552 use crate::columnar::updates::UpdatesTyped;
553 use crate::trace::chunk::merge_chains;
554 use crate::trace::Navigable;
555
556 type Upd = (u64, u64, u64, i64);
557
558 fn chunk(updates: Vec<Upd>) -> ColChunk<Upd> {
560 let mut u = UpdatesTyped::<Upd>::default();
561 for (k, v, t, d) in updates { u.push((&k, &v, &t, &d)); }
562 ColChunk::from_trie(u.consolidate())
563 }
564
565 fn flat<I: IntoIterator<Item = ColChunk<Upd>>>(chunks: I) -> Vec<Upd> {
567 chunks.into_iter().flat_map(|c| c.trie().iter().map(|(k, v, t, d)| (*k, *v, *t, *d)).collect::<Vec<_>>()).collect()
568 }
569
570 fn chain(updates: &[Upd], sz: usize) -> Vec<ColChunk<Upd>> {
572 updates.chunks(sz).map(|c| chunk(c.to_vec())).collect()
573 }
574
575 #[test]
580 fn merge_matches_reference() {
581 use crate::consolidation::consolidate_updates;
582
583 let mut seed = 0x2545F4914F6CDD1Du64;
584 let mut rng = move || { seed ^= seed << 13; seed ^= seed >> 7; seed ^= seed << 17; seed };
585
586 fn gen(rng: &mut impl FnMut() -> u64, n: usize) -> Vec<Upd> {
587 let mut v: Vec<Upd> = (0..n).map(|_| {
588 let k = rng() % 20; let val = rng() % 3; let t = rng() % 8;
589 let d = if rng() % 4 == 0 { -1 } else { 1 };
590 (k, val, t, d)
591 }).collect();
592 let mut rows: Vec<((u64, u64), u64, i64)> = v.drain(..).map(|(k, val, t, d)| ((k, val), t, d)).collect();
593 consolidate_updates(&mut rows);
594 rows.into_iter().map(|((k, val), t, d)| (k, val, t, d)).collect()
595 }
596
597 for _ in 0..300 {
598 let (n1, n2) = ((rng() as usize % 60) + 1, (rng() as usize % 60) + 1);
599 let u1 = gen(&mut rng, n1);
600 let u2 = gen(&mut rng, n2);
601 if u1.is_empty() || u2.is_empty() { continue; }
602 let sz = (rng() as usize % 5) + 1;
603
604 let mut out = VecDeque::new();
605 merge_chains(chain(&u1, sz), chain(&u2, sz), &mut out);
606 let merged = flat(out);
607
608 let mut reference: Vec<((u64, u64), u64, i64)> =
609 u1.iter().chain(u2.iter()).map(|&(k, v, t, d)| ((k, v), t, d)).collect();
610 consolidate_updates(&mut reference);
611 let reference: Vec<Upd> = reference.into_iter().map(|((k, v), t, d)| (k, v, t, d)).collect();
612
613 assert_eq!(merged, reference, "chunk size {sz}\n u1={u1:?}\n u2={u2:?}");
614 }
615 }
616
617 #[test]
620 fn settle_maximal_packing() {
621 use super::TARGET;
622 use crate::trace::chunk::is_graded;
623
624 let t = TARGET;
625 let sizes = [t / 3, t / 3, t / 3, t, t / 2, t / 2, t, 1, t - 1];
626 let total: usize = sizes.iter().sum();
627 let mut key = 0u64;
628 let mut input = VecDeque::new();
629 let mut output = VecDeque::new();
630 for &s in &sizes {
631 let updates: Vec<Upd> = (0..s).map(|_| { let k = key; key += 1; (k, 0, 0, 1) }).collect();
632 input.push_back(chunk(updates));
633 ColChunk::settle(&mut input, false, &mut output);
634 }
635 ColChunk::settle(&mut input, true, &mut output);
636 let chunks: Vec<_> = output.into();
637
638 assert!(is_graded(&chunks), "not graded: {:?}", chunks.iter().map(Chunk::len).collect::<Vec<_>>());
639 let got = flat(chunks);
640 assert_eq!(got.len(), total);
641 assert!(got.windows(2).all(|w| w[0].0 < w[1].0));
642 }
643
644 #[test]
648 fn cursor_handles_straddle() {
649 use crate::trace::cursor::Cursor;
650 use crate::trace::Description;
651 use crate::trace::chunk::ChunkBatch;
652 use timely::progress::Antichain;
653
654 let chunks = vec![
655 chunk(vec![(0, 0, 0, 1), (1, 0, 0, 1), (1, 1, 0, 1)]),
656 chunk(vec![(1, 1, 1, 1), (1, 2, 0, 1)]),
657 chunk(vec![(2, 0, 0, 1)]),
658 ];
659 let desc = Description::new(
660 Antichain::from_elem(0u64), Antichain::from_elem(2u64), Antichain::from_elem(0u64));
661 let batch = ChunkBatch::new(chunks, desc);
662
663 let mut cursor = batch.cursor();
664 let got = cursor.to_vec(&batch, |k| *k, |v| *v);
665 let want = vec![
666 ((0u64, 0u64), vec![(0u64, 1i64)]),
667 ((1, 0), vec![(0, 1)]),
668 ((1, 1), vec![(0, 1), (1, 1)]),
669 ((1, 2), vec![(0, 1)]),
670 ((2, 0), vec![(0, 1)]),
671 ];
672 assert_eq!(got, want);
673 }
674
675 #[test]
680 fn batch_merger_resumable_matches_reference() {
681 use crate::trace::{Description, Merger};
682 use crate::trace::chunk::{ChunkBatch, ChunkBatchMerger, is_graded};
683 use crate::trace::cursor::Cursor;
684 use crate::consolidation::consolidate_updates;
685 use timely::progress::Antichain;
686
687 let mut seed = 0x9E3779B97F4A7C15u64;
688 let mut rng = move || { seed ^= seed << 13; seed ^= seed >> 7; seed ^= seed << 17; seed };
689
690 fn gen(rng: &mut impl FnMut() -> u64) -> Vec<Upd> {
691 let n = rng() as usize % 40 + 1;
692 let mut rows: Vec<((u64, u64), u64, i64)> = (0..n).map(|_| {
693 let k = rng() % 10; let val = rng() % 3; let t = rng() % 6;
694 let d = if rng() % 4 == 0 { -1 } else { 1 };
695 ((k, val), t, d)
696 }).collect();
697 consolidate_updates(&mut rows);
698 rows.into_iter().map(|((k, v), t, d)| (k, v, t, d)).collect()
699 }
700 fn batch(updates: &[Upd], sz: usize) -> ChunkBatch<ColChunk<Upd>> {
701 let chunks: Vec<_> = updates.chunks(sz).map(|c| chunk(c.to_vec())).collect();
702 let desc = Description::new(
703 Antichain::from_elem(0u64), Antichain::from_elem(10u64), Antichain::from_elem(0u64));
704 ChunkBatch::new(chunks, desc)
705 }
706 fn read(b: &ChunkBatch<ColChunk<Upd>>) -> Vec<Upd> {
707 let mut out = Vec::new();
708 let mut c = b.cursor();
709 while c.key_valid(b) {
710 let k = *c.key(b);
711 while c.val_valid(b) {
712 let v = *c.val(b);
713 c.map_times(b, |t, d| out.push(((k, v), *t, *d)));
714 c.step_val(b);
715 }
716 c.step_key(b);
717 }
718 consolidate_updates(&mut out);
719 out.into_iter().map(|((k, v), t, d)| (k, v, t, d)).collect()
720 }
721
722 for _ in 0..200 {
723 let u1 = gen(&mut rng);
724 let u2 = gen(&mut rng);
725 if u1.is_empty() || u2.is_empty() { continue; }
726 let sz = (rng() as usize % 4) + 1;
727 let f = rng() % 6;
728 let (s1, s2) = (batch(&u1, sz), batch(&u2, sz));
729 let frontier = Antichain::from_elem(f);
730
731 let mut merger = ChunkBatchMerger::new(&s1, &s2, frontier.borrow());
732 loop {
733 let mut fuel = 1isize;
734 merger.work(&s1, &s2, &mut fuel);
735 if fuel > 0 { break; }
736 }
737 let result = merger.done();
738
739 assert!(is_graded(&result.chunks), "ungraded result: {:?}",
740 result.chunks.iter().map(Chunk::len).collect::<Vec<_>>());
741 let got = read(&result);
742 let mut want: Vec<((u64, u64), u64, i64)> =
743 u1.iter().chain(u2.iter()).map(|&(k, v, t, d)| ((k, v), t.max(f), d)).collect();
744 consolidate_updates(&mut want);
745 let want: Vec<Upd> = want.into_iter().map(|((k, v), t, d)| (k, v, t, d)).collect();
746 assert_eq!(got, want, "fuel-driven merge mismatch\n u1={u1:?}\n u2={u2:?}\n f={f}");
747 }
748 }
749
750 #[test]
754 fn settle_pages_and_round_trips() {
755 use std::cell::RefCell;
756 use std::rc::Rc;
757 use std::sync::Arc;
758 use std::sync::atomic::Ordering::Relaxed;
759 use crate::columnar::trace::spill::{self, BytesSource, BytesStore, SpillStats};
760
761 struct MemStore(Rc<RefCell<Vec<Vec<u8>>>>);
763 struct MemSource(Rc<RefCell<Vec<Vec<u8>>>>, usize);
764 impl BytesStore for MemStore {
765 fn store(&mut self, bytes: &[u8]) -> Box<dyn BytesSource> {
766 let mut a = self.0.borrow_mut();
767 let id = a.len();
768 a.push(bytes.to_vec());
769 Box::new(MemSource(self.0.clone(), id))
770 }
771 }
772 impl BytesSource for MemSource { fn load(&self) -> Vec<u8> { self.0.borrow()[self.1].clone() } }
773
774 let arena = Rc::new(RefCell::new(Vec::new()));
775 let stats = Arc::new(SpillStats::default());
776 spill::install(1, Box::new(MemStore(arena)), stats.clone()); let n = 5 * super::TARGET as u64;
780 let mut input: VecDeque<_> = (0..n).map(|k| chunk(vec![(k, 0, 0, 1)])).collect();
781 let mut out = VecDeque::new();
782 ColChunk::settle(&mut input, true, &mut out);
783
784 assert!(stats.spilled_chunks.load(Relaxed) > 0, "nothing was paged");
785 assert!(out.iter().any(|c| matches!(c, ColChunk::Paged(_))), "no paged chunk in output");
786 let got = flat(out);
788 let want: Vec<Upd> = (0..n).map(|k| (k, 0, 0, 1)).collect();
789 assert_eq!(got, want);
790 assert!(stats.fetched_chunks.load(Relaxed) > 0, "nothing was fetched back");
791
792 spill::uninstall();
793 }
794
795 #[test]
799 fn advance_single_key_spanning_pushes() {
800 use timely::progress::Antichain;
801 let frontier = Antichain::from_elem(100u64);
802 let n = 50u64;
803 let mut q = VecDeque::new();
804 let mut out = VecDeque::new();
805 for t in 0..n {
806 q.push_back(chunk(vec![(7, 0, t, 1)]));
807 ColChunk::advance(&mut q, frontier.borrow(), false, &mut out);
808 }
809 ColChunk::advance(&mut q, frontier.borrow(), true, &mut out);
810 assert_eq!(flat(out), vec![(7, 0, 100, n as i64)]);
811 }
812
813 #[test]
816 fn advance_emits_complete_groups_eagerly() {
817 use timely::progress::Antichain;
818 let frontier = Antichain::from_elem(5u64);
819 let mut q = VecDeque::from([chunk(vec![(0, 0, 0, 1), (0, 0, 1, 1), (1, 0, 0, 1)])]);
821 let mut out = VecDeque::new();
822 ColChunk::advance(&mut q, frontier.borrow(), false, &mut out);
823 assert_eq!(q.len(), 1);
825 assert_eq!(Chunk::len(&q[0]), 1);
826 assert_eq!(flat(out), vec![(0, 0, 5, 2)]);
828 }
829
830 #[test]
834 fn advance_resumable_matches_oneshot() {
835 use timely::progress::Antichain;
836 let frontier = Antichain::from_elem(3u64);
837 let input = || vec![
839 chunk(vec![(0, 0, 0, 1), (0, 0, 1, 1), (1, 0, 0, 1)]),
840 chunk(vec![(1, 0, 5, 1), (1, 1, 0, 1), (2, 0, 0, 1)]),
841 chunk(vec![(2, 0, 2, 1), (2, 0, 9, 1)]),
842 ];
843 let oneshot = {
844 let mut q: VecDeque<_> = input().into();
845 let mut out = VecDeque::new();
846 ColChunk::advance(&mut q, frontier.borrow(), false, &mut out);
847 ColChunk::advance(&mut q, frontier.borrow(), true, &mut out);
848 flat(out)
849 };
850 let incremental = {
851 let mut q = VecDeque::new();
852 let mut out = VecDeque::new();
853 for c in input() { q.push_back(c); ColChunk::advance(&mut q, frontier.borrow(), false, &mut out); }
854 ColChunk::advance(&mut q, frontier.borrow(), true, &mut out);
855 flat(out)
856 };
857 assert_eq!(oneshot, incremental);
858 for u in &oneshot { assert!(u.2 >= 3); }
859 }
860
861 #[test]
866 fn advance_matches_row_reference() {
867 use timely::progress::Antichain;
868 use crate::consolidation::consolidate_updates;
869
870 let mut seed = 0x2545F4914F6CDD1Du64;
871 let mut rng = move || { seed ^= seed << 13; seed ^= seed >> 7; seed ^= seed << 17; seed };
872
873 for _ in 0..200 {
874 let n = rng() as usize % 60 + 1;
877 let mut rows: Vec<((u64, u64), u64, i64)> = (0..n).map(|_| {
878 let k = rng() % 8; let v = rng() % 3; let t = rng() % 6;
879 let d = if rng() % 4 == 0 { -1 } else { 1 };
880 ((k, v), t, d)
881 }).collect();
882 consolidate_updates(&mut rows);
883 if rows.is_empty() { continue; }
884 let f = rng() % 6;
885 let frontier = Antichain::from_elem(f);
886 let sz = rng() as usize % 5 + 1; let mut q = VecDeque::new();
890 let mut out = VecDeque::new();
891 for c in rows.chunks(sz) {
892 q.push_back(chunk(c.iter().map(|&((k, v), t, d)| (k, v, t, d)).collect()));
893 ColChunk::advance(&mut q, frontier.borrow(), false, &mut out);
894 }
895 ColChunk::advance(&mut q, frontier.borrow(), true, &mut out);
896 let got = flat(out);
897
898 let mut want: Vec<((u64, u64), u64, i64)> =
900 rows.iter().map(|&((k, v), t, d)| ((k, v), t.max(f), d)).collect();
901 consolidate_updates(&mut want);
902 let want: Vec<Upd> = want.into_iter().map(|((k, v), t, d)| (k, v, t, d)).collect();
903
904 assert_eq!(got, want, "frontier {f}, chunk size {sz}, rows {rows:?}");
905 }
906 }
907}