diff --git a/Cargo.toml b/Cargo.toml index a5b050c78..21521b9ac 100644 --- a/Cargo.toml +++ b/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] diff --git a/columnar/benches/bench_optional_index.rs b/columnar/benches/bench_optional_index.rs index 03ff1df97..223514821 100644 --- a/columnar/benches/bench_optional_index.rs +++ b/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, +} + +fn rank_if_exists_batch_over_chunks(input: &BatchInput) -> usize { + let mut doc_ids_out: Vec = Vec::with_capacity(BATCH_LEN); + let mut row_ids_out: Vec = 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 = 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 = (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 = 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(); } diff --git a/columnar/src/column_index/optional_index/mod.rs b/columnar/src/column_index/optional_index/mod.rs index 72fe102e8..590108b89 100644 --- a/columnar/src/column_index/optional_index/mod.rs +++ b/columnar/src/column_index/optional_index/mod.rs @@ -330,6 +330,10 @@ impl OptionalIndex { doc_ids_out: &mut Vec, row_ids_out: &mut Vec, ) { + // 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); 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 6b2b89df7..945dfed04 100644 --- a/columnar/src/column_index/optional_index/set_block/sparse.rs +++ b/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 for SparseBlock<'_> { #[inline(always)] fn rank_if_exists(&self, el: u16) -> Option { - 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 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, + mut sorted_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; + 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 { - self.binary_search_from(target, 0) + fn range(&self, els_range: Range) -> 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 { - 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 { + 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 { + // If you are wondering, the alignment is guaranteed by Rust's reference: + // > An array of [T; N] has a size of size_of::() * 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 { + let mut buffer: Vec = 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 = (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 = (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 = 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, u16)> { + proptest::collection::btree_set(any::(), 0..1_000).prop_flat_map(|vals| { + let vals: Vec = 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::().boxed() + } else { + prop_oneof![any::(), 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, Vec)> { + // 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 = 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) } }