Better rank for sparse columns (#3159)

* Faster lookups in sparse optional index blocks

- Branchless binary search in SparseBlock, using
  std::hint::select_unpredictable. Indices are usize: with u16, LLVM
  emits a branch on x86_64, even with the hint.
- SparseBlock::rank_if_exists_batch checks the next value first, then
  uses an exponential search. The search is bounded by the doc id gap:
  values are distinct and sorted, so vals[c + k] >= vals[c] + k.
- OptionalIndex::rank_if_exists_batch reserves its output vectors.
- MSRV 1.88, for select_unpredictable.
- Add a rank_if_exists_batch group to bench_optional_index, unit tests
  and proptests.

* CR comments

---------

Co-authored-by: Paul Masurel <paul.masurel@datadoghq.com>

Cargo.toml

@@ -11,7 +11,7 @@ repository = "https://github.com/quickwit-oss/tantivy"
 readme = "README.md"
 keywords = ["search", "information", "retrieval"]
 edition = "2021"
-rust-version = "1.87"
+rust-version = "1.88"
 exclude = ["benches/*.json", "benches/*.txt"]
 
 [dependencies]

columnar/benches/bench_optional_index.rs

@@ -55,6 +55,31 @@ fn walk_over_data_from_positions(
     dense_idx
 }
 
+/// Number of doc ids per call to `rank_if_exists_batch`, matching tantivy's
+/// `COLLECT_BLOCK_BUFFER_LEN`.
+const BATCH_LEN: usize = 64;
+
+struct BatchInput {
+    optional_index: OptionalIndex,
+    /// Sorted doc ids, queried in chunks of `BATCH_LEN`.
+    doc_ids: Vec<u32>,
+}
+
+fn rank_if_exists_batch_over_chunks(input: &BatchInput) -> usize {
+    let mut doc_ids_out: Vec<u32> = Vec::with_capacity(BATCH_LEN);
+    let mut row_ids_out: Vec<u32> = Vec::with_capacity(BATCH_LEN);
+    let mut num_found = 0;
+    for doc_ids in input.doc_ids.chunks(BATCH_LEN) {
+        doc_ids_out.clear();
+        row_ids_out.clear();
+        input
+            .optional_index
+            .rank_if_exists_batch(doc_ids, &mut doc_ids_out, &mut row_ids_out);
+        num_found += row_ids_out.len();
+    }
+    num_found
+}
+
 fn main() {
     // Build separate inputs for each fill ratio.
     let inputs: Vec<(String, OptionalIndex)> = vec![
@@ -103,4 +128,41 @@ fn main() {
     });
 
     group.run();
+
+    // `rank_if_exists_batch` on sorted batches of doc ids. fill=1% and fill=5% only produce
+    // sparse blocks, fill=50% only dense blocks. `step` is the average distance between
+    // queried doc ids.
+    let mut batch_inputs: Vec<(String, BatchInput)> = Vec::new();
+    for (fill_name, fill_ratio) in [("1%", 0.01), ("5%", 0.05), ("50%", 0.50)] {
+        for avg_step_size in [1u32, 10, 100] {
+            let doc_ids: Vec<u32> = if avg_step_size == 1 {
+                (0..TOTAL_NUM_VALUES).collect()
+            } else {
+                random_range_iterator(0, TOTAL_NUM_VALUES, avg_step_size, avg_step_size - 1)
+                    .collect()
+            };
+            batch_inputs.push((
+                format!("fill={fill_name} step={avg_step_size}"),
+                BatchInput {
+                    optional_index: gen_optional_index(fill_ratio),
+                    doc_ids,
+                },
+            ));
+        }
+    }
+    // Mirrors the `sparse` runner of `agg_bench` filtered on `single_term`: one doc in 21 has a
+    // value (sparse blocks), and exactly those docs are queried, so every lookup is a hit.
+    let every_21st_doc: Vec<u32> = (0..TOTAL_NUM_VALUES).step_by(21).collect();
+    batch_inputs.push((
+        "every 21st doc, all hits".to_string(),
+        BatchInput {
+            optional_index: OptionalIndex::for_test(TOTAL_NUM_VALUES, &every_21st_doc),
+            doc_ids: every_21st_doc,
+        },
+    ));
+    let mut batch_group: InputGroup<BatchInput> = InputGroup::new_with_inputs(batch_inputs);
+    batch_group.register("rank_if_exists_batch", |input: &BatchInput| {
+        black_box(rank_if_exists_batch_over_chunks(input));
+    });
+    batch_group.run();
 }

columnar/src/column_index/optional_index/mod.rs

@@ -330,6 +330,10 @@ impl OptionalIndex {
         doc_ids_out: &mut Vec<DocId>,
         row_ids_out: &mut Vec<RowId>,
     ) {
+        // At most one output per queried doc id: reserving upfront avoids reallocations
+        // in the middle of the batch.
+        doc_ids_out.reserve(doc_ids.len());
+        row_ids_out.reserve(doc_ids.len());
         // TODO: change the API to require sorted input, and remove the fallback.
         if !doc_ids.is_sorted() {
             for &doc_id in doc_ids {
@@ -357,7 +361,7 @@ impl OptionalIndex {
                 self.block(block_meta).rank_if_exists_batch(
                     doc_ids[block_doc_start..block_doc_end]
                         .iter()
-                        .map(|doc_id| (doc_id - block_doc_id_start) as u16),
+                        .map(|doc_id| row_addr_from_row_id(*doc_id).in_block_row_id),
                     |in_block_doc_id, in_block_row_id| {
                         doc_ids_out.push(block_doc_id_start + in_block_doc_id as u32);
                         row_ids_out.push(row_id_start + in_block_row_id as u32);

columnar/src/column_index/optional_index/set_block/sparse.rs

@@ -1,3 +1,5 @@
+use std::ops::Range;
+
 use crate::column_index::optional_index::{SelectCursor, Set, SetCodec};
 
 pub struct SparseBlockCodec;
@@ -43,12 +45,12 @@ impl Set<u16> for SparseBlock<'_> {
 
     #[inline(always)]
     fn rank_if_exists(&self, el: u16) -> Option<u16> {
-        self.binary_search(el).ok()
+        self.binary_search(el).ok().map(|el| el as u16)
     }
 
     #[inline(always)]
     fn rank(&self, el: u16) -> u16 {
-        self.binary_search(el).unwrap_or_else(|el| el)
+        self.binary_search(el).unwrap_or_else(|el| el) as u16
     }
 
     #[inline(always)]
@@ -63,105 +65,358 @@ impl Set<u16> for SparseBlock<'_> {
     }
 }
 
-#[inline(always)]
-fn get_u16(data: &[u8], byte_position: usize) -> u16 {
-    let bytes: [u8; 2] = data[byte_position..byte_position + 2].try_into().unwrap();
-    u16::from_le_bytes(bytes)
-}
-
 impl SparseBlock<'_> {
-    #[inline(always)]
-    fn value_at_idx(&self, data: &[u8], idx: u16) -> u16 {
-        let start_offset: usize = idx as usize * 2;
-        get_u16(data, start_offset)
-    }
-
     #[inline]
     fn num_vals(&self) -> u16 {
         (self.0.len() / 2) as u16
     }
 
+    /// Returns the value at `idx`, without bounds checking.
+    ///
+    /// # Safety
+    ///
+    /// `idx < self.num_vals()` is required.
+    #[inline(always)]
+    unsafe fn value_at_idx_unchecked(&self, idx: usize) -> u16 {
+        debug_assert!(idx < self.num_vals() as usize);
+        let val_ptr: *const u16 = self.0.as_ptr() as *const u16;
+        u16::from_le(unsafe { val_ptr.add(idx).read_unaligned() })
+    }
+
     /// Looks up several sorted elements while keeping a cursor into the sparse block.
     ///
     /// The callback receives `(element, rank)` for each element present in the block. Once a
     /// lookup has passed a sparse value, later lookups never search the prefix containing that
     /// value again. Queries below the next sparse value only require a comparison with that value.
+    ///
+    /// els are required to be sorted, or the results could be incomplete.
     #[inline]
     pub(crate) fn rank_if_exists_batch(
         &self,
-        els: impl Iterator<Item = u16>,
+        mut sorted_els: impl Iterator<Item = u16>,
         mut collect: impl FnMut(u16, u16),
     ) {
-        let data = self.0;
-        let num_vals = self.num_vals();
-        let mut candidate_rank = 0u16;
-        let mut previous_el = None;
+        let num_vals = self.num_vals() as usize;
+        let mut candidate_rank: usize = 0;
 
-        for el in els {
-            debug_assert!(previous_el.is_none_or(|previous_el| previous_el <= el));
-            previous_el = Some(el);
+        while candidate_rank < num_vals {
+            let candidate = unsafe { self.value_at_idx_unchecked(candidate_rank) };
+            let Some(needle) = sorted_els.find(|needle| *needle >= candidate) else {
+                return;
+            };
 
+            if candidate == needle {
+                // The cursor stays on the matched value, so that a duplicate needle matches too.
+                collect(needle, candidate_rank as u16);
+                continue;
+            }
+
+            // Check the immediately following value first. This makes querying consecutive
+            // sparse values linear instead of performing a binary search for every value.
+            candidate_rank += 1;
             if candidate_rank >= num_vals {
+                // `needle` and the remaining needles are larger than every value.
                 return;
             }
-            let mut candidate = self.value_at_idx(data, candidate_rank);
-            if candidate < el {
-                // Check the immediately following value first. This makes querying consecutive
-                // sparse values linear instead of performing a binary search for every value.
-                candidate_rank += 1;
-                if candidate_rank >= num_vals {
-                    return;
-                }
-                candidate = self.value_at_idx(data, candidate_rank);
-
-                if candidate < el {
-                    candidate_rank = self
-                        .binary_search_from(el, candidate_rank + 1)
-                        .unwrap_or_else(|rank| rank);
-                    if candidate_rank >= num_vals {
-                        return;
-                    }
-                    candidate = self.value_at_idx(data, candidate_rank);
+            let candidate = unsafe { self.value_at_idx_unchecked(candidate_rank) };
+            if candidate >= needle {
+                if candidate == needle {
+                    collect(needle, candidate_rank as u16);
                 }
+                // Otherwise `needle` is absent. Either way, the cursor now points to `candidate`.
+                continue;
             }
 
-            if candidate == el {
-                collect(el, candidate_rank);
+            // `candidate_rank < num_vals`, so `left <= num_vals`.
+            let left = candidate_rank + 1;
+            // No need to always search the rest of the block here!
+            // Values are distinct and sorted, so `vals[candidate_rank + k] >= candidate + k`.
+            // With `k = needle - candidate`, the value at `candidate_rank + k` is already
+            // `>= needle`: the search can stop right after that index.
+            // `needle > candidate`, so `gap_bound >= left + 1`, hence `left <= right`.
+            let gap_bound = candidate_rank + (needle - candidate) as usize + 1;
+            let right = gap_bound.min(num_vals);
+            // The value at `left - 1` is lower than `needle`, and the value at `right`, if any,
+            // is greater than `needle`, so the insertion point found in `left..right` is also
+            // the insertion point in the whole block.
+            // The needle is usually a few values after the cursor, hence the exponential search.
+            match self.range(left..right).exponential_search(needle) {
+                Ok(rel_rank) => {
+                    candidate_rank = left + rel_rank;
+                    collect(needle, candidate_rank as u16);
+                }
+                Err(rel_rank) => {
+                    // `left + rel_rank` may be `num_vals`: the next iteration then returns.
+                    candidate_rank = left + rel_rank;
+                }
             }
         }
     }
 
+    /// Returns the sub-block made of the values at indices `els_range`.
+    ///
+    /// Indices in the returned block are relative to `els_range.start`: callers need to add it
+    /// back to get indices in `self`.
+    ///
+    /// Panics if `els_range.start > els_range.end` or `els_range.end > self.num_vals()`.
     #[inline]
-    // Looks for the element in the block. Returns the position if found.
-    // TODO try different implem.
-    // e.g. exponential search into binary search
-    fn binary_search(&self, target: u16) -> Result<u16, u16> {
-        self.binary_search_from(target, 0)
+    fn range(&self, els_range: Range<usize>) -> Self {
+        let Range { start, end } = els_range;
+        SparseBlock(&self.0[start * 2..end * 2])
     }
 
+    /// Searches `target` among all the values of the block, starting from its beginning.
+    ///
+    /// Returns the same result as `binary_search`: `Ok(idx)` if found, `Err(idx)` with the
+    /// insertion point otherwise. It is faster when `target` is close to the start of the block.
     #[inline]
-    #[expect(clippy::comparison_chain)]
-    fn binary_search_from(&self, target: u16, mut left: u16) -> Result<u16, u16> {
-        let data = &self.0;
-        let mut right = self.num_vals();
-        let mut size = right - left;
-        while left < right {
-            let mid = left + size / 2;
-
-            // TODO do boundary check only once, and then use an
-            // unsafe `value_at_idx`
-            let mid_val = self.value_at_idx(data, mid);
-
-            if target > mid_val {
-                left = mid + 1;
-            } else if target < mid_val {
-                right = mid;
-            } else {
-                return Ok(mid);
+    fn exponential_search(&self, target: u16) -> Result<usize, usize> {
+        let num_vals = self.num_vals() as usize;
+        // Invariant: the values at indices `< left` are lower than `target`, and if `target` is
+        // present, it is in `left..right`.
+        let mut left = 0;
+        let mut right = num_vals;
+        // Probes indices `0, 2, 6, 14, 30`... (i.e. `2^(k+1) - 2`)
+        let mut step: usize = 1;
+        loop {
+            let probe = left + step - 1;
+            if probe >= right {
+                break;
             }
+            // SAFETY: `probe < right <= num_vals`.
+            let probe_val = unsafe { self.value_at_idx_unchecked(probe) };
+            if probe_val >= target {
+                right = probe + 1;
+                break;
+            }
+            left = probe + 1;
+            step *= 2;
+        }
+        // Either `right == num_vals`, or the value at `right - 1` is `>= target`: the insertion
+        // point found in `left..right` is also the insertion point in the whole block.
+        match self.range(left..right).binary_search(target) {
+            Ok(rel_idx) => Ok(left + rel_idx),
+            Err(rel_idx) => Err(left + rel_idx),
+        }
+    }
 
-            size = right - left;
+    /// Searches `target` among all the values of the block.
+    ///
+    /// Returns `Ok(idx)` if found, `Err(idx)` with the insertion point otherwise.
+    ///
+    /// The bytes cannot be cast to `&[u16]`: the block has no alignment guarantee, and values are
+    /// little endian. They are viewed as `&[[u8; 2]]` instead, which is always valid, so that the
+    /// std binary search can be used.
+    #[inline]
+    fn binary_search(&self, target: u16) -> Result<usize, usize> {
+        // If you are wondering, the alignment is guaranteed by Rust's reference:
+        // > An array of [T; N] has a size of size_of::<T>() * N and the same alignment of T.
+        let (vals, remainder): (&[[u8; 2]], &[u8]) = self.0.as_chunks::<2>();
+        // The block length is even, so the remainder is empty.
+        debug_assert!(remainder.is_empty());
+        // Values are distinct, so the returned index is unique.
+        vals.binary_search_by_key(&target, |val_bytes| u16::from_le_bytes(*val_bytes))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use proptest::prelude::*;
+
+    use super::*;
+
+    fn serialize_sparse(vals: &[u16]) -> Vec<u8> {
+        let mut buffer: Vec<u8> = Vec::with_capacity(vals.len() * 2);
+        SparseBlockCodec::serialize(vals.iter().copied(), &mut buffer).unwrap();
+        buffer
+    }
+
+    #[test]
+    fn test_binary_search_empty() {
+        let buffer = serialize_sparse(&[]);
+        let block = SparseBlockCodec::open(&buffer);
+        assert_eq!(block.binary_search(0), Err(0));
+        assert_eq!(block.binary_search(17), Err(0));
+        assert_eq!(block.binary_search(u16::MAX), Err(0));
+    }
+
+    #[test]
+    fn test_binary_search_single_value() {
+        let buffer = serialize_sparse(&[10]);
+        let block = SparseBlockCodec::open(&buffer);
+        assert_eq!(block.binary_search(9), Err(0));
+        assert_eq!(block.binary_search(10), Ok(0));
+        assert_eq!(block.binary_search(11), Err(1));
+    }
+
+    #[test]
+    fn test_binary_search_extreme_values() {
+        let buffer = serialize_sparse(&[0, u16::MAX]);
+        let block = SparseBlockCodec::open(&buffer);
+        assert_eq!(block.binary_search(0), Ok(0));
+        assert_eq!(block.binary_search(1), Err(1));
+        assert_eq!(block.binary_search(u16::MAX - 1), Err(1));
+        assert_eq!(block.binary_search(u16::MAX), Ok(1));
+    }
+
+    #[test]
+    fn test_binary_search_every_position() {
+        let vals: Vec<u16> = (0..100u16).map(|val| val * 3 + 1).collect();
+        let buffer = serialize_sparse(&vals);
+        let block = SparseBlockCodec::open(&buffer);
+        for (idx, &val) in vals.iter().enumerate() {
+            assert_eq!(block.binary_search(val), Ok(idx));
+            assert_eq!(block.binary_search(val - 1), Err(idx));
+            assert_eq!(block.binary_search(val + 1), Err(idx + 1));
+        }
+    }
+
+    #[test]
+    fn test_exponential_search_every_position() {
+        let empty_buffer = serialize_sparse(&[]);
+        let empty_block = SparseBlockCodec::open(&empty_buffer);
+        assert_eq!(empty_block.exponential_search(0), Err(0));
+        assert_eq!(empty_block.exponential_search(u16::MAX), Err(0));
+
+        let vals: Vec<u16> = (0..100u16).map(|val| val * 3 + 1).collect();
+        let buffer = serialize_sparse(&vals);
+        let block = SparseBlockCodec::open(&buffer);
+        for (idx, &val) in vals.iter().enumerate() {
+            assert_eq!(block.exponential_search(val), Ok(idx));
+            assert_eq!(block.exponential_search(val - 1), Err(idx));
+            assert_eq!(block.exponential_search(val + 1), Err(idx + 1));
+        }
+    }
+
+    #[test]
+    fn test_range_binary_search() {
+        let vals: [u16; 5] = [1, 3, 5, 7, 9];
+        let buffer = serialize_sparse(&vals);
+        let block = SparseBlockCodec::open(&buffer);
+        // Values before index 2 are ignored, and indices are relative to 2.
+        let tail = block.range(2..5);
+        assert_eq!(tail.binary_search(1), Err(0));
+        assert_eq!(tail.binary_search(3), Err(0));
+        assert_eq!(tail.binary_search(5), Ok(0));
+        assert_eq!(tail.binary_search(6), Err(1));
+        assert_eq!(tail.binary_search(9), Ok(2));
+        assert_eq!(tail.binary_search(10), Err(3));
+        // An empty range is allowed, including at `num_vals`.
+        let empty = block.range(5..5);
+        assert_eq!(empty.binary_search(9), Err(0));
+        assert_eq!(empty.binary_search(10), Err(0));
+    }
+
+    #[test]
+    fn test_binary_search_unaligned_data() {
+        // The block data has no alignment guarantee: make sure we read correctly from an odd
+        // address.
+        let vals: [u16; 4] = [2, 0x0100, 0x1234, u16::MAX];
+        let mut buffer: Vec<u8> = vec![0u8];
+        buffer.extend_from_slice(&serialize_sparse(&vals));
+        let block = SparseBlockCodec::open(&buffer[1..]);
+        for (idx, &val) in vals.iter().enumerate() {
+            assert_eq!(block.binary_search(val), Ok(idx));
+        }
+        assert_eq!(block.binary_search(1), Err(0));
+        assert_eq!(block.binary_search(0x1235), Err(3));
+    }
+
+    fn vals_target_strategy() -> impl Strategy<Value = (Vec<u16>, u16)> {
+        proptest::collection::btree_set(any::<u16>(), 0..1_000).prop_flat_map(|vals| {
+            let vals: Vec<u16> = vals.into_iter().collect();
+            // Pick targets among existing values half of the time, to exercise the `Ok` path.
+            let target_strategy = if vals.is_empty() {
+                any::<u16>().boxed()
+            } else {
+                prop_oneof![any::<u16>(), proptest::sample::select(vals.clone())].boxed()
+            };
+            (Just(vals), target_strategy)
+        })
+    }
+
+    proptest! {
+        #[test]
+        fn test_proptest_binary_search((vals, target) in vals_target_strategy()) {
+            let buffer = serialize_sparse(&vals);
+            let block = SparseBlockCodec::open(&buffer);
+            // `vals` is strictly increasing, so both the `Ok` and the `Err` result are uniquely
+            // defined, and `slice::binary_search` is a valid reference.
+            prop_assert_eq!(block.binary_search(target), vals.binary_search(&target));
+        }
+
+        #[test]
+        fn test_proptest_exponential_search((vals, target) in vals_target_strategy()) {
+            let buffer = serialize_sparse(&vals);
+            let block = SparseBlockCodec::open(&buffer);
+            prop_assert_eq!(block.exponential_search(target), vals.binary_search(&target));
+        }
+    }
+
+    fn rank_if_exists_batch_to_vec(block: &SparseBlock, sorted_needles: &[u16]) -> Vec<(u16, u16)> {
+        let mut found: Vec<(u16, u16)> = Vec::with_capacity(sorted_needles.len());
+        block.rank_if_exists_batch(sorted_needles.iter().copied(), |needle, rank| {
+            found.push((needle, rank))
+        });
+        found
+    }
+
+    /// Naive reference for `rank_if_exists_batch`: one independent lookup per needle.
+    fn expected_rank_if_exists_batch(vals: &[u16], sorted_needles: &[u16]) -> Vec<(u16, u16)> {
+        let mut found: Vec<(u16, u16)> = Vec::with_capacity(sorted_needles.len());
+        for &needle in sorted_needles {
+            if let Ok(rank) = vals.binary_search(&needle) {
+                found.push((needle, rank as u16));
+            }
+        }
+        found
+    }
+
+    #[test]
+    fn test_rank_if_exists_batch_match_on_gap_bound() {
+        // Consecutive values make the gap bound tight: when searching for 14 from rank 1
+        // (value 11), the bound stops the search right after rank 1 + (14 - 11) = 4, which is
+        // exactly where 14 is.
+        let vals: [u16; 6] = [10, 11, 12, 13, 14, 20];
+        let buffer = serialize_sparse(&vals);
+        let block = SparseBlockCodec::open(&buffer);
+        assert_eq!(
+            rank_if_exists_batch_to_vec(&block, &[10, 14, 15, 20, 21]),
+            vec![(10, 0), (14, 4), (20, 5)]
+        );
+    }
+
+    fn vals_needles_strategy() -> impl Strategy<Value = (Vec<u16>, Vec<u16>)> {
+        // A small value domain gives small gaps between values, so the gap bound often
+        // matters. A large domain exercises the case where the bound exceeds the block.
+        prop_oneof![Just(200u16), Just(2_000u16), Just(u16::MAX)].prop_flat_map(|max_val| {
+            proptest::collection::btree_set(0..=max_val, 0..1_000).prop_flat_map(move |vals| {
+                let vals: Vec<u16> = vals.into_iter().collect();
+                // Pick needles among existing values half of the time, to exercise matches.
+                let needle_strategy = if vals.is_empty() {
+                    (0..=max_val).boxed()
+                } else {
+                    prop_oneof![0..=max_val, proptest::sample::select(vals.clone())].boxed()
+                };
+                let needles_strategy =
+                    proptest::collection::vec(needle_strategy, 0..100).prop_map(|mut needles| {
+                        needles.sort_unstable();
+                        needles
+                    });
+                (Just(vals), needles_strategy)
+            })
+        })
+    }
+
+    proptest! {
+        #[test]
+        fn test_proptest_rank_if_exists_batch((vals, sorted_needles) in vals_needles_strategy()) {
+            let buffer = serialize_sparse(&vals);
+            let block = SparseBlockCodec::open(&buffer);
+            prop_assert_eq!(
+                rank_if_exists_batch_to_vec(&block, &sorted_needles),
+                expected_rank_if_exists_batch(&vals, &sorted_needles)
+            );
         }
-        Err(left)
     }
 }