diff --git a/crates/service/src/algorithms/diskann.rs b/crates/service/src/algorithms/diskann.rs new file mode 100644 index 000000000..973878430 --- /dev/null +++ b/crates/service/src/algorithms/diskann.rs @@ -0,0 +1,544 @@ +use super::raw::Raw; +use crate::index::segments::growing::GrowingSegment; +use crate::index::segments::sealed::SealedSegment; +use crate::index::{IndexOptions, SearchOptions, VectorOptions}; +use crate::prelude::*; +use crate::utils::dir_ops::sync_dir; +use crate::utils::element_heap::ElementHeap; +use crate::utils::mmap_array::MmapArray; +use parking_lot::{RwLock, RwLockWriteGuard}; +use rand::distributions::Uniform; +use rand::prelude::SliceRandom; +use rand::Rng; +use rayon::prelude::{IntoParallelIterator, ParallelIterator}; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::collections::{BTreeMap, HashSet}; +use std::fs::create_dir; +use std::path::PathBuf; +use std::sync::Arc; + +pub struct DiskANN { + mmap: DiskANNMmap, +} + +impl DiskANN { + pub fn create( + path: PathBuf, + options: IndexOptions, + sealed: Vec>>, + growing: Vec>>, + ) -> Self { + create_dir(&path).unwrap(); + let ram = make(path.clone(), sealed, growing, options.clone()); + let mmap = save(ram, path.clone()); + sync_dir(&path); + Self { mmap } + } + pub fn open(path: PathBuf, options: IndexOptions) -> Self { + let mmap = load(path, options.clone()); + Self { mmap } + } + + pub fn len(&self) -> u32 { + self.mmap.raw.len() + } + + pub fn vector(&self, i: u32) -> &[S::Scalar] { + self.mmap.raw.vector(i) + } + + pub fn payload(&self, i: u32) -> Payload { + self.mmap.raw.payload(i) + } + + pub fn basic( + &self, + vector: &[S::Scalar], + opts: &SearchOptions, + filter: impl Filter, + ) -> BinaryHeap> { + basic(&self.mmap, vector, opts.disk_ann_k, filter) + } +} + +unsafe impl Send for DiskANN {} +unsafe impl Sync for DiskANN {} + +pub struct VertexWithDistance { + pub id: u32, + pub distance: F32, +} + +impl VertexWithDistance { + pub fn new(id: u32, distance: F32) -> Self { + Self { id, distance } + } +} + +impl PartialEq for VertexWithDistance { + fn eq(&self, other: &Self) -> bool { + self.distance.eq(&other.distance) + } +} + +impl Eq for VertexWithDistance {} + +impl PartialOrd for VertexWithDistance { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.distance.cmp(&other.distance)) + } +} + +impl Ord for VertexWithDistance { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.distance.cmp(&other.distance) + } +} + +pub struct SearchState { + pub visited: HashSet, + candidates: BTreeMap, + heap: BinaryHeap>, + heap_visited: HashSet, + l: usize, +} + +impl SearchState { + /// Creates a new search state. + pub(crate) fn new(l: usize) -> Self { + Self { + visited: HashSet::new(), + candidates: BTreeMap::new(), + heap: BinaryHeap::new(), + heap_visited: HashSet::new(), + l, + } + } + + /// Return the next unvisited vertex. + fn pop(&mut self) -> Option { + while let Some(vertex) = self.heap.pop() { + if !self.candidates.contains_key(&vertex.0.distance) { + // The vertex has been removed from the candidate lists, + // from [`push()`]. + continue; + } + + self.visited.insert(vertex.0.id); + return Some(vertex.0.id); + } + + None + } + + /// Push a new (unvisited) vertex into the search state. + fn push(&mut self, vertex_id: u32, distance: F32) { + assert!(!self.visited.contains(&vertex_id)); + self.heap_visited.insert(vertex_id); + self.heap + .push(Reverse(VertexWithDistance::new(vertex_id, distance))); + self.candidates.insert(distance, vertex_id); + if self.candidates.len() > self.l { + self.candidates.pop_last(); + } + } + + /// Mark a vertex as visited. + fn visit(&mut self, vertex_id: u32) { + self.visited.insert(vertex_id); + } + + // Returns true if the vertex has been visited. + fn is_visited(&self, vertex_id: u32) -> bool { + self.visited.contains(&vertex_id) || self.heap_visited.contains(&vertex_id) + } +} + +struct VertexNeighbor { + neighbors: Vec, +} + +// DiskANNRam is for constructing the index +// it stores the intermediate structure when constructing +// the index and these data are stored in memory +pub struct DiskANNRam { + raw: Arc>, + // quantization: Quantization, + vertexs: Vec>, + /// the entry for the entire graph, the closet vector to centroid + medoid: u32, + dims: u16, + max_degree: u32, + l_build: u32, +} + +pub struct DiskANNMmap { + raw: Arc>, + neighbors: MmapArray, + neighbor_offset: MmapArray, + medoid: MmapArray, + l: u32, +} + +impl DiskANNRam { + fn _init_graph(&self, n: u32, mut rng: impl Rng) { + let distribution = Uniform::new(0, n); + for i in 0..n { + let mut neighbor_ids: HashSet = HashSet::new(); + if self.max_degree < n { + while neighbor_ids.len() < self.max_degree as usize { + let neighbor_id = rng.sample(distribution); + if neighbor_id != i { + neighbor_ids.insert(neighbor_id); + } + } + } else { + neighbor_ids = (0..n).collect(); + } + + self._set_neighbors(i, &neighbor_ids); + } + } + + fn _set_neighbors(&self, vertex_index: u32, neighbor_ids: &HashSet) { + assert!(neighbor_ids.len() <= self.max_degree as usize); + assert!((vertex_index as usize) < self.vertexs.len()); + + let mut vertex = self.vertexs[vertex_index as usize].write(); + vertex.neighbors.clear(); + for item in neighbor_ids.iter() { + vertex.neighbors.push(*item); + } + } + + fn _set_neighbors_with_write_guard( + &self, + neighbor_ids: &HashSet, + guard: &mut RwLockWriteGuard, + ) { + assert!(neighbor_ids.len() <= self.max_degree as usize); + guard.neighbors.clear(); + for item in neighbor_ids.iter() { + guard.neighbors.push(*item); + } + } + + fn _get_neighbors(&self, vertex_index: u32) -> VertexNeighbor { + let vertex = self.vertexs[vertex_index as usize].read(); + VertexNeighbor { + neighbors: vertex.neighbors.clone(), + } + } + + fn _find_medoid(&self, n: u32) -> u32 { + let centroid = self._compute_centroid(n); + let centroid_arr: &[S::Scalar] = ¢roid; + + let mut medoid_index = 0; + let mut min_dis = F32::infinity(); + for i in 0..n { + let dis = S::distance(centroid_arr, self.raw.vector(i)); + if dis < min_dis { + min_dis = dis; + medoid_index = i; + } + } + medoid_index + } + + fn _compute_centroid(&self, n: u32) -> Vec { + let dim = self.dims as usize; + let mut sum = vec![0_f32; dim]; + for i in 0..n { + let vec = self.raw.vector(i); + for j in 0..dim { + sum[j] += vec[j].to_f32(); + } + } + + let collection: Vec = sum + .iter() + .map(|v| S::Scalar::from_f32((*v / n as f32) as f32)) + .collect(); + collection + } + + // r and l leave here for multiple pass extension + fn _one_pass(&self, n: u32, alpha: f32, r: u32, l: u32, mut rng: impl Rng) { + let mut ids = (0..n).collect::>(); + ids.shuffle(&mut rng); + + ids.into_par_iter() + .for_each(|id| self.search_and_prune_for_one_vertex(id, alpha, r, l)); + } + + #[allow(unused_assignments)] + fn search_and_prune_for_one_vertex(&self, id: u32, alpha: f32, r: u32, l: u32) { + let query = self.raw.vector(id); + let mut state = self._greedy_search(self.medoid, query, l as usize); + state.visited.remove(&id); // in case visited has id itself + let mut new_neighbor_ids: HashSet = HashSet::new(); + { + let mut guard = self.vertexs[id as usize].write(); + let neighbor_ids = &guard.neighbors; + state.visited.extend(neighbor_ids.iter().copied()); + let neighbor_ids = self._robust_prune(id, state.visited, alpha, r); + let neighbor_ids: HashSet = neighbor_ids.into_iter().collect(); + self._set_neighbors_with_write_guard(&neighbor_ids, &mut guard); + new_neighbor_ids = neighbor_ids; + } + + for &neighbor_id in new_neighbor_ids.iter() { + { + let mut guard = self.vertexs[neighbor_id as usize].write(); + let old_neighbors = &guard.neighbors; + let mut old_neighbors: HashSet = old_neighbors.iter().copied().collect(); + old_neighbors.insert(id); + if old_neighbors.len() > r as usize { + // need robust prune + let new_neighbors = self._robust_prune(neighbor_id, old_neighbors, alpha, r); + let new_neighbors: HashSet = new_neighbors.into_iter().collect(); + self._set_neighbors_with_write_guard(&new_neighbors, &mut guard); + } else { + self._set_neighbors_with_write_guard(&old_neighbors, &mut guard); + } + } + } + } + + fn _greedy_search(&self, start: u32, query: &[S::Scalar], search_size: usize) -> SearchState { + let mut state = SearchState::new(search_size); + + let dist = S::distance(query, self.raw.vector(start)); + state.push(start, dist); + while let Some(id) = state.pop() { + // only pop id in the search list but not visited + state.visit(id); + { + let neighbor_ids = self._get_neighbors(id).neighbors; + for neighbor_id in neighbor_ids { + if state.is_visited(neighbor_id) { + continue; + } + + let dist = S::distance(query, self.raw.vector(neighbor_id)); + state.push(neighbor_id, dist); // push and retain closet l nodes + } + } + } + + state + } + + fn _robust_prune(&self, id: u32, mut visited: HashSet, alpha: f32, r: u32) -> Vec { + let mut heap: BinaryHeap = visited + .iter() + .map(|v| { + let dist = S::distance(self.raw.vector(id), self.raw.vector(*v)); + VertexWithDistance { + id: *v, + distance: dist, + } + }) + .collect(); + + let mut new_neighbor_ids: Vec = vec![]; + while !visited.is_empty() { + if let Some(mut p) = heap.pop() { + while !visited.contains(&p.id) { + match heap.pop() { + Some(value) => { + p = value; + } + None => { + return new_neighbor_ids; + } + } + } + new_neighbor_ids.push(p.id); + if new_neighbor_ids.len() >= r as usize { + break; + } + let mut to_remove: HashSet = HashSet::new(); + for pv in visited.iter() { + let dist_prime = S::distance(self.raw.vector(p.id), self.raw.vector(*pv)); + let dist_query = S::distance(self.raw.vector(id), self.raw.vector(*pv)); + if F32::from(alpha) * dist_prime <= dist_query { + to_remove.insert(*pv); + } + } + for pv in to_remove.iter() { + visited.remove(pv); + } + } else { + return new_neighbor_ids; + } + } + new_neighbor_ids + } +} + +impl DiskANNMmap { + fn _get_neighbors(&self, id: u32) -> Vec { + let start = self.neighbor_offset[id as usize]; + let end = self.neighbor_offset[id as usize + 1]; + self.neighbors[start..end].to_vec() + } +} + +unsafe impl Send for DiskANNMmap {} +unsafe impl Sync for DiskANNMmap {} + +fn calculate_offsets(iter: impl Iterator) -> impl Iterator { + let mut offset = 0usize; + let mut iter = std::iter::once(0).chain(iter); + std::iter::from_fn(move || { + let x = iter.next()?; + offset += x; + Some(offset) + }) +} + +pub fn make( + path: PathBuf, + sealed: Vec>>, + growing: Vec>>, + options: IndexOptions, +) -> DiskANNRam { + let idx_opts = options.indexing.clone().unwrap_diskann(); + let raw = Arc::new(Raw::create( + path.join("raw"), + options.clone(), + sealed, + growing, + )); + + let n = raw.len(); + let r = idx_opts.max_degree; + let VectorOptions { dims, .. } = options.vector; + + let vertexs: Vec> = (0..n) + .map(|_| { + RwLock::new(VertexNeighbor { + neighbors: Vec::new(), + }) + }) + .collect(); + + let medoid = 0; + + let mut new_vamana = DiskANNRam:: { + raw, + vertexs, + medoid, + dims, + max_degree: idx_opts.max_degree, + l_build: idx_opts.l_build, + }; + + // 1. init graph with r random neighbors for each node + let rng = rand::thread_rng(); + new_vamana._init_graph(n, rng.clone()); + + // 2. find medoid + new_vamana.medoid = new_vamana._find_medoid(n); + + // 3. iterate pass + new_vamana._one_pass(n, 1.0, r, idx_opts.l_build, rng.clone()); + + new_vamana._one_pass(n, idx_opts.alpha, r, idx_opts.max_degree, rng.clone()); + + new_vamana +} + +pub fn save(ram: DiskANNRam, path: PathBuf) -> DiskANNMmap { + let neighbors_iter = ram.vertexs.iter().flat_map(|vertex| { + let vertex = vertex.read(); + vertex.neighbors.to_vec().into_iter() + }); + + // Create the neighbors array using MmapArray::create. + let neighbors = MmapArray::create(path.join("neighbors"), neighbors_iter); + + // Create an iterator for the size of each neighbor list. + let neighbor_offset_iter = { + let iter = ram.vertexs.iter().map(|vertex| { + let vertex = vertex.read(); + vertex.neighbors.len() + }); + calculate_offsets(iter) + }; + + // Create the neighbor_size array using MmapArray::create. + let neighbor_offset = MmapArray::create(path.join("neighbor_offset"), neighbor_offset_iter); + + let medoid_vec = vec![ram.medoid]; + let medoid = MmapArray::create(path.join("medoid"), medoid_vec.into_iter()); + + DiskANNMmap { + raw: ram.raw, + neighbors, + neighbor_offset, + medoid, + l: ram.l_build, + } +} + +pub fn load(path: PathBuf, options: IndexOptions) -> DiskANNMmap { + let idx_opts = options.indexing.clone().unwrap_diskann(); + let raw = Arc::new(Raw::open(path.join("raw"), options.clone())); + let neighbors = MmapArray::open(path.join("neighbors")); + let neighbor_offset = MmapArray::open(path.join("neighbor_offset")); + let medoid = MmapArray::open(path.join("medoid")); + assert!(medoid.len() == 1); + + DiskANNMmap { + raw, + neighbors, + neighbor_offset, + medoid, + l: idx_opts.l_build, + } +} + +pub fn basic( + mmap: &DiskANNMmap, + vector: &[S::Scalar], + k: u32, + mut filter: impl Filter, +) -> BinaryHeap> { + let mut state = SearchState::new(mmap.l as usize); + + let start = mmap.medoid[0]; + let dist = S::distance(vector, mmap.raw.vector(start)); + state.push(start, dist); + while let Some(id) = state.pop() { + // only pop id in the search list but not visited + state.visit(id); + { + let neighbor_ids = mmap._get_neighbors(id); + for neighbor_id in neighbor_ids { + if state.is_visited(neighbor_id) { + continue; + } + + let payload = mmap.raw.payload(neighbor_id); + + if filter.check(payload) { + let dist = S::distance(vector, mmap.raw.vector(neighbor_id)); + state.push(neighbor_id, dist); // push and retain closet l nodes + } + } + } + } + + let mut results = ElementHeap::new(k as usize); + for (distance, id) in state.candidates { + results.push(Element { + distance, + payload: mmap.raw.payload(id), + }); + } + results.into_reversed_heap() +} diff --git a/crates/service/src/algorithms/mod.rs b/crates/service/src/algorithms/mod.rs index a3c5ffd52..9c20f5655 100644 --- a/crates/service/src/algorithms/mod.rs +++ b/crates/service/src/algorithms/mod.rs @@ -1,4 +1,5 @@ pub mod clustering; +pub mod diskann; pub mod flat; pub mod hnsw; pub mod ivf; diff --git a/crates/service/src/index/indexing/diskann.rs b/crates/service/src/index/indexing/diskann.rs new file mode 100644 index 000000000..15651c82a --- /dev/null +++ b/crates/service/src/index/indexing/diskann.rs @@ -0,0 +1,185 @@ +use super::AbstractIndexing; +use crate::algorithms::diskann::DiskANN; +use crate::algorithms::quantization::QuantizationOptions; +use crate::index::segments::growing::GrowingSegment; +use crate::index::segments::sealed::SealedSegment; +use crate::index::IndexOptions; +use crate::index::SearchOptions; +use crate::prelude::*; +use serde::{Deserialize, Serialize}; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::path::PathBuf; +use std::sync::Arc; +use validator::Validate; + +#[derive(Debug, Clone, Serialize, Deserialize, Validate)] +#[serde(deny_unknown_fields)] +pub struct DiskANNIndexingOptions { + // TODO(avery): Referenced from Microsoft.DiskANN algorithm, One design is to + // leave the definition of memory usage to users and estimate the required + // memory and the given memory to decide the quantization options. + // + // Current design is to let users define the ratio of PQ for in memory index. + // + // Besides, it is hard to estimate current memory usage as sealed segment and + // growing segments are passed to RawMmap and RawRam. Different from the direct + // calculation of the vector layout. + + // #[serde(default = "DiskANNIndexingOptions::default_index_path_prefix")] + // pub index_path_prefix: PathBuf, + + // #[serde(default = "DiskANNIndexingOptions::default_data_path")] + // pub data_path: PathBuf, + + // // DRAM budget in GB for searching the index to set the compressed level + // // for data while search happens + + // // bound on the memory footprint of the index at search time in GB. Once built, + // // the index will use up only the specified RAM limit, the rest will reside on disk. + // // This will dictate how aggressively we compress the data vectors to store in memory. + // // Larger will yield better performance at search time. For an n point index, to use + // // b byte PQ compressed representation in memory, use `B = ((n * b) / 2^30 + (250000*(4*R + sizeof(T)*ndim)) / 2^30)`. + // // The second term in the summation is to allow some buffer for caching about 250,000 nodes from the graph in memory while serving. + // // If you are not sure about this term, add 0.25GB to the first term. + // #[serde(default = "DiskANNIndexingOptions::default_search_DRAM_budget")] + // pub search_DRAM_budget: u32, + + // // DRAM budget in GB for building the index + // // Limit on the memory allowed for building the index in GB. + // // If you specify a value less than what is required to build the index + // // in one pass, the index is built using a divide and conquer approach so + // // that sub-graphs will fit in the RAM budget. The sub-graphs are overlaid + // // to build the overall index. This approach can be upto 1.5 times slower than + // // building the index in one shot. Allocate as much memory as your RAM allows. + // #[serde(default = "DiskANNIndexingOptions::default_build_DRAM_budget")] + // pub build_DRAM_budget: u32, + #[serde(default = "DiskANNIndexingOptions::default_num_threads")] + pub num_threads: u32, + + // R in the paper + #[serde(default = "DiskANNIndexingOptions::default_max_degree")] + pub max_degree: u32, + + // L in the paper + #[serde(default = "DiskANNIndexingOptions::default_l_build")] + pub l_build: u32, + + // alpha in the paper, slack factor + #[serde(default = "DiskANNIndexingOptions::default_alpha")] + pub alpha: f32, + + // TODO: QD (quantized dimension) + // TODO: codebook prefix + // TODO: PQ disk bytes (compressed bytes on SSD; 0 for no compression) + // TODO: append reorder data (include full precision data in the index; use only in conjunction with compressed data on SSD) + // TODO: build_PQ_bytes + // TODO: use opq + // TODO: label file (for filtered diskANN) + // TODO: universal label (for filtered diskANN) + // TODO: filtered Lbuild (for filtered diskANN) + // TODO: filter threshold (for filtered diskANN) + // TODO: label type (for filtered diskANN) + #[serde(default)] + #[validate] + pub quantization: QuantizationOptions, +} + +impl DiskANNIndexingOptions { + // fn default_index_path_prefix() -> PathBuf { + // "DiskANN_index".to_string().into() + // } + // fn default_data_path() -> PathBuf { + // "DiskANN_data".to_string().into() + // } + // fn default_search_DRAM_budget() -> u32 { + // 1 + // } + // fn default_build_DRAM_budget() -> u32 { + // 1 + // } + + fn default_num_threads() -> u32 { + match std::thread::available_parallelism() { + Ok(threads) => (threads.get() as f64).sqrt() as _, + Err(_) => 1, + } + } + fn default_max_degree() -> u32 { + 64 + } + fn default_l_build() -> u32 { + 100 + } + fn default_alpha() -> f32 { + 1.2 + } +} + +impl Default for DiskANNIndexingOptions { + fn default() -> Self { + Self { + // index_path_prefix: Self::default_index_path_prefix(), + // data_path: Self::default_data_path(), + // search_DRAM_budget: Self::default_search_DRAM_budget(), + // build_DRAM_budget: Self::default_build_DRAM_budget(), + num_threads: Self::default_num_threads(), + max_degree: Self::default_max_degree(), + l_build: Self::default_l_build(), + alpha: Self::default_alpha(), + quantization: Default::default(), + } + } +} + +pub struct DiskANNIndexing { + raw: DiskANN, +} + +impl AbstractIndexing for DiskANNIndexing { + fn create( + path: PathBuf, + options: IndexOptions, + sealed: Vec>>, + growing: Vec>>, + ) -> Self { + let raw = DiskANN::create(path, options, sealed, growing); + Self { raw } + } + + fn open(path: PathBuf, options: IndexOptions) -> Self { + let raw = DiskANN::open(path, options); + Self { raw } + } + + fn len(&self) -> u32 { + self.raw.len() + } + + fn vector(&self, i: u32) -> &[S::Scalar] { + self.raw.vector(i) + } + + fn payload(&self, i: u32) -> Payload { + self.raw.payload(i) + } + + fn basic( + &self, + vector: &[S::Scalar], + opts: &SearchOptions, + filter: impl Filter, + ) -> BinaryHeap> { + self.raw.basic(vector, opts, filter) + } + + #[allow(unused_variables)] + fn vbase<'a>( + &'a self, + vector: &'a [S::Scalar], + opts: &'a SearchOptions, + filter: impl Filter + 'a, + ) -> (Vec, Box<(dyn Iterator + 'a)>) { + unimplemented!("DiskANN does not support vbase mode") + } +} diff --git a/crates/service/src/index/indexing/mod.rs b/crates/service/src/index/indexing/mod.rs index 9b2468d87..e3b1e3de5 100644 --- a/crates/service/src/index/indexing/mod.rs +++ b/crates/service/src/index/indexing/mod.rs @@ -1,7 +1,9 @@ +pub mod diskann; pub mod flat; pub mod hnsw; pub mod ivf; +use self::diskann::{DiskANNIndexing, DiskANNIndexingOptions}; use self::flat::{FlatIndexing, FlatIndexingOptions}; use self::hnsw::{HnswIndexing, HnswIndexingOptions}; use self::ivf::{IvfIndexing, IvfIndexingOptions}; @@ -24,6 +26,7 @@ pub enum IndexingOptions { Flat(FlatIndexingOptions), Ivf(IvfIndexingOptions), Hnsw(HnswIndexingOptions), + DiskANN(DiskANNIndexingOptions), } impl IndexingOptions { @@ -45,6 +48,12 @@ impl IndexingOptions { }; x } + pub fn unwrap_diskann(self) -> DiskANNIndexingOptions { + let IndexingOptions::DiskANN(x) = self else { + unreachable!() + }; + x + } } impl Default for IndexingOptions { @@ -59,6 +68,7 @@ impl Validate for IndexingOptions { Self::Flat(x) => x.validate(), Self::Ivf(x) => x.validate(), Self::Hnsw(x) => x.validate(), + Self::DiskANN(x) => x.validate(), } } } @@ -92,6 +102,7 @@ pub enum DynamicIndexing { Flat(FlatIndexing), Ivf(IvfIndexing), Hnsw(HnswIndexing), + DiskANN(DiskANNIndexing), } impl DynamicIndexing { @@ -111,6 +122,9 @@ impl DynamicIndexing { IndexingOptions::Hnsw(_) => { Self::Hnsw(HnswIndexing::create(path, options, sealed, growing)) } + IndexingOptions::DiskANN(_) => { + Self::DiskANN(DiskANNIndexing::create(path, options, sealed, growing)) + } } } @@ -119,6 +133,7 @@ impl DynamicIndexing { IndexingOptions::Flat(_) => Self::Flat(FlatIndexing::open(path, options)), IndexingOptions::Ivf(_) => Self::Ivf(IvfIndexing::open(path, options)), IndexingOptions::Hnsw(_) => Self::Hnsw(HnswIndexing::open(path, options)), + IndexingOptions::DiskANN(_) => Self::DiskANN(DiskANNIndexing::open(path, options)), } } @@ -127,6 +142,7 @@ impl DynamicIndexing { DynamicIndexing::Flat(x) => x.len(), DynamicIndexing::Ivf(x) => x.len(), DynamicIndexing::Hnsw(x) => x.len(), + DynamicIndexing::DiskANN(x) => x.len(), } } @@ -135,6 +151,7 @@ impl DynamicIndexing { DynamicIndexing::Flat(x) => x.vector(i), DynamicIndexing::Ivf(x) => x.vector(i), DynamicIndexing::Hnsw(x) => x.vector(i), + DynamicIndexing::DiskANN(x) => x.vector(i), } } @@ -143,6 +160,7 @@ impl DynamicIndexing { DynamicIndexing::Flat(x) => x.payload(i), DynamicIndexing::Ivf(x) => x.payload(i), DynamicIndexing::Hnsw(x) => x.payload(i), + DynamicIndexing::DiskANN(x) => x.payload(i), } } @@ -156,6 +174,7 @@ impl DynamicIndexing { DynamicIndexing::Flat(x) => x.basic(vector, opts, filter), DynamicIndexing::Ivf(x) => x.basic(vector, opts, filter), DynamicIndexing::Hnsw(x) => x.basic(vector, opts, filter), + DynamicIndexing::DiskANN(x) => x.basic(vector, opts, filter), } } @@ -169,6 +188,7 @@ impl DynamicIndexing { DynamicIndexing::Flat(x) => x.vbase(vector, opts, filter), DynamicIndexing::Ivf(x) => x.vbase(vector, opts, filter), DynamicIndexing::Hnsw(x) => x.vbase(vector, opts, filter), + DynamicIndexing::DiskANN(x) => x.vbase(vector, opts, filter), } } } diff --git a/crates/service/src/index/mod.rs b/crates/service/src/index/mod.rs index 0baad255f..c7f0b3562 100644 --- a/crates/service/src/index/mod.rs +++ b/crates/service/src/index/mod.rs @@ -68,6 +68,8 @@ pub struct SearchOptions { pub hnsw_ef_search: usize, #[validate(range(min = 1, max = 1_000_000))] pub ivf_nprobe: u32, + #[validate(range(min = 1, max = 65535))] + pub disk_ann_k: u32, } #[derive(Debug, Serialize, Deserialize)]