Skip to content

Commit d37d9eb

Browse files
frankmcsherryclaude
andcommitted
compute: do not pre-sort chunks bound for the snapshot batcher
The chunkers sort and consolidate every input container before the batcher sees it. With the snapshot batcher sorting everything it holds at `seal`, that work is wasted on its fast path. `UnsortedChunker` packs updates into chunks in arrival order, and the batcher sorts and consolidates each chunk itself only when it hands chunks to the merge batcher on the fallback path, which still receives what it requires. Wired at every Row-keyed arrangement site that already uses the snapshot batcher: the reduce inputs, the TopK stages, `FormArrangementKey` and the join's arrangements. Error arrangements keep their chunker. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1 parent cffe59b commit d37d9eb

5 files changed

Lines changed: 209 additions & 17 deletions

File tree

src/compute/src/render/context.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1282,7 +1282,7 @@ impl<'scope, T: RenderTimestamp> CollectionBundle<'scope, T> {
12821282
>(exchange, name),
12831283
ArrangementBatcher::Columnation => ok_stream.mz_arrange_core::<
12841284
_,
1285-
batcher::Chunker<_>,
1285+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
12861286
mz_row_spine::RowRowBatcher<_, _>,
12871287
RowRowBuilder<_, _>,
12881288
RowRowSpine<_, _>,

src/compute/src/render/join/linear_join.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -405,7 +405,7 @@ where
405405
>(exchange, "JoinStage"),
406406
ArrangementBatcher::Columnation => keyed.mz_arrange_core::<
407407
_,
408-
batcher::Chunker<_>,
408+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
409409
mz_row_spine::RowRowBatcher<_, _>,
410410
RowRowBuilder<_, _>,
411411
RowRowSpine<_, _>,

src/compute/src/render/reduce.rs

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -301,7 +301,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
301301

302302
let arranged = collection
303303
.mz_arrange::<
304-
ColumnationChunker<_>,
304+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
305305
RowRowBatcher<_, _>,
306306
RowRowBuilder<_, _>,
307307
RowRowSpine<_, _>,
@@ -411,7 +411,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
411411

412412
let arranged = differential_dataflow::collection::concatenate(input.scope(), to_collect)
413413
.mz_arrange::<
414-
ColumnationChunker<_>,
414+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
415415
RowValBatcher<_, _, _>,
416416
RowValBuilder<_, _, _>,
417417
RowValSpine<_, _, _>,
@@ -577,7 +577,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
577577
};
578578
let arranged = partial
579579
.mz_arrange::<
580-
ColumnationChunker<_>,
580+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
581581
RowRowBatcher<_, _>,
582582
RowRowBuilder<_, _>,
583583
RowRowSpine<_, _>,
@@ -806,7 +806,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
806806

