columnar storage - add batched rank

columnar/src/column_index/mod.rs

@@ -107,12 +107,7 @@ impl ColumnIndex {
                 row_ids.extend_from_slice(doc_ids);
             }
             ColumnIndex::Optional(optional_index) => {
-                for doc_id in doc_ids {
-                    if let Some(row_id) = optional_index.rank_if_exists(*doc_id) {
-                        doc_ids_out.push(*doc_id);
-                        row_ids.push(row_id);
-                    }
-                }
+                optional_index.rank_if_exists_batch(doc_ids, doc_ids_out, row_ids);
             }
             ColumnIndex::Multivalued(multivalued_index) => {
                 for doc_id in doc_ids {

columnar/src/column_index/optional_index/mod.rs

@@ -280,6 +280,55 @@ impl OptionalIndex {
         self.num_non_null_docs
     }
 
+    /// Appends the queried doc ids that exist in this index and their corresponding ranks.
+    ///
+    /// Sorted input is processed one encoded block at a time. Unsorted input is supported too,
+    /// but falls back to independent lookups in order to preserve its input order.
+    #[inline]
+    pub fn rank_if_exists_batch(
+        &self,
+        doc_ids: &[DocId],
+        doc_ids_out: &mut Vec<DocId>,
+        row_ids_out: &mut Vec<RowId>,
+    ) {
+        if !doc_ids.is_sorted() {
+            for &doc_id in doc_ids {
+                if let Some(row_id) = self.rank_if_exists(doc_id) {
+                    doc_ids_out.push(doc_id);
+                    row_ids_out.push(row_id);
+                }
+            }
+            return;
+        }
+
+        let mut block_doc_start = 0usize;
+        while block_doc_start < doc_ids.len() {
+            let block_id = row_addr_from_row_id(doc_ids[block_doc_start]).block_id;
+            let mut block_doc_end = block_doc_start + 1;
+            while block_doc_end < doc_ids.len()
+                && row_addr_from_row_id(doc_ids[block_doc_end]).block_id == block_id
+            {
+                block_doc_end += 1;
+            }
+
+            if let Some(&block_meta) = self.block_metas.get(block_id as usize) {
+                let block_doc_id_start = block_id as u32 * ELEMENTS_PER_BLOCK;
+                let row_id_start = block_meta.non_null_rows_before_block;
+                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),
+                    |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);
+                    },
+                );
+            }
+
+            block_doc_start = block_doc_end;
+        }
+    }
+
     pub fn iter_non_null_docs(&self) -> impl Iterator<Item = RowId> + '_ {
         // TODO optimize. We could iterate over the blocks directly.
         // We use the dense value ids and retrieve the doc ids via select.
@@ -335,6 +384,28 @@ enum Block<'a> {
     Sparse(SparseBlock<'a>),
 }
 
+impl Block<'_> {
+    #[inline]
+    fn rank_if_exists_batch(
+        self,
+        doc_ids: impl Iterator<Item = u16>,
+        mut collect: impl FnMut(u16, u16),
+    ) {
+        match self {
+            Block::Dense(dense_block) => {
+                for doc_id in doc_ids {
+                    if let Some(row_id) = dense_block.rank_if_exists(doc_id) {
+                        collect(doc_id, row_id);
+                    }
+                }
+            }
+            Block::Sparse(sparse_block) => {
+                sparse_block.rank_if_exists_batch(doc_ids, collect);
+            }
+        }
+    }
+}
+
 fn serialize_optional_index_block(block_els: &[u16], out: &mut impl io::Write) -> io::Result<()> {
     let is_sparse = is_sparse(block_els.len() as u32);
     if is_sparse {

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

@@ -81,16 +81,68 @@ impl SparseBlock<'_> {
         (self.0.len() / 2) as u16
     }
 
+    /// 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.
+    #[inline]
+    pub(crate) fn rank_if_exists_batch(
+        &self,
+        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;
+
+        for el in els {
+            debug_assert!(previous_el.is_none_or(|previous_el| previous_el <= el));
+            previous_el = Some(el);
+
+            if candidate_rank >= num_vals {
+                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);
+                }
+            }
+
+            if candidate == el {
+                collect(el, candidate_rank);
+            }
+        }
+    }
+
+    #[inline]
+    // Looks for the element in the block. Returns the position if found.
+    fn binary_search(&self, target: u16) -> Result<u16, u16> {
+        self.binary_search_from(target, 0)
+    }
+
     #[inline]
     #[expect(clippy::comparison_chain)]
