vector index as enum

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::*;

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<PathBuf>;
 }
 
-pub type VectorIndexSS = dyn VectorIndex + Sync + Send;
+pub enum VectorIndexEnum {
+    Plain(PlainIndex),
+    HnswRam(HNSWIndex<GraphLinksRam>),
+    HnswMmap(HNSWIndex<GraphLinksMmap>),
+}
+
+impl VectorIndex for VectorIndexEnum {
+    fn search(
+        &self,
+        vectors: &[&[VectorElementType]],
+        filter: Option<&Filter>,
+        top: usize,
+        params: Option<&SearchParams>,
+    ) -> Vec<Vec<ScoredPointOffset>> {
+        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<PathBuf> {
+        match self {
+            VectorIndexEnum::Plain(index) => index.files(),
+            VectorIndexEnum::HnswRam(index) => index.files(),
+            VectorIndexEnum::HnswMmap(index) => index.files(),
+        }
+    }
+}

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<AtomicRefCell<VectorIndexSS>>,
+    pub vector_index: Arc<AtomicRefCell<VectorIndexEnum>>,
     pub vector_storage: Arc<AtomicRefCell<VectorStorageSS>>,
 }
 

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};

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<AtomicRefCell<VectorIndexSS>> = match config.index {
-            Indexes::Plain { .. } => sp(PlainIndex::new(
+        let vector_index: Arc<AtomicRefCell<VectorIndexEnum>> = 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::<GraphLinksMmap>::open(
-                        &vector_index_path,
-                        vector_storage.clone(),
-                        payload_index.clone(),
-                        hnsw_config,
-                    )?)
+                    sp(VectorIndexEnum::HnswMmap(
+                        HNSWIndex::<GraphLinksMmap>::open(
+                            &vector_index_path,
+                            vector_storage.clone(),
+                            payload_index.clone(),
+                            hnsw_config,
+                        )?,
+                    ))
                 } else {
-                    sp(HNSWIndex::<GraphLinksRam>::open(
+                    sp(VectorIndexEnum::HnswRam(HNSWIndex::<GraphLinksRam>::open(
                         &vector_index_path,
                         vector_storage.clone(),
                         payload_index.clone(),
                         hnsw_config,
-                    )?)
+                    )?))
                 }
             }
         };