807807
let input: KeyCollection<_, _, _> = input.into();
808808
let arranged = input.mz_arrange::<
809-
ColumnationChunker<_>,
809+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
810810
RowBatcher<_, _>,
811811
RowBuilder<_, _>,
812812
RowSpine<_, _>,
@@ -941,7 +941,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
941941
// view mz_introspection.mz_expected_group_size_advice.
942942
let arranged = partial
943943
.mz_arrange::<
944-
ColumnationChunker<_>,
944+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
945945
RowRowBatcher<_, _>,
946946
RowRowBuilder<_, _>,
947947
RowRowSpine<_, _>,
@@ -1152,7 +1152,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
11521152
// view mz_introspection.mz_expected_group_size_advice.
11531153
let arranged_input = input
11541154
.mz_arrange::<
1155-
ColumnationChunker<_>,
1155+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
11561156
RowRowBatcher<_, _>,
11571157
RowRowBuilder<_, _>,
11581158
RowRowSpine<_, _>,
@@ -1285,7 +1285,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
12851285
let partial: KeyCollection<_, _, _> = partial.into();
12861286
let arranged = partial
12871287
.mz_arrange::<
1288-
ColumnationChunker<_>,
1288+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
12891289
RowBatcher<_, _>,
12901290
RowBuilder<_, _>,
12911291
RowSpine<_, Vec<ReductionMonoid>>,
@@ -1441,7 +1441,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
14411441
(pairer.merge(&key, std::iter::once(value)), ())
14421442
})
14431443
.mz_arrange::<
1444-
ColumnationChunker<_>,
1444+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
14451445
RowBatcher<_, _>,
14461446
RowBuilder<_, _>,
14471447
RowSpine<_, _>,
@@ -1484,7 +1484,7 @@ impl<'scope, T: RenderTimestamp> Context<'scope, T> {
14841484
let err_full_aggrs = full_aggrs.clone();
14851485
let arranged = collection
14861486
.mz_arrange::<
1487-
ColumnationChunker<_>,
1487+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
14881488
RowBatcher<_, _>,
14891489
RowBuilder<_, _>,
14901490
RowSpine<_, (Vec<Accum>, Diff)>,

src/compute/src/render/top_k.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -581,7 +581,7 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim
581581
.into();
582582
let result = partial
583583
.mz_arrange::<
584-
ColumnationChunker<_>,
584+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
585585
RowBatcher<_, _>,
586586
RowBuilder<_, _>,
587587
RowSpine<_, _>,
@@ -647,7 +647,7 @@ where
647647
let arranged = input
648648
.clone()
649649
.mz_arrange::<
650-
ColumnationChunker<_>,
650+
mz_row_spine::snapshot_batcher::UnsortedChunker<_, _, _>,
651651
RowRowBatcher<_, _>,
652652
RowRowBuilder<_, _>,
653653
RowRowSpine<_, _>,

src/row-spine/src/snapshot_batcher.rs

Lines changed: 196 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,18 +20,25 @@
2020
//! on equal prefixes, and emits the sorted, consolidated chain in one copy pass. The first chunk
2121
//! carrying a second time hands everything held to a [`MergeBatcher`], and the batcher stays on
2222
//! that path from then on, so steady-state behaviour is the merge batcher's.
23+
//!
24+
//! [`UnsortedChunker`] pairs with it: chunks arrive in arrival order, since sorting them ahead
25+
//! of a batcher that sorts everything at `seal` is wasted, and the batcher sorts and
26+
//! consolidates each chunk itself before handing it to the merge batcher on the fallback path.
2327
2428
use std::cmp::Ordering;
29+
use std::collections::VecDeque;
2530

31+
use columnar::{Columnar, Index, Len};
2632
use columnation::Columnation;
2733
use differential_dataflow::difference::Semigroup;
2834
use differential_dataflow::lattice::Lattice;
2935
use differential_dataflow::logging::Logger;
3036
use differential_dataflow::trace::implementations::merge_batcher::MergeBatcher;
3137
use differential_dataflow::trace::{Batcher, Description};
3238
use mz_repr::Row;
39+
use mz_timely_util::columnar::Column;
3340
use mz_timely_util::columnation::{ColInternalMerger, ColumnationStack};
34-
use timely::container::PushInto;
41+
use timely::container::{ContainerBuilder, PushInto};
3542
use timely::progress::Timestamp;
3643
use timely::progress::frontier::{Antichain, AntichainRef};
3744

@@ -103,6 +110,35 @@ where
103110
}
104111
}
105112

