diff --git a/lib/sparse/src/index/search_context.rs b/lib/sparse/src/index/search_context.rs index da3a85c018..310a8e0883 100644 --- a/lib/sparse/src/index/search_context.rs +++ b/lib/sparse/src/index/search_context.rs @@ -1,9 +1,9 @@ -use std::cmp::Ordering; +use std::cmp::{max, min, Ordering}; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering::Relaxed; use common::fixed_length_priority_queue::FixedLengthPriorityQueue; -use common::types::{PointOffsetType, ScoredPointOffset}; +use common::types::{PointOffsetType, ScoreType, ScoredPointOffset}; use crate::common::sparse_vector::SparseVector; use crate::common::types::{DimId, DimWeight}; @@ -17,14 +17,18 @@ pub struct IndexedPostingListIterator<'a> { query_weight: DimWeight, } +/// Making this larger makes the search faster but uses more memory +const ADVANCE_BATCH_SIZE: usize = 1_000; + pub struct SearchContext<'a> { postings_iterators: Vec>, query: SparseVector, top: usize, is_stopped: &'a AtomicBool, result_queue: FixedLengthPriorityQueue, // keep the largest elements and peek smallest - min_record_id: Option, // min record ids across all posting lists - contains_empty_posting: bool, // whether the posting list iterators contain exhausted posting lists + min_record_id: Option, // min_record_id ids across all posting lists + max_record_id: u32, // max_record_id ids across all posting lists + batch_scores: Vec, // scores for the current batch use_pruning: bool, } @@ -36,16 +40,32 @@ impl<'a> SearchContext<'a> { is_stopped: &'a AtomicBool, ) -> SearchContext<'a> { let mut postings_iterators = Vec::new(); - + // track min and max record ids across all posting lists + let mut max_record_id = 0; + let mut min_record_id = u32::MAX; + // iterate over query indices for (query_weight_offset, id) in query.indices.iter().enumerate() { if let Some(posting_list_iterator) = inverted_index.get(id) { - let query_index = *id; - let query_weight = query.values[query_weight_offset]; - postings_iterators.push(IndexedPostingListIterator { - posting_list_iterator, - query_index, - query_weight, - }); + let posting_elements = posting_list_iterator.elements; + if !posting_elements.is_empty() { + // check if new min + let min_record_id_posting = posting_elements[0].record_id; + min_record_id = min(min_record_id, min_record_id_posting); + + // check if new max + let max_record_id_posting = posting_elements.last().unwrap().record_id; + max_record_id = max(max_record_id, max_record_id_posting); + + // capture query info + let query_index = *id; + let query_weight = query.values[query_weight_offset]; + + postings_iterators.push(IndexedPostingListIterator { + posting_list_iterator, + query_index, + query_weight, + }); + } } } let result_queue = FixedLengthPriorityQueue::new(top); @@ -53,10 +73,9 @@ impl<'a> SearchContext<'a> { // The max contribution per posting list that we calculate is not made to compute the max value of two negative numbers. // This is a limitation of the current pruning implementation. let use_pruning = query.values.iter().all(|v| *v >= 0.0); - // find min record id across all posting lists - let min_record_id = Self::next_min_id(&postings_iterators); - // no empty posting lists at the beginning - let contains_empty_posting = false; + // TODO pool this Vec to reuse memory across searches + let batch_scores = Vec::with_capacity(ADVANCE_BATCH_SIZE); + let min_record_id = Some(min_record_id); SearchContext { postings_iterators, query, @@ -64,7 +83,8 @@ impl<'a> SearchContext<'a> { is_stopped, result_queue, min_record_id, - contains_empty_posting, + max_record_id, + batch_scores, use_pruning, } } @@ -106,72 +126,49 @@ impl<'a> SearchContext<'a> { queue.into_vec() } - /// Advance posting lists iterators and return the next candidate by increasing ids. - /// - /// Example - /// - /// postings_iterators: - /// - /// 1, 30, 34, 60, 230 - /// 10, 30, 35, 51, 230 - /// 2, 21, 34, 60, 200 - /// 2, 30, 34, 60, 230 - /// - /// Next: - /// - /// a, 30, 34, 60, 230 - /// 10, 30, 35, 51, 230 - /// 2, 21, 34, 60, 200 - /// 2, 30, 34, 60, 230 - /// - /// Next: - /// - /// a, 30, 34, 60, 230 - /// 10, 30, 35, 51, 230 - /// b, 21, 34, 60, 200 - /// b, 30, 34, 60, 230 - /// - /// Next: - /// - /// a, 30, 34, 60, 230 - /// c, 30, 35, 51, 230 - /// b, 21, 34, 60, 200 - /// b, 30, 34, 60, 230 - fn advance(&mut self) -> Option { - // Get current min record id from all posting list iterators - let current_min_record_id = self.min_record_id?; - let mut score = 0.0; - let mut found = false; - // Iterate to advance matching posting iterators - for posting_iterator in self.postings_iterators.iter_mut() { - if let Some(element) = posting_iterator.posting_list_iterator.peek() { - // accumulate score for the current record id - if element.record_id == current_min_record_id { - found = true; - score += element.weight * posting_iterator.query_weight; - // advance posting list iterator to next element - posting_iterator.posting_list_iterator.advance(); + /// Advance posting lists iterators in a batch fashion. + fn advance_batch bool>( + &mut self, + batch_start_id: PointOffsetType, + batch_last_id: PointOffsetType, + filter_condition: &F, + ) { + for posting in self.postings_iterators.iter_mut() { + for element in posting.posting_list_iterator.remaining_elements() { + let element_id = element.record_id; + if element_id > batch_last_id { + // reaching end of the batch + break; } - } else { - // mark search context for cleanup if any posting list iterator is exhausted - self.contains_empty_posting = true; + let element_score = element.weight * posting.query_weight; + // update score for id + let local_id = (element_id - batch_start_id) as usize; + self.batch_scores[local_id] += element_score; + } + // advance posting to the batch last id + posting.posting_list_iterator.skip_to(batch_last_id + 1); + } + + // publish only the non-zero scores above the current min + let min_score_to_beat = if self.result_queue.len() == self.top { + self.result_queue.top().map(|e| e.score) + } else { + None + }; + for (local_index, &score) in self.batch_scores.iter().enumerate() { + if score != 0.0 && Some(score) > min_score_to_beat { + 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 { + score, + idx: real_id, + }; + self.result_queue.push(score_point_offset); } } - - // update min record id for next iteration - if found { - // assume the next min record id is the current one + 1 - self.min_record_id = self.min_record_id.map(|min_id| min_id + 1); - } else { - // no match found, compute next min record id from all posting list iterators - self.min_record_id = Self::next_min_id(&self.postings_iterators); - return self.advance(); - } - - Some(ScoredPointOffset { - score, - idx: current_min_record_id, - }) } /// Compute scores for the last posting list quickly @@ -247,30 +244,49 @@ impl<'a> SearchContext<'a> { return Vec::new(); } let mut best_min_score = f32::MIN; - while let Some(candidate) = self.advance() { - // check filter condition - if !filter_condition(candidate.idx) { - continue; - } - // check for cancellation + loop { + // check for cancellation (atomic amortized by batch) if self.is_stopped.load(Relaxed) { break; } - // push candidate to result queue - self.result_queue.push(candidate); + + // prepare next iterator of batched ids + let start_batch_id = match self.min_record_id { + Some(min_id) => min_id, + None => break, // all posting lists exhausted + }; + + // compute batch range of contiguous ids for the next batch + let last_batch_id = min( + start_batch_id + ADVANCE_BATCH_SIZE as u32, + self.max_record_id, + ); + let batch_len = last_batch_id - start_batch_id + 1; + + // init batch scores + self.batch_scores.clear(); + self.batch_scores.resize(batch_len as usize, 0.0); + + // advance and score posting lists iterators + self.advance_batch(start_batch_id, last_batch_id, filter_condition); // remove empty posting lists if necessary - if self.contains_empty_posting { - self.postings_iterators.retain(|posting_iterator| { - posting_iterator.posting_list_iterator.len_to_end() != 0 - }); - // reset flag - self.contains_empty_posting = false; - // if only one posting list left, we can score it quickly - if self.postings_iterators.len() == 1 { - self.process_last_posting_list(filter_condition); - break; - } + self.postings_iterators.retain(|posting_iterator| { + posting_iterator.posting_list_iterator.len_to_end() != 0 + }); + + // update min_record_id + self.min_record_id = Self::next_min_id(&self.postings_iterators); + + // check if all posting lists are exhausted + if self.postings_iterators.is_empty() { + break; + } + + // if only one posting list left, we can score it quickly + if self.postings_iterators.len() == 1 { + self.process_last_posting_list(filter_condition); + break; } // we potentially have enough results to prune low performing posting lists @@ -289,8 +305,7 @@ impl<'a> SearchContext<'a> { // prune posting list that cannot possibly contribute to the top results let pruned = self.prune_longest_posting_list(new_min_score); if pruned { - // recompute new min record id for next iteration - // the pruned posting list is always at the head and with the lowest record_id + // update min_record_id self.min_record_id = Self::next_min_id(&self.postings_iterators); } } @@ -367,10 +382,7 @@ impl<'a> SearchContext<'a> { #[cfg(test)] mod tests { - use std::collections::HashSet; - - use rand::rngs::StdRng; - use rand::{Rng, SeedableRng}; + use rand::Rng; use super::*; use crate::common::sparse_vector_fixture::random_sparse_vector; @@ -398,62 +410,6 @@ mod tests { true } - fn _advance_test(inverted_index: &impl InvertedIndex) { - let is_stopped = AtomicBool::new(false); - let mut search_context = SearchContext::new( - SparseVector { - indices: vec![1, 2, 3], - values: vec![1.0, 1.0, 1.0], - }, - 10, - inverted_index, - &is_stopped, - ); - - assert_eq!( - search_context.advance(), - Some(ScoredPointOffset { - score: 30.0, - idx: 1 - }) - ); - assert_eq!( - search_context.advance(), - Some(ScoredPointOffset { - score: 60.0, - idx: 2 - }) - ); - assert_eq!( - search_context.advance(), - Some(ScoredPointOffset { - score: 90.0, - idx: 3 - }) - ); - } - - #[test] - fn advance_test() { - let inverted_index_ram = InvertedIndexBuilder::new() - .add(1, PostingList::from(vec![(1, 10.0), (2, 20.0), (3, 30.0)])) - .add(2, PostingList::from(vec![(1, 10.0), (2, 20.0), (3, 30.0)])) - .add(3, PostingList::from(vec![(1, 10.0), (2, 20.0), (3, 30.0)])) - .build(); - - // test with ram index - _advance_test(&inverted_index_ram); - - // test with mmap index - let tmp_dir_path = tempfile::Builder::new() - .prefix("test_index_dir") - .tempdir() - .unwrap(); - let inverted_index_mmap = - InvertedIndexMmap::convert_and_save(&inverted_index_ram, &tmp_dir_path).unwrap(); - _advance_test(&inverted_index_mmap); - } - fn _search_test(inverted_index: &impl InvertedIndex) { let is_stopped = AtomicBool::new(false); let mut search_context = SearchContext::new( @@ -678,124 +634,6 @@ mod tests { _search_with_hot_key_test(&inverted_index_mmap); } - fn _prune_test(inverted_index: &impl InvertedIndex) { - let is_stopped = AtomicBool::new(false); - let mut search_context = SearchContext::new( - SparseVector { - indices: vec![1, 2, 3], - values: vec![1.0, 1.0, 1.0], - }, - 3, - inverted_index, - &is_stopped, - ); - - // initial state - assert_eq!( - search_context.postings_iterators[0] - .posting_list_iterator - .len_to_end(), - 9 - ); - assert_eq!( - search_context.advance(), - Some(ScoredPointOffset { - score: 30.0, - idx: 1 - }) - ); - assert_eq!( - search_context.postings_iterators[0] - .posting_list_iterator - .len_to_end(), - 8 - ); - assert!(!search_context.prune_longest_posting_list(30.0)); - assert_eq!( - search_context.postings_iterators[0] - .posting_list_iterator - .len_to_end(), - 8 - ); - - assert_eq!( - search_context.advance(), - Some(ScoredPointOffset { - score: 60.0, - idx: 2 - }) - ); - assert_eq!( - search_context.postings_iterators[0] - .posting_list_iterator - .len_to_end(), - 7 - ); - assert!(!search_context.prune_longest_posting_list(30.0)); - assert_eq!( - search_context.postings_iterators[0] - .posting_list_iterator - .len_to_end(), - 7 - ); - - assert_eq!( - search_context.advance(), - Some(ScoredPointOffset { - score: 90.0, - idx: 3 - }) - ); - // pruning can take place - assert_eq!( - search_context.postings_iterators[0] - .posting_list_iterator - .len_to_end(), - 6 - ); - assert!(search_context.prune_longest_posting_list(30.0)); - assert_eq!( - search_context.postings_iterators[0] - .posting_list_iterator - .len_to_end(), - 0 - ); - } - - #[test] - fn prune_test() { - let inverted_index_ram = InvertedIndexBuilder::new() - .add( - 1, - PostingList::from(vec![ - (1, 10.0), - (2, 20.0), - (3, 30.0), - (4, 1.0), - (5, 2.0), - (6, 3.0), - (7, 4.0), - (8, 5.0), - (9, 6.0), - ]), - ) - .add(2, PostingList::from(vec![(1, 10.0), (2, 20.0), (3, 30.0)])) - .add(3, PostingList::from(vec![(1, 10.0), (2, 20.0), (3, 30.0)])) - .build(); - - // test with ram index - _prune_test(&inverted_index_ram); - - // test with mmap index - let tmp_dir_path = tempfile::Builder::new() - .prefix("test_index_dir") - .tempdir() - .unwrap(); - let inverted_index_mmap = - InvertedIndexMmap::convert_and_save(&inverted_index_ram, &tmp_dir_path).unwrap(); - _prune_test(&inverted_index_mmap); - } - #[test] fn pruning_single_to_end_test() { let inverted_index_ram = InvertedIndexBuilder::new() @@ -900,71 +738,6 @@ mod tests { ); } - #[test] - fn pruning_does_not_skip_negative_score_test() { - let inverted_index_ram = InvertedIndexBuilder::new() - .add( - 1, - PostingList::from(vec![(1, 1.0), (2, 2.0), (3, 3.0), (4, 1.0), (5, -40.0)]), - ) - .build(); - - let is_stopped = AtomicBool::new(false); - let mut search_context = SearchContext::new( - SparseVector { - indices: vec![1, 2, 3], - values: vec![-1.0, 1.0, 1.0], - }, - 2, - &inverted_index_ram, - &is_stopped, - ); - - // pruning is automatically deactivated because the query vector contains negative values - assert!(!search_context.use_pruning); - assert_eq!( - search_context.search(&match_all), - vec![ - ScoredPointOffset { - score: 40.0, - idx: 5 - }, - ScoredPointOffset { - score: -1.0, - idx: 1 - }, - ] - ); - - // try again with pruning to show the problem - let mut search_context = SearchContext::new( - SparseVector { - indices: vec![1, 2, 3], - values: vec![-1.0, 1.0, 1.0], - }, - 2, - &inverted_index_ram, - &is_stopped, - ); - search_context.use_pruning = true; - assert!(search_context.use_pruning); - - // the last value has been pruned although it could have contributed a high score -1 * -40 = 40 - assert_eq!( - search_context.search(&match_all), - vec![ - ScoredPointOffset { - score: -1.0, - idx: 1 - }, - ScoredPointOffset { - score: -2.0, - idx: 2 - } - ] - ); - } - /// Generates a random inverted index with `num_vectors` vectors fn random_inverted_index( rnd_gen: &mut R, @@ -980,72 +753,6 @@ mod tests { inverted_index_ram } - #[test] - fn next_min_partial_scan_test() { - let num_vectors = 100; - let max_sparse_dimension = 25; - let mut rnd = StdRng::seed_from_u64(42); - let is_stopped = AtomicBool::new(false); - let inverted_index_ram = random_inverted_index(&mut rnd, num_vectors, max_sparse_dimension); - let mut search_context = SearchContext::new( - SparseVector { - indices: vec![1, 2, 3], - values: vec![1.0, 1.0, 1.0], - }, - 3, - &inverted_index_ram, - &is_stopped, - ); - - let mut all_next_min_observed = HashSet::new(); - - while let Some(next_min) = - SearchContext::next_min_id(search_context.postings_iterators.as_slice()) - { - all_next_min_observed.insert(next_min); - let next_candidate_id = search_context.advance().map(|s| s.idx); - assert_eq!(next_candidate_id, Some(next_min)); - } - - // Not all vectors are observed because only the indices of the query vector are explored. - assert!(all_next_min_observed.len() < num_vectors as usize); - } - - #[test] - fn next_min_full_scan_test() { - let num_vectors = 100; - let max_sparse_dimension = 25; - let mut rnd = StdRng::seed_from_u64(42); - let is_stopped = AtomicBool::new(false); - let inverted_index_ram = random_inverted_index(&mut rnd, num_vectors, max_sparse_dimension); - let mut search_context = SearchContext::new( - SparseVector { - indices: (1..=max_sparse_dimension as u32).collect(), - values: vec![1.0; max_sparse_dimension], - }, - 3, - &inverted_index_ram, - &is_stopped, - ); - - // initial state - let min = SearchContext::next_min_id(search_context.postings_iterators.as_slice()); - // no side effect - assert_eq!(min, Some(1)); - assert_eq!(min, Some(1)); - - // Complete scan over all vectors because the query vector contains all dimensions in the index. - for i in 1..num_vectors { - let before_min = - SearchContext::next_min_id(search_context.postings_iterators.as_slice()); - assert_eq!(before_min, Some(i)); - let next = search_context.advance().map(|s| s.idx); - assert_eq!(next, Some(i)); - let new_min = SearchContext::next_min_id(search_context.postings_iterators.as_slice()); - assert_eq!(new_min, Some(i + 1)); - } - } - #[test] fn promote_longest_test() { let is_stopped = AtomicBool::new(false);