Corgi: search few keys in a long chunk in lockstep

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<T, Diff>, 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;
         }

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<T> CorgiReduceBackend<T> {
 
 /// 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<T: ColTime>(chunks: &[&CorgiChunk<T, Diff>], 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()
 }
 

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<usize>,
+    out: &mut Vec<(usize, Range<usize>)>,
+) {
+    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::<Vec<_>>(),
                 expected
             );
+            let (mut scratch, mut out) = (Vec::new(), Vec::new());
+            matching_ranges(&needles, &haystack, &mut scratch, &mut out);
+            assert_eq!(out, expected);
         }
     }
 }