use crate::common::bitvec::BitSlice;
use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::types::ScoredPointOffset;
use itertools::Itertools;
use crate::segment::common::operation_error::OperationResult;
use crate::segment::data_types::vectors::QueryVector;
use crate::segment::index::hnsw_index::point_scorer::FilteredScorer;
use crate::segment::types::{
SearchParams, default_quantization_ignore_value, default_quantization_oversampling_value,
};
use crate::segment::vector_storage::quantized::quantized_vectors::QuantizedVectorsRead;
use crate::segment::vector_storage::{RawScorerBuilder, VectorStorageRead};
pub fn is_quantized_search<Q: QuantizedVectorsRead>(
quantized_storage: Option<&Q>,
params: Option<&SearchParams>,
) -> bool {
let ignore_quantization = params
.and_then(|p| p.quantization)
.map(|q| q.ignore)
.unwrap_or(default_quantization_ignore_value());
let exact = params.is_some_and(|p| p.exact);
quantized_storage.is_some() && !ignore_quantization && !exact
}
pub fn get_oversampled_top<Q: QuantizedVectorsRead>(
quantized_storage: Option<&Q>,
params: Option<&SearchParams>,
top: usize,
) -> usize {
let quantization_enabled = is_quantized_search(quantized_storage, params);
let oversampling_value = params
.and_then(|p| p.quantization)
.map(|q| q.oversampling)
.unwrap_or(default_quantization_oversampling_value());
match oversampling_value {
Some(oversampling) if quantization_enabled && oversampling > 1.0 => {
(oversampling * top as f64) as usize
}
_ => top,
}
}
#[allow(clippy::too_many_arguments)]
pub fn postprocess_search_result<V, Q>(
mut search_result: Vec<ScoredPointOffset>,
point_deleted: &BitSlice,
vector_storage: &V,
quantized_vectors: Option<&Q>,
vector: &QueryVector,
params: Option<&SearchParams>,
top: usize,
hardware_counter: HardwareCounterCell,
) -> OperationResult<Vec<ScoredPointOffset>>
where
V: VectorStorageRead + RawScorerBuilder,
Q: QuantizedVectorsRead,
{
let quantization_enabled = is_quantized_search(quantized_vectors, params);
let default_rescoring = quantized_vectors
.as_ref()
.map(|q| q.default_rescoring())
.unwrap_or(false);
let rescore = quantization_enabled
&& params
.and_then(|p| p.quantization)
.and_then(|q| q.rescore)
.unwrap_or(default_rescoring);
if rescore {
let mut scorer = FilteredScorer::new(
vector.to_owned(),
vector_storage,
None::<&Q>,
None,
point_deleted,
hardware_counter,
)?;
search_result = scorer
.score_points(&mut search_result.iter().map(|x| x.idx).collect_vec(), 0)
.collect();
search_result.sort_unstable();
search_result.reverse();
}
search_result.truncate(top);
Ok(search_result)
}