diff --git a/columnar/src/column_index/mod.rs b/columnar/src/column_index/mod.rs index 0003c1a3b..bf6b5b3df 100644 --- a/columnar/src/column_index/mod.rs +++ b/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 { diff --git a/columnar/src/column_index/optional_index/mod.rs b/columnar/src/column_index/optional_index/mod.rs index a21991177..8dc7086d4 100644 --- a/columnar/src/column_index/optional_index/mod.rs +++ b/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, + row_ids_out: &mut Vec, + ) { + 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 + '_ { // 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, + 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 { diff --git a/columnar/src/column_index/optional_index/set_block/sparse.rs b/columnar/src/column_index/optional_index/set_block/sparse.rs index 6fc0e3469..bc1fed0af 100644 --- a/columnar/src/column_index/optional_index/set_block/sparse.rs +++ b/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, + 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 { + 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 { + fn binary_search_from(&self, target: u16, mut left: u16) -> Result { 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; diff --git a/columnar/src/column_index/optional_index/tests.rs b/columnar/src/column_index/optional_index/tests.rs index 170f8fa72..16d61be4e 100644 --- a/columnar/src/column_index/optional_index/tests.rs +++ b/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 = (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 = (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]);