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

ConditionChecker::check_batched: use in sparse

Browse Source
This commit is contained in:
xzfc
2026-06-30 11:31:37 +00:00
parent d234d02b7e
commit c19266e559
8 changed files with 110 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
lib/segment/src/index/sparse_index/sparse_vector_index/read_view/search.rs
View File

@@ -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(&not_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(

6
lib/sparse/benches/search.rs
View File

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

100
lib/sparse/src/index/search_context.rs
View File

@@ -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<T: PostingListIter> {
/// 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<IndexedPostingListIterator<T>>,
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<ScoreType>,
/// Buffer for batched filtering from [`SearchScratch`].
candidates: &'a mut Vec<ScoredPointOffset>,
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<F: Fn(PointOffsetType) -> bool>(
fn advance_batch<C: ConditionChecker>(
&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<F: Fn(PointOffsetType) -> bool>(&mut self, filter_condition: &F) {
fn process_last_posting_list<C: ConditionChecker>(&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<F: Fn(PointOffsetType) -> bool>(
&mut self,
filter_condition: &F,
) -> Vec<ScoredPointOffset> {
pub fn search<C: ConditionChecker>(&mut self, checker: &C) -> Vec<ScoredPointOffset> {
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<C: ConditionChecker>(
candidates: &mut Vec<ScoredPointOffset>,
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();
}

5
lib/sparse/src/index/tests/common.rs
View File

@@ -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::<W, _>(&mut rnd_gen, count, density, vocab1, vocab2)
}
pub fn match_all(_p: PointOffsetType) -> bool {
true
pub fn match_all() -> ConstantConditionChecker<std::convert::Infallible> {
ConstantConditionChecker::MATCH_ALL
}

2
lib/sparse/src/index/tests/hw_counter_test.rs
View File

@@ -26,7 +26,7 @@ fn do_search<I: InvertedIndex>(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);

2
lib/sparse/src/index/tests/indexed_vs_plain_test.rs
View File

@@ -24,7 +24,7 @@ fn query<I: InvertedIndex>(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);

19
lib/sparse/src/index/tests/search_context_tests.rs
View File

@@ -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<std::convert::Infallible> {
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,

14
lib/sparse/src/search_scratch.rs
View File

@@ -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<Vec<(SearchScratchScores, Blink)>>,
pool: Mutex<Vec<(SearchScratchScores, SearchScratchCandidates, Blink)>>,
}
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<ScoreType>;
type SearchScratchCandidates = Vec<ScoredPointOffset>;
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));
}
}
}