From f9a729aad034d63d91236b8976823a644d2fc40e Mon Sep 17 00:00:00 2001 From: Frank McSherry Date: Wed, 23 Sep 2026 07:12:13 -0400 Subject: [PATCH] Corgi: search few keys in a long chunk in lockstep Reduce presentation and the join's probe side find each key's rows in every chunk with `MatchingRanges`, which gallops from the previous match. Each search then depends on the last, and in an incremental round (tens of keys against a chunk of 10^5 rows) each lands on a cold line, so the misses are paid one after another. In SCC that loop was ~13% of churn time. Fewer probes did not help: interpolating the start halved them and saved 4%. `matching_ranges` searches the keys in lockstep instead: branch-free binary searches that advance one level at a time together, so each level issues one independent load per key and the misses overlap. It keeps the merge when the keys are dense relative to the chunk (4 x keys x log2(rows) > rows + keys), where sequential steps are cheaper. Corgi, 1 worker, 5 interleaved runs, medians (before -> after): scc 200k edges, 200 x 50 churn 4.49 -> 4.06 s reach 2M edges, 100 x 100 churn 691 -> 480 ms kcore 200k edges, 200 x 50 churn 836 -> 829 ms (noise) count 1M rows, 200 x 500 churn 2.27 -> 2.29 s Loads within 1.5%. Outputs match; interactive tests pass; AoC 2023 corgi 33/33. Co-Authored-By: Claude Opus 5.5 (1M context) --- interactive/src/corgi/join.rs | 6 ++-- interactive/src/corgi/reduce.rs | 14 +++++---- interactive/src/corgi/search.rs | 53 +++++++++++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 8 deletions(-) diff --git a/interactive/src/corgi/join.rs b/interactive/src/corgi/join.rs index 0f559472a..cbb6e33b2 100644 --- a/interactive/src/corgi/join.rs +++ b/interactive/src/corgi/join.rs @@ -38,7 +38,7 @@ use differential_dataflow::operators::int_proxy::join::JoinMatches; use differential_dataflow::trace::chunk::{Chunk, ChunkBatch}; use corgi::arrange::{compare_at, gather, gather_lanes, leaf_slice}; -use crate::corgi::search::MatchingRanges; +use crate::corgi::search::matching_ranges; use corgi::{shape_of_value, Shape, Value as CValue}; use crate::corgi::chunk::{key_is_hashed, key_lane, recover_key, CorgiChunk}; @@ -523,7 +523,9 @@ impl<'a, T: ColTime> Probe<'a, T> { fn new(chunk: &'a CorgiChunk, cid: usize, needles: &[u64], leaf_vals: bool) -> Self { let keys = corgi::arrange::leaf_slice(key_lane(chunk.keys())).expect("identifier lane is a u64 leaf"); let (mut lo, mut hi) = (vec![0; needles.len()], vec![0; needles.len()]); - for (j, range) in MatchingRanges::new(needles, keys) { + let mut found = Vec::new(); + matching_ranges(needles, keys, &mut Vec::new(), &mut found); + for (j, range) in found { lo[j] = range.start; hi[j] = range.end; } diff --git a/interactive/src/corgi/reduce.rs b/interactive/src/corgi/reduce.rs index 9fc83478f..ddf335f49 100644 --- a/interactive/src/corgi/reduce.rs +++ b/interactive/src/corgi/reduce.rs @@ -35,7 +35,7 @@ use corgi::arrange::{compare_at, gather_lanes, sort_blocks}; use corgi::{ArithOp, Bounds, NumOp, OpLike, Value as CValue}; use crate::corgi::col_times::{ColTime, ColTimes}; -use crate::corgi::search::MatchingRanges; +use crate::corgi::search::matching_ranges; use crate::corgi::chunk::{columns_to_batch, key_ids, key_lane, CorgiChunk}; use crate::ir::{Diff, Reducer}; @@ -218,14 +218,16 @@ impl CorgiReduceBackend { /// Where `chunks` hold the ascending `keys`. /// -/// Both the keys and stored identifier lane are sorted. Match them with -/// monotone positions, galloping over long gaps and stepping through adjacent -/// keys. The same compiled search covers narrow updates and broad cascades. +/// Both the keys and stored identifier lane are sorted; [`matching_ranges`] chooses between +/// merging them and searching the keys in lockstep. fn search(chunks: &[&CorgiChunk], keys: &[u64]) -> Matches { + let mut scratch = Vec::new(); chunks.iter().map(|chunk| { - if chunk.diffs().is_empty() { return Vec::new(); } + let mut found = Vec::new(); + if chunk.diffs().is_empty() { return found; } let lane = corgi::arrange::leaf_slice(key_lane(chunk.keys())).expect("the identifier lane is a u64 leaf"); - MatchingRanges::new(keys, lane).collect() + matching_ranges(keys, lane, &mut scratch, &mut found); + found }).collect() } diff --git a/interactive/src/corgi/search.rs b/interactive/src/corgi/search.rs index db8d205ca..8b603b683 100644 --- a/interactive/src/corgi/search.rs +++ b/interactive/src/corgi/search.rs @@ -52,6 +52,56 @@ impl Iterator for MatchingRanges<'_> { } } +/// The items of [`MatchingRanges`], written to `out`: for each present needle, its index and the +/// range of `haystack` equal to it, in needle order. +/// +/// Few needles in a long haystack are searched in lockstep by branch-free binary search, so each +/// level issues one independent load per needle and their cache misses overlap. Galloping from the +/// previous match makes each search depend on the last and serializes the misses. Otherwise the +/// two sorted lists are merged by [`MatchingRanges`]. `scratch` holds the search positions. +pub(crate) fn matching_ranges( + needles: &[u64], + haystack: &[u64], + scratch: &mut Vec, + out: &mut Vec<(usize, Range)>, +) { + out.clear(); + let n = haystack.len(); + if n == 0 || needles.is_empty() { + return; + } + let depth = (usize::BITS - n.leading_zeros()) as usize; + // A random probe costs several sequential merge steps; at four, loads that present most of + // a chunk keep the merge, and incremental rounds that present few keys take the search. + if 4 * needles.len() * depth > n + needles.len() { + out.extend(MatchingRanges::new(needles, haystack)); + return; + } + debug_assert!(needles.windows(2).all(|w| w[0] < w[1])); + debug_assert!(haystack.windows(2).all(|w| w[0] <= w[1])); + scratch.clear(); + scratch.resize(needles.len(), 0); + let mut size = n; + while size > 1 { + let half = size / 2; + for (base, &needle) in scratch.iter_mut().zip(needles) { + let mid = *base + half; + *base = if haystack[mid] < needle { mid } else { *base }; + } + size -= half; + } + for (j, (&base, &needle)) in scratch.iter().zip(needles).enumerate() { + let start = base + (haystack[base] < needle) as usize; + if start < n && haystack[start] == needle { + let mut end = start + 1; + while end < n && haystack[end] == needle { + end += 1; + } + out.push((j, start..end)); + } + } +} + /// First index at or after `start` outside a predicate's prefix. fn gallop(xs: &[u64], start: usize, predicate: impl Fn(&u64) -> bool) -> usize { let mut pos = start; @@ -109,6 +159,9 @@ mod tests { MatchingRanges::new(&needles, &haystack).collect::>(), expected ); + let (mut scratch, mut out) = (Vec::new(), Vec::new()); + matching_ranges(&needles, &haystack, &mut scratch, &mut out); + assert_eq!(out, expected); } } }