From 428f1e76d43324ef2b186a55d48e861dbfdebfc4 Mon Sep 17 00:00:00 2001 From: Ivan Boldyrev Date: Thu, 20 Nov 2025 17:34:02 +0700 Subject: [PATCH] Batched scans for dense vectors (#7514) * Batched iteration for plain `HNSWIndex` searches * Batched iteration for `PlainVectorIndex` search --- lib/segment/benches/scorer_mmap.rs | 49 ++++- lib/segment/benches/vector_search.rs | 48 ++++- lib/segment/src/index/hnsw_index/hnsw.rs | 105 ++++++----- .../src/index/hnsw_index/point_scorer.rs | 171 ++++++++++++++++++ lib/segment/src/index/plain_vector_index.rs | 135 +++++--------- .../dense/memmap_dense_vector_storage.rs | 73 +++++--- .../test_appendable_dense_vector_storage.rs | 73 +++++--- ...t_appendable_multi_dense_vector_storage.rs | 61 +++++-- 8 files changed, 504 insertions(+), 211 deletions(-) diff --git a/lib/segment/benches/scorer_mmap.rs b/lib/segment/benches/scorer_mmap.rs index 5755d89fa3..a26d4740cd 100644 --- a/lib/segment/benches/scorer_mmap.rs +++ b/lib/segment/benches/scorer_mmap.rs @@ -10,7 +10,7 @@ use segment::data_types::named_vectors::CowVector; use segment::data_types::vectors::{DenseVector, QueryVector}; use segment::fixtures::payload_context_fixture::FixtureIdTracker; use segment::id_tracker::IdTrackerSS; -use segment::index::hnsw_index::point_scorer::FilteredScorer; +use segment::index::hnsw_index::point_scorer::BatchFilteredSearcher; use segment::types::Distance; use segment::vector_storage::dense::memmap_dense_vector_storage::open_memmap_vector_storage; use segment::vector_storage::{DEFAULT_STOPPED, VectorStorage, VectorStorageEnum}; @@ -63,12 +63,13 @@ fn benchmark_scorer_mmap(c: &mut Criterion) { b.iter_batched( || QueryVector::from(random_vector(DIM)), |vector| { - FilteredScorer::new_for_test( - vector, + BatchFilteredSearcher::new_for_test( + &[vector], &storage, borrowed_id_tracker.deleted_point_bitslice(), + 10, ) - .peek_top_all(10, &DEFAULT_STOPPED) + .peek_top_all(&DEFAULT_STOPPED) .unwrap() }, BatchSize::SmallInput, @@ -76,18 +77,54 @@ fn benchmark_scorer_mmap(c: &mut Criterion) { }); } +// Batched search gives performance benefit only when memory is contended. +// For a single-threaded criterion run, it only shows that batching penalty is relatively small. +// We might run a thread pool explicitly, though. +fn benchmark_scorer_mmap_4(c: &mut Criterion) { + let dir = Builder::new().prefix("storage_dir").tempdir().unwrap(); + + let dist = Distance::Dot; + let (storage, id_tracker) = init_mmap_vector_storage(dir.path(), DIM, NUM_VECTORS, dist); + let borrowed_id_tracker = id_tracker.borrow(); + + let mut group = c.benchmark_group("storage-score-all"); + + group.bench_function("storage batched vector scoring, 4 vectors batch", |b| { + b.iter_batched( + || { + [ + QueryVector::from(random_vector(DIM)), + QueryVector::from(random_vector(DIM)), + QueryVector::from(random_vector(DIM)), + QueryVector::from(random_vector(DIM)), + ] + }, + |vecs| { + BatchFilteredSearcher::new_for_test( + &vecs, + &storage, + borrowed_id_tracker.deleted_point_bitslice(), + 10, + ) + .peek_top_all(&DEFAULT_STOPPED) + .unwrap() + }, + BatchSize::SmallInput, + ) + }); +} #[cfg(not(target_os = "windows"))] criterion_group! { name = benches; config = Criterion::default().with_profiler(prof::FlamegraphProfiler::new(100)); - targets = benchmark_scorer_mmap + targets = benchmark_scorer_mmap, benchmark_scorer_mmap_4, } #[cfg(target_os = "windows")] criterion_group! { name = benches; config = Criterion::default(); - targets = benchmark_scorer_mmap, + targets = benchmark_scorer_mmap, benchmark_scorer_mmap_4, } criterion_main!(benches); diff --git a/lib/segment/benches/vector_search.rs b/lib/segment/benches/vector_search.rs index f7aa0fd6da..0fdc7d1f62 100644 --- a/lib/segment/benches/vector_search.rs +++ b/lib/segment/benches/vector_search.rs @@ -12,7 +12,7 @@ use segment::common::rocksdb_wrapper::{DB_VECTOR_CF, open_db}; use segment::data_types::vectors::{DenseVector, VectorInternal, VectorRef}; use segment::fixtures::payload_context_fixture::FixtureIdTracker; use segment::id_tracker::IdTrackerSS; -use segment::index::hnsw_index::point_scorer::FilteredScorer; +use segment::index::hnsw_index::point_scorer::{BatchFilteredSearcher, FilteredScorer}; use segment::types::{Distance, VectorStorageDatatype}; use segment::vector_storage::dense::simple_dense_vector_storage::open_simple_dense_vector_storage; use segment::vector_storage::{DEFAULT_STOPPED, VectorStorage, VectorStorageEnum}; @@ -72,12 +72,45 @@ fn benchmark_naive(c: &mut Criterion) { b.iter(|| { let vector = random_vector(DIM); let vector = vector.as_slice().into(); - FilteredScorer::new_for_test( - vector, + BatchFilteredSearcher::new_for_test( + &[vector], &storage, borrowed_id_tracker.deleted_point_bitslice(), + 10, ) - .peek_top_all(10, &DEFAULT_STOPPED) + .peek_top_all(&DEFAULT_STOPPED) + .unwrap(); + }) + }); +} + +// Batched search gives performance benefit only when memory is contended. +// For a single-threaded criterion run, it only shows that batching penalty is relatively small. +// We might run a thread pool explicitly, though. +fn benchmark_naive_4(c: &mut Criterion) { + let dir = Builder::new().prefix("storage_dir").tempdir().unwrap(); + + let dist = Distance::Dot; + let (storage, id_tracker) = init_vector_storage(dir.path(), DIM, NUM_VECTORS, dist); + let borrowed_id_tracker = id_tracker.borrow(); + + let mut group = c.benchmark_group("storage-score-all"); + + group.bench_function("storage vector search, 4 vectors batch", |b| { + b.iter(|| { + let vectors = [ + random_vector(DIM).into(), + random_vector(DIM).into(), + random_vector(DIM).into(), + random_vector(DIM).into(), + ]; + BatchFilteredSearcher::new_for_test( + &vectors, + &storage, + borrowed_id_tracker.deleted_point_bitslice(), + 10, + ) + .peek_top_all(&DEFAULT_STOPPED) .unwrap(); }) }); @@ -111,5 +144,10 @@ fn random_access_benchmark(c: &mut Criterion) { eprintln!("total_score = {total_score:?}"); } -criterion_group!(benches, benchmark_naive, random_access_benchmark); +criterion_group!( + benches, + benchmark_naive, + benchmark_naive_4, + random_access_benchmark +); criterion_main!(benches); diff --git a/lib/segment/src/index/hnsw_index/hnsw.rs b/lib/segment/src/index/hnsw_index/hnsw.rs index aaa555c966..b452dede1d 100644 --- a/lib/segment/src/index/hnsw_index/hnsw.rs +++ b/lib/segment/src/index/hnsw_index/hnsw.rs @@ -28,6 +28,7 @@ use super::gpu::gpu_devices_manager::LockedGpuDevice; use super::gpu::gpu_insert_context::GpuInsertContext; #[cfg(feature = "gpu")] use super::gpu::gpu_vector_storage::GpuVectorStorage; +use super::point_scorer::BatchFilteredSearcher; use crate::common::BYTES_IN_KB; use crate::common::operation_error::{OperationError, OperationResult, check_process_stopped}; use crate::common::operation_time_statistics::{ @@ -1123,14 +1124,14 @@ impl HNSWIndex { .collect() } - fn search_plain_iterator( + fn search_plain_iterator_batched( &self, - vector: &QueryVector, - points: &mut dyn Iterator, + query_vectors: &[&QueryVector], + points: impl Iterator, top: usize, params: Option<&SearchParams>, vector_query_context: &VectorQueryContext, - ) -> OperationResult> { + ) -> OperationResult>> { let id_tracker = self.id_tracker.borrow(); let vector_storage = self.vector_storage.borrow(); let quantized_vectors = self.quantized_vectors.borrow(); @@ -1140,60 +1141,61 @@ impl HNSWIndex { .unwrap_or_else(|| id_tracker.deleted_point_bitslice()); let is_stopped = vector_query_context.is_stopped(); + let oversampled_top = get_oversampled_top(quantized_vectors.as_ref(), params, top); - let points_scorer = Self::construct_search_scorer( - vector, + let batch_filtered_searcher = Self::construct_batch_searcher( + query_vectors, &vector_storage, quantized_vectors.as_ref(), + oversampled_top, deleted_points, params, vector_query_context.hardware_counter(), None, )?; - let oversampled_top = get_oversampled_top(quantized_vectors.as_ref(), params, top); - - let search_result = points_scorer.peek_top_iter(points, oversampled_top, &is_stopped)?; - - let res = postprocess_search_result( - search_result, - id_tracker.deleted_point_bitslice(), - &vector_storage, - quantized_vectors.as_ref(), - vector, - params, - top, - vector_query_context.hardware_counter(), - )?; - Ok(res) + let mut search_results = batch_filtered_searcher.peek_top_iter(points, &is_stopped)?; + for (search_result, query_vector) in search_results.iter_mut().zip(query_vectors) { + *search_result = postprocess_search_result( + std::mem::take(search_result), + id_tracker.deleted_point_bitslice(), + &vector_storage, + quantized_vectors.as_ref(), + query_vector, + params, + top, + vector_query_context.hardware_counter(), + )?; + } + Ok(search_results) } - fn search_plain( + fn search_plain_batched( &self, - vector: &QueryVector, + vectors: &[&QueryVector], filtered_points: &[PointOffsetType], top: usize, params: Option<&SearchParams>, vector_query_context: &VectorQueryContext, - ) -> OperationResult> { - self.search_plain_iterator( - vector, - &mut filtered_points.iter().copied(), + ) -> OperationResult>> { + self.search_plain_iterator_batched( + vectors, + filtered_points.iter().copied(), top, params, vector_query_context, ) } - fn search_plain_unfiltered( + fn search_plain_unfiltered_batched( &self, - vector: &QueryVector, + vectors: &[&QueryVector], top: usize, params: Option<&SearchParams>, vector_query_context: &VectorQueryContext, - ) -> OperationResult> { + ) -> OperationResult>> { let id_tracker = self.id_tracker.borrow(); - let mut ids_iterator = id_tracker.iter_internal(); - self.search_plain_iterator(vector, &mut ids_iterator, top, params, vector_query_context) + let ids_iterator = id_tracker.iter_internal(); + self.search_plain_iterator_batched(vectors, ids_iterator, top, params, vector_query_context) } fn search_vectors_plain( @@ -1205,18 +1207,12 @@ impl HNSWIndex { vector_query_context: &VectorQueryContext, ) -> OperationResult>> { let payload_index = self.payload_index.borrow(); - // share filtered points for all query vectors let filtered_points = payload_index.query_points( filter, &vector_query_context.hardware_counter(), &vector_query_context.is_stopped(), ); - vectors - .iter() - .map(|vector| { - self.search_plain(vector, &filtered_points, top, params, vector_query_context) - }) - .collect() + self.search_plain_batched(vectors, &filtered_points, top, params, vector_query_context) } fn discovery_search_with_graph( @@ -1276,6 +1272,29 @@ impl HNSWIndex { ) } + #[allow(clippy::too_many_arguments)] + fn construct_batch_searcher<'a>( + vectors: &[&QueryVector], + vector_storage: &'a VectorStorageEnum, + quantized_storage: Option<&'a QuantizedVectors>, + top: usize, + deleted_points: &'a BitSlice, + params: Option<&SearchParams>, + hardware_counter: HardwareCounterCell, + filter_context: Option>, + ) -> OperationResult> { + let quantization_enabled = is_quantized_search(quantized_storage, params); + BatchFilteredSearcher::new( + vectors, + vector_storage, + quantization_enabled.then_some(quantized_storage).flatten(), + filter_context.map(BoxCow::Owned), + top, + deleted_points, + hardware_counter, + ) + } + /// Read underlying data from disk into disk cache. pub fn populate(&self) -> OperationResult<()> { self.graph.populate() @@ -1340,13 +1359,7 @@ impl VectorIndex for HNSWIndex { }); let params_ref = if exact { exact_params.as_ref() } else { params }; - - vectors - .iter() - .map(|&vector| { - self.search_plain_unfiltered(vector, top, params_ref, query_context) - }) - .collect() + self.search_plain_unfiltered_batched(vectors, top, params_ref, query_context) } else { let _timer = ScopeDurationMeasurer::new(&self.searches_telemetry.unfiltered_hnsw); diff --git a/lib/segment/src/index/hnsw_index/point_scorer.rs b/lib/segment/src/index/hnsw_index/point_scorer.rs index 13ae96b980..cce63a9053 100644 --- a/lib/segment/src/index/hnsw_index/point_scorer.rs +++ b/lib/segment/src/index/hnsw_index/point_scorer.rs @@ -5,6 +5,7 @@ use common::counter::hardware_counter::HardwareCounterCell; use common::cow::BoxCow; use common::fixed_length_priority_queue::FixedLengthPriorityQueue; use common::types::{PointOffsetType, ScoreType, ScoredPointOffset}; +use smallvec::SmallVec; use crate::common::operation_error::{CancellableResult, OperationResult, check_process_stopped}; use crate::data_types::vectors::QueryVector; @@ -37,6 +38,19 @@ use crate::vector_storage::{ /// - Scoring logic /// - Complex queries /// ``` +/// +/// The `BatchFilteredSearcher` contains an array of `RawScorer`s, a common filter and certain parameters. +/// +/// ```plaintext +/// BatchFilteredSearcher RawScorer +/// ┌─────────────────┐ ┌───────────────┐ +/// │ [RawScorer] ◄───┼──┤ QueryScorer ◄─┼── (ditto) +/// │ │ └───────────────┘ +/// │ FilterContext │ +/// │ │ +/// │ top │ +/// └─────────────────┘ +/// ``` pub struct FilteredScorer<'a> { raw_scorer: Box, filters: ScorerFilters<'a>, @@ -309,3 +323,160 @@ impl<'a> FilteredScorer<'a> { Ok(pq.into_sorted_vec()) } } + +// We keep each scorer with its queue to reduce allocations and improve data locality. +struct BatchSearch<'a> { + raw_scorer: Box, + pq: FixedLengthPriorityQueue, +} + +pub struct BatchFilteredSearcher<'a> { + scorer_batch: SmallVec<[BatchSearch<'a>; 1]>, + filters: ScorerFilters<'a>, + top: usize, +} + +impl<'a> BatchFilteredSearcher<'a> { + /// Create a new batch filtered searcher. + /// + /// If present, `quantized_vectors` will be used for scoring, otherwise `vectors` will be used. + pub fn new( + queries: &[&QueryVector], + vectors: &'a VectorStorageEnum, + quantized_vectors: Option<&'a QuantizedVectors>, + filter_context: Option>, + top: usize, + point_deleted: &'a BitSlice, + hardware_counter: HardwareCounterCell, + ) -> OperationResult { + let scorer_batch = queries + .iter() + .map(|&query| { + let query = query.to_owned(); + let hardware_counter = hardware_counter.fork(); + let raw_scorer = match quantized_vectors { + Some(quantized_vectors) => { + quantized_vectors.raw_scorer(query, hardware_counter) + } + None => new_raw_scorer(query, vectors, hardware_counter), + }; + let pq = FixedLengthPriorityQueue::new(top); + raw_scorer.map(|raw_scorer| BatchSearch { raw_scorer, pq }) + }) + .collect::>()?; + let filters = ScorerFilters { + filter_context, + point_deleted, + vec_deleted: vectors.deleted_vector_bitslice(), + }; + Ok(Self { + scorer_batch, + filters, + top, + }) + } + + /// Create a new batched filtered searcher for testing purposes. + /// + /// # Panics + /// + /// Panics if [`new_raw_scorer`] fails. + #[cfg(feature = "testing")] + pub fn new_for_test( + vectors: &[QueryVector], + vector_storage: &'a VectorStorageEnum, + point_deleted: &'a BitSlice, + top: usize, + ) -> Self { + let scorer_batch = vectors + .iter() + .map(|vector| { + let raw_scorer = new_raw_scorer( + vector.to_owned(), + vector_storage, + HardwareCounterCell::new(), + ) + .unwrap(); + BatchSearch { + raw_scorer, + pq: FixedLengthPriorityQueue::new(top), + } + }) + .collect(); + Self { + scorer_batch, + filters: ScorerFilters { + filter_context: None, + point_deleted, + vec_deleted: vector_storage.deleted_vector_bitslice(), + }, + top, + } + } + + pub fn peek_top_all( + self, + is_stopped: &AtomicBool, + ) -> CancellableResult>> { + let iter = self + .filters + .point_deleted + .iter_zeros() + .map(|p| p as PointOffsetType); + self.peek_top_iter(iter, is_stopped) + } + + pub fn peek_top_iter( + mut self, + mut points: impl Iterator, + is_stopped: &AtomicBool, + ) -> CancellableResult>> { + if self.top == 0 { + return Ok(vec![vec![]; self.scorer_batch.len()]); + } + + // Reuse the same buffer for all chunks, to avoid reallocation + let mut chunk = [0; VECTOR_READ_BATCH_SIZE]; + let mut scores_buffer = [0.0; VECTOR_READ_BATCH_SIZE]; + + loop { + check_process_stopped(is_stopped)?; + + let mut chunk_size = 0; + for point_id in &mut points { + check_process_stopped(is_stopped)?; + if !self.filters.check_vector(point_id) { + continue; + } + chunk[chunk_size] = point_id; + chunk_size += 1; + if chunk_size == VECTOR_READ_BATCH_SIZE { + break; + } + } + + if chunk_size == 0 { + break; + } + + // Switching the loops improves batching performance, but slightly degrades single-query performance. + for BatchSearch { raw_scorer, pq } in &mut self.scorer_batch { + raw_scorer.score_points(&chunk[..chunk_size], &mut scores_buffer[..chunk_size]); + + for i in 0..chunk_size { + pq.push(ScoredPointOffset { + idx: chunk[i], + score: scores_buffer[i], + }); + } + } + } + + let results = self + .scorer_batch + .into_iter() + .map(|BatchSearch { pq, .. }| pq.into_sorted_vec()) + .collect(); + Ok(results) + } +} diff --git a/lib/segment/src/index/plain_vector_index.rs b/lib/segment/src/index/plain_vector_index.rs index 47bc0565b6..4e9c18c383 100644 --- a/lib/segment/src/index/plain_vector_index.rs +++ b/lib/segment/src/index/plain_vector_index.rs @@ -6,7 +6,7 @@ use common::counter::hardware_counter::HardwareCounterCell; use common::types::{PointOffsetType, ScoredPointOffset, TelemetryDetail}; use parking_lot::Mutex; -use super::hnsw_index::point_scorer::FilteredScorer; +use super::hnsw_index::point_scorer::BatchFilteredSearcher; use crate::common::BYTES_IN_KB; use crate::common::operation_error::OperationResult; use crate::common::operation_time_statistics::{ @@ -83,7 +83,7 @@ impl PlainVectorIndex { impl VectorIndex for PlainVectorIndex { fn search( &self, - vectors: &[&QueryVector], + query_vectors: &[&QueryVector], filter: Option<&Filter>, top: usize, params: Option<&SearchParams>, @@ -97,101 +97,64 @@ impl VectorIndex for PlainVectorIndex { &query_context.hardware_counter(), ) { - return Ok(vec![vec![]; vectors.len()]); + return Ok(vec![vec![]; query_vectors.len()]); + } + if top == 0 { + return Ok(vec![vec![]; query_vectors.len()]); } let is_stopped = query_context.is_stopped(); let hw_counter = query_context.hardware_counter(); - match filter { + let _timer = ScopeDurationMeasurer::new(if filter.is_some() { + &self.filtered_searches_telemetry + } else { + &self.unfiltered_searches_telemetry + }); + let vector_storage = self.vector_storage.borrow(); + let quantized_storage = self.quantized_vectors.borrow(); + let id_tracker = self.id_tracker.borrow(); + let deleted_points = query_context + .deleted_points() + .unwrap_or_else(|| id_tracker.deleted_point_bitslice()); + let quantization_enabled = is_quantized_search(quantized_storage.as_ref(), params); + let quantized_vectors = quantization_enabled + .then_some(quantized_storage.as_ref()) + .flatten(); + let oversampled_top = get_oversampled_top(quantized_storage.as_ref(), params, top); + let batch_searcher = BatchFilteredSearcher::new( + query_vectors, + &vector_storage, + quantized_vectors, + None, + oversampled_top, + deleted_points, + query_context.hardware_counter(), + )?; + + let mut search_results = match filter { Some(filter) => { - let _timer = ScopeDurationMeasurer::new(&self.filtered_searches_telemetry); - let id_tracker = self.id_tracker.borrow(); let payload_index = self.payload_index.borrow(); - let vector_storage = self.vector_storage.borrow(); - let quantized_storage = self.quantized_vectors.borrow(); let filtered_ids_vec = payload_index.query_points(filter, &hw_counter, &is_stopped); - let deleted_points = query_context - .deleted_points() - .unwrap_or_else(|| id_tracker.deleted_point_bitslice()); - vectors - .iter() - .map(|&vector| { - let quantization_enabled = - is_quantized_search(quantized_storage.as_ref(), params); - let scorer = FilteredScorer::new( - vector.to_owned(), - &vector_storage, - quantization_enabled - .then_some(quantized_storage.as_ref()) - .flatten(), - None, - deleted_points, - query_context.hardware_counter(), - )?; - let oversampled_top = - get_oversampled_top(quantized_storage.as_ref(), params, top); - let search_result = scorer.peek_top_iter( - &mut filtered_ids_vec.iter().copied(), - oversampled_top, - &is_stopped, - )?; - let res = postprocess_search_result( - search_result, - id_tracker.deleted_point_bitslice(), - &vector_storage, - quantized_storage.as_ref(), - vector, - params, - top, - query_context.hardware_counter(), - )?; - Ok(res) - }) - .collect() - } - None => { - let _timer = ScopeDurationMeasurer::new(&self.unfiltered_searches_telemetry); - let vector_storage = self.vector_storage.borrow(); - let quantized_storage = self.quantized_vectors.borrow(); - let id_tracker = self.id_tracker.borrow(); - let deleted_points = query_context - .deleted_points() - .unwrap_or_else(|| id_tracker.deleted_point_bitslice()); - vectors - .iter() - .map(|&vector| { - let quantization_enabled = - is_quantized_search(quantized_storage.as_ref(), params); - let scorer = FilteredScorer::new( - vector.to_owned(), - &vector_storage, - quantization_enabled - .then_some(quantized_storage.as_ref()) - .flatten(), - None, - deleted_points, - query_context.hardware_counter(), - )?; - let oversampled_top = - get_oversampled_top(quantized_storage.as_ref(), params, top); - let search_result = scorer.peek_top_all(oversampled_top, &is_stopped)?; - let res = postprocess_search_result( - search_result, - id_tracker.deleted_point_bitslice(), - &vector_storage, - quantized_storage.as_ref(), - vector, - params, - top, - query_context.hardware_counter(), - )?; - Ok(res) - }) - .collect() + batch_searcher.peek_top_iter(&mut filtered_ids_vec.iter().copied(), &is_stopped)? } + None => batch_searcher.peek_top_all(&is_stopped)?, + }; + + for (search_result, query_vector) in search_results.iter_mut().zip(query_vectors) { + *search_result = postprocess_search_result( + std::mem::take(search_result), + deleted_points, + &vector_storage, + quantized_storage.as_ref(), + query_vector, + params, + top, + query_context.hardware_counter(), + )?; } + Ok(search_results) } fn get_telemetry_data(&self, detail: TelemetryDetail) -> VectorIndexSearchesTelemetry { diff --git a/lib/segment/src/vector_storage/dense/memmap_dense_vector_storage.rs b/lib/segment/src/vector_storage/dense/memmap_dense_vector_storage.rs index 731a97161a..835b94d59c 100644 --- a/lib/segment/src/vector_storage/dense/memmap_dense_vector_storage.rs +++ b/lib/segment/src/vector_storage/dense/memmap_dense_vector_storage.rs @@ -323,7 +323,7 @@ mod tests { use crate::data_types::vectors::{DenseVector, QueryVector}; use crate::fixtures::payload_context_fixture::FixtureIdTracker; use crate::id_tracker::id_tracker_base::IdTracker; - use crate::index::hnsw_index::point_scorer::FilteredScorer; + use crate::index::hnsw_index::point_scorer::{BatchFilteredSearcher, FilteredScorer}; use crate::types::{PointIdType, QuantizationConfig, ScalarQuantizationConfig}; use crate::vector_storage::dense::volatile_dense_vector_storage::new_volatile_dense_vector_storage; use crate::vector_storage::quantized::quantized_vectors::{ @@ -413,19 +413,32 @@ mod tests { assert_eq!(stored_ids, [0, 1, 3, 4]); - let scorer = FilteredScorer::new_for_test( - points[2].as_slice().into(), + let searcher = BatchFilteredSearcher::new_for_test( + &[points[2].as_slice().into()], &storage, borrowed_id_tracker.deleted_point_bitslice(), + 2, ); - let res = scorer.peek_top_all(2, &DEFAULT_STOPPED).unwrap(); + let res = searcher + .peek_top_all(&DEFAULT_STOPPED) + .unwrap() + .pop() + .unwrap(); assert_eq!(res.len(), 2); assert_ne!(res[0].idx, 2); - let res = scorer - .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), 2, &DEFAULT_STOPPED) + let searcher = BatchFilteredSearcher::new_for_test( + &[points[2].as_slice().into()], + &storage, + borrowed_id_tracker.deleted_point_bitslice(), + 2, + ); + let res = searcher + .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), &DEFAULT_STOPPED) + .unwrap() + .pop() .unwrap(); assert_eq!(res.len(), 2); @@ -487,20 +500,22 @@ mod tests { let vector = vec![0.0, 1.0, 1.1, 1.0]; let query = vector.as_slice().into(); - let scorer = FilteredScorer::new_for_test( - query, + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), &storage, borrowed_id_tracker.deleted_point_bitslice(), + 5, ); - let closest = scorer - .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), 5, &DEFAULT_STOPPED) + let closest = searcher + .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), &DEFAULT_STOPPED) + .unwrap() + .pop() .unwrap(); assert_eq!(closest.len(), 3, "must have 3 vectors, 2 are deleted"); assert_eq!(closest[0].idx, 0); assert_eq!(closest[1].idx, 1); assert_eq!(closest[2].idx, 4); - drop(scorer); // Delete 1, redelete 2 storage.delete_vector(1 as PointOffsetType).unwrap(); @@ -514,18 +529,20 @@ mod tests { let vector = vec![1.0, 0.0, 0.0, 0.0]; let query = vector.as_slice().into(); - let scorer = FilteredScorer::new_for_test( - query, + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), &storage, borrowed_id_tracker.deleted_point_bitslice(), + 5, ); - let closest = scorer - .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), 5, &DEFAULT_STOPPED) + let closest = searcher + .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), &DEFAULT_STOPPED) + .unwrap() + .pop() .unwrap(); assert_eq!(closest.len(), 2, "must have 2 vectors, 3 are deleted"); assert_eq!(closest[0].idx, 4); assert_eq!(closest[1].idx, 0); - drop(scorer); // Delete all storage.delete_vector(0 as PointOffsetType).unwrap(); @@ -538,12 +555,17 @@ mod tests { let vector = vec![1.0, 0.0, 0.0, 0.0]; let query = vector.as_slice().into(); - let scorer = FilteredScorer::new_for_test( - query, + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), &storage, borrowed_id_tracker.deleted_point_bitslice(), + 5, ); - let closest = scorer.peek_top_all(5, &DEFAULT_STOPPED).unwrap(); + let closest = searcher + .peek_top_all(&DEFAULT_STOPPED) + .unwrap() + .pop() + .unwrap(); assert!(closest.is_empty(), "must have no results, all deleted"); } @@ -595,17 +617,18 @@ mod tests { let vector = vec![0.0, 1.0, 1.1, 1.0]; let query = vector.as_slice().into(); - let scorer = FilteredScorer::new_for_test( - query, + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), &storage, borrowed_id_tracker.deleted_point_bitslice(), + 5, ); - let closest = scorer - .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), 5, &DEFAULT_STOPPED) + let closest = searcher + .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), &DEFAULT_STOPPED) + .unwrap() + .pop() .unwrap(); - drop(scorer); - assert_eq!(closest.len(), 3, "must have 3 vectors, 2 are deleted"); assert_eq!(closest[0].idx, 0); assert_eq!(closest[1].idx, 1); diff --git a/lib/segment/src/vector_storage/tests/test_appendable_dense_vector_storage.rs b/lib/segment/src/vector_storage/tests/test_appendable_dense_vector_storage.rs index 3ac73d91b6..66d7bf3796 100644 --- a/lib/segment/src/vector_storage/tests/test_appendable_dense_vector_storage.rs +++ b/lib/segment/src/vector_storage/tests/test_appendable_dense_vector_storage.rs @@ -11,7 +11,7 @@ use crate::common::rocksdb_wrapper::{DB_VECTOR_CF, open_db}; use crate::data_types::vectors::QueryVector; use crate::fixtures::payload_context_fixture::FixtureIdTracker; use crate::id_tracker::IdTrackerSS; -use crate::index::hnsw_index::point_scorer::FilteredScorer; +use crate::index::hnsw_index::point_scorer::{BatchFilteredSearcher, FilteredScorer}; use crate::types::{Distance, PointIdType, QuantizationConfig, ScalarQuantizationConfig}; use crate::vector_storage::dense::appendable_dense_vector_storage::open_appendable_memmap_vector_storage; #[cfg(feature = "rocksdb")] @@ -62,16 +62,21 @@ fn do_test_delete_points(storage: &mut VectorStorageEnum) { let vector = vec![0.0, 1.0, 1.1, 1.0]; let query = vector.as_slice().into(); - let scorer = - FilteredScorer::new_for_test(query, storage, borrowed_id_tracker.deleted_point_bitslice()); - let closest = scorer - .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), 5, &DEFAULT_STOPPED) + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), + storage, + borrowed_id_tracker.deleted_point_bitslice(), + 5, + ); + let closest = searcher + .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), &DEFAULT_STOPPED) + .unwrap() + .pop() .unwrap(); assert_eq!(closest.len(), 3, "must have 3 vectors, 2 are deleted"); assert_eq!(closest[0].idx, 0); assert_eq!(closest[1].idx, 1); assert_eq!(closest[2].idx, 4); - drop(scorer); // Delete 1, redelete 2 storage.delete_vector(1 as PointOffsetType).unwrap(); @@ -84,15 +89,20 @@ fn do_test_delete_points(storage: &mut VectorStorageEnum) { let vector = vec![1.0, 0.0, 0.0, 0.0]; let query = vector.as_slice().into(); - let scorer = - FilteredScorer::new_for_test(query, storage, borrowed_id_tracker.deleted_point_bitslice()); - let closest = scorer - .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), 5, &DEFAULT_STOPPED) + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), + storage, + borrowed_id_tracker.deleted_point_bitslice(), + 5, + ); + let closest = searcher + .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), &DEFAULT_STOPPED) + .unwrap() + .pop() .unwrap(); assert_eq!(closest.len(), 2, "must have 2 vectors, 3 are deleted"); assert_eq!(closest[0].idx, 4); assert_eq!(closest[1].idx, 0); - drop(scorer); // Delete all storage.delete_vector(0 as PointOffsetType).unwrap(); @@ -105,9 +115,17 @@ fn do_test_delete_points(storage: &mut VectorStorageEnum) { let vector = vec![1.0, 0.0, 0.0, 0.0]; let query = vector.as_slice().into(); - let scorer = - FilteredScorer::new_for_test(query, storage, borrowed_id_tracker.deleted_point_bitslice()); - let closest = scorer.peek_top_all(5, &DEFAULT_STOPPED).unwrap(); + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), + storage, + borrowed_id_tracker.deleted_point_bitslice(), + 5, + ); + let closest = searcher + .peek_top_all(&DEFAULT_STOPPED) + .unwrap() + .pop() + .unwrap(); assert!(closest.is_empty(), "must have no results, all deleted"); } @@ -156,12 +174,17 @@ fn do_test_update_from_delete_points(storage: &mut VectorStorageEnum) { let vector = vec![0.0, 1.0, 1.1, 1.0]; let query = vector.as_slice().into(); - let scorer = - FilteredScorer::new_for_test(query, storage, borrowed_id_tracker.deleted_point_bitslice()); - let closest = scorer - .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), 5, &DEFAULT_STOPPED) + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), + storage, + borrowed_id_tracker.deleted_point_bitslice(), + 5, + ); + let closest = searcher + .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), &DEFAULT_STOPPED) + .unwrap() + .pop() .unwrap(); - drop(scorer); assert_eq!(closest.len(), 3, "must have 3 vectors, 2 are deleted"); assert_eq!(closest[0].idx, 0); assert_eq!(closest[1].idx, 1); @@ -200,15 +223,17 @@ fn do_test_score_points(storage: &mut VectorStorageEnum) { let query: QueryVector = [0.0, 1.0, 1.1, 1.0].into(); - let scorer = FilteredScorer::new_for_test( - query.clone(), + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), storage, borrowed_id_tracker.deleted_point_bitslice(), + 2, ); - let closest = scorer - .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), 2, &DEFAULT_STOPPED) + let closest = searcher + .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), &DEFAULT_STOPPED) + .unwrap() + .pop() .unwrap(); - drop(scorer); let top_idx = match closest.first() { Some(scored_point) => { diff --git a/lib/segment/src/vector_storage/tests/test_appendable_multi_dense_vector_storage.rs b/lib/segment/src/vector_storage/tests/test_appendable_multi_dense_vector_storage.rs index 7acd8aaf3d..0c53b9e306 100644 --- a/lib/segment/src/vector_storage/tests/test_appendable_multi_dense_vector_storage.rs +++ b/lib/segment/src/vector_storage/tests/test_appendable_multi_dense_vector_storage.rs @@ -13,7 +13,7 @@ use crate::data_types::vectors::{ }; use crate::fixtures::payload_context_fixture::FixtureIdTracker; use crate::id_tracker::IdTrackerSS; -use crate::index::hnsw_index::point_scorer::FilteredScorer; +use crate::index::hnsw_index::point_scorer::BatchFilteredSearcher; use crate::types::{Distance, MultiVectorConfig}; use crate::vector_storage::common::CHUNK_SIZE; use crate::vector_storage::multi_dense::appendable_mmap_multi_dense_vector_storage::open_appendable_memmap_multi_vector_storage_full; @@ -141,16 +141,21 @@ fn do_test_delete_points(vector_dim: usize, vec_count: usize, storage: &mut Vect ); let vector: Vec> = vec![vec![2.0; vector_dim]]; let query = QueryVector::Nearest(vector.try_into().unwrap()); - let scorer = - FilteredScorer::new_for_test(query, storage, borrowed_id_tracker.deleted_point_bitslice()); - let closest = scorer - .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), 5, &DEFAULT_STOPPED) + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), + storage, + borrowed_id_tracker.deleted_point_bitslice(), + 5, + ); + let closest = searcher + .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), &DEFAULT_STOPPED) + .unwrap() + .pop() .unwrap(); assert_eq!(closest.len(), 3, "must have 3 vectors, 2 are deleted"); assert_eq!(closest[0].idx, 4); assert_eq!(closest[1].idx, 1); assert_eq!(closest[2].idx, 0); - drop(scorer); // Delete 1, redelete 2 storage.delete_vector(1 as PointOffsetType).unwrap(); @@ -163,15 +168,20 @@ fn do_test_delete_points(vector_dim: usize, vec_count: usize, storage: &mut Vect let vector: Vec> = vec![vec![1.0; vector_dim]]; let query = QueryVector::Nearest(vector.try_into().unwrap()); - let scorer = - FilteredScorer::new_for_test(query, storage, borrowed_id_tracker.deleted_point_bitslice()); - let closest = scorer - .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), 5, &DEFAULT_STOPPED) + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), + storage, + borrowed_id_tracker.deleted_point_bitslice(), + 5, + ); + let closest = searcher + .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), &DEFAULT_STOPPED) + .unwrap() + .pop() .unwrap(); assert_eq!(closest.len(), 2, "must have 2 vectors, 3 are deleted"); assert_eq!(closest[0].idx, 4); assert_eq!(closest[1].idx, 0); - drop(scorer); // Delete all storage.delete_vector(0 as PointOffsetType).unwrap(); @@ -184,9 +194,17 @@ fn do_test_delete_points(vector_dim: usize, vec_count: usize, storage: &mut Vect let vector: Vec> = vec![vec![1.0; vector_dim]]; let query = QueryVector::Nearest(vector.try_into().unwrap()); - let scorer = - FilteredScorer::new_for_test(query, storage, borrowed_id_tracker.deleted_point_bitslice()); - let closest = scorer.peek_top_all(5, &DEFAULT_STOPPED).unwrap(); + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), + storage, + borrowed_id_tracker.deleted_point_bitslice(), + 5, + ); + let closest = searcher + .peek_top_all(&DEFAULT_STOPPED) + .unwrap() + .pop() + .unwrap(); assert!(closest.is_empty(), "must have no results, all deleted"); } @@ -240,12 +258,17 @@ fn do_test_update_from_delete_points( let query = QueryVector::Nearest(vector.try_into().unwrap()); - let scorer = - FilteredScorer::new_for_test(query, storage, borrowed_id_tracker.deleted_point_bitslice()); - let closest = scorer - .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), 5, &DEFAULT_STOPPED) + let searcher = BatchFilteredSearcher::new_for_test( + std::slice::from_ref(&query), + storage, + borrowed_id_tracker.deleted_point_bitslice(), + 5, + ); + let closest = searcher + .peek_top_iter(&mut [0, 1, 2, 3, 4].iter().cloned(), &DEFAULT_STOPPED) + .unwrap() + .pop() .unwrap(); - drop(scorer); assert_eq!(closest.len(), 3, "must have 3 vectors, 2 are deleted"); assert_eq!(closest[0].idx, 4); assert_eq!(closest[1].idx, 1);