113+
/// Sort and consolidate one chunk, for handing to the general path, which requires it.
114+
fn sort_chunk(chunk: Chunk<D, T, R>) -> Chunk<D, T, R> {
115+
let mut index: Vec<usize> = (0..chunk.len()).collect();
116+
index.sort_unstable_by(|&a, &b| {
117+
let (da, ta, _) = &chunk[a];
118+
let (db, tb, _) = &chunk[b];
119+
(da, ta).cmp(&(db, tb))
120+
});
121+
let mut out: Chunk<D, T, R> = ColumnationStack::with_capacity(chunk.len());
122+
let mut iter = index.iter().peekable();
123+
while let Some(&i) = iter.next() {
124+
let (d, t, r) = &chunk[i];
125+
let mut diff = r.clone();
126+
while let Some(&&j) = iter.peek() {
127+
let (d2, t2, r2) = &chunk[j];
128+
if (d2, t2).cmp(&(d, t)) == Ordering::Equal {
129+
diff.plus_equals(r2);
130+
iter.next();
131+
} else {
132+
break;
133+
}
134+
}
135+
if !diff.is_zero() {
136+
out.copy_destructured(d, t, &diff);
137+
}
138+
}
139+
out
140+
}
141+
106142
/// Sort and consolidate everything pending into a chain of chunks, in one copy pass.
107143
fn sort_pending(&mut self) -> Vec<Chunk<D, T, R>> {
108144
let mut index: Vec<(u64, u32, u32)> = Vec::with_capacity(self.pending_len);
@@ -243,15 +279,132 @@ where
243279
self.pending.push(chunk);
244280
return;
245281
}
246-
// A second time: hand everything held to the general path, for good.
282+
// A second time: hand everything held to the general path, for good. The merge
283+
// batcher needs sorted, consolidated chunks, which the chunker did not provide.
247284
self.general = true;
248285
self.time = None;
249286
self.pending_len = 0;
250287
for held in self.pending.drain(..) {
251-
self.inner.push_into(held);
288+
self.inner.push_into(Self::sort_chunk(held));
289+
}
290+
}
291+
self.inner.push_into(Self::sort_chunk(chunk));
292+
}
293+
}
294+
295+
/// A chunker that packs incoming updates into [`ColumnationStack`] chunks in arrival order.
296+
///
297+
/// The sorting chunkers sort and consolidate every input container before the batcher sees
298+
/// it. A [`SnapshotBatcher`] sorts everything it holds at `seal`, so that work is wasted on
299+
/// its fast path; on the fallback path it sorts the held chunks itself (see
300+
/// [`SnapshotBatcher::sort_chunk`]). Only pair this chunker with that batcher.
301+
pub struct UnsortedChunker<D, T, R>
302+
where
303+
D: Columnation,
304+
T: Columnation,
305+
R: Columnation,
306+
{
307+
pending: Vec<(D, T, R)>,
308+
ready: VecDeque<ColumnationStack<(D, T, R)>>,
309+
empty: Option<ColumnationStack<(D, T, R)>>,
310+
}
311+
312+
impl<D, T, R> Default for UnsortedChunker<D, T, R>
313+
where
314+
D: Columnation,
315+
T: Columnation,
316+
R: Columnation,
317+
{
318+
fn default() -> Self {
319+
Self {
320+
pending: Vec::new(),
321+
ready: VecDeque::new(),
322+
empty: None,
323+
}
324+
}
325+
}
326+
327+
impl<D, T, R> UnsortedChunker<D, T, R>
328+
where
329+
D: Columnation,
330+
T: Columnation,
331+
R: Columnation,
332+
{
333+
/// Records per chunk, matching the merge batcher's 64 KiB chunks.
334+
fn chunk_capacity() -> usize {
335+
const BUFFER_SIZE_BYTES: usize = 64 << 10;
336+
let size = std::mem::size_of::<(D, T, R)>();
337+
if size == 0 {
338+
BUFFER_SIZE_BYTES
339+
} else if size <= BUFFER_SIZE_BYTES {
340+
BUFFER_SIZE_BYTES / size
341+
} else {
342+
1
343+
}
344+
}
345+
346+
fn form_chunks(&mut self, all: bool) {
347+
let cap = Self::chunk_capacity();
348+
while self.pending.len() >= cap || (all && !self.pending.is_empty()) {
349+
let take = std::cmp::min(self.pending.len(), cap);
350+
let mut chunk = ColumnationStack::with_capacity(cap);
351+
for item in self.pending.drain(..take) {
352+
chunk.copy(&item);
252353
}
354+
self.ready.push_back(chunk);
355+
}
356+
}
357+
}
358+
359+
impl<'a, D, T, R> PushInto<&'a mut Vec<(D, T, R)>> for UnsortedChunker<D, T, R>
360+
where
361+
D: Columnation,
362+
T: Columnation,
363+
R: Columnation,
364+
{
365+
fn push_into(&mut self, container: &'a mut Vec<(D, T, R)>) {
366+
self.pending.append(container);
367+
self.form_chunks(false);
368+
}
369+
}
370+
371+
impl<'a, D, T, R> PushInto<&'a mut Column<(D, T, R)>> for UnsortedChunker<D, T, R>
372+
where
373+
D: Columnar + Columnation,
374+
T: Columnar + Columnation,
375+
R: Columnar + Columnation,
376+
{
377+
fn push_into(&mut self, container: &'a mut Column<(D, T, R)>) {
378+
let borrowed = container.borrow();
379+
self.pending.reserve(borrowed.len());
380+
for (d, t, r) in borrowed.into_index_iter() {
381+
self.pending
382+
.push((D::into_owned(d), T::into_owned(t), R::into_owned(r)));
253383
}
254-
self.inner.push_into(chunk);
384+
self.form_chunks(false);
385+
}
386+
}
387+
388+
impl<D, T, R> ContainerBuilder for UnsortedChunker<D, T, R>
389+
where
390+
D: Columnation + Clone + 'static,
391+
T: Columnation + Clone + 'static,
392+
R: Columnation + Clone + 'static,
393+
{
394+
type Container = ColumnationStack<(D, T, R)>;
395+
396+
fn extract(&mut self) -> Option<&mut Self::Container> {
397+
if let Some(ready) = self.ready.pop_front() {
398+
self.empty = Some(ready);
399+
self.empty.as_mut()
400+
} else {
401+
None
402+
}
403+
}
404+
405+
fn finish(&mut self) -> Option<&mut Self::Container> {
406+
self.form_chunks(true);
407+
self.extract()
255408
}
256409
}
257410

@@ -350,6 +503,45 @@ mod tests {
350503
assert!(!b.general);
351504
}
352505

506+
/// A chunk in arrival order, as [`UnsortedChunker`] produces.
507+
fn unsorted(updates: &[(i64, u64, i64)]) -> Chunk<(Row, ()), u64, i64> {
508+
let mut out = ColumnationStack::with_capacity(updates.len());
509+
for (k, t, r) in updates {
510+
out.copy(&((Row::pack_slice(&[Datum::Int64(*k)]), ()), *t, *r));
511+
}
512+
out
513+
}
514+
515+
#[mz_ore::test]
516+
fn unsorted_chunks_on_both_paths() {
517+
let mut b = B::new(None, 0);
518+
b.push_into(unsorted(&[(3, 1, 1), (1, 1, 1), (3, 1, 1)]));
519+
let (chain, _) = b.seal(upper(2));
520+
assert_eq!(collect(&chain), vec![(1, 1, 1), (3, 1, 2)]);
521+
b.push_into(unsorted(&[(9, 3, 1), (2, 3, 1), (9, 3, -1)]));
522+
b.push_into(unsorted(&[(5, 4, 1), (2, 3, 1)]));
523+
let (chain, _) = b.seal(upper(5));
524+
assert!(b.general);
525+
assert_eq!(collect(&chain), vec![(2, 3, 2), (5, 4, 1)]);
526+
}
527+
528+
#[mz_ore::test]
529+
fn unsorted_chunker_keeps_arrival_order_and_chunk_size() {
530+
let mut c: UnsortedChunker<(Row, ()), u64, i64> = UnsortedChunker::default();
531+
let mut input: Vec<((Row, ()), u64, i64)> = (0..5)
532+
.rev()
533+
.map(|k| ((Row::pack_slice(&[Datum::Int64(k)]), ()), 1, 1))
534+
.collect();
535+
c.push_into(&mut input);
536+
assert!(c.extract().is_none(), "below chunk capacity, nothing ready");
537+
let chunk = c.finish().expect("finish flushes");
538+
let keys: Vec<i64> = chunk
539+
.iter()
540+
.map(|((k, ()), _, _)| k.iter().next().unwrap().unwrap_int64())
541+
.collect();
542+
assert_eq!(keys, vec![4, 3, 2, 1, 0]);
543+
}
544+
353545
#[mz_ore::test]
354546
fn sort_prefix_agrees_with_row_order() {
355547
let rows: Vec<Row> = [

0 commit comments

Comments
 (0)