use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use atomic_refcell::AtomicRefCell;
use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::types::{DeferredBehavior, PointOffsetType, ScoredPointOffset, TelemetryDetail};
use parking_lot::Mutex;
use crate::sparse::common::types::DimId;
use super::hnsw_index::point_scorer::BatchFilteredSearcher;
use crate::segment::common::BYTES_IN_KB;
use crate::segment::common::operation_error::OperationResult;
use crate::segment::common::operation_time_statistics::{
OperationDurationStatistics, OperationDurationsAggregator, ScopeDurationMeasurer,
};
use crate::segment::data_types::query_context::VectorQueryContext;
use crate::segment::data_types::vectors::{QueryVector, VectorRef};
use crate::segment::id_tracker::{IdTrackerEnum, IdTrackerRead};
use crate::segment::index::struct_payload_index::StructPayloadIndex;
use crate::segment::index::vector_index_search_common::{
get_oversampled_top, is_quantized_search, postprocess_search_result,
};
use crate::segment::index::{PayloadIndexRead, VectorIndex, VectorIndexRead};
use crate::segment::telemetry::VectorIndexSearchesTelemetry;
use crate::segment::types::{Filter, SearchParams};
use crate::segment::vector_storage::quantized::quantized_vectors::QuantizedVectors;
use crate::segment::vector_storage::{VectorStorage, VectorStorageEnum, VectorStorageRead};
#[derive(Debug)]
pub struct PlainVectorIndex {
id_tracker: Arc<AtomicRefCell<IdTrackerEnum>>,
vector_storage: Arc<AtomicRefCell<VectorStorageEnum>>,
quantized_vectors: Arc<AtomicRefCell<Option<QuantizedVectors>>>,
payload_index: Arc<AtomicRefCell<StructPayloadIndex>>,
filtered_searches_telemetry: Arc<Mutex<OperationDurationsAggregator>>,
unfiltered_searches_telemetry: Arc<Mutex<OperationDurationsAggregator>>,
}
impl PlainVectorIndex {
pub fn new(
id_tracker: Arc<AtomicRefCell<IdTrackerEnum>>,
vector_storage: Arc<AtomicRefCell<VectorStorageEnum>>,
quantized_vectors: Arc<AtomicRefCell<Option<QuantizedVectors>>>,
payload_index: Arc<AtomicRefCell<StructPayloadIndex>>,
) -> PlainVectorIndex {
PlainVectorIndex {
id_tracker,
vector_storage,
quantized_vectors,
payload_index,
filtered_searches_telemetry: OperationDurationsAggregator::new(),
unfiltered_searches_telemetry: OperationDurationsAggregator::new(),
}
}
pub fn is_small_enough_for_unindexed_search(
&self,
search_optimized_threshold_kb: usize,
filter: Option<&Filter>,
hw_counter: &HardwareCounterCell,
) -> OperationResult<bool> {
let vector_storage = self.vector_storage.borrow();
let available_vector_count = vector_storage.available_vector_count();
if available_vector_count == 0 {
return Ok(true);
}
let vector_size_bytes =
vector_storage.size_of_available_vectors_in_bytes() / available_vector_count;
let indexing_threshold_bytes = search_optimized_threshold_kb * BYTES_IN_KB;
if let Some(payload_filter) = filter {
let cardinality = self
.payload_index
.borrow()
.with_view(|v| v.estimate_cardinality(payload_filter, hw_counter))?;
let scan_size = vector_size_bytes.saturating_mul(cardinality.max);
Ok(scan_size <= indexing_threshold_bytes)
} else {
let vector_storage_size = vector_size_bytes.saturating_mul(available_vector_count);
Ok(vector_storage_size <= indexing_threshold_bytes)
}
}
}
impl VectorIndexRead for PlainVectorIndex {
fn search(
&self,
query_vectors: &[&QueryVector],
filter: Option<&Filter>,
top: usize,
params: Option<&SearchParams>,
query_context: &VectorQueryContext,
) -> OperationResult<Vec<Vec<ScoredPointOffset>>> {
let is_indexed_only = params.is_some_and(|p| p.indexed_only);
if is_indexed_only
&& !self.is_small_enough_for_unindexed_search(
query_context.search_optimized_threshold_kb(),
filter,
&query_context.hardware_counter(),
)?
{
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();
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 filtered_ids_vec = self
.payload_index
.borrow()
.with_view(|v| v.query_points(filter, &hw_counter, &is_stopped))?;
batch_searcher.peek_top_iter(filtered_ids_vec.iter().copied(), &is_stopped)?
}
None => {
let iter = id_tracker.point_mappings().filter_deferred_and_deleted(
batch_searcher.iter_not_deleted(),
DeferredBehavior::Exclude,
);
batch_searcher.peek_top_iter(iter, &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 {
VectorIndexSearchesTelemetry {
index_name: None,
unfiltered_plain: self
.unfiltered_searches_telemetry
.lock()
.get_statistics(detail),
filtered_plain: self
.filtered_searches_telemetry
.lock()
.get_statistics(detail),
unfiltered_hnsw: OperationDurationStatistics::default(),
filtered_small_cardinality: OperationDurationStatistics::default(),
filtered_large_cardinality: OperationDurationStatistics::default(),
filtered_exact: OperationDurationStatistics::default(),
filtered_sparse: Default::default(),
unfiltered_exact: OperationDurationStatistics::default(),
unfiltered_sparse: OperationDurationStatistics::default(),
}
}
fn indexed_vector_count(&self) -> usize {
0
}
fn size_of_searchable_vectors_in_bytes(&self) -> usize {
self.vector_storage
.borrow()
.size_of_available_vectors_in_bytes()
}
fn fill_idf_statistics(
&self,
_idf: &mut HashMap<DimId, usize>,
_hw_counter: &HardwareCounterCell,
) -> OperationResult<()> {
Ok(())
}
fn is_index(&self) -> bool {
false
}
}
impl VectorIndex for PlainVectorIndex {
fn files(&self) -> Vec<PathBuf> {
vec![]
}
fn update_vector(
&mut self,
id: PointOffsetType,
vector: Option<VectorRef>,
hw_counter: &HardwareCounterCell,
) -> OperationResult<()> {
let mut vector_storage = self.vector_storage.borrow_mut();
if let Some(vector) = vector {
vector_storage.insert_vector(id, vector, hw_counter)?;
let mut quantized_vectors = self.quantized_vectors.borrow_mut();
if let Some(quantized_vectors) = quantized_vectors.as_mut() {
quantized_vectors.upsert_vector(id, vector, hw_counter)?;
}
} else {
if id as usize >= vector_storage.total_vector_count() {
debug_assert!(id as usize == vector_storage.total_vector_count());
let default_vector = vector_storage.default_vector();
vector_storage.insert_vector(id, VectorRef::from(&default_vector), hw_counter)?;
}
vector_storage.delete_vector(id)?;
}
Ok(())
}
}