upstream-archive/qdrant
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant.git synced 2026-08-05 09:31:05 -05:00
Code Issues Packages Projects Releases Wiki Activity

Batched scans for dense vectors (#7514)

Browse Source 
* Batched iteration for plain `HNSWIndex` searches
* Batched iteration for `PlainVectorIndex` search
This commit is contained in:
Ivan Boldyrev
2025-11-20 17:34:02 +07:00
committed by GitHub
parent 113e8aab54
commit 428f1e76d4
8 changed files with 504 additions and 211 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
lib/segment/benches/scorer_mmap.rs
View File

@@ -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);

48
lib/segment/benches/vector_search.rs
View File

@@ -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);

105
lib/segment/src/index/hnsw_index/hnsw.rs
View File

@@ -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<Item = PointOffsetType>,
query_vectors: &[&QueryVector],
points: impl Iterator<Item = PointOffsetType>,
top: usize,
params: Option<&SearchParams>,
vector_query_context: &VectorQueryContext,
) -> OperationResult<Vec<ScoredPointOffset>> {
) -> OperationResult<Vec<Vec<ScoredPointOffset>>> {
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<Vec<ScoredPointOffset>> {
self.search_plain_iterator(
vector,
&mut filtered_points.iter().copied(),
) -> OperationResult<Vec<Vec<ScoredPointOffset>>> {
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<Vec<ScoredPointOffset>> {
) -> OperationResult<Vec<Vec<ScoredPointOffset>>> {
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<Vec<Vec<ScoredPointOffset>>> {
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<Box<dyn FilterContext + 'a>>,
) -> OperationResult<BatchFilteredSearcher<'a>> {
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);

171
lib/segment/src/index/hnsw_index/point_scorer.rs
View File

@@ -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<dyn RawScorer + 'a>,
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<dyn RawScorer + 'a>,
pq: FixedLengthPriorityQueue<ScoredPointOffset>,
}
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<BoxCow<'a, dyn FilterContext + 'a>>,
top: usize,
point_deleted: &'a BitSlice,
hardware_counter: HardwareCounterCell,
) -> OperationResult<Self> {
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::<Result<_, _>>()?;
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<Vec<Vec<ScoredPointOffset>>> {
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<Item = PointOffsetType>,
is_stopped: &AtomicBool,
) -> CancellableResult<Vec<Vec<ScoredPointOffset>>> {
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)
}
}

135
lib/segment/src/index/plain_vector_index.rs
View File

@@ -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 {

73
lib/segment/src/vector_storage/dense/memmap_dense_vector_storage.rs
View File

@@ -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);

73
lib/segment/src/vector_storage/tests/test_appendable_dense_vector_storage.rs
View File

@@ -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) => {

61
lib/segment/src/vector_storage/tests/test_appendable_multi_dense_vector_storage.rs
View File

@@ -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<f32>> = 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<f32>> = 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<f32>> = 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);