diff --git a/diskann-benchmark-core/src/internal/buffer.rs b/diskann-benchmark-core/src/internal/buffer.rs index 29cfb9eb7..36d7e7264 100644 --- a/diskann-benchmark-core/src/internal/buffer.rs +++ b/diskann-benchmark-core/src/internal/buffer.rs @@ -3,7 +3,10 @@ * Licensed under the MIT license. */ -use diskann::graph::{SearchOutputBuffer, search_output_buffer::BufferState}; +use diskann::{ + graph::{SearchOutputBuffer, search_output_buffer::BufferState}, + neighbor::Neighbor, +}; /// A [`SearchOutputBuffer`] implementation that either references a slice in-place for /// fixed sized outputs or references a growable vector. @@ -58,7 +61,8 @@ impl SearchOutputBuffer for Buffer<'_, I> { >::current_len(self) } - fn push(&mut self, id: I, _distance: D) -> BufferState { + fn push(&mut self, neighbor: Neighbor) -> BufferState { + let (id, _) = neighbor.as_tuple(); match &mut self.0 { Inner::Slice { slice, written } => match slice.get_mut(*written) { Some(slot) => { @@ -81,14 +85,14 @@ impl SearchOutputBuffer for Buffer<'_, I> { fn extend(&mut self, itr: Itr) -> usize where - Itr: IntoIterator, + Itr: IntoIterator>, { match &mut self.0 { Inner::Slice { slice, written } => match slice.get_mut(*written..) { Some(left) => { let count = std::iter::zip(left.iter_mut(), itr) .map(|(dst, src)| { - *dst = src.0; + *dst = src.as_tuple().0; }) .count(); *written += count; @@ -98,7 +102,7 @@ impl SearchOutputBuffer for Buffer<'_, I> { }, Inner::Vec(vec) => { let before = vec.len(); - vec.extend(itr.into_iter().map(|i| i.0)); + vec.extend(itr.into_iter().map(|i| i.as_tuple().0)); vec.len() - before } } @@ -120,6 +124,10 @@ mod tests { buffer.size_hint() } + fn from_tuple((id, distance): (I, D)) -> Neighbor { + Neighbor::new(id, distance) + } + #[test] fn test_slice_buffer_creation() { let mut data = [0u32; 5]; @@ -144,7 +152,7 @@ mod tests { let mut data = [0u32; 5]; let mut buffer = Buffer::slice(&mut data); - assert_eq!(buffer.push(42, 0.0), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(42, 0.0)), BufferState::Available); assert_eq!(buffer.current_len(), 1); assert_eq!(data, [42, 0, 0, 0, 0]); } @@ -157,27 +165,27 @@ mod tests { assert_eq!(buffer.current_len(), 0); assert_eq!(size_hint(&buffer), Some(5)); - assert_eq!(buffer.push(100, 0.0), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(100, 0.0)), BufferState::Available); assert_eq!(buffer.current_len(), 1); assert_eq!(size_hint(&buffer), Some(4)); - assert_eq!(buffer.push(200, 0.0), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(200, 0.0)), BufferState::Available); assert_eq!(buffer.current_len(), 2); assert_eq!(size_hint(&buffer), Some(3)); - assert_eq!(buffer.push(300, 0.0), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(300, 0.0)), BufferState::Available); assert_eq!(buffer.current_len(), 3); assert_eq!(size_hint(&buffer), Some(2)); - assert_eq!(buffer.push(400, 0.0), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(400, 0.0)), BufferState::Available); assert_eq!(buffer.current_len(), 4); assert_eq!(size_hint(&buffer), Some(1)); - assert_eq!(buffer.push(500, 0.0), BufferState::Full); + assert_eq!(buffer.push(Neighbor::new(500, 0.0)), BufferState::Full); assert_eq!(buffer.current_len(), 5); assert_eq!(size_hint(&buffer), Some(0)); - assert_eq!(buffer.push(600, 0.0), BufferState::Full); + assert_eq!(buffer.push(Neighbor::new(600, 0.0)), BufferState::Full); assert_eq!(buffer.current_len(), 5); assert_eq!(size_hint(&buffer), Some(0)); @@ -190,7 +198,7 @@ mod tests { let mut data = [0u32; 0]; let mut buffer = Buffer::slice(&mut data); - assert_eq!(buffer.push(42, 0.0), BufferState::Full); + assert_eq!(buffer.push(Neighbor::new(42, 0.0)), BufferState::Full); } #[test] @@ -203,13 +211,13 @@ mod tests { "vector-type buffers have no upper bound" ); - assert_eq!(buffer.push(42, 0.0), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(42, 0.0)), BufferState::Available); assert_eq!(buffer.current_len(), 1); - assert_eq!(buffer.push(50, 0.0), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(50, 0.0)), BufferState::Available); assert_eq!(buffer.current_len(), 2); - assert_eq!(buffer.push(3, 0.0), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(3, 0.0)), BufferState::Available); assert_eq!(buffer.current_len(), 3); assert_eq!(&vec, &[42, 50, 3]); @@ -220,7 +228,7 @@ mod tests { let mut data = [0u32; 5]; let mut buffer = Buffer::slice(&mut data); - let items = vec![(10, 1.0), (20, 2.0), (30, 3.0)]; + let items = [(10, 1.0), (20, 2.0), (30, 3.0)].map(from_tuple); let count = buffer.extend(items); assert_eq!(count, 3); @@ -233,7 +241,7 @@ mod tests { let mut data = [0u32; 3]; let mut buffer = Buffer::slice(&mut data); - let items = vec![(10, 1.0), (20, 2.0), (30, 3.0), (40, 4.0), (50, 5.0)]; + let items = [(10, 1.0), (20, 2.0), (30, 3.0), (40, 4.0), (50, 5.0)].map(from_tuple); let count = buffer.extend(items); // Only first 3 items should be written @@ -247,7 +255,7 @@ mod tests { let mut vec = Vec::::new(); let mut buffer = Buffer::vector(&mut vec); - let items = vec![(100, 1.0), (200, 2.0)]; + let items = [(100, 1.0), (200, 2.0)].map(from_tuple); let count = buffer.extend(items); assert_eq!(count, 2); @@ -260,14 +268,14 @@ mod tests { let mut vec = vec![1u32, 2, 3, 4, 5]; let mut buffer = Buffer::vector(&mut vec); - let items = vec![(10, 1.0), (20, 2.0)]; + let items = [(10, 1.0), (20, 2.0)].map(from_tuple); assert_eq!(buffer.extend(items), 2); - let items = vec![(21, 1.0), (22, 2.0)]; + let items = [(21, 1.0), (22, 2.0)].map(from_tuple); assert_eq!(buffer.extend(items), 2); assert_eq!( - buffer.extend::<[(u32, f32); 0]>([]), + buffer.extend::<[Neighbor; 0]>([]), 0, "empty iterator should add nothing" ); @@ -281,27 +289,30 @@ mod tests { // Push then extend let mut data = [0u32; 5]; let mut buffer = Buffer::slice(&mut data); - assert_eq!(buffer.push(1, 0.0), BufferState::Available); - assert_eq!(buffer.push(2, 0.0), BufferState::Available); - assert_eq!(buffer.extend(vec![(3, 0.0), (4, 0.0)]), 2); + assert_eq!(buffer.push(Neighbor::new(1, 0.0)), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(2, 0.0)), BufferState::Available); + assert_eq!(buffer.extend([(3, 0.0), (4, 0.0)].map(from_tuple)), 2); assert_eq!(data, [1, 2, 3, 4, 0]); // Extend then push to fill let mut data = [0u32; 5]; let mut buffer = Buffer::slice(&mut data); - buffer.extend(vec![(10, 0.0), (20, 0.0)]); - assert_eq!(buffer.push(30, 0.0), BufferState::Available); - assert_eq!(buffer.push(40, 0.0), BufferState::Available); - assert_eq!(buffer.push(50, 0.0), BufferState::Full); + buffer.extend([(10, 0.0), (20, 0.0)].map(from_tuple)); + assert_eq!(buffer.push(Neighbor::new(30, 0.0)), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(40, 0.0)), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(50, 0.0)), BufferState::Full); assert_eq!(data, [10, 20, 30, 40, 50]); // Interleaved operations let mut data = [0u32; 6]; let mut buffer = Buffer::slice(&mut data); - assert_eq!(buffer.push(1, 0.0), BufferState::Available); - assert_eq!(buffer.extend(vec![(2, 0.0), (3, 0.0)]), 2); - assert_eq!(buffer.push(4, 0.0), BufferState::Available); - assert_eq!(buffer.extend(vec![(5, 0.0), (6, 0.0), (7, 0.0)]), 2); + assert_eq!(buffer.push(Neighbor::new(1, 0.0)), BufferState::Available); + assert_eq!(buffer.extend([(2, 0.0), (3, 0.0)].map(from_tuple)), 2); + assert_eq!(buffer.push(Neighbor::new(4, 0.0)), BufferState::Available); + assert_eq!( + buffer.extend([(5, 0.0), (6, 0.0), (7, 0.0)].map(from_tuple)), + 2 + ); assert_eq!(data, [1, 2, 3, 4, 5, 6]); } @@ -310,25 +321,25 @@ mod tests { // Extend fills buffer, push returns Full let mut data = [0u32; 2]; let mut buffer = Buffer::slice(&mut data); - buffer.extend(vec![(1, 0.0), (2, 0.0)]); - assert_eq!(buffer.push(99, 0.0), BufferState::Full); + buffer.extend([(1, 0.0), (2, 0.0)].map(from_tuple)); + assert_eq!(buffer.push(Neighbor::new(99, 0.0)), BufferState::Full); assert_eq!(data, [1, 2]); // Push fills buffer, extend returns 0 let mut data = [0u32; 2]; let mut buffer = Buffer::slice(&mut data); - assert_eq!(buffer.push(1, 0.0), BufferState::Available); - assert_eq!(buffer.push(2, 0.0), BufferState::Full); - assert_eq!(buffer.extend(vec![(99, 0.0)]), 0); + assert_eq!(buffer.push(Neighbor::new(1, 0.0)), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(2, 0.0)), BufferState::Full); + assert_eq!(buffer.extend([Neighbor::new(99, 0.0)]), 0); assert_eq!(data, [1, 2]); // Extend truncates when exceeding remaining capacity after push let mut data = [0u32; 4]; let mut buffer = Buffer::slice(&mut data); - assert_eq!(buffer.push(1, 0.0), BufferState::Available); - assert_eq!(buffer.push(2, 0.0), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(1, 0.0)), BufferState::Available); + assert_eq!(buffer.push(Neighbor::new(2, 0.0)), BufferState::Available); assert_eq!( - buffer.extend(vec![(3, 0.0), (4, 0.0), (5, 0.0), (6, 0.0)]), + buffer.extend([(3, 0.0), (4, 0.0), (5, 0.0), (6, 0.0)].map(from_tuple)), 2 ); assert_eq!(data, [1, 2, 3, 4]); @@ -340,11 +351,11 @@ mod tests { let mut buffer = Buffer::vector(&mut vec); // Interleave push and extend - vec has no capacity limit - assert_eq!(buffer.push(1, 0.0), BufferState::Available); - assert_eq!(buffer.extend(vec![(2, 0.0), (3, 0.0)]), 2); - assert_eq!(buffer.push(4, 0.0), BufferState::Available); - assert_eq!(buffer.extend::<[(u32, f32); 0]>([]), 0); // empty extend - assert_eq!(buffer.extend(vec![(5, 0.0)]), 1); + assert_eq!(buffer.push(Neighbor::new(1, 0.0)), BufferState::Available); + assert_eq!(buffer.extend([(2, 0.0), (3, 0.0)].map(from_tuple)), 2); + assert_eq!(buffer.push(Neighbor::new(4, 0.0)), BufferState::Available); + assert_eq!(buffer.extend::<[Neighbor; 0]>([]), 0); // empty extend + assert_eq!(buffer.extend(vec![Neighbor::new(5, 0.0)]), 1); assert_eq!(&vec, &[1, 2, 3, 4, 5]); } diff --git a/diskann-benchmark-core/src/search/api.rs b/diskann-benchmark-core/src/search/api.rs index 636a9ab19..811f02b0e 100644 --- a/diskann-benchmark-core/src/search/api.rs +++ b/diskann-benchmark-core/src/search/api.rs @@ -584,6 +584,8 @@ mod tests { use std::hash::{self, Hash, Hasher}; + use diskann::neighbor::Neighbor; + // We intentionally do not derive `Clone` to ensure that it is not needed // in the implementations. #[derive(Debug)] @@ -675,7 +677,8 @@ mod tests { O: graph::SearchOutputBuffer + Send, { let count = self.count(index, params); - let set = buffer.extend((0..count).map(|i| (self.format(index, i), i as f32))); + let set = + buffer.extend((0..count).map(|i| Neighbor::new(self.format(index, i), i as f32))); assert_eq!(set, count); Ok(count) } diff --git a/diskann-benchmark-core/src/search/mod.rs b/diskann-benchmark-core/src/search/mod.rs index 36b4be3ad..1eecd8f3f 100644 --- a/diskann-benchmark-core/src/search/mod.rs +++ b/diskann-benchmark-core/src/search/mod.rs @@ -39,6 +39,7 @@ //! //! ```rust //! use std::{sync::Arc, num::NonZeroUsize}; +//! use diskann::neighbor::Neighbor; //! use diskann_benchmark_core::search; //! //! /// A simple example implementation of the `Search` trait. @@ -96,7 +97,7 @@ //! use diskann::graph::SearchOutputBuffer; //! //! // Fill the buffer with `index`. -//! buffer.extend((0..parameters.num_ids).map(|_| (index, 0.0f32))); +//! buffer.extend((0..parameters.num_ids).map(|_| Neighbor::new(index, 0.0f32))); //! Ok(Output::new(index)) //! } //! } diff --git a/diskann-bftree/src/provider.rs b/diskann-bftree/src/provider.rs index 7e9e53d63..eefc5bf2b 100644 --- a/diskann-bftree/src/provider.rs +++ b/diskann-bftree/src/provider.rs @@ -30,7 +30,7 @@ use diskann::{ workingset::map, AdjacencyList, SearchOutputBuffer, }, - neighbor::Neighbor, + neighbor::{self, Neighbor}, provider::{DataProvider, DefaultContext, Delete, ElementStatus, HasId, NoopGuard, SetElement}, utils::VectorRepr, ANNError, ANNResult, @@ -1588,7 +1588,7 @@ where let provider = accessor.provider; let f = T::distance(provider.metric, Some(provider.full_vectors.dim())); - let mut reranked: Vec<(I, f32)> = Vec::new(); + let mut reranked = Vec::new(); for n in candidates { match provider .full_vectors @@ -1596,7 +1596,7 @@ where .allow_transient("stale candidate during rerank") { Ok(Some(vec)) => { - reranked.push((*n.id(), f.evaluate_similarity(query, &vec))); + reranked.push(Neighbor::new(*n.id(), f.evaluate_similarity(query, &vec))); } Ok(None) => { // Transient (deleted/missing) — skip this candidate. @@ -1605,8 +1605,7 @@ where } } - reranked - .sort_unstable_by(|a, b| (a.1).partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); + reranked.sort_unstable_by(neighbor::ord::fast_distance); std::future::ready(Ok(output.extend(reranked))) } } diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index c36e8ffd7..a4986729d 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -23,7 +23,7 @@ use diskann::{ search::{AdaptiveL, InlineFilterSearch, Knn}, search_output_buffer, DiskANNIndex, }, - neighbor::{Neighbor, NeighborPriorityQueue}, + neighbor::{self, Neighbor, NeighborPriorityQueue}, provider::{DataProvider, DefaultContext, HasId, NoopGuard}, utils::{IntoUsize, VectorRepr}, ANNError, ANNResult, @@ -350,10 +350,10 @@ where let provider = accessor.provider; let mut uncached_ids = Vec::new(); - let mut reranked = { + let mut reranked: Vec<_> = { let mut process = |n: u32| { if let Some(entry) = accessor.scratch.distance_cache.get(&n) { - Some(Ok::<((u32, _), f32), ANNError>(((n, entry.1), entry.0))) + Some(Neighbor::new((n, entry.1), entry.0)) } else { uncached_ids.push(n); None @@ -363,12 +363,12 @@ where PostprocessStrategy::AcceptAll => candidates .map(|n| *n.id()) .filter_map(&mut process) - .collect::, _>>()?, + .collect(), PostprocessStrategy::Apply(f) => candidates .map(|n| *n.id()) .filter(|id| f(id)) .filter_map(&mut process) - .collect::, _>>()?, + .collect(), } }; if !uncached_ids.is_empty() { @@ -377,13 +377,13 @@ where let v = accessor.scratch.vertex_provider.get_vector(n)?; let d = provider.distance_comparer.evaluate_similarity(query, v); let a = accessor.scratch.vertex_provider.get_associated_data(n)?; - reranked.push(((*n, *a), d)); + reranked.push(Neighbor::new((*n, *a), d)); } } // Sort the full precision distances. - reranked - .sort_unstable_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); + reranked.sort_unstable_by(neighbor::ord::fast_distance); + // Store the reranked results. Ok(output.extend(reranked)) } @@ -463,7 +463,7 @@ where Ok(output.extend(reranked.into_iter().map(|idx| { let id = candidate_ids[idx]; let distance = candidate_distances[idx]; - ((id, associated_data[idx]), distance) + Neighbor::new((id, associated_data[idx]), distance) }))) } } diff --git a/diskann-garnet/src/lib.rs b/diskann-garnet/src/lib.rs index cc9dc92d8..41ba5d468 100644 --- a/diskann-garnet/src/lib.rs +++ b/diskann-garnet/src/lib.rs @@ -20,6 +20,7 @@ use diskann::{ config::{self, defaults::GRAPH_SLACK_FACTOR}, search::{self, AdaptiveL}, }, + neighbor::Neighbor, utils::VectorRepr, }; use diskann_providers::index::wrapped_async::DiskANNIndex; @@ -126,7 +127,9 @@ impl SearchOutputBuffer for SearchResults<'_> { Some(self.dists.len() - self.index) } - fn push(&mut self, id: GarnetId, distance: f32) -> diskann::graph::BufferState { + fn push(&mut self, neighbor: Neighbor) -> diskann::graph::BufferState { + let (id, distance) = neighbor.as_tuple(); + if self.index >= self.dists.len() || self.id_index + mem::size_of::() + id.len() > self.ids.len() { @@ -156,12 +159,12 @@ impl SearchOutputBuffer for SearchResults<'_> { fn extend(&mut self, itr: Itr) -> usize where - Itr: IntoIterator, + Itr: IntoIterator>, { let initial = self.current_len(); - for (id, dist) in itr { - if self.push(id, dist).is_full() { + for neighbor in itr { + if self.push(neighbor).is_full() { break; } } @@ -829,7 +832,10 @@ pub unsafe extern "C" fn check_external_id_valid( mod tests { use std::{mem, ptr}; - use diskann::graph::{BufferState, SearchOutputBuffer}; + use diskann::{ + graph::{BufferState, SearchOutputBuffer}, + neighbor::Neighbor, + }; use diskann_vector::distance::Metric; use crate::{ @@ -858,11 +864,11 @@ mod tests { assert_eq!(sr.size_hint(), Some(5)); let test_data = [ - (GarnetId::from(bytemuck::bytes_of(&1u32)), 1.1f32), - (GarnetId::from(bytemuck::bytes_of(&2u32)), 2.1), - (GarnetId::from(bytemuck::bytes_of(&3u32)), 3.1), - (GarnetId::from(bytemuck::bytes_of(&4u32)), 4.1), - (GarnetId::from(bytemuck::bytes_of(&5u32)), 5.1), + Neighbor::new(GarnetId::from(bytemuck::bytes_of(&1u32)), 1.1f32), + Neighbor::new(GarnetId::from(bytemuck::bytes_of(&2u32)), 2.1), + Neighbor::new(GarnetId::from(bytemuck::bytes_of(&3u32)), 3.1), + Neighbor::new(GarnetId::from(bytemuck::bytes_of(&4u32)), 4.1), + Neighbor::new(GarnetId::from(bytemuck::bytes_of(&5u32)), 5.1), ]; assert_eq!(sr.current_len(), 0); @@ -889,7 +895,10 @@ mod tests { } assert_eq!( - sr.push(GarnetId::from(bytemuck::bytes_of(&6u32)), 6.1f32), + sr.push(Neighbor::new( + GarnetId::from(bytemuck::bytes_of(&6u32)), + 6.1f32 + )), BufferState::Full ); } diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index f4e713867..b989373c2 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -1247,7 +1247,7 @@ impl<'a, T: VectorRepr> SearchPostProcess, &[T], GarnetId Err(_) => continue, // Can't read the mapping; skip. }; - if output.push(id, n.distance()).is_full() { + if output.push(Neighbor::new(id, *n.distance())).is_full() { break; } } @@ -1297,7 +1297,7 @@ impl<'a, 'b, T: VectorRepr> SearchPostProcessStep, &'b [T let mut v = Poly::broadcast(0u8, provider.dim * mem::size_of::(), AlignToEight)?; // Filter before computing the full precision distances. - let mut reranked: Vec<(u32, f32)> = candidates + let mut reranked: Vec<_> = candidates .filter_map(|n| { if !provider.vector_iid_exists(accessor.context, *n.id()) { None @@ -1306,7 +1306,7 @@ impl<'a, 'b, T: VectorRepr> SearchPostProcessStep, &'b [T *n.id(), &mut v, ) { - Some(( + Some(Neighbor::new( *n.id(), f.evaluate_similarity(query, bytemuck::cast_slice::(&v)), )) @@ -1317,17 +1317,11 @@ impl<'a, 'b, T: VectorRepr> SearchPostProcessStep, &'b [T .collect(); // Sort the full precision distances. - reranked - .sort_unstable_by(|a, b| (a.1).partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); - - next.post_process( - accessor, - query, - reranked.into_iter().map(|(id, d)| Neighbor::new(id, d)), - output, - ) - .await - .map_err(|e| GarnetProviderError::PostProcessing(Box::new(e))) + reranked.sort_unstable_by(diskann::neighbor::ord::fast_distance); + + next.post_process(accessor, query, reranked.into_iter(), output) + .await + .map_err(|e| GarnetProviderError::PostProcessing(Box::new(e))) } } diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs index d355379b5..0a0ccc310 100644 --- a/diskann-inmem/src/provider.rs +++ b/diskann-inmem/src/provider.rs @@ -932,7 +932,10 @@ where let mut count = 0; for c in candidates { if let Some(ext) = provider.mapping.to_external(*c.id()) { - if output.push(ext, c.distance()).is_available() { + if output + .push(Neighbor::new(ext, *c.distance())) + .is_available() + { count += 1; } else { break; diff --git a/diskann-providers/src/index/diskann_async.rs b/diskann-providers/src/index/diskann_async.rs index 4acf5458f..26705bcaa 100644 --- a/diskann-providers/src/index/diskann_async.rs +++ b/diskann-providers/src/index/diskann_async.rs @@ -917,7 +917,7 @@ pub(crate) mod tests { let expected: Neighbor = gt[gt.len() - 1 - position]; if id != *expected.id() { // We can allow it if the distance is the same. - if distance == expected.distance() { + if distance == *expected.distance() { Ok(()) } else { Err(Box::new(format!( @@ -925,7 +925,7 @@ pub(crate) mod tests { expected, id ))) } - } else if distance != expected.distance() { + } else if distance != *expected.distance() { Err(Box::new(format!( "expected neighbor {:?}, but found {}", expected, distance @@ -1085,7 +1085,7 @@ pub(crate) mod tests { let mut gt = groundtruth(corpus.as_view(), &query, |a, b| SquaredL2::evaluate(a, b)); for n in gt.iter_mut() { if filter.is_match(*n.id()) { - *n = Neighbor::new(*n.id(), n.distance() * beta); + *n = Neighbor::new(*n.id(), *n.distance() * beta); } } gt.sort_unstable_by(neighbor::ord::reverse(neighbor::ord::fast_distance)); diff --git a/diskann-providers/src/index/wrapped_async.rs b/diskann-providers/src/index/wrapped_async.rs index 527c072d1..092ea3bc2 100644 --- a/diskann-providers/src/index/wrapped_async.rs +++ b/diskann-providers/src/index/wrapped_async.rs @@ -839,7 +839,7 @@ mod tests { "monotonicity should at least hold for the 1d grid" ); assert_eq!( - neighbor.distance(), + *neighbor.distance(), (i as f32) * (i as f32), "distance was computed incorrectly!", ); diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs b/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs index 8437e5e90..b33e73a86 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs +++ b/diskann-providers/src/model/graph/provider/async_/inmem/full_precision.rs @@ -17,7 +17,7 @@ use diskann::{ }, workingset, }, - neighbor::Neighbor, + neighbor::{self, Neighbor}, provider::{DefaultContext, ExecutionContext, HasId}, utils::{IntoUsize, VectorRepr}, }; @@ -379,12 +379,12 @@ where let f = full.distance(); // Filter before computing the full precision distances. - let mut reranked: Vec<(u32, f32)> = candidates + let mut reranked: Vec<_> = candidates .filter_map(|n| { if checker.deletion_check(*n.id()) { None } else { - Some(( + Some(Neighbor::new( *n.id(), f.evaluate_similarity(query, unsafe { full.get_vector_sync(n.id().into_usize()) @@ -395,8 +395,7 @@ where .collect(); // Sort the full precision distances. - reranked - .sort_unstable_by(|a, b| (a.1).partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); + reranked.sort_unstable_by(neighbor::ord::fast_distance); // Store the reranked results. std::future::ready(Ok(output.extend(reranked))) @@ -447,7 +446,7 @@ where // pattern used by `Rerank` above. let vector = unsafe { store.get_vector_sync(candidate.id().into_usize()) }; ids.push(*candidate.id()); - distances.push(candidate.distance()); + distances.push(*candidate.distance()); vectors.row_mut(i).copy_from_slice(vector); } @@ -462,7 +461,9 @@ where Err(error) => return std::future::ready(Err(error.into())), }; - let reranked = indices.into_iter().map(|idx| (ids[idx], distances[idx])); + let reranked = indices + .into_iter() + .map(|idx| Neighbor::new(ids[idx], distances[idx])); std::future::ready(Ok(output.extend(reranked))) } diff --git a/diskann-providers/src/model/graph/provider/async_/postprocess.rs b/diskann-providers/src/model/graph/provider/async_/postprocess.rs index 75505fe2e..e605f4c5d 100644 --- a/diskann-providers/src/model/graph/provider/async_/postprocess.rs +++ b/diskann-providers/src/model/graph/provider/async_/postprocess.rs @@ -55,13 +55,7 @@ where B: SearchOutputBuffer + Send + ?Sized, { let checker = accessor.as_deletion_check(); - let count = output.extend(candidates.filter_map(|n| { - if checker.deletion_check(*n.id()) { - None - } else { - Some(n.as_tuple()) - } - })); + let count = output.extend(candidates.filter(|n| !checker.deletion_check(*n.id()))); std::future::ready(Ok(count)) } } diff --git a/diskann-providers/src/test_utils/search_utils.rs b/diskann-providers/src/test_utils/search_utils.rs index 1ce24bbd5..6265b3cc6 100644 --- a/diskann-providers/src/test_utils/search_utils.rs +++ b/diskann-providers/src/test_utils/search_utils.rs @@ -70,7 +70,7 @@ pub fn assert_top_k_exactly_match( for i in 0..top_k { let neighbor = gt[gt.len() - 1 - i]; assert_eq!( - neighbor.distance(), + *neighbor.distance(), distances[i], "failed on query {} for result {}", query_id, @@ -100,12 +100,12 @@ pub fn assert_range_results_exactly_match( ) { let gt_ids = if let Some(inner_radius) = inner_radius { gt.iter() - .filter(|nbh| nbh.distance() >= inner_radius && nbh.distance() <= radius) + .filter(|nbh| *nbh.distance() >= inner_radius && *nbh.distance() <= radius) .map(|nbh| *nbh.id()) .collect::>() } else { gt.iter() - .filter(|nbh| nbh.distance() <= radius) + .filter(|nbh| *nbh.distance() <= radius) .map(|nbh| *nbh.id()) .collect::>() }; diff --git a/diskann-tools/src/utils/ground_truth.rs b/diskann-tools/src/utils/ground_truth.rs index 0ed2094ed..c083e9ddd 100644 --- a/diskann-tools/src/utils/ground_truth.rs +++ b/diskann-tools/src/utils/ground_truth.rs @@ -539,7 +539,7 @@ fn write_ground_truth( for mut query_neighbors in ground_truth { while let Some(closest_node) = query_neighbors.closest_notvisited() { gt_ids.push(*closest_node.id()); - gt_distances.push(closest_node.distance()); + gt_distances.push(*closest_node.distance()); } } diff --git a/diskann/src/flat/test/harness.rs b/diskann/src/flat/test/harness.rs index 30a300d65..553b9fab8 100644 --- a/diskann/src/flat/test/harness.rs +++ b/diskann/src/flat/test/harness.rs @@ -88,7 +88,7 @@ impl KnnOracleRun { .take(stats.result_count as usize) .collect(); sort_neighbors(&mut top_k); - let top_k_distances = top_k.iter().map(|n| n.distance()).collect(); + let top_k_distances = top_k.iter().map(|n| *n.distance()).collect(); let ground_truth = oracle.expected(brute_force_topk(index.provider(), Metric::L2, query, k)); @@ -154,11 +154,7 @@ where I: Iterator> + Send, B: SearchOutputBuffer + Send + ?Sized, { - let count = output.extend( - candidates - .filter(|n| *n.id() % 2 == 0) - .map(|n| n.as_tuple()), - ); + let count = output.extend(candidates.filter(|n| *n.id() % 2 == 0)); std::future::ready(Ok(count)) } } diff --git a/diskann/src/graph/glue.rs b/diskann/src/graph/glue.rs index c92e151a6..31d74bc8b 100644 --- a/diskann/src/graph/glue.rs +++ b/diskann/src/graph/glue.rs @@ -688,7 +688,7 @@ where I: Iterator> + Send, B: SearchOutputBuffer + Send + ?Sized, { - let count = output.extend(candidates.map(|n| n.as_tuple())); + let count = output.extend(candidates); std::future::ready(Ok(count)) } } diff --git a/diskann/src/graph/index.rs b/diskann/src/graph/index.rs index 83134d984..7a4d09064 100644 --- a/diskann/src/graph/index.rs +++ b/diskann/src/graph/index.rs @@ -2610,9 +2610,9 @@ where // Filter out self loops. let id = neighbor.id(); if exclude(*id) { - (neighbor.distance(), None) + (*neighbor.distance(), None) } else { - (neighbor.distance(), map.get(*id)) + (*neighbor.distance(), map.get(*id)) } }) .collect(); diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index 3b3df35b5..ba04c357a 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -203,7 +203,7 @@ where let max_returned = self.max_returned().unwrap_or(usize::MAX); for neighbor in scratch.best.iter().take(starting_l) { - if neighbor.distance() <= self.radius() { + if *neighbor.distance() <= self.radius() { in_range.push(neighbor); } } @@ -291,9 +291,9 @@ where self.inner.size_hint() } - fn push(&mut self, id: I, distance: f32) -> search_output_buffer::BufferState { - if (self.predicate)(distance) { - self.inner.push(id, distance) + fn push(&mut self, neighbor: Neighbor) -> search_output_buffer::BufferState { + if (self.predicate)(*neighbor.distance()) { + self.inner.push(neighbor) } else { match self.inner.size_hint() { Some(0) => search_output_buffer::BufferState::Full, @@ -308,10 +308,10 @@ where fn extend(&mut self, itr: Itr) -> usize where - Itr: IntoIterator, + Itr: IntoIterator>, { self.inner - .extend(itr.into_iter().filter(|(_, dist)| (self.predicate)(*dist))) + .extend(itr.into_iter().filter(|n| (self.predicate)(*n.distance()))) } } @@ -365,7 +365,7 @@ where // The predicate ensures that the contents of `neighbors` are unique. for neighbor in neighbors.iter() { - if neighbor.distance() <= search_params.radius() * search_params.range_slack() + if *neighbor.distance() <= search_params.radius() * search_params.range_slack() && scratch.in_range.len() < max_returned { scratch.in_range.push(*neighbor); @@ -417,10 +417,10 @@ mod tests { let mut inner: Vec> = Vec::new(); let mut filtered = DistanceFiltered::new(&mut inner, |d| d < 1.0); - assert_eq!(filtered.push(1, 0.5), BufferState::Available); + assert_eq!(filtered.push(Neighbor::new(1, 0.5)), BufferState::Available); assert_eq!(filtered.current_len(), 1); assert_eq!(*inner[0].id(), 1); - assert_eq!(inner[0].distance(), 0.5); + assert_eq!(*inner[0].distance(), 0.5); } #[test] @@ -428,7 +428,7 @@ mod tests { let mut inner: Vec> = Vec::new(); let mut filtered = DistanceFiltered::new(&mut inner, |d| d < 1.0); - assert_eq!(filtered.push(1, 1.5), BufferState::Available); + assert_eq!(filtered.push(Neighbor::new(1, 1.5)), BufferState::Available); assert_eq!(filtered.current_len(), 0); } @@ -438,7 +438,7 @@ mod tests { let mut filtered = DistanceFiltered::new(&mut inner, |d| d < 1.0); assert!(filtered.size_hint().is_none()); - let items = vec![(1u32, 0.3), (2, 1.5), (3, 0.7), (4, 2.0), (5, 0.9)]; + let items = [(1u32, 0.3), (2, 1.5), (3, 0.7), (4, 2.0), (5, 0.9)].map(Neighbor::from_tuple); let count = filtered.extend(items); assert_eq!(count, 3); @@ -456,7 +456,7 @@ mod tests { let mut filtered = DistanceFiltered::new(&mut inner, |d| d < 1.0); assert_eq!(filtered.size_hint(), Some(2)); - let items = vec![(1u32, 0.1), (2, 0.2), (3, 0.3)]; + let items = [(1u32, 0.1), (2, 0.2), (3, 0.3)].map(Neighbor::from_tuple); let count = filtered.extend(items); assert_eq!(count, 2); @@ -477,7 +477,7 @@ mod tests { dist < radius }); - let items = vec![(1u32, 0.1), (2, 0.5), (3, 0.3), (4, 1.0), (5, 0.8)]; + let items = [(1u32, 0.1), (2, 0.5), (3, 0.3), (4, 1.0), (5, 0.8)].map(Neighbor::from_tuple); let count = filtered.extend(items); // 0.1 and 0.3 are <= inner_radius, 1.0 is not < radius diff --git a/diskann/src/graph/search_output_buffer.rs b/diskann/src/graph/search_output_buffer.rs index fddd7c73b..03c3b2504 100644 --- a/diskann/src/graph/search_output_buffer.rs +++ b/diskann/src/graph/search_output_buffer.rs @@ -3,6 +3,8 @@ * Licensed under the MIT license. */ +use crate::neighbor::Neighbor; + /// A generalized trait for extracting search results. /// /// Putting this behind a trait allows multiple different output containers to use common @@ -17,13 +19,13 @@ pub trait SearchOutputBuffer { /// `None` may be returned in instances where the output has unknown or unbounded size. fn size_hint(&self) -> Option; - /// Push an `id` and `distance` pair to the next position in the buffer. + /// Push a [`Neighbor`] to the next position in the buffer. /// /// Returns a [`BufferState`] to indicate whether future insertions will succeed. /// /// Unlike the iterator interface, implementations should return [`BufferState::Full`] /// if **future** insertions will fail to prevent unnecessary work. - fn push(&mut self, id: I, distance: D) -> BufferState; + fn push(&mut self, neighbor: Neighbor) -> BufferState; /// Return the number of items pushed into the buffer. fn current_len(&self) -> usize; @@ -33,7 +35,7 @@ pub trait SearchOutputBuffer { /// The entire iterator may not be consumed if there is insufficient capacity in `self`. fn extend(&mut self, itr: Itr) -> usize where - Itr: IntoIterator; + Itr: IntoIterator>; } /// Indicate whether future calls to [`SearchOutputBuffer::push`] will succeed or not. @@ -102,11 +104,13 @@ impl SearchOutputBuffer for IdDistance<'_, I> { Some(self.capacity() - self.position) } - fn push(&mut self, id: I, distance: f32) -> BufferState { + fn push(&mut self, neighbor: Neighbor) -> BufferState { if self.position == self.capacity() { return BufferState::Full; } + let (id, distance) = neighbor.as_tuple(); + self.ids[self.position] = id; self.distances[self.position] = distance; self.position += 1; @@ -125,7 +129,7 @@ impl SearchOutputBuffer for IdDistance<'_, I> { fn extend(&mut self, itr: Itr) -> usize where - Itr: IntoIterator, + Itr: IntoIterator>, { let mut i = 0; let p = self.position; @@ -134,7 +138,8 @@ impl SearchOutputBuffer for IdDistance<'_, I> { self.distances.iter_mut().skip(p), ) .zip(itr) - .for_each(|((i_out, d_out), (i_in, d_in))| { + .for_each(|((i_out, d_out), neighbor)| { + let (i_in, d_in) = neighbor.as_tuple(); i += 1; *i_out = i_in; *d_out = d_in; @@ -184,12 +189,12 @@ impl SearchOutputBuffer<(I, A), f32> Some(self.capacity() - self.position) } - fn push(&mut self, item: (I, A), distance: f32) -> BufferState { + fn push(&mut self, neighbor: Neighbor<(I, A)>) -> BufferState { if self.position == self.capacity() { return BufferState::Full; } - let (id, assoc) = item; + let ((id, assoc), distance) = neighbor.as_tuple(); self.ids[self.position] = id; self.distances[self.position] = distance; self.associated_data[self.position] = assoc; @@ -209,11 +214,11 @@ impl SearchOutputBuffer<(I, A), f32> fn extend(&mut self, itr: Itr) -> usize where - Itr: IntoIterator, + Itr: IntoIterator>, { let mut i = 0; let p = self.position; - for (((id_out, dist_out), assoc_out), ((id, assoc), dist)) in self + for (((id_out, dist_out), assoc_out), neighbor) in self .ids .iter_mut() .skip(p) @@ -221,6 +226,7 @@ impl SearchOutputBuffer<(I, A), f32> .zip(self.associated_data.iter_mut().skip(p)) .zip(itr) { + let ((id, assoc), dist) = neighbor.as_tuple(); *id_out = id; *dist_out = dist; *assoc_out = assoc; @@ -262,28 +268,28 @@ mod tests { assert_eq!(buffer.size_hint(), Some(MAX_LENGTH)); assert_eq!(buffer.current_len(), 0); - assert!(buffer.push(1, 1.0).is_available()); + assert!(buffer.push(Neighbor::new(1, 1.0)).is_available()); assert_eq!(buffer.current_len(), 1); assert_eq!(buffer.size_hint(), Some(MAX_LENGTH - 1)); - assert!(buffer.push(2, 2.0).is_available()); + assert!(buffer.push(Neighbor::new(2, 2.0)).is_available()); assert_eq!(buffer.current_len(), 2); assert_eq!(buffer.size_hint(), Some(MAX_LENGTH - 2)); - assert!(buffer.push(3, 3.0).is_available()); + assert!(buffer.push(Neighbor::new(3, 3.0)).is_available()); assert_eq!(buffer.current_len(), 3); assert_eq!(buffer.size_hint(), Some(MAX_LENGTH - 3)); - assert!(buffer.push(4, 4.0).is_available()); + assert!(buffer.push(Neighbor::new(4, 4.0)).is_available()); assert_eq!(buffer.current_len(), 4); assert_eq!(buffer.size_hint(), Some(MAX_LENGTH - 4)); // This should error since further attempts will not work. - assert!(buffer.push(5, 5.0).is_full()); + assert!(buffer.push(Neighbor::new(5, 5.0)).is_full()); assert_eq!(buffer.current_len(), 5); assert_eq!(buffer.size_hint(), Some(0)); - assert!(buffer.push(6, 6.0).is_full()); + assert!(buffer.push(Neighbor::new(6, 6.0)).is_full()); assert_eq!(buffer.current_len(), 5); assert_eq!(buffer.size_hint(), Some(0)); @@ -301,15 +307,18 @@ mod tests { assert_eq!(buffer.size_hint(), Some(MAX_LENGTH)); assert_eq!(buffer.current_len(), 0); - let set = buffer.extend([(1, 1.0), (2, 2.0), (3, 3.0), (4, 4.0), (5, 5.0), (6, 6.0)]); + let set = buffer.extend( + [(1, 1.0), (2, 2.0), (3, 3.0), (4, 4.0), (5, 5.0), (6, 6.0)] + .map(Neighbor::from_tuple), + ); assert_eq!(set, MAX_LENGTH); assert_eq!(buffer.current_len(), MAX_LENGTH); assert_eq!(buffer.size_hint(), Some(0)); // Ensure that `pushing` respects the limit. - assert!(buffer.push(7, 7.0).is_full()); + assert!(buffer.push(Neighbor::new(7, 7.0)).is_full()); - let set = buffer.extend([(10, 10.0), (20, 20.0)]); + let set = buffer.extend([(10, 10.0), (20, 20.0)].map(Neighbor::from_tuple)); assert_eq!(set, 0, "no more items can be added"); assert_eq!(&ids, &[1, 2, 3, 4, 5]); @@ -322,19 +331,19 @@ mod tests { let mut distances = [0.0f32; MAX_LENGTH]; let mut buffer = IdDistance::new(&mut ids, &mut distances); - assert!(buffer.push(1, 1.0).is_available()); + assert!(buffer.push(Neighbor::new(1, 1.0)).is_available()); - let set = buffer.extend([(2, 2.0), (3, 3.0)]); + let set = buffer.extend([(2, 2.0), (3, 3.0)].map(Neighbor::from_tuple)); assert_eq!(set, 2, "only two items were pushed"); assert_eq!(buffer.current_len(), 3); assert_eq!(buffer.size_hint(), Some(2)); - assert!(buffer.push(4, 4.0).is_available()); + assert!(buffer.push(Neighbor::new(4, 4.0)).is_available()); assert_eq!(buffer.current_len(), 4); assert_eq!(buffer.size_hint(), Some(1)); - let set = buffer.extend([(5, 5.0), (6, 6.0)]); + let set = buffer.extend([(5, 5.0), (6, 6.0)].map(Neighbor::from_tuple)); assert_eq!( set, 1, "there should only be room for one more item in the buffer" @@ -342,7 +351,7 @@ mod tests { assert_eq!(buffer.current_len(), 5); assert_eq!(buffer.size_hint(), Some(0)); - assert_eq!(&ids, &[1, 2, 3, 4, 5],); + assert_eq!(&ids, &[1, 2, 3, 4, 5]); assert_eq!(&distances, &[1.0, 2.0, 3.0, 4.0, 5.0]); } } @@ -372,28 +381,28 @@ mod tests { assert_eq!(buffer.size_hint(), Some(MAX_LENGTH)); assert_eq!(buffer.current_len(), 0); - assert!(buffer.push((1, 10), 1.0).is_available()); + assert!(buffer.push(Neighbor::new((1, 10), 1.0)).is_available()); assert_eq!(buffer.current_len(), 1); assert_eq!(buffer.size_hint(), Some(MAX_LENGTH - 1)); - assert!(buffer.push((2, 20), 2.0).is_available()); + assert!(buffer.push(Neighbor::new((2, 20), 2.0)).is_available()); assert_eq!(buffer.current_len(), 2); assert_eq!(buffer.size_hint(), Some(MAX_LENGTH - 2)); - assert!(buffer.push((3, 30), 3.0).is_available()); + assert!(buffer.push(Neighbor::new((3, 30), 3.0)).is_available()); assert_eq!(buffer.current_len(), 3); assert_eq!(buffer.size_hint(), Some(MAX_LENGTH - 3)); - assert!(buffer.push((4, 40), 4.0).is_available()); + assert!(buffer.push(Neighbor::new((4, 40), 4.0)).is_available()); assert_eq!(buffer.current_len(), 4); assert_eq!(buffer.size_hint(), Some(MAX_LENGTH - 4)); // This should error since further attempts will not work. - assert!(buffer.push((5, 50), 5.0).is_full()); + assert!(buffer.push(Neighbor::new((5, 50), 5.0)).is_full()); assert_eq!(buffer.current_len(), 5); assert_eq!(buffer.size_hint(), Some(0)); - assert!(buffer.push((6, 60), 6.0).is_full()); + assert!(buffer.push(Neighbor::new((6, 60), 6.0)).is_full()); assert_eq!(buffer.current_len(), 5); assert_eq!(buffer.size_hint(), Some(0)); @@ -414,22 +423,26 @@ mod tests { assert_eq!(buffer.size_hint(), Some(MAX_LENGTH)); assert_eq!(buffer.current_len(), 0); - let set = buffer.extend([ - ((1, 10), 1.0), - ((2, 20), 2.0), - ((3, 30), 3.0), - ((4, 40), 4.0), - ((5, 50), 5.0), - ((6, 60), 6.0), - ]); + let set = buffer.extend( + [ + ((1, 10), 1.0), + ((2, 20), 2.0), + ((3, 30), 3.0), + ((4, 40), 4.0), + ((5, 50), 5.0), + ((6, 60), 6.0), + ] + .map(Neighbor::from_tuple), + ); assert_eq!(set, MAX_LENGTH); assert_eq!(buffer.current_len(), MAX_LENGTH); assert_eq!(buffer.size_hint(), Some(0)); // Ensure that `pushing` respects the limit. - assert!(buffer.push((7, 70), 7.0).is_full()); + assert!(buffer.push(Neighbor::new((7, 70), 7.0)).is_full()); - let set = buffer.extend([((10, 100), 10.0), ((20, 200), 20.0)]); + let set = + buffer.extend([((10, 100), 10.0), ((20, 200), 20.0)].map(Neighbor::from_tuple)); assert_eq!(set, 0, "no more items can be added"); assert_eq!(&ids, &[1, 2, 3, 4, 5]); @@ -445,19 +458,19 @@ mod tests { let mut buffer = IdDistanceAssociatedData::new(&mut ids, &mut distances, &mut associated); - assert!(buffer.push((1, 10), 1.0).is_available()); + assert!(buffer.push(Neighbor::new((1, 10), 1.0)).is_available()); - let set = buffer.extend([((2, 20), 2.0), ((3, 30), 3.0)]); + let set = buffer.extend([((2, 20), 2.0), ((3, 30), 3.0)].map(Neighbor::from_tuple)); assert_eq!(set, 2, "only two items were pushed"); assert_eq!(buffer.current_len(), 3); assert_eq!(buffer.size_hint(), Some(2)); - assert!(buffer.push((4, 40), 4.0).is_available()); + assert!(buffer.push(Neighbor::new((4, 40), 4.0)).is_available()); assert_eq!(buffer.current_len(), 4); assert_eq!(buffer.size_hint(), Some(1)); - let set = buffer.extend([((5, 50), 5.0), ((6, 60), 6.0)]); + let set = buffer.extend([((5, 50), 5.0), ((6, 60), 6.0)].map(Neighbor::from_tuple)); assert_eq!( set, 1, "there should only be room for one more item in the buffer" @@ -465,9 +478,9 @@ mod tests { assert_eq!(buffer.current_len(), 5); assert_eq!(buffer.size_hint(), Some(0)); - assert_eq!(&ids, &[1, 2, 3, 4, 5],); + assert_eq!(&ids, &[1, 2, 3, 4, 5]); assert_eq!(&distances, &[1.0, 2.0, 3.0, 4.0, 5.0]); - assert_eq!(&associated, &[10, 20, 30, 40, 50],); + assert_eq!(&associated, &[10, 20, 30, 40, 50]); } } } diff --git a/diskann/src/graph/test/cases/range_search.rs b/diskann/src/graph/test/cases/range_search.rs index bfb539e60..2917205eb 100644 --- a/diskann/src/graph/test/cases/range_search.rs +++ b/diskann/src/graph/test/cases/range_search.rs @@ -93,7 +93,7 @@ fn assert_no_duplicates(results: &[Neighbor]) { fn assert_range_invariants(results: &[Neighbor], radius: f32, inner_radius: Option) { for n in results { assert!( - n.distance() <= radius, + *n.distance() <= radius, "result {} distance {} exceeds radius {}", n.id(), n.distance(), @@ -101,7 +101,7 @@ fn assert_range_invariants(results: &[Neighbor], radius: f32, inner_radius: ); if let Some(inner) = inner_radius { assert!( - n.distance() > inner, + *n.distance() > inner, "result {} distance {} is within inner radius {}", n.id(), n.distance(), diff --git a/diskann/src/neighbor/mod.rs b/diskann/src/neighbor/mod.rs index 911deb3ed..3359ae1e6 100644 --- a/diskann/src/neighbor/mod.rs +++ b/diskann/src/neighbor/mod.rs @@ -65,28 +65,28 @@ pub use diverse_priority_queue::{ /// ); /// ``` #[derive(Debug, Default, Clone, Copy)] -pub struct Neighbor { +pub struct Neighbor { id: I, - distance: f32, + distance: D, } -impl Neighbor { +impl Neighbor { /// Create a [`Neighbor`] with `id` and `distance`. #[inline] - pub fn new(id: I, distance: f32) -> Self { + pub fn new(id: I, distance: D) -> Self { Self { id, distance } } /// Return the ID and distance in `self` as a tuple. #[inline] - pub fn as_tuple(self) -> (I, f32) { + pub fn as_tuple(self) -> (I, D) { (self.id, self.distance) } /// Return the distance. #[inline] - pub fn distance(&self) -> f32 { - self.distance + pub fn distance(&self) -> &D { + &self.distance } /// Return the ID. @@ -94,12 +94,18 @@ impl Neighbor { pub fn id(&self) -> &I { &self.id } + + #[cfg(test)] + pub(crate) fn from_tuple((id, distance): (I, D)) -> Self { + Self::new(id, distance) + } } #[cfg(test)] -impl crate::test::cmp::VerboseEq for Neighbor +impl crate::test::cmp::VerboseEq for Neighbor where I: crate::test::cmp::VerboseEq, + D: crate::test::cmp::VerboseEq, { #[inline(never)] #[track_caller] @@ -143,7 +149,7 @@ pub mod ord { /// ``` pub fn fast_distance(x: &Neighbor, y: &Neighbor) -> std::cmp::Ordering { x.distance() - .partial_cmp(&y.distance()) + .partial_cmp(y.distance()) .unwrap_or(std::cmp::Ordering::Equal) } @@ -203,19 +209,19 @@ pub mod ord { } } -/// A [`SearchOutputBuffer`] wrapper around `&mut [Neighbor]`. This can be used to +/// A [`SearchOutputBuffer`] wrapper around `&mut [Neighbor]`. This can be used to /// populate such a mutable slice as the result of [`crate::graph::DiskANNIndex::search`]. #[derive(Debug)] -pub struct BackInserter<'a, I> { - buffer: &'a mut [Neighbor], +pub struct BackInserter<'a, I, D = f32> { + buffer: &'a mut [Neighbor], position: usize, } -impl<'a, I> BackInserter<'a, I> { +impl<'a, I, D> BackInserter<'a, I, D> { /// Construct a new [`BackInserter`] around the provided slice. /// /// The buffer will have a capacity equal to the length of `buffer`. - pub fn new(buffer: &'a mut [Neighbor]) -> Self { + pub fn new(buffer: &'a mut [Neighbor]) -> Self { Self { buffer, position: 0, @@ -228,19 +234,19 @@ impl<'a, I> BackInserter<'a, I> { } } -impl SearchOutputBuffer for BackInserter<'_, I> { +impl SearchOutputBuffer for BackInserter<'_, I, D> { fn size_hint(&self) -> Option { // We maintain the invariant that `self.position <= self.buffer.len()`, so this // subtraction should not underflow. Some(self.buffer.len() - self.position) } - fn push(&mut self, id: I, distance: f32) -> search_output_buffer::BufferState { + fn push(&mut self, neighbor: Neighbor) -> search_output_buffer::BufferState { if self.position == self.buffer.len() { return search_output_buffer::BufferState::Full; } - self.buffer[self.position] = Neighbor::new(id, distance); + self.buffer[self.position] = neighbor; self.position += 1; // Return `Full` if we added the last item. @@ -257,15 +263,13 @@ impl SearchOutputBuffer for BackInserter<'_, I> { fn extend(&mut self, itr: Itr) -> usize where - Itr: IntoIterator, + Itr: IntoIterator>, { let mut i = 0; - std::iter::zip(self.buffer.iter_mut().skip(self.position), itr).for_each( - |(neighbor, (id, distance))| { - i += 1; - *neighbor = Neighbor::new(id, distance); - }, - ); + std::iter::zip(self.buffer.iter_mut().skip(self.position), itr).for_each(|(dst, src)| { + i += 1; + *dst = src; + }); self.position += i; @@ -273,13 +277,13 @@ impl SearchOutputBuffer for BackInserter<'_, I> { } } -impl SearchOutputBuffer for Vec> { +impl SearchOutputBuffer for Vec> { fn size_hint(&self) -> Option { None } - fn push(&mut self, id: I, distance: f32) -> search_output_buffer::BufferState { - self.push(Neighbor::new(id, distance)); + fn push(&mut self, neighbor: Neighbor) -> search_output_buffer::BufferState { + self.push(neighbor); search_output_buffer::BufferState::Available } @@ -289,13 +293,10 @@ impl SearchOutputBuffer for Vec> { fn extend(&mut self, itr: Itr) -> usize where - Itr: IntoIterator, + Itr: IntoIterator>, { let before = self.len(); - Extend::extend( - self, - itr.into_iter().map(|(id, dist)| Neighbor::new(id, dist)), - ); + Extend::extend(self, itr); self.len() - before } } @@ -381,28 +382,28 @@ mod neighbor_test { assert_eq!(inserter.size_hint(), Some(MAX_LENGTH)); assert_eq!(inserter.current_len(), 0); - assert!(inserter.push(1, 1.0).is_available()); + assert!(inserter.push(Neighbor::new(1, 1.0)).is_available()); assert_eq!(inserter.current_len(), 1); assert_eq!(inserter.size_hint(), Some(MAX_LENGTH - 1)); - assert!(inserter.push(2, 2.0).is_available()); + assert!(inserter.push(Neighbor::new(2, 2.0)).is_available()); assert_eq!(inserter.current_len(), 2); assert_eq!(inserter.size_hint(), Some(MAX_LENGTH - 2)); - assert!(inserter.push(3, 3.0).is_available()); + assert!(inserter.push(Neighbor::new(3, 3.0)).is_available()); assert_eq!(inserter.current_len(), 3); assert_eq!(inserter.size_hint(), Some(MAX_LENGTH - 3)); - assert!(inserter.push(4, 4.0).is_available()); + assert!(inserter.push(Neighbor::new(4, 4.0)).is_available()); assert_eq!(inserter.current_len(), 4); assert_eq!(inserter.size_hint(), Some(MAX_LENGTH - 4)); // This should error since further attempts will not work. - assert!(inserter.push(5, 5.0).is_full()); + assert!(inserter.push(Neighbor::new(5, 5.0)).is_full()); assert_eq!(inserter.current_len(), 5); assert_eq!(inserter.size_hint(), Some(0)); - assert!(inserter.push(6, 6.0).is_full()); + assert!(inserter.push(Neighbor::new(6, 6.0)).is_full()); assert_eq!(inserter.current_len(), 5); assert_eq!(inserter.size_hint(), Some(0)); @@ -417,15 +418,18 @@ mod neighbor_test { assert_eq!(inserter.size_hint(), Some(MAX_LENGTH)); assert_eq!(inserter.current_len(), 0); - let set = inserter.extend([(1, 1.0), (2, 2.0), (3, 3.0), (4, 4.0), (5, 5.0), (6, 6.0)]); + let set = inserter.extend( + [(1, 1.0), (2, 2.0), (3, 3.0), (4, 4.0), (5, 5.0), (6, 6.0)] + .map(Neighbor::from_tuple), + ); assert_eq!(set, MAX_LENGTH); assert_eq!(inserter.current_len(), MAX_LENGTH); assert_eq!(inserter.size_hint(), Some(0)); // Ensure that `pushing` respects the limit. - assert!(inserter.push(7, 7.0).is_full()); + assert!(inserter.push(Neighbor::new(7, 7.0)).is_full()); - let set = inserter.extend([(10, 10.0), (20, 20.0)]); + let set = inserter.extend([(10, 10.0), (20, 20.0)].map(Neighbor::from_tuple)); assert_eq!(set, 0, "no more items can be added"); assert_eq_verbose!(buffer, [f(1), f(2), f(3), f(4), f(5)]); @@ -436,19 +440,19 @@ mod neighbor_test { let mut buffer = [Neighbor::::default(); MAX_LENGTH]; let mut inserter = BackInserter::new(&mut buffer); - assert!(inserter.push(1, 1.0).is_available()); + assert!(inserter.push(Neighbor::new(1, 1.0)).is_available()); - let set = inserter.extend([(2, 2.0), (3, 3.0)]); + let set = inserter.extend([(2, 2.0), (3, 3.0)].map(Neighbor::from_tuple)); assert_eq!(set, 2, "only two items were pushed"); assert_eq!(inserter.current_len(), 3); assert_eq!(inserter.size_hint(), Some(2)); - assert!(inserter.push(4, 4.0).is_available()); + assert!(inserter.push(Neighbor::new(4, 4.0)).is_available()); assert_eq!(inserter.current_len(), 4); assert_eq!(inserter.size_hint(), Some(1)); - let set = inserter.extend([(5, 5.0), (6, 6.0)]); + let set = inserter.extend([(5, 5.0), (6, 6.0)].map(Neighbor::from_tuple)); assert_eq!( set, 1, "there should only be room for one more item in the buffer" @@ -469,14 +473,17 @@ mod neighbor_test { assert_eq!(SearchOutputBuffer::::current_len(&buf), 0); // push grows unboundedly - assert!(SearchOutputBuffer::push(&mut buf, 1, 0.5).is_available()); - assert!(SearchOutputBuffer::push(&mut buf, 2, 1.0).is_available()); + assert!(SearchOutputBuffer::push(&mut buf, Neighbor::new(1, 0.5)).is_available()); + assert!(SearchOutputBuffer::push(&mut buf, Neighbor::new(2, 1.0)).is_available()); assert_eq!(SearchOutputBuffer::::current_len(&buf), 2); assert_eq_verbose!(buf[0], Neighbor::new(1, 0.5)); assert_eq_verbose!(buf[1], Neighbor::new(2, 1.0)); // extend appends and returns count - let count = SearchOutputBuffer::extend(&mut buf, vec![(3u32, 1.5), (4, 2.0), (5, 2.5)]); + let count = SearchOutputBuffer::extend( + &mut buf, + [(3u32, 1.5), (4, 2.0), (5, 2.5)].map(Neighbor::from_tuple), + ); assert_eq!(count, 3); assert_eq!(SearchOutputBuffer::::current_len(&buf), 5); assert_eq_verbose!(buf[4], Neighbor::new(5, 2.5));