-    // Looks for the element in the block. Returns the positions if found.
-    fn binary_search(&self, target: u16) -> Result<u16, u16> {
+    fn binary_search_from(&self, target: u16, mut left: u16) -> Result<u16, u16> {
         let data = &self.0;
-        let mut size = self.num_vals();
-        let mut left = 0;
-        let mut right = size;
-        // TODO try different implem.
-        //  e.g. exponential search into binary search
+        let mut right = self.num_vals();
+        let mut size = right - left;
         while left < right {
             let mid = left + size / 2;
 

columnar/src/column_index/optional_index/tests.rs

@@ -2,7 +2,7 @@ use proptest::prelude::*;
 use proptest::{prop_oneof, proptest};
 
 use super::*;
-use crate::{ColumnarReader, ColumnarWriter, DynamicColumnHandle};
+use crate::{ColumnIndex, ColumnarReader, ColumnarWriter, DynamicColumnHandle};
 
 #[test]
 fn test_optional_index_bug_2293() {
@@ -123,6 +123,11 @@ fn test_null_index(data: &[bool]) {
 
     // 100 samples
     let step_size = (data.len() / 100).max(1);
+    let sampled_doc_ids: Vec<DocId> = (0..data.len())
+        .step_by(step_size)
+        .map(|doc_id| doc_id as DocId)
+        .collect();
+    assert_rank_if_exists_batch_matches_scalar(&null_index, &sampled_doc_ids);
     for (pos, value) in data.iter().enumerate().step_by(step_size) {
         assert_eq!(null_index.contains(pos as u32), *value);
     }
@@ -143,6 +148,98 @@ fn test_optional_index_translate() {
     assert_eq!(optional_index.rank_if_exists(2), Some(1));
 }
 
+fn assert_rank_if_exists_batch_matches_scalar(optional_index: &OptionalIndex, doc_ids: &[DocId]) {
+    let mut expected_doc_ids = vec![u32::MAX];
+    let mut expected_row_ids = vec![u32::MAX];
+    for &doc_id in doc_ids {
+        if let Some(row_id) = optional_index.rank_if_exists(doc_id) {
+            expected_doc_ids.push(doc_id);
+            expected_row_ids.push(row_id);
+        }
+    }
+
+    let mut actual_doc_ids = vec![u32::MAX];
+    let mut actual_row_ids = vec![u32::MAX];
+    optional_index.rank_if_exists_batch(doc_ids, &mut actual_doc_ids, &mut actual_row_ids);
+    assert_eq!(actual_doc_ids, expected_doc_ids);
+    assert_eq!(actual_row_ids, expected_row_ids);
+
+    let mut column_doc_ids = vec![u32::MAX];
+    let mut column_row_ids = vec![u32::MAX];
+    ColumnIndex::Optional(optional_index.clone()).docids_to_rowids(
+        doc_ids,
+        &mut column_doc_ids,
+        &mut column_row_ids,
+    );
+    assert_eq!(column_doc_ids, expected_doc_ids);
+    assert_eq!(column_row_ids, expected_row_ids);
+}
+
+#[test]
+fn test_optional_index_rank_if_exists_batch_sparse_blocks() {
+    let row_ids = [
+        1,
+        3,
+        10,
+        20,
+        30,
+        40,
+        50,
+        ELEMENTS_PER_BLOCK - 1,
+        2 * ELEMENTS_PER_BLOCK + 2,
+        2 * ELEMENTS_PER_BLOCK + 200,
+    ];
+    let optional_index = OptionalIndex::for_test(3 * ELEMENTS_PER_BLOCK, &row_ids);
+
+    let sorted_doc_ids = [
+        0,
+        1,
+        2,
+        3,
+        3,
+        49,
+        50,
+        ELEMENTS_PER_BLOCK - 1,
+        ELEMENTS_PER_BLOCK,
+        ELEMENTS_PER_BLOCK + 1,
+        2 * ELEMENTS_PER_BLOCK + 1,
+        2 * ELEMENTS_PER_BLOCK + 2,
+        2 * ELEMENTS_PER_BLOCK + 199,
+        2 * ELEMENTS_PER_BLOCK + 200,
+        3 * ELEMENTS_PER_BLOCK,
+    ];
+    assert_rank_if_exists_batch_matches_scalar(&optional_index, &sorted_doc_ids);
+
+    // Arbitrary input order and duplicate doc ids retain the scalar lookup semantics.
+    let unsorted_doc_ids = [
+        2 * ELEMENTS_PER_BLOCK + 200,
+        3,
+        ELEMENTS_PER_BLOCK,
+        1,
+        3,
+        u32::MAX,
+        50,
+    ];
+    assert_rank_if_exists_batch_matches_scalar(&optional_index, &unsorted_doc_ids);
+}
+
+#[test]
+fn test_optional_index_rank_if_exists_batch_dense_and_sparse_blocks() {
+    let mut row_ids: Vec<RowId> = (0..6_000).collect();
+    row_ids.extend([ELEMENTS_PER_BLOCK + 5, 2 * ELEMENTS_PER_BLOCK + 7]);
+    let optional_index = OptionalIndex::for_test(2 * ELEMENTS_PER_BLOCK + 8, row_ids.as_slice());
+    let doc_ids = [
+        0,
+        1,
+        5_999,
+        6_000,
+        ELEMENTS_PER_BLOCK - 1,
+        ELEMENTS_PER_BLOCK + 5,
+        2 * ELEMENTS_PER_BLOCK + 7,
+    ];
+    assert_rank_if_exists_batch_matches_scalar(&optional_index, &doc_ids);
+}
+
 #[test]
 fn test_optional_index_small() {
     let optional_index = OptionalIndex::for_test(4, &[0, 2]);