From f0b77b1532dbebd9b788a0a4cc43200405c7d933 Mon Sep 17 00:00:00 2001 From: Frank McSherry Date: Tue, 10 Nov 2026 10:24:02 -0500 Subject: [PATCH] Chunk: merge and advance as one step A spine merge of chunked batches merges two chunks, then advances and consolidates the result. Each step writes its output, so an implementor for whom writing a chunk is the cost pays it twice. `Chunk` gains `merge_advance`, which does both. Its default is `merge` followed by `advance`, so no implementor has to change. `ChunkBatchMerger` calls it while both inputs have data, and calls `advance` as before for the tails and the final flush. The merger now counts a step's fuel by the updates it takes from the inputs, where it counted the updates `merge` produced. DDIR's `CorgiChunk` implements it. The body of its `merge` moves to `merge_fronts`, which takes an optional frontier and carry. With them, the merged order stays as indices into the source columns. The times of the complete groups are advanced, each group of several rows is sorted and summed on its times, and only the surviving rows are gathered. The last group may continue in a later chunk, so it is held back, and joins the next step as a third source. On a merge-heavy DDIR workload the second gather was 15% of merge time, and this removes it. The transaction benchmark's maintained program runs 2 to 4% more transactions a second, and the graph examples are unchanged. Co-Authored-By: Claude Fable 5.2 Claude-Session: https://claude.ai/code/session_01H6wB4EnzBnYpvwz5E2ZGmD --- differential-dataflow/src/trace/chunk/mod.rs | 43 +++- interactive/src/runtime/chunk.rs | 227 +++++++++++++++---- 2 files changed, 219 insertions(+), 51 deletions(-) diff --git a/differential-dataflow/src/trace/chunk/mod.rs b/differential-dataflow/src/trace/chunk/mod.rs index 5907753cf..fad3ae5b2 100644 --- a/differential-dataflow/src/trace/chunk/mod.rs +++ b/differential-dataflow/src/trace/chunk/mod.rs @@ -88,6 +88,29 @@ pub trait Chunk: Sized + Clone { /// point it appends the queue to the output without the method's assistance. fn merge(in1: &mut VecDeque, in2: &mut VecDeque, out: &mut VecDeque); + /// Merge the fronts of two input deques and advance the result by `frontier`, as one step. + /// + /// This is [`merge`](Chunk::merge) followed by [`advance`](Chunk::advance), and that is its + /// default. An implementor for whom writing a chunk is the expensive part can override it to + /// write each merged update once, where the two steps write it twice. + /// + /// `carry` holds merged updates that are not yet advanced, as `advance` leaves them at the + /// front of its input: a trailing `(key, val)` group that may continue. The method may leave + /// such a group there in turn. Advanced, consolidated chunks go to `out`. + /// The obligations of `merge` hold: an invocation consumes at least one of its inputs. + fn merge_advance( + in1: &mut VecDeque, + in2: &mut VecDeque, + frontier: AntichainRef, + carry: &mut VecDeque, + out: &mut VecDeque, + ) { + let mut produced = VecDeque::new(); + Self::merge(in1, in2, &mut produced); + carry.extend(produced); + Self::advance(carry, frontier, false, out); + } + /// Partition `input` updates into `keep` (greater or equal `frontier`) or not (`ship`). /// /// An implementation should yield with some frequency to allow the output to "settle". @@ -378,15 +401,20 @@ where self.idx2 += 1; } - // Merge's per-tick output (a burst's worth, or one tail chunk), measured - // for fuel before it joins the carry already in `merged`. - let mut produced = VecDeque::new(); + // A step's work, for fuel, is the updates it takes from the inputs: a burst's + // worth merged, or one tail chunk. + let work: usize; if !self.in1.is_empty() && !self.in2.is_empty() { - // Both sides have data: drain the loaded burst. - C::merge(&mut self.in1, &mut self.in2, &mut produced); + // Both sides have data: merge and advance the fronts of the loaded burst. + let loaded = |in1: &VecDeque, in2: &VecDeque| in1.iter().chain(in2.iter()).map(C::len).sum::(); + let before = loaded(&self.in1, &self.in2); + C::merge_advance(&mut self.in1, &mut self.in2, self.frontier.borrow(), &mut self.merged, &mut self.advanced); + work = before - loaded(&self.in1, &self.in2); } else if let Some(chunk) = self.in1.pop_front().or_else(|| self.in2.pop_front()) { // Exactly one side has data: flush its verbatim tail, one chunk a step. - produced.push_back(chunk); + work = chunk.len(); + self.merged.push_back(chunk); + C::advance(&mut self.merged, self.frontier.borrow(), false, &mut self.advanced); } else { // Both sources drained: final flush of advance's and settle's carries. C::advance(&mut self.merged, self.frontier.borrow(), true, &mut self.advanced); @@ -395,9 +423,6 @@ where break; } - let work: usize = produced.iter().map(C::len).sum(); - self.merged.extend(produced); - C::advance(&mut self.merged, self.frontier.borrow(), false, &mut self.advanced); // Maintain grading at the yield boundary: this step may exhaust `fuel` and // suspend with `advanced` held, and held chunk state must be graded. C::settle(&mut self.advanced, false, &mut self.settled); diff --git a/interactive/src/runtime/chunk.rs b/interactive/src/runtime/chunk.rs index 6b3597841..ecb45bbce 100644 --- a/interactive/src/runtime/chunk.rs +++ b/interactive/src/runtime/chunk.rs @@ -136,56 +136,59 @@ where } } - /// Concatenate a run of (globally-sorted) chunks into one combined `(kv, times, diffs)`, the - /// key/val columns a range at a time ([`concat_blocks`]). - fn concat(chunks: &[Self]) -> (CValue, ColTimes, Vec) { - let kvs: Vec = chunks.iter().map(Self::kv).collect(); - let total: usize = chunks.iter().map(Self::len_).sum(); - let (mut times, mut diffs) = (ColTimes::new(), Vec::with_capacity(total)); - for ch in chunks { - times.push_range(ch.times(), 0, ch.len_()); - diffs.extend_from_slice(ch.diffs()); - } - let kv = if total == 0 { CValue::Unit(0) } else { concat_blocks(&kvs) }; - (kv, times, diffs) - } - - fn len_(&self) -> usize { self.0.times.len() } -} - -impl Chunk for CorgiChunk -where - T: ColTime, - R: Semigroup + Clone + 'static, -{ - type Time = T; - const TARGET: usize = TARGET; - - fn len(&self) -> usize { self.0.times.len() } - - /// Merge of the two front chunks through their shared horizon, FULLY consolidating equal - /// `(key, val, time)` triples and pushing back the survivor's suffix (the fueled-merger - /// contract). Batched: corgi's `survey_groups` reports the interleaving of the two `(key, val)` - /// columns as maximal ranges exclusive to one side and equal classes as their ranges on BOTH - /// sides, so rows only one side holds are copied by range, and only the classes both sides - /// hold — where consolidation can happen — are merged row by row, on their times, which corgi - /// does not own. The shape is walked once per level per chunk, never once per row. - fn merge(in1: &mut VecDeque, in2: &mut VecDeque, out: &mut VecDeque) { + /// The merge of the two front chunks through their shared horizon, as [`Chunk::merge`] + /// describes; and with `compact`, a frontier and a carry, the advancing and consolidating of + /// the merged rows as well, as [`Chunk::advance`] would do next. + /// + /// With `compact` the merged order stays as indices into the sources. Complete `(key, val)` + /// groups have their times advanced, each group of several rows is sorted and summed on its + /// times, and only the surviving rows are gathered. The key and value columns are so written + /// once, where merging and then advancing writes them twice. + /// The last group may continue in a later chunk. It is held back in the carry, unadvanced, + /// and a group the step before held back joins this step as a third source. + fn merge_fronts( + in1: &mut VecDeque, + in2: &mut VecDeque, + mut compact: Option<(AntichainRef, &mut VecDeque)>, + out: &mut VecDeque, + ) { let c1 = in1.pop_front().unwrap(); let c2 = in2.pop_front().unwrap(); let (kv1, kv2) = (c1.kv(), c2.kv()); let (n1, n2) = (c1.len_(), c2.len_()); let (t1, d1) = (c1.times(), c1.diffs()); let (t2, d2) = (c2.times(), c2.diffs()); + // A group the step before held back: one `(key, val)`, whose rows come first. + let held = compact.as_mut().and_then(|(_, carry)| carry.pop_front()); + let kv0 = held.as_ref().map(Self::kv); + let n0 = held.as_ref().map_or(0, Self::len_); let runs = survey_groups(&kv1, &kv2); - let (mut tags, mut offs) = (Vec::with_capacity(n1 + n2), Vec::with_capacity(n1 + n2)); + let (mut tags, mut offs) = (Vec::with_capacity(n0 + n1 + n2), Vec::with_capacity(n0 + n1 + n2)); // Emitted chunks own this allocation. Size it by surviving rows so cancellation // does not leave a small result holding storage for both full inputs. Times are // reserved for both inputs, to copy ranges without regrowing, and shrunk below if // fewer than half the rows survive. let (mut times, mut diffs): (ColTimes, Vec) = (ColTimes::new(), Vec::new()); - times.reserve(t1.width().max(t2.width()), n1 + n2); + times.reserve(t1.width().max(t2.width()), n0 + n1 + n2); + if let Some(held) = &held { + tags.resize(n0, 2); + offs.extend(0..n0); + times.push_range(held.times(), 0, n0); + diffs.extend_from_slice(held.diffs()); + } + // When compacting: where each group starts in the merged order. A run exclusive to one + // side starts a group, as does each group within it, read off that side's own group ends. + // A class both sides hold is one group. + let ends = compact.as_ref().map(|_| (group_bounds(&kv1), group_bounds(&kv2))); + let (mut e1, mut e2) = (0, 0); + let mut starts: Vec = Vec::new(); + if n0 > 0 { starts.push(0); } + let mark = |starts: &mut Vec, ends: &[usize], e: &mut usize, base: usize, lo: usize, hi: usize| { + starts.push(base); + while *e < ends.len() && ends[*e] <= lo { *e += 1; } + while *e < ends.len() && ends[*e] < hi { starts.push(base + ends[*e] - lo); *e += 1; } + }; // Where the survivor's pushed-back suffix starts: the last run, if it is exclusive. let (mut p1, mut p2) = (n1, n2); let copy = |tags: &mut Vec, offs: &mut Vec, times: &mut ColTimes, diffs: &mut Vec, side: usize, lo: usize, hi: usize| { @@ -199,14 +202,21 @@ where let last = r + 1 == runs.len(); match *run { GroupRun::A(lo, hi) => { - if last { p1 = lo; } else { copy(&mut tags, &mut offs, &mut times, &mut diffs, 0, lo, hi); } + if last { p1 = lo; } else { + if let Some((ends1, _)) = &ends { mark(&mut starts, ends1, &mut e1, tags.len(), lo, hi); } + copy(&mut tags, &mut offs, &mut times, &mut diffs, 0, lo, hi); + } } GroupRun::B(lo, hi) => { - if last { p2 = lo; } else { copy(&mut tags, &mut offs, &mut times, &mut diffs, 1, lo, hi); } + if last { p2 = lo; } else { + if let Some((_, ends2)) = &ends { mark(&mut starts, ends2, &mut e2, tags.len(), lo, hi); } + copy(&mut tags, &mut offs, &mut times, &mut diffs, 1, lo, hi); + } } GroupRun::Both(a_lo, a_hi, b_lo, b_hi) => { // Both classes hold one `(key, val)`, sorted by time: merge on time, summing the // diffs of equal times, which is the consolidation. + if ends.is_some() { starts.push(tags.len()); } let (mut i, mut j) = (a_lo, b_lo); while i < a_hi && j < b_hi { match t1.cmp_cross(i, t2, j) { @@ -238,9 +248,62 @@ where // The survey can be large. Its last use precedes materializing the // output columns, so do not retain it across that allocation peak. drop(runs); - if times.len() * 2 < n1 + n2 { times.shrink_to_fit(); } - let srcs = [Some(&kv1), Some(&kv2)]; - Self::emit(&srcs, &tags, &offs, times, diffs, out); + let mut srcs = vec![Some(&kv1), Some(&kv2)]; + if let Some(kv0) = &kv0 { srcs.push(Some(kv0)); } + match compact { + None => { + if times.len() * 2 < n1 + n2 { times.shrink_to_fit(); } + Self::emit(&srcs, &tags, &offs, times, diffs, out); + } + Some((frontier, carry)) => { + let m = tags.len(); + // A class whose rows all cancelled leaves a start with no rows. + starts.dedup(); + if starts.last() == Some(&m) { starts.pop(); } + // The held group continues if the first merged row has its `(key, val)`. + if n0 > 0 && starts.get(1) == Some(&n0) { + let first = if tags[n0] == 0 { &kv1 } else { &kv2 }; + if equal_idx(kv0.as_ref().unwrap(), first, &[0], &[offs[n0]])[0] == 0 { starts.remove(1); } + } + if let Some(&g) = starts.last() { + // The last group may continue in a later chunk: hold it back, as it is. + let mut last = ColTimes::new(); + last.push_range(×, g, m); + carry.push_front(Self::from_kv(gather_lanes(&srcs, &tags[g..], &offs[g..]), last, diffs[g..].to_vec())); + // Advance the complete groups' times in one lane-wise pass. Advancing can reorder + // a group's times and make some equal, so each group of several rows is sorted + // on its times and equal times sum. All rows of a group share `(key, val)`, so + // any of them stands for the group in the gather. + times.advance_by(frontier, g); + let (mut otags, mut ooffs, mut rows) = (Vec::with_capacity(g), Vec::with_capacity(g), Vec::with_capacity(g)); + let mut odiffs: Vec = Vec::with_capacity(g); + let mut order: Vec = Vec::new(); + for group in starts.windows(2) { + let (lo, hi) = (group[0], group[1]); + if hi - lo == 1 { + if !diffs[lo].is_zero() { otags.push(tags[lo]); ooffs.push(offs[lo]); rows.push(lo); odiffs.push(diffs[lo].clone()); } + continue; + } + order.clear(); + order.extend(lo..hi); + order.sort_by(|&a, &b| times.cmp(a, b)); + let mut k = 0; + while k < order.len() { + let rep = order[k]; + let mut d = diffs[rep].clone(); + k += 1; + while k < order.len() && times.cmp(order[k], rep) == Ordering::Equal { + d.plus_equals(&diffs[order[k]]); + k += 1; + } + if !d.is_zero() { otags.push(tags[rep]); ooffs.push(offs[rep]); rows.push(rep); odiffs.push(d); } + } + } + // The one gather of the key and value columns. + if !rows.is_empty() { Self::emit(&srcs, &otags, &ooffs, times.gather(&rows), odiffs, out); } + } + } + } drop(tags); drop(offs); @@ -259,6 +322,63 @@ where } } + /// Concatenate a run of (globally-sorted) chunks into one combined `(kv, times, diffs)`, the + /// key/val columns a range at a time ([`concat_blocks`]). + fn concat(chunks: &[Self]) -> (CValue, ColTimes, Vec) { + let kvs: Vec = chunks.iter().map(Self::kv).collect(); + let total: usize = chunks.iter().map(Self::len_).sum(); + let (mut times, mut diffs) = (ColTimes::new(), Vec::with_capacity(total)); + for ch in chunks { + times.push_range(ch.times(), 0, ch.len_()); + diffs.extend_from_slice(ch.diffs()); + } + let kv = if total == 0 { CValue::Unit(0) } else { concat_blocks(&kvs) }; + (kv, times, diffs) + } + + fn len_(&self) -> usize { self.0.times.len() } +} + +impl Chunk for CorgiChunk +where + T: ColTime, + R: Semigroup + Clone + 'static, +{ + type Time = T; + const TARGET: usize = TARGET; + + fn len(&self) -> usize { self.0.times.len() } + + /// Merge of the two front chunks through their shared horizon, FULLY consolidating equal + /// `(key, val, time)` triples and pushing back the survivor's suffix (the fueled-merger + /// contract). Batched: corgi's `survey_groups` reports the interleaving of the two `(key, val)` + /// columns as maximal ranges exclusive to one side and equal classes as their ranges on BOTH + /// sides, so rows only one side holds are copied by range, and only the classes both sides + /// hold — where consolidation can happen — are merged row by row, on their times, which corgi + /// does not own. The shape is walked once per level per chunk, never once per row. + fn merge(in1: &mut VecDeque, in2: &mut VecDeque, out: &mut VecDeque) { + Self::merge_fronts(in1, in2, None, out); + } + + /// Merge and advance as one step ([`merge_fronts`](Self::merge_fronts)), which gathers the + /// key and value columns once. With more than a single group held over, the two steps. + fn merge_advance( + in1: &mut VecDeque, + in2: &mut VecDeque, + frontier: AntichainRef, + carry: &mut VecDeque, + out: &mut VecDeque, + ) { + if carry.len() > 1 { + let mut produced = VecDeque::new(); + Self::merge(in1, in2, &mut produced); + carry.extend(produced); + Self::advance(carry, frontier, false, out); + } else { + Self::merge_fronts(in1, in2, Some((frontier, carry)), out); + } + } + fn extract( input: &mut VecDeque, frontier: AntichainRef, @@ -1181,4 +1301,27 @@ mod test { assert_eq!(result.as_ref().map_or_else(BTreeMap::new, read_batch), want, "u1={u1:?}\nu2={u2:?}\nf={f}"); } } + + // A `(key, val)` group that spans every chunk of both inputs is held back from step to step, + // and advanced once, whole. Compacted to a frontier past all of its times it is one update. + #[test] + fn merge_advance_carries_a_group_across_chunks() { + let mut u1: Vec<((u64, u64), u64, i64)> = (0..9).map(|t| ((7, 7), 2 * t, 1)).collect(); + let mut u2: Vec<((u64, u64), u64, i64)> = (0..9).map(|t| ((7, 7), 2 * t + 1, 1)).collect(); + u1.push(((1, 1), 0, 1)); + u2.push(((9, 9), 0, 1)); + for sz in 1..5 { + let (s1, s2) = (batch(&u1, sz), batch(&u2, sz)); + let frontier = Antichain::from_elem(100); + let mut merger = ChunkBatchMerger::new(&s1, &s2, frontier.borrow()); + loop { + let mut fuel = 1isize; + merger.work(&s1, &s2, &mut fuel); + if fuel > 0 { break; } + } + let result = merger.done().expect("three updates remain"); + assert_eq!(read_batch(&result), reference(&u1, &u2, 100), "chunks of {sz}"); + assert_eq!(result.chunks.iter().map(Chunk::len).sum::(), 3, "chunks of {sz}"); + } + } }