diff --git a/lib/segment/src/index/sparse_index/sparse_vector_index/read_view/search.rs b/lib/segment/src/index/sparse_index/sparse_vector_index/read_view/search.rs index 8c0a4f33bf..2b7505c7c2 100644 --- a/lib/segment/src/index/sparse_index/sparse_vector_index/read_view/search.rs +++ b/lib/segment/src/index/sparse_index/sparse_vector_index/read_view/search.rs @@ -14,7 +14,7 @@ use crate::data_types::vectors::{QueryVector, VectorInternal}; use crate::id_tracker::IdTrackerRead; use crate::index::PayloadIndexRead; use crate::index::field_index::CardinalityEstimation; -use crate::index::hnsw_index::point_scorer::BatchFilteredSearcher; +use crate::index::hnsw_index::point_scorer::{BatchFilteredSearcher, ScorerFilters}; use crate::index::query_estimator::adjust_to_available_vectors; use crate::telemetry::VectorIndexSearchesTelemetry; use crate::types::{DEFAULT_SPARSE_FULL_SCAN_THRESHOLD, Filter}; @@ -219,9 +219,6 @@ where .unwrap_or(self.id_tracker.deleted_point_bitslice()); let not_deleted = self.vector_storage.not_deleted_checker(point_deleted); - let not_deleted_condition = - |idx: PointOffsetType| -> bool { not_deleted.check_infallible(idx) }; - let is_stopped = vector_query_context.is_stopped(); let sparse_vector = self.indices_tracker.remap_vector(sparse_vector.clone()); @@ -243,16 +240,11 @@ where &hw_counter, )?; - match filter { - Some(filter) => { - let filter_context = self.payload_index.filter_context(filter, &hw_counter)?; - let matches_filter_condition = |idx: PointOffsetType| -> bool { - not_deleted_condition(idx) && filter_context.check_infallible(idx) - }; - Ok(search_context.search(&matches_filter_condition)) - } - None => Ok(search_context.search(¬_deleted_condition)), - } + let filter_context = filter + .map(|filter| self.payload_index.filter_context(filter, &hw_counter)) + .transpose()?; + let filter_condition = ScorerFilters::new(filter_context, not_deleted); + Ok(search_context.search(&filter_condition)) } fn search_nearest_query( diff --git a/lib/sparse/benches/search.rs b/lib/sparse/benches/search.rs index dd1c57c196..2d8ad90bab 100644 --- a/lib/sparse/benches/search.rs +++ b/lib/sparse/benches/search.rs @@ -2,6 +2,7 @@ use std::borrow::Cow; use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicBool; +use common::condition_checker::ConstantConditionChecker; use common::counter::hardware_counter::HardwareCounterCell; use common::types::PointOffsetType; use common::universal_io::{MmapFile, MmapFs}; @@ -190,6 +191,7 @@ fn run_bench2( let mut it = query_vectors.iter().cycle(); let hardware_counter = HardwareCounterCell::new(); + let checker = ConstantConditionChecker::<()>::MATCH_ALL; group.bench_function("basic", |b| { b.iter_batched( @@ -198,7 +200,7 @@ fn run_bench2( let mut scratch = pool.get(); SearchContext::new(vec, TOP, index, &mut scratch, &stopped, &hardware_counter) .unwrap() - .search(&|_| true) + .search(&checker) }, criterion::BatchSize::SmallInput, ) @@ -214,7 +216,7 @@ fn run_bench2( let mut scratch = pool.get(); SearchContext::new(vec, TOP, index, &mut scratch, &stopped, &hardware_counter) .unwrap() - .search(&|_| true) + .search(&checker) }, criterion::BatchSize::SmallInput, ) diff --git a/lib/sparse/src/index/search_context.rs b/lib/sparse/src/index/search_context.rs index 16b1844554..e358ccaa1e 100644 --- a/lib/sparse/src/index/search_context.rs +++ b/lib/sparse/src/index/search_context.rs @@ -2,6 +2,7 @@ use std::cmp::{Ordering, max, min}; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering::Relaxed; +use common::condition_checker::{ConditionChecker, Rest, Select}; use common::counter::hardware_counter::HardwareCounterCell; use common::top_k::TopK; use common::types::{PointOffsetType, ScoreType, ScoredPointOffset}; @@ -24,6 +25,13 @@ pub struct IndexedPostingListIterator { /// Making this larger makes the search faster but uses more (pooled) memory const ADVANCE_BATCH_SIZE: usize = 10_000; +/// How many candidates to accumulate before flushing them through the filter. +/// +/// Each flush updates the "min score to beat" ([`TopK::threshold`]), which +/// gates further candidates, so smaller batches mean fewer wasted filter +/// checks, while larger batches amortize the filter better. +const FILTER_BATCH_SIZE: usize = 128; + pub struct SearchContext<'a, T: PostingListIter = PostingListIterator<'a>> { postings_iterators: Vec>, query: RemappedSparseVector, @@ -34,6 +42,8 @@ pub struct SearchContext<'a, T: PostingListIter = PostingListIterator<'a>> { max_record_id: PointOffsetType, // max_record_id ids across all posting lists /// Scores buffer from [`SearchScratch`]. scores: &'a mut Vec, + /// Buffer for batched filtering from [`SearchScratch`]. + candidates: &'a mut Vec, use_pruning: bool, hardware_counter: &'a HardwareCounterCell, } @@ -81,6 +91,7 @@ impl<'a, T: PostingListIter> SearchContext<'a, T> { min_record_id, max_record_id, scores: &mut scratch.scores, + candidates: &mut scratch.candidates, use_pruning, hardware_counter, }) @@ -143,11 +154,11 @@ impl<'a, T: PostingListIter> SearchContext<'a, T> { } /// Advance posting lists iterators in a batch fashion. - fn advance_batch bool>( + fn advance_batch( &mut self, batch_start_id: PointOffsetType, batch_last_id: PointOffsetType, - filter_condition: &F, + checker: &C, ) { // init batch scores let batch_len = batch_last_id - batch_start_id + 1; @@ -169,39 +180,64 @@ impl<'a, T: PostingListIter> SearchContext<'a, T> { ); } - for (local_index, &score) in self.scores.iter().enumerate() { - // publish only the non-zero scores above the current min to beat - if score != 0.0 && score > self.top_results.threshold() { + // Collect only the non-zero scores above the current min to beat, + // then check them against the filter batch by batch. + let Self { + scores, + candidates, + top_results, + .. + } = self; + candidates.clear(); + // Hoisted out of the hot loop: only a flush can move the threshold. + let mut threshold = top_results.threshold(); + for (local_index, &score) in scores.iter().enumerate() { + if score != 0.0 && score > threshold { let real_id = batch_start_id + local_index as PointOffsetType; - // do not score if filter condition is not satisfied - if !filter_condition(real_id) { - continue; - } - let score_point_offset = ScoredPointOffset { + candidates.push(ScoredPointOffset { score, idx: real_id, - }; - self.top_results.push(score_point_offset); + }); + if candidates.len() == FILTER_BATCH_SIZE { + push_filtered(candidates, top_results, checker); + threshold = top_results.threshold(); + } } } + push_filtered(candidates, top_results, checker); } /// Compute scores for the last posting list quickly - fn process_last_posting_list bool>(&mut self, filter_condition: &F) { + fn process_last_posting_list(&mut self, checker: &C) { debug_assert_eq!(self.postings_iterators.len(), 1); - let posting = &mut self.postings_iterators[0]; + let Self { + postings_iterators, + top_results, + candidates, + .. + } = self; + let posting = &mut postings_iterators[0]; + let query_weight = posting.query_weight; + candidates.clear(); + // Hoisted out of the hot loop: only a flush can move the threshold. + let mut threshold = top_results.threshold(); posting.posting_list_iterator.for_each_till_id( PointOffsetType::MAX, &mut (), |_, id, weight| { - // do not score if filter condition is not satisfied - if !filter_condition(id) { - return; + let score = weight * query_weight; + // The same "min to beat" check as in `TopK::push`, but before + // paying for the filter check. + if score > threshold { + candidates.push(ScoredPointOffset { score, idx: id }); + if candidates.len() == FILTER_BATCH_SIZE { + push_filtered(candidates, top_results, checker); + threshold = top_results.threshold(); + } } - let score = weight * posting.query_weight; - self.top_results.push(ScoredPointOffset { score, idx: id }); }, ); + push_filtered(candidates, top_results, checker); } /// Returns the next min record id from all posting list iterators @@ -260,10 +296,7 @@ impl<'a, T: PostingListIter> SearchContext<'a, T> { } /// Search for the top k results that satisfy the filter condition - pub fn search bool>( - &mut self, - filter_condition: &F, - ) -> Vec { + pub fn search(&mut self, checker: &C) -> Vec { if self.postings_iterators.is_empty() { return Vec::new(); } @@ -300,7 +333,7 @@ impl<'a, T: PostingListIter> SearchContext<'a, T> { ); // advance and score posting lists iterators - self.advance_batch(start_batch_id, last_batch_id, filter_condition); + self.advance_batch(start_batch_id, last_batch_id, checker); // remove empty posting lists if necessary self.postings_iterators.retain(|posting_iterator| { @@ -317,7 +350,7 @@ impl<'a, T: PostingListIter> SearchContext<'a, T> { // if only one posting list left, we can score it quickly if self.postings_iterators.len() == 1 { - self.process_last_posting_list(filter_condition); + self.process_last_posting_list(checker); break; } @@ -413,3 +446,20 @@ impl<'a, T: PostingListIter> SearchContext<'a, T> { false } } + +/// Batch-check `candidates` against the filter and push the survivors into +/// `top_results`. Drains `candidates`. +fn push_filtered( + candidates: &mut Vec, + top_results: &mut TopK, + checker: &C, +) { + // Errors are swallowed the same way as in `ConditionChecker::check_infallible`. + let n = checker + .check_batched(candidates, Select::Matches, Rest::Discard) + .unwrap_or(0); + for &score_point_offset in &candidates[..n] { + top_results.push(score_point_offset); + } + candidates.clear(); +} diff --git a/lib/sparse/src/index/tests/common.rs b/lib/sparse/src/index/tests/common.rs index a29d6ba1d7..1f4ad1e2e1 100644 --- a/lib/sparse/src/index/tests/common.rs +++ b/lib/sparse/src/index/tests/common.rs @@ -1,5 +1,6 @@ use std::borrow::Cow; +use common::condition_checker::ConstantConditionChecker; use common::types::PointOffsetType; use common::universal_io::{MmapFile, MmapFs}; use rand::{RngExt, SeedableRng}; @@ -102,6 +103,6 @@ where generate_sparse_index::(&mut rnd_gen, count, density, vocab1, vocab2) } -pub fn match_all(_p: PointOffsetType) -> bool { - true +pub fn match_all() -> ConstantConditionChecker { + ConstantConditionChecker::MATCH_ALL } diff --git a/lib/sparse/src/index/tests/hw_counter_test.rs b/lib/sparse/src/index/tests/hw_counter_test.rs index ccafb030bb..9435dd8852 100644 --- a/lib/sparse/src/index/tests/hw_counter_test.rs +++ b/lib/sparse/src/index/tests/hw_counter_test.rs @@ -26,7 +26,7 @@ fn do_search(index: &I, query: RemappedSparseVector) -> HwMeas ) .unwrap(); - let result = search_context.search(&match_all); + let result = search_context.search(&match_all()); // there might be less than `top` result // happens if index contains less than `top` sparse vectors with indices overlapping the query indices assert!(result.len() <= top); diff --git a/lib/sparse/src/index/tests/indexed_vs_plain_test.rs b/lib/sparse/src/index/tests/indexed_vs_plain_test.rs index 049cf309e7..ea6ecd2e4c 100644 --- a/lib/sparse/src/index/tests/indexed_vs_plain_test.rs +++ b/lib/sparse/src/index/tests/indexed_vs_plain_test.rs @@ -24,7 +24,7 @@ fn query(index: &I, query: RemappedSparseVector) { ) .unwrap(); - let result = search_context.search(&match_all); + let result = search_context.search(&match_all()); let docs: Vec<_> = result.iter().map(|x| x.idx).collect(); drop(search_context); diff --git a/lib/sparse/src/index/tests/search_context_tests.rs b/lib/sparse/src/index/tests/search_context_tests.rs index 44b3d65150..941e9c5452 100644 --- a/lib/sparse/src/index/tests/search_context_tests.rs +++ b/lib/sparse/src/index/tests/search_context_tests.rs @@ -2,9 +2,10 @@ use std::any::TypeId; use std::borrow::Cow; use std::sync::atomic::AtomicBool; +use common::condition_checker::ConstantConditionChecker; use common::counter::hardware_accumulator::HwMeasurementAcc; use common::counter::hardware_counter::HardwareCounterCell; -use common::types::{PointOffsetType, ScoredPointOffset}; +use common::types::ScoredPointOffset; #[cfg(target_os = "linux")] use common::universal_io::{IoUringFile, IoUringFs}; use common::universal_io::{MmapFile, MmapFs}; @@ -22,8 +23,8 @@ use crate::index::posting_list_common::PostingListIter; use crate::index::search_context::SearchContext; /// Match all filter condition for testing -fn match_all(_p: PointOffsetType) -> bool { - true +fn match_all() -> ConstantConditionChecker { + ConstantConditionChecker::MATCH_ALL } #[duplicate::duplicate_item( @@ -103,7 +104,7 @@ mod test_mod { &hw_counter, ) .unwrap(); - assert_eq!(search_context.search(&match_all), Vec::new()); + assert_eq!(search_context.search(&match_all()), Vec::new()); } #[test] @@ -134,7 +135,7 @@ mod test_mod { .unwrap(); assert_eq!( - round_scores(search_context.search(&match_all)), + round_scores(search_context.search(&match_all())), vec![ ScoredPointOffset { score: 90.0, @@ -194,7 +195,7 @@ mod test_mod { .unwrap(); assert_eq!( - round_scores(search_context.search(&match_all)), + round_scores(search_context.search(&match_all())), vec![ ScoredPointOffset { score: 90.0, @@ -238,7 +239,7 @@ mod test_mod { .unwrap(); assert_eq!( - search_context.search(&match_all), + search_context.search(&match_all()), vec![ ScoredPointOffset { score: 120.0, @@ -294,7 +295,7 @@ mod test_mod { .unwrap(); assert_eq!( - round_scores(search_context.search(&match_all)), + round_scores(search_context.search(&match_all())), vec![ ScoredPointOffset { score: 90.0, @@ -337,7 +338,7 @@ mod test_mod { .unwrap(); assert_eq!( - round_scores(search_context.search(&match_all)), + round_scores(search_context.search(&match_all())), vec![ ScoredPointOffset { score: 90.0, diff --git a/lib/sparse/src/search_scratch.rs b/lib/sparse/src/search_scratch.rs index adaee45e45..80fc9a3f9c 100644 --- a/lib/sparse/src/search_scratch.rs +++ b/lib/sparse/src/search_scratch.rs @@ -2,11 +2,11 @@ use std::fmt::Debug; use blink_alloc::Blink; use common::defaults::POOL_KEEP_LIMIT; -use common::types::ScoreType; +use common::types::{ScoreType, ScoredPointOffset}; use parking_lot::Mutex; pub struct SearchScratchPool { - pool: Mutex>, + pool: Mutex>, } impl Debug for SearchScratchPool { @@ -25,10 +25,11 @@ impl SearchScratchPool { /// Take a single [`SearchScratch`] from the pool. pub fn get(&self) -> SearchScratch<'_> { - let (scores, arena) = self.pool.lock().pop().unwrap_or_default(); + let (scores, candidates, arena) = self.pool.lock().pop().unwrap_or_default(); SearchScratch { pool: self, scores, + candidates, arena, } } @@ -40,6 +41,8 @@ pub struct SearchScratch<'a> { pool: &'a SearchScratchPool, /// Used for batched scoring. pub(crate) scores: SearchScratchScores, + /// Used for batched filtering. + pub(crate) candidates: SearchScratchCandidates, /// Used to own/store posting list bytes while reading them from the file. pub(crate) arena: Blink, } @@ -54,20 +57,23 @@ impl SearchScratch<'_> { } type SearchScratchScores = Vec; +type SearchScratchCandidates = Vec; impl Drop for SearchScratch<'_> { fn drop(&mut self) { let SearchScratch { pool: SearchScratchPool { pool }, scores, + candidates, arena, } = self; let mut pool = pool.lock(); if pool.len() < *POOL_KEEP_LIMIT { let scores = std::mem::take(scores); + let candidates = std::mem::take(candidates); let mut arena = std::mem::take(arena); arena.reset(); - pool.push((scores, arena)); + pool.push((scores, candidates, arena)); } } }