diff --git a/lib/segment/src/index/mod.rs b/lib/segment/src/index/mod.rs index 5f742d3746..d6bc579ae4 100644 --- a/lib/segment/src/index/mod.rs +++ b/lib/segment/src/index/mod.rs @@ -1,15 +1,15 @@ pub mod field_index; pub mod hnsw_index; -mod vector_index_base; -mod payload_index_base; mod key_encoding; mod payload_config; +mod payload_index_base; pub mod plain_payload_index; pub mod query_estimator; mod query_optimization; mod sample_estimation; mod struct_filter_context; pub mod struct_payload_index; +mod vector_index_base; mod visited_pool; -pub use vector_index_base::*; pub use payload_index_base::*; +pub use vector_index_base::*; diff --git a/lib/segment/src/index/vector_index_base.rs b/lib/segment/src/index/vector_index_base.rs index a00feac143..55c419bf83 100644 --- a/lib/segment/src/index/vector_index_base.rs +++ b/lib/segment/src/index/vector_index_base.rs @@ -1,12 +1,13 @@ use std::path::PathBuf; use std::sync::atomic::AtomicBool; +use super::hnsw_index::graph_links::{GraphLinksMmap, GraphLinksRam}; +use super::hnsw_index::hnsw::HNSWIndex; +use super::plain_payload_index::PlainIndex; use crate::data_types::vectors::VectorElementType; use crate::entry::entry_point::OperationResult; use crate::telemetry::VectorIndexSearchesTelemetry; -use crate::types::{ - Filter, SearchParams, -}; +use crate::types::{Filter, SearchParams}; use crate::vector_storage::ScoredPointOffset; /// Trait for vector searching @@ -28,4 +29,48 @@ pub trait VectorIndex { fn files(&self) -> Vec; } -pub type VectorIndexSS = dyn VectorIndex + Sync + Send; +pub enum VectorIndexEnum { + Plain(PlainIndex), + HnswRam(HNSWIndex), + HnswMmap(HNSWIndex), +} + +impl VectorIndex for VectorIndexEnum { + fn search( + &self, + vectors: &[&[VectorElementType]], + filter: Option<&Filter>, + top: usize, + params: Option<&SearchParams>, + ) -> Vec> { + match self { + VectorIndexEnum::Plain(index) => index.search(vectors, filter, top, params), + VectorIndexEnum::HnswRam(index) => index.search(vectors, filter, top, params), + VectorIndexEnum::HnswMmap(index) => index.search(vectors, filter, top, params), + } + } + + fn build_index(&mut self, stopped: &AtomicBool) -> OperationResult<()> { + match self { + VectorIndexEnum::Plain(index) => index.build_index(stopped), + VectorIndexEnum::HnswRam(index) => index.build_index(stopped), + VectorIndexEnum::HnswMmap(index) => index.build_index(stopped), + } + } + + fn get_telemetry_data(&self) -> VectorIndexSearchesTelemetry { + match self { + VectorIndexEnum::Plain(index) => index.get_telemetry_data(), + VectorIndexEnum::HnswRam(index) => index.get_telemetry_data(), + VectorIndexEnum::HnswMmap(index) => index.get_telemetry_data(), + } + } + + fn files(&self) -> Vec { + match self { + VectorIndexEnum::Plain(index) => index.files(), + VectorIndexEnum::HnswRam(index) => index.files(), + VectorIndexEnum::HnswMmap(index) => index.files(), + } + } +} diff --git a/lib/segment/src/segment.rs b/lib/segment/src/segment.rs index 4702f311dd..fa62c630f9 100644 --- a/lib/segment/src/segment.rs +++ b/lib/segment/src/segment.rs @@ -23,7 +23,7 @@ use crate::entry::entry_point::{ use crate::id_tracker::IdTrackerSS; use crate::index::field_index::CardinalityEstimation; use crate::index::struct_payload_index::StructPayloadIndex; -use crate::index::{PayloadIndex, VectorIndexSS}; +use crate::index::{PayloadIndex, VectorIndex, VectorIndexEnum}; use crate::spaces::tools::peek_top_smallest_iterable; use crate::telemetry::SegmentTelemetry; use crate::types::{ @@ -81,7 +81,7 @@ pub struct Segment { } pub struct VectorData { - pub vector_index: Arc>, + pub vector_index: Arc>, pub vector_storage: Arc>, } diff --git a/lib/segment/src/segment_constructor/segment_builder.rs b/lib/segment/src/segment_constructor/segment_builder.rs index 9eb2ca08c1..107bfdca98 100644 --- a/lib/segment/src/segment_constructor/segment_builder.rs +++ b/lib/segment/src/segment_constructor/segment_builder.rs @@ -8,7 +8,7 @@ use crate::common::error_logging::LogError; use crate::entry::entry_point::{ check_process_stopped, OperationError, OperationResult, SegmentEntry, }; -use crate::index::PayloadIndex; +use crate::index::{PayloadIndex, VectorIndex}; use crate::segment::Segment; use crate::segment_constructor::{build_segment, load_segment}; use crate::types::{PayloadFieldSchema, PayloadKeyType, SegmentConfig}; diff --git a/lib/segment/src/segment_constructor/segment_constructor_base.rs b/lib/segment/src/segment_constructor/segment_constructor_base.rs index 91ac5fdf21..82c558585c 100644 --- a/lib/segment/src/segment_constructor/segment_constructor_base.rs +++ b/lib/segment/src/segment_constructor/segment_constructor_base.rs @@ -20,7 +20,7 @@ use crate::index::hnsw_index::graph_links::{GraphLinksMmap, GraphLinksRam}; use crate::index::hnsw_index::hnsw::HNSWIndex; use crate::index::plain_payload_index::PlainIndex; use crate::index::struct_payload_index::StructPayloadIndex; -use crate::index::VectorIndexSS; +use crate::index::VectorIndexEnum; use crate::payload_storage::on_disk_payload_storage::OnDiskPayloadStorage; use crate::payload_storage::simple_payload_storage::SimplePayloadStorage; use crate::segment::{Segment, SegmentVersion, VectorData, SEGMENT_STATE_FILE}; @@ -117,26 +117,28 @@ fn create_segment( .load_quantization(&quantized_data_path)?; } - let vector_index: Arc> = match config.index { - Indexes::Plain { .. } => sp(PlainIndex::new( + let vector_index: Arc> = match config.index { + Indexes::Plain { .. } => sp(VectorIndexEnum::Plain(PlainIndex::new( vector_storage.clone(), payload_index.clone(), - )), + ))), Indexes::Hnsw(hnsw_config) => { if hnsw_config.on_disk.unwrap_or(false) { - sp(HNSWIndex::::open( - &vector_index_path, - vector_storage.clone(), - payload_index.clone(), - hnsw_config, - )?) + sp(VectorIndexEnum::HnswMmap( + HNSWIndex::::open( + &vector_index_path, + vector_storage.clone(), + payload_index.clone(), + hnsw_config, + )?, + )) } else { - sp(HNSWIndex::::open( + sp(VectorIndexEnum::HnswRam(HNSWIndex::::open( &vector_index_path, vector_storage.clone(), payload_index.clone(), hnsw_config, - )?) + )?)) } } };