From 9a2ba74d1e0dd401770f605de1bb7ac0f7cd4e91 Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Thu, 23 Jul 2026 17:32:12 -0700 Subject: [PATCH 1/2] Initial cleanups. --- diskann-benchmark/src/exhaustive/algos.rs | 3 +- diskann-bftree/src/provider.rs | 14 +- .../src/search/provider/disk_provider.rs | 12 +- diskann-garnet/src/provider.rs | 10 +- .../inline_beta_search/inline_beta_filter.rs | 19 +- diskann-providers/src/index/diskann_async.rs | 22 +- diskann-providers/src/index/wrapped_async.rs | 7 +- .../provider/async_/inmem/full_precision.rs | 12 +- .../graph/provider/async_/postprocess.rs | 4 +- .../src/test_utils/search_utils.rs | 28 +- diskann-tools/src/utils/ground_truth.rs | 14 +- diskann/src/flat/test/harness.rs | 17 +- diskann/src/graph/glue.rs | 19 +- diskann/src/graph/index.rs | 26 +- .../src/graph/internal/sorted_neighbors.rs | 75 +-- .../src/graph/search/inline_filter_search.rs | 10 +- .../graph/search/multihop_filter_search.rs | 18 +- diskann/src/graph/search/paged.rs | 2 +- diskann/src/graph/search/range_search.rs | 24 +- diskann/src/graph/search/record.rs | 4 +- diskann/src/graph/test/cases/multihop.rs | 10 +- diskann/src/graph/test/cases/paged_search.rs | 16 +- diskann/src/graph/test/cases/range_search.rs | 24 +- diskann/src/graph/test/provider.rs | 2 +- diskann/src/neighbor/mod.rs | 515 ++++++++++-------- diskann/src/neighbor/queue.rs | 4 +- 26 files changed, 482 insertions(+), 429 deletions(-) diff --git a/diskann-benchmark/src/exhaustive/algos.rs b/diskann-benchmark/src/exhaustive/algos.rs index 1d7d56989..2a973a4d0 100644 --- a/diskann-benchmark/src/exhaustive/algos.rs +++ b/diskann-benchmark/src/exhaustive/algos.rs @@ -80,7 +80,8 @@ where let search = start.elapsed().into(); - std::iter::zip(o.iter_mut(), queue.iter()).for_each(|(o, neighbor)| *o = neighbor.id); + std::iter::zip(o.iter_mut(), queue.iter()) + .for_each(|(o, neighbor)| *o = *neighbor.id()); progress.inc(1); Ok(Times { preprocess, search }) }) diff --git a/diskann-bftree/src/provider.rs b/diskann-bftree/src/provider.rs index 1a9f773ca..4fd90a79f 100644 --- a/diskann-bftree/src/provider.rs +++ b/diskann-bftree/src/provider.rs @@ -1540,11 +1540,11 @@ where for n in candidates { match provider .full_vectors - .get_vector_sync(n.id.into_usize()) + .get_vector_sync(n.id().into_usize()) .allow_transient("stale candidate during rerank") { Ok(Some(vec)) => { - reranked.push((n.id, f.evaluate_similarity(query, &vec))); + reranked.push((*n.id(), f.evaluate_similarity(query, &vec))); } Ok(None) => { // Transient (deleted/missing) — skip this candidate. @@ -2092,7 +2092,7 @@ mod tests { res.result_count, 5, "there are 15 points and we're asking for 5, we expect 5" ); - assert_eq!(neighbors[0].id, 3); + assert_eq!(*neighbors[0].id(), 3); } #[tokio::test] @@ -2143,7 +2143,7 @@ mod tests { res.result_count, 5, "there are 15 points and we're asking for 5, we expect 5" ); - let neighbor_ids: Vec = neighbors.iter().map(|n| n.id).collect(); + let neighbor_ids: Vec = neighbors.iter().map(|n| *n.id()).collect(); for expected in 1u32..=5 { assert!( neighbor_ids.contains(&expected), @@ -2190,7 +2190,7 @@ mod tests { .unwrap(); assert_eq!(res.result_count, 5); - let neighbor_ids: Vec = neighbors.iter().map(|n| n.id).collect(); + let neighbor_ids: Vec = neighbors.iter().map(|n| *n.id()).collect(); assert!(!neighbor_ids.contains(&2u32)); assert!(!neighbor_ids.contains(&4u32)); } @@ -2264,7 +2264,7 @@ mod tests { res.result_count, 5, "there are 15 points and we're asking for 5, we expect 5" ); - assert_eq!(neighbors[0].id, 3); + assert_eq!(*neighbors[0].id(), 3); } #[tokio::test] @@ -2317,7 +2317,7 @@ mod tests { .unwrap(); assert_eq!(res.result_count, 5); - let neighbor_ids: Vec = neighbors.iter().map(|n| n.id).collect(); + let neighbor_ids: Vec = neighbors.iter().map(|n| *n.id()).collect(); assert!(!neighbor_ids.contains(&2u32)); assert!(!neighbor_ids.contains(&4u32)); } diff --git a/diskann-disk/src/search/provider/disk_provider.rs b/diskann-disk/src/search/provider/disk_provider.rs index c497414a6..a5ebf2549 100644 --- a/diskann-disk/src/search/provider/disk_provider.rs +++ b/diskann-disk/src/search/provider/disk_provider.rs @@ -360,11 +360,11 @@ where }; match self.filter { PostprocessStrategy::AcceptAll => candidates - .map(|n| n.id) + .map(|n| *n.id()) .filter_map(&mut process) .collect::, _>>()?, PostprocessStrategy::Apply(f) => candidates - .map(|n| n.id) + .map(|n| *n.id()) .filter(|id| f(id)) .filter_map(&mut process) .collect::, _>>()?, @@ -419,9 +419,9 @@ where let query_f32 = Data::VectorDataType::as_f32(query).map_err(Into::into)?; let candidate_ids: Vec = match self.filter { - PostprocessStrategy::AcceptAll => candidates.map(|candidate| candidate.id).collect(), + PostprocessStrategy::AcceptAll => candidates.map(|candidate| *candidate.id()).collect(), PostprocessStrategy::Apply(f) => candidates - .map(|candidate| candidate.id) + .map(|candidate| *candidate.id()) .filter(|id| f(id)) .collect(), }; @@ -1805,7 +1805,7 @@ mod disk_provider_tests { let ids = search_record .visited .iter() - .map(|n| n.id) + .map(|n| *n.id()) .collect::>(); const EXPECTED_NODES: [u32; 18] = [ @@ -2503,7 +2503,7 @@ mod disk_provider_tests { let visited_ids = search_record .visited .iter() - .map(|n| n.id) + .map(|n| *n.id()) .collect::>(); let query_stats = strategy.io_tracker; diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index 5024b7b9d..3468c2c90 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -1227,12 +1227,12 @@ impl<'a, T: VectorRepr> SearchPostProcess, &[T], GarnetId { let initial = output.current_len(); for n in candidates { - let id = match accessor.provider.to_external_id(accessor.context, n.id) { + let id = match accessor.provider.to_external_id(accessor.context, *n.id()) { Ok(id) => id, Err(_) => continue, // Can't read the mapping; skip. }; - if output.push(id, n.distance).is_full() { + if output.push(id, n.distance()).is_full() { break; } } @@ -1284,15 +1284,15 @@ impl<'a, 'b, T: VectorRepr> SearchPostProcessStep, &'b [T // Filter before computing the full precision distances. let mut reranked: Vec<(u32, f32)> = candidates .filter_map(|n| { - if !provider.vector_iid_exists(accessor.context, n.id) { + if !provider.vector_iid_exists(accessor.context, *n.id()) { None } else if provider.callbacks.read_single_iid( &accessor.context.term(Term::Vector), - n.id, + *n.id(), &mut v, ) { Some(( - n.id, + *n.id(), f.evaluate_similarity(query, bytemuck::cast_slice::(&v)), )) } else { diff --git a/diskann-label-filter/src/inline_beta_search/inline_beta_filter.rs b/diskann-label-filter/src/inline_beta_search/inline_beta_filter.rs index 2e3c7497a..c9c67c371 100644 --- a/diskann-label-filter/src/inline_beta_search/inline_beta_filter.rs +++ b/diskann-label-filter/src/inline_beta_search/inline_beta_filter.rs @@ -156,14 +156,17 @@ where // TODO: Fix for performance. let mut filtered_candidates = Vec::>::new(); for candidate in candidates { - accessor.attributes_for(candidate.id, |computer, attributes| -> ANNResult<()> { - let pe = PredicateEvaluator::new(&*attributes); - - if computer.filter_expr().encoded_filter_expr().accept(&pe)? { - filtered_candidates.push(Neighbor::new(candidate.id, candidate.distance)); - } - Ok(()) - })??; + accessor.attributes_for( + *candidate.id(), + |computer, attributes| -> ANNResult<()> { + let pe = PredicateEvaluator::new(&*attributes); + + if computer.filter_expr().encoded_filter_expr().accept(&pe)? { + filtered_candidates.push(candidate); + } + Ok(()) + }, + )??; } // Assuming that the job of the post processor is to only forward the right set of diff --git a/diskann-providers/src/index/diskann_async.rs b/diskann-providers/src/index/diskann_async.rs index 41389910e..f22fb889f 100644 --- a/diskann-providers/src/index/diskann_async.rs +++ b/diskann-providers/src/index/diskann_async.rs @@ -171,7 +171,7 @@ pub(crate) mod tests { search::Range, search_output_buffer, }, - neighbor::Neighbor, + neighbor::{self, Neighbor}, provider::{ DataProvider, DefaultContext, Delete, ExecutionContext, Guard, NeighborAccessor, NeighborAccessorMut, SetElement, @@ -916,9 +916,9 @@ pub(crate) mod tests { let checker = |position, (id, distance)| -> Result<(), Box> { let expected: Neighbor = gt[gt.len() - 1 - position]; - if id != expected.id { + 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!( @@ -926,7 +926,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,11 +1085,11 @@ pub(crate) mod tests { let gt = { 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.distance *= beta; + if filter.is_match(*n.id()) { + *n = Neighbor::new(*n.id(), n.distance() * beta); } } - gt.sort_unstable_by(|a, b| a.cmp(b).reverse()); + gt.sort_unstable_by(neighbor::ord::reverse(neighbor::ord::fast_distance)); gt }; @@ -1528,7 +1528,7 @@ pub(crate) mod tests { .await .unwrap(); - let ids: Vec = results.iter().map(|n| n.id).collect(); + let ids: Vec = results.iter().map(|n| *n.id()).collect(); assert_range_results_exactly_match(q, >, &ids, radius, None); } @@ -1541,7 +1541,7 @@ pub(crate) mod tests { .await .unwrap(); - let ids: Vec = results.iter().map(|n| n.id).collect(); + let ids: Vec = results.iter().map(|n| *n.id()).collect(); assert_range_results_exactly_match(q, >, &ids, radius, None); } @@ -1565,7 +1565,7 @@ pub(crate) mod tests { .await .unwrap(); - let ids: Vec = results.iter().map(|n| n.id).collect(); + let ids: Vec = results.iter().map(|n| *n.id()).collect(); assert_range_results_exactly_match(q, >, &ids, radius, Some(inner_radius)); } @@ -1582,7 +1582,7 @@ pub(crate) mod tests { // check that ids don't have duplicates let mut ids_set = std::collections::HashSet::new(); for n in &results { - assert!(ids_set.insert(n.id)); + assert!(ids_set.insert(*n.id())); } } } diff --git a/diskann-providers/src/index/wrapped_async.rs b/diskann-providers/src/index/wrapped_async.rs index 46084b6b1..78ba0e9a4 100644 --- a/diskann-providers/src/index/wrapped_async.rs +++ b/diskann-providers/src/index/wrapped_async.rs @@ -836,16 +836,17 @@ mod tests { for neighbor in v { assert_ne!( - neighbor.id, + *neighbor.id(), u32::MAX, "paged search should not return start point", ); assert_eq!( - neighbor.id, i, + *neighbor.id(), + i, "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 15c2d6af1..8437e5e90 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 @@ -381,13 +381,13 @@ where // Filter before computing the full precision distances. let mut reranked: Vec<(u32, f32)> = candidates .filter_map(|n| { - if checker.deletion_check(n.id) { + if checker.deletion_check(*n.id()) { None } else { Some(( - n.id, + *n.id(), f.evaluate_similarity(query, unsafe { - full.get_vector_sync(n.id.into_usize()) + full.get_vector_sync(n.id().into_usize()) }), )) } @@ -445,9 +445,9 @@ where for (i, candidate) in candidates.into_iter().enumerate() { // SAFETY: We accept potential unsynchronized concurrent mutation, matching the // pattern used by `Rerank` above. - let vector = unsafe { store.get_vector_sync(candidate.id.into_usize()) }; - ids.push(candidate.id); - distances.push(candidate.distance); + let vector = unsafe { store.get_vector_sync(candidate.id().into_usize()) }; + ids.push(*candidate.id()); + distances.push(candidate.distance()); vectors.row_mut(i).copy_from_slice(vector); } diff --git a/diskann-providers/src/model/graph/provider/async_/postprocess.rs b/diskann-providers/src/model/graph/provider/async_/postprocess.rs index 94dce287e..75505fe2e 100644 --- a/diskann-providers/src/model/graph/provider/async_/postprocess.rs +++ b/diskann-providers/src/model/graph/provider/async_/postprocess.rs @@ -56,10 +56,10 @@ where { let checker = accessor.as_deletion_check(); let count = output.extend(candidates.filter_map(|n| { - if checker.deletion_check(n.id) { + if checker.deletion_check(*n.id()) { None } else { - Some((n.id, n.distance)) + Some(n.as_tuple()) } })); 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 1cf74fd22..1ce24bbd5 100644 --- a/diskann-providers/src/test_utils/search_utils.rs +++ b/diskann-providers/src/test_utils/search_utils.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use diskann::neighbor::Neighbor; +use diskann::neighbor::{self, Neighbor}; use diskann_utils::views::MatrixView; /// Compute the ground truth for a small dataset. @@ -23,7 +23,7 @@ where .map(|(i, row)| Neighbor::new(i as u32, f(row, query))) .collect(); - results.sort_unstable_by(|a, b| a.cmp(b).reverse()); + results.sort_unstable_by(neighbor::ord::reverse(neighbor::ord::fast_distance)); results } @@ -43,10 +43,10 @@ pub fn is_match( for i in (0..groundtruth.len()).rev() { // Check if the distance matches. let gt = groundtruth[i]; - if (gt.distance - neighbor.distance).abs() > margin { + if (gt.distance() - neighbor.distance()).abs() > margin { return None; } - if gt.id == neighbor.id { + if gt.id() == neighbor.id() { return Some(i); } } @@ -70,14 +70,18 @@ pub fn assert_top_k_exactly_match( for i in 0..top_k { let neighbor = gt[gt.len() - 1 - i]; assert_eq!( - neighbor.distance, distances[i], + neighbor.distance(), + distances[i], "failed on query {} for result {}", - query_id, i + query_id, + i ); assert_eq!( - neighbor.id, ids[i], + *neighbor.id(), + ids[i], "failed on query {} for result {}", - query_id, i + query_id, + i ); } } @@ -96,13 +100,13 @@ 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) - .map(|nbh| nbh.id) + .filter(|nbh| nbh.distance() >= inner_radius && nbh.distance() <= radius) + .map(|nbh| *nbh.id()) .collect::>() } else { gt.iter() - .filter(|nbh| nbh.distance <= radius) - .map(|nbh| nbh.id) + .filter(|nbh| nbh.distance() <= radius) + .map(|nbh| *nbh.id()) .collect::>() }; if ids.iter().any(|id| !gt_ids.contains(id)) { diff --git a/diskann-tools/src/utils/ground_truth.rs b/diskann-tools/src/utils/ground_truth.rs index 77c2f663c..0ed2094ed 100644 --- a/diskann-tools/src/utils/ground_truth.rs +++ b/diskann-tools/src/utils/ground_truth.rs @@ -340,7 +340,7 @@ where if allowed_by_bitmap { let distance = distance_comparer.evaluate_similarity(data, query); if distance <= radius { - query_results.push(Neighbor { id: idx, distance }); + query_results.push(Neighbor::new(idx, distance)); } } } @@ -502,7 +502,7 @@ fn write_range_search_ground_truth( // In the file, we write the neighbor IDs array first, then write the distances array. 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_ids.push(*closest_node.id()); + gt_distances.push(closest_node.distance()); } } @@ -654,7 +654,7 @@ where if allowed_by_bitmap { let distance = distance_comparer.evaluate_similarity(data, query); - neighbor_queue.insert(Neighbor { id: idx, distance }); + neighbor_queue.insert(Neighbor::new(idx, distance)); } } }, @@ -683,7 +683,7 @@ where if allowed_by_bitmap { let distance = distance_comparer.evaluate_similarity(&data_vector, query); - neighbor_queue.insert(Neighbor { id: idx, distance }) + neighbor_queue.insert(Neighbor::new(idx, distance)) } } } @@ -777,7 +777,7 @@ where }; // insert into neighbor queue let idx = idx_base as u32; - neighbor_queue.insert(Neighbor { id: idx, distance }); + neighbor_queue.insert(Neighbor::new(idx, distance)); } } }, diff --git a/diskann/src/flat/test/harness.rs b/diskann/src/flat/test/harness.rs index 2f1efb4cb..6eaaa40ae 100644 --- a/diskann/src/flat/test/harness.rs +++ b/diskann/src/flat/test/harness.rs @@ -22,7 +22,7 @@ use crate::{ SearchOutputBuffer, glue::{CopyIds, SearchPostProcess}, }, - neighbor::{BackInserter, Neighbor}, + neighbor::{self, BackInserter, Neighbor}, provider::HasId, test::tokio::current_thread_runtime, utils::VectorRepr, @@ -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)); @@ -156,8 +156,8 @@ where { let count = output.extend( candidates - .filter(|n| n.id % 2 == 0) - .map(|n| (n.id, n.distance)), + .filter(|n| *n.id() % 2 == 0) + .map(|n| n.as_tuple()), ); std::future::ready(Ok(count)) } @@ -175,7 +175,7 @@ impl OracleProcessor for EvenIdsOnlyOracle { } fn expected(&self, gt: Vec>) -> Vec> { - gt.into_iter().filter(|n| n.id % 2 == 0).collect() + gt.into_iter().filter(|n| *n.id() % 2 == 0).collect() } } @@ -203,10 +203,5 @@ pub(crate) fn brute_force_topk( /// Sort a slice of [`Neighbor`] by `(distance asc, id asc)`. fn sort_neighbors(neighbors: &mut [Neighbor]) { - neighbors.sort_by(|a, b| { - a.distance - .partial_cmp(&b.distance) - .unwrap_or(Ordering::Equal) - .then(a.id.cmp(&b.id)) - }); + neighbors.sort_by(neighbor::ord::fast_distance_total); } diff --git a/diskann/src/graph/glue.rs b/diskann/src/graph/glue.rs index 91cd7643d..c92e151a6 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.id, n.distance))); + let count = output.extend(candidates.map(|n| n.as_tuple())); std::future::ready(Ok(count)) } } @@ -758,12 +758,17 @@ where Next: SearchPostProcess + Sync, { let filter = accessor.is_not_start_point().await?; - next.post_process(accessor, query, candidates.filter(|n| filter(n.id)), output) - .await - .map_err(|err| { - let err = err.into(); - err.context("after filtering start points") - }) + next.post_process( + accessor, + query, + candidates.filter(|n| filter(*n.id())), + output, + ) + .await + .map_err(|err| { + let err = err.into(); + err.context("after filtering start points") + }) } } diff --git a/diskann/src/graph/index.rs b/diskann/src/graph/index.rs index 37d5bc946..1df48072f 100644 --- a/diskann/src/graph/index.rs +++ b/diskann/src/graph/index.rs @@ -1218,7 +1218,7 @@ where let mut undeleted_ids: Vec<_> = output .iter() .take(num_results) - .map(|neighbor| neighbor.id) + .map(|neighbor| *neighbor.id()) .collect(); // Collect IDs whose adjacency lists need to be updated. @@ -1710,8 +1710,8 @@ where )); } - pool.sort_unstable(); - let best = pool.iter().take(num_to_replace).map(|x| x.id).collect(); + pool.sort_unstable_by(neighbor::ord::fast_distance); + let best = pool.iter().take(num_to_replace).map(|x| *x.id()).collect(); edges_to_add.insert(*neighbor, best); } @@ -1740,10 +1740,10 @@ where )); } - pool.sort_unstable(); + pool.sort_unstable_by(neighbor::ord::fast_distance); pool.iter().take(num_to_replace).for_each(|n| { edges_to_add - .entry(n.id) + .entry(*n.id()) .or_insert_with(Vec::new) .push(neighbor); }); @@ -1972,7 +1972,7 @@ where && let Some(closest_node) = scratch.best.closest_notvisited() { search_record.record(closest_node, scratch.hops, scratch.cmps); - scratch.beam_nodes.push(closest_node.id); + scratch.beam_nodes.push(*closest_node.id()); } neighbors.clear(); @@ -2368,7 +2368,7 @@ where } let (view, computer) = accessor - .fill(context.pool.iter().map(|n| n.id)) + .fill(context.pool.iter().map(|n| *n.id())) .send() .await?; @@ -2613,11 +2613,11 @@ where .iter() .map(|neighbor| { // Filter out self loops. - let id = &neighbor.id; + 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(); @@ -2759,7 +2759,7 @@ where let mut guard = neighbors.resize(found); std::iter::zip(guard.iter_mut(), states.iter()).for_each(|(d, s)| { - *d = pool[s.neighbor.into_usize()].id; + *d = *pool[s.neighbor.into_usize()].id(); }); guard.finish(found); @@ -2772,10 +2772,10 @@ where break; } - if !exclude(neighbor.id) { + if !exclude(*neighbor.id()) { // `AdjacencyList` filters out duplicates. No need to explicitly // check. - neighbors.push(neighbor.id); + neighbors.push(*neighbor.id()); } } } diff --git a/diskann/src/graph/internal/sorted_neighbors.rs b/diskann/src/graph/internal/sorted_neighbors.rs index dd6dc97f2..c40677115 100644 --- a/diskann/src/graph/internal/sorted_neighbors.rs +++ b/diskann/src/graph/internal/sorted_neighbors.rs @@ -7,7 +7,7 @@ use std::ops::Deref; -use crate::neighbor::Neighbor; +use crate::neighbor::{self, Neighbor}; /// A utility that asserts the contained neighbors are sorted by distance. #[derive(Debug)] @@ -34,8 +34,9 @@ where // `neighbors.len() == 0`. In either case, the resulting slice will be empty // and there's no actual work to be done. if let Some(position) = max.min(neighbors.len()).checked_sub(1) { - let (prefix, _, _) = neighbors.select_nth_unstable(position); - prefix.sort_unstable() + let (prefix, _, _) = + neighbors.select_nth_unstable_by(position, neighbor::ord::fast_distance); + prefix.sort_unstable_by(neighbor::ord::fast_distance) } neighbors.truncate(max); @@ -72,38 +73,38 @@ mod tests { } } - #[test] - fn test_sorted_neighbors() { - let reference = [ - Neighbor::new(1, 0.1), - Neighbor::new(2, 0.2), - Neighbor::new(3, 0.3), - Neighbor::new(4, 0.4), - Neighbor::new(5, 0.5), - Neighbor::new(6, 0.6), - Neighbor::new(7, 0.7), - Neighbor::new(8, 0.8), - Neighbor::new(9, 0.9), - Neighbor::new(10, 1.0), - ]; - - let mut rng = StdRng::seed_from_u64(0xd6152fb91c744f54); - - let ntrials = 10; - for max in 0..reference.len() + 2 { - for _ in 0..ntrials { - let mut shuffled = reference.to_vec(); - shuffled.shuffle(&mut rng); - - let sorted = SortedNeighbors::new(&mut shuffled, max); - - let expected_len = reference.len().min(max); - assert_eq!(sorted.len(), expected_len); - assert_eq!(&sorted[..expected_len], &reference[..expected_len],); - - // Changes are visible on the taken vector. - assert_eq!(shuffled.len(), expected_len) - } - } - } + // #[test] + // fn test_sorted_neighbors() { + // let reference = [ + // Neighbor::new(1, 0.1), + // Neighbor::new(2, 0.2), + // Neighbor::new(3, 0.3), + // Neighbor::new(4, 0.4), + // Neighbor::new(5, 0.5), + // Neighbor::new(6, 0.6), + // Neighbor::new(7, 0.7), + // Neighbor::new(8, 0.8), + // Neighbor::new(9, 0.9), + // Neighbor::new(10, 1.0), + // ]; + + // let mut rng = StdRng::seed_from_u64(0xd6152fb91c744f54); + + // let ntrials = 10; + // for max in 0..reference.len() + 2 { + // for _ in 0..ntrials { + // let mut shuffled = reference.to_vec(); + // shuffled.shuffle(&mut rng); + + // let sorted = SortedNeighbors::new(&mut shuffled, max); + + // let expected_len = reference.len().min(max); + // assert_eq!(sorted.len(), expected_len); + // assert_eq!(&sorted[..expected_len], &reference[..expected_len],); + + // // Changes are visible on the taken vector. + // assert_eq!(shuffled.len(), expected_len) + // } + // } + // } } diff --git a/diskann/src/graph/search/inline_filter_search.rs b/diskann/src/graph/search/inline_filter_search.rs index c535c87b6..d2e84bb08 100644 --- a/diskann/src/graph/search/inline_filter_search.rs +++ b/diskann/src/graph/search/inline_filter_search.rs @@ -18,7 +18,7 @@ use crate::{ search::record::NoopSearchRecord, search_output_buffer::SearchOutputBuffer, }, - neighbor::Neighbor, + neighbor::{self, Neighbor}, provider::DataProvider, utils::VectorId, }; @@ -223,7 +223,7 @@ where break; }; search_record.record(closest_node, scratch.hops, scratch.cmps); - scratch.beam_nodes.push(closest_node.id); + scratch.beam_nodes.push(*closest_node.id()); } // Exit if no nodes to process @@ -276,11 +276,7 @@ where } } - matched_results.sort_unstable_by(|a, b| { - a.distance - .partial_cmp(&b.distance) - .unwrap_or(std::cmp::Ordering::Equal) - }); + matched_results.sort_unstable_by(neighbor::ord::fast_distance); Ok(Ret { cmps: scratch.cmps, diff --git a/diskann/src/graph/search/multihop_filter_search.rs b/diskann/src/graph/search/multihop_filter_search.rs index f5c32d207..f8d9f28e3 100644 --- a/diskann/src/graph/search/multihop_filter_search.rs +++ b/diskann/src/graph/search/multihop_filter_search.rs @@ -19,7 +19,7 @@ use crate::{ search::record::NoopSearchRecord, search_output_buffer::SearchOutputBuffer, }, - neighbor::Neighbor, + neighbor::{self, Neighbor}, provider::DataProvider, utils::VectorId, }; @@ -91,7 +91,7 @@ where scratch .best .iter() - .filter(|n| !ret.rejected_start_points.contains(&n.id)) + .filter(|n| !ret.rejected_start_points.contains(n.id())) .take(self.inner.l_value().get()), output, ) @@ -174,7 +174,7 @@ where && let Some(closest_node) = scratch.best.closest_notvisited() { search_record.record(closest_node, scratch.hops, scratch.cmps); - scratch.beam_nodes.push(closest_node.id); + scratch.beam_nodes.push(*closest_node.id()); } // compute distances from query to one-hop neighbors, and mark them visited @@ -204,11 +204,7 @@ where scratch.hops += scratch.beam_nodes.len() as u32; // sort the candidates for two-hop expansion by distance to query point - candidates_two_hop_expansion.sort_unstable_by(|a, b| { - a.distance - .partial_cmp(&b.distance) - .unwrap_or(std::cmp::Ordering::Equal) - }); + candidates_two_hop_expansion.sort_unstable_by(neighbor::ord::fast_distance); // limit the number of two-hop candidates to avoid too many expansions candidates_two_hop_expansion.truncate(max_degree_with_slack / 2); @@ -216,8 +212,10 @@ where // Expand each two-hop candidate: if its neighbor is a match, compute its distance // to the query and insert into `scratch.visited` // If it is not a match, do nothing - let two_hop_expansion_candidate_ids: Vec = - candidates_two_hop_expansion.iter().map(|n| n.id).collect(); + let two_hop_expansion_candidate_ids: Vec = candidates_two_hop_expansion + .iter() + .map(|n| *n.id()) + .collect(); accessor .expand_beam_accept_only( diff --git a/diskann/src/graph/search/paged.rs b/diskann/src/graph/search/paged.rs index 791cd7acb..e2cff5464 100644 --- a/diskann/src/graph/search/paged.rs +++ b/diskann/src/graph/search/paged.rs @@ -134,7 +134,7 @@ where let mut candidates = Vec::with_capacity(page_size); for n in best.iter() { total += 1; - if !start_points.contains(&n.id) { + if !start_points.contains(n.id()) { candidates.push(n); if candidates.len() >= page_size { break; diff --git a/diskann/src/graph/search/range_search.rs b/diskann/src/graph/search/range_search.rs index 01c859ca5..dbe5fec78 100644 --- a/diskann/src/graph/search/range_search.rs +++ b/diskann/src/graph/search/range_search.rs @@ -208,7 +208,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); } } @@ -216,7 +216,7 @@ where // clear the visited set and repopulate it with just the in-range points scratch.visited.clear(); for neighbor in in_range.iter() { - scratch.visited.insert(neighbor.id); + scratch.visited.insert(*neighbor.id()); } scratch.in_range = in_range; @@ -340,7 +340,7 @@ where let beam_width = search_params.beam_width().unwrap_or(1); for neighbor in &scratch.in_range { - scratch.range_frontier.push_back(neighbor.id); + scratch.range_frontier.push_back(*neighbor.id()); } let mut neighbors = Vec::with_capacity(max_degree_with_slack); @@ -370,11 +370,11 @@ 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); - scratch.range_frontier.push_back(neighbor.id); + scratch.range_frontier.push_back(*neighbor.id()); } } scratch.cmps += neighbors.len() as u32; @@ -424,8 +424,8 @@ mod tests { assert_eq!(filtered.push(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].id(), 1); + assert_eq!(inner[0].distance(), 0.5); } #[test] @@ -448,9 +448,9 @@ mod tests { assert_eq!(count, 3); assert_eq!(inner.len(), 3); - assert_eq!(inner[0].id, 1); - assert_eq!(inner[1].id, 3); - assert_eq!(inner[2].id, 5); + assert_eq!(*inner[0].id(), 1); + assert_eq!(*inner[1].id(), 3); + assert_eq!(*inner[2].id(), 5); } #[test] @@ -487,7 +487,7 @@ mod tests { // 0.1 and 0.3 are <= inner_radius, 1.0 is not < radius assert_eq!(count, 2); - assert_eq!(inner[0].id, 2); - assert_eq!(inner[1].id, 5); + assert_eq!(*inner[0].id(), 2); + assert_eq!(*inner[1].id(), 5); } } diff --git a/diskann/src/graph/search/record.rs b/diskann/src/graph/search/record.rs index cd8d8f97d..783bfde1f 100644 --- a/diskann/src/graph/search/record.rs +++ b/diskann/src/graph/search/record.rs @@ -86,7 +86,7 @@ where } pub fn ids(&self) -> impl ExactSizeIterator + Clone + Send + Sync { - self.visited.iter().map(|n| n.id.clone()) + self.visited.iter().map(|n| n.id().clone()) } } @@ -134,7 +134,7 @@ where pub fn push(&mut self, neighbor: Neighbor, hops: u32) { self.hops.push(hops); - if self.groundtruth.contains(&neighbor.id) { + if self.groundtruth.contains(neighbor.id()) { self.running_recall += 1; } self.recall.push(self.running_recall); diff --git a/diskann/src/graph/test/cases/multihop.rs b/diskann/src/graph/test/cases/multihop.rs index 1020964ff..04b272c06 100644 --- a/diskann/src/graph/test/cases/multihop.rs +++ b/diskann/src/graph/test/cases/multihop.rs @@ -165,7 +165,7 @@ fn accept_all_finds_all_nodes() { let (stats, results) = run(&index, &[1.5], 3, 10, &AcceptAll); - let ids: Vec = results.iter().map(|n| n.id).collect(); + let ids: Vec = results.iter().map(|n| *n.id()).collect(); assert!(ids.contains(&0), "node 0 should be found"); assert!(ids.contains(&1), "node 1 should be found"); assert!(ids.contains(&2), "node 2 should be found"); @@ -206,7 +206,7 @@ fn reject_triggers_two_hop_expansion() { let filter = EvenFilter; let (stats, results) = run(&index, &[2.0], 5, 20, &filter); - let ids: Vec = results.iter().map(|n| n.id).collect(); + let ids: Vec = results.iter().map(|n| *n.id()).collect(); // Even nodes reachable only via two-hop through odd nodes. assert!( @@ -221,13 +221,13 @@ fn reject_triggers_two_hop_expansion() { // All results in the best set should be even (matching). for n in &results { - if n.id == start_id { + if *n.id() == start_id { continue; } assert!( - n.id.is_multiple_of(2), + n.id().is_multiple_of(2), "non-matching node {} should not be in best set", - n.id + n.id() ); } diff --git a/diskann/src/graph/test/cases/paged_search.rs b/diskann/src/graph/test/cases/paged_search.rs index 6de4e7205..e2b676f5d 100644 --- a/diskann/src/graph/test/cases/paged_search.rs +++ b/diskann/src/graph/test/cases/paged_search.rs @@ -80,9 +80,9 @@ fn assert_no_duplicates_across_pages(pages: &[Vec>]) { for (page_idx, page) in pages.iter().enumerate() { for n in page { assert!( - seen.insert(n.id), + seen.insert(*n.id()), "duplicate id {} found on page {}", - n.id, + n.id(), page_idx ); } @@ -94,13 +94,13 @@ fn assert_non_decreasing_distances(pages: &[Vec>]) { for (page_idx, page) in pages.iter().enumerate() { for window in page.windows(2) { assert!( - window[0].distance <= window[1].distance, + window[0].distance() <= window[1].distance(), "page {}: distances not non-decreasing: id {} dist {} followed by id {} dist {}", page_idx, - window[0].id, - window[0].distance, - window[1].id, - window[1].distance, + window[0].id(), + window[0].distance(), + window[1].id(), + window[1].distance(), ); } } @@ -135,7 +135,7 @@ fn build_baseline( page_size, pages: pages .iter() - .map(|p| p.iter().map(|n| (n.id, n.distance)).collect()) + .map(|p| p.iter().map(|n| n.as_tuple()).collect()) .collect(), total_results: pages.iter().map(|p| p.len()).sum(), } diff --git a/diskann/src/graph/test/cases/range_search.rs b/diskann/src/graph/test/cases/range_search.rs index 6e2f4dce2..bfb539e60 100644 --- a/diskann/src/graph/test/cases/range_search.rs +++ b/diskann/src/graph/test/cases/range_search.rs @@ -86,25 +86,25 @@ verbose_eq!(RangeSearchBaseline { fn assert_no_duplicates(results: &[Neighbor]) { let mut seen = std::collections::HashSet::new(); for n in results { - assert!(seen.insert(n.id), "duplicate result id {}", n.id); + assert!(seen.insert(*n.id()), "duplicate result id {}", n.id()); } } 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, + n.id(), + n.distance(), radius ); if let Some(inner) = inner_radius { assert!( - n.distance > inner, + n.distance() > inner, "result {} distance {} is within inner radius {}", - n.id, - n.distance, + n.id(), + n.distance(), inner ); } @@ -142,7 +142,7 @@ fn basic_range_search() { radius, inner_radius: None, starting_l, - results: results.iter().map(|n| (n.id, n.distance)).collect(), + results: results.iter().map(|n| n.as_tuple()).collect(), comparisons: stats.cmps as usize, hops: stats.hops as usize, result_count: results.len(), @@ -189,7 +189,7 @@ fn inner_radius_filtering() { radius, inner_radius: Some(inner_radius), starting_l, - results: results.iter().map(|n| (n.id, n.distance)).collect(), + results: results.iter().map(|n| n.as_tuple()).collect(), comparisons: stats.cmps as usize, hops: stats.hops as usize, result_count: results.len(), @@ -234,7 +234,7 @@ fn two_round_search() { radius, inner_radius: None, starting_l, - results: results.iter().map(|n| (n.id, n.distance)).collect(), + results: results.iter().map(|n| n.as_tuple()).collect(), comparisons: stats.cmps as usize, hops: stats.hops as usize, result_count: results.len(), @@ -318,7 +318,7 @@ fn max_results_respected_means_no_second_round() { radius, inner_radius: None, starting_l, - results: results.iter().map(|n| (n.id, n.distance)).collect(), + results: results.iter().map(|n| n.as_tuple()).collect(), comparisons: stats.cmps as usize, hops: stats.hops as usize, result_count: results.len(), @@ -376,7 +376,7 @@ fn max_results_respected_and_second_round_triggered() { radius, inner_radius: None, starting_l, - results: results.iter().map(|n| (n.id, n.distance)).collect(), + results: results.iter().map(|n| n.as_tuple()).collect(), comparisons: stats.cmps as usize, hops: stats.hops as usize, result_count: results.len(), diff --git a/diskann/src/graph/test/provider.rs b/diskann/src/graph/test/provider.rs index 56f1c23aa..dd0e7082c 100644 --- a/diskann/src/graph/test/provider.rs +++ b/diskann/src/graph/test/provider.rs @@ -1437,7 +1437,7 @@ impl<'a, 'b, O> glue::SearchPostProcessStep, &'b [f32], O> for Filt next.post_process( accessor, query, - candidates.filter(|n| !provider.is_deleted(n.id).unwrap_or(true)), + candidates.filter(|n| !provider.is_deleted(*n.id()).unwrap_or(true)), output, ) } diff --git a/diskann/src/neighbor/mod.rs b/diskann/src/neighbor/mod.rs index 391a5435e..ee4f8d742 100644 --- a/diskann/src/neighbor/mod.rs +++ b/diskann/src/neighbor/mod.rs @@ -4,7 +4,7 @@ */ // Imports -use std::{cmp::Ordering, fmt::Debug}; +use std::fmt::Debug; use crate::graph::{SearchOutputBuffer, search_output_buffer}; @@ -23,99 +23,146 @@ pub use diverse_priority_queue::{ // Neighbor // ////////////// -/// Neighbor node #[derive(Debug, Clone, Copy)] -pub struct Neighbor -where - VectorIdType: Eq, -{ +pub struct Neighbor { /// The id of the node - pub id: VectorIdType, + id: I, /// The distance from the query node to current node - pub distance: f32, + distance: f32, } -impl Neighbor -where - VectorIdType: Eq, -{ +impl Neighbor { /// Create the neighbor node and it has not been visited - pub fn new(id: VectorIdType, distance: f32) -> Self { + #[inline] + pub fn new(id: I, distance: f32) -> Self { Self { id, distance } } /// Return the contents of `self` as a tuple. - pub fn as_tuple(self) -> (VectorIdType, f32) { + #[inline] + pub fn as_tuple(self) -> (I, f32) { (self.id, self.distance) } + + #[inline] + pub fn distance(&self) -> f32 { + self.distance + } + + #[inline] + pub fn id(&self) -> &I { + &self.id + } } -impl Default for Neighbor +impl Default for Neighbor where - VectorIdType: Default + Eq, + I: Default, { fn default() -> Self { Self { - id: VectorIdType::default(), - distance: 0.0_f32, + id: I::default(), + distance: 0.0, } } } -impl PartialEq for Neighbor -where - VectorIdType: Eq, -{ - #[inline] - fn eq(&self, other: &Self) -> bool { - self.id == other.id +// impl PartialEq for Neighbor +// where +// VectorIdType: Eq, +// { +// #[inline] +// fn eq(&self, other: &Self) -> bool { +// self.id == other.id +// } +// } +// +// impl Eq for Neighbor where VectorIdType: Eq {} + +// /// PERF SENSITIVE: does not do well with comparing item with self. +// /// Not doing so, allows for a 1% gain. So use it with care. +// impl Ord for Neighbor +// where +// VectorIdType: Eq + Debug, +// { +// fn cmp(&self, other: &Self) -> Ordering { +// debug_assert!( +// self.id.ne(&other.id), +// "Neighbor id should not be equal: {:?}, {:?}", +// self.id, +// other.id +// ); +// self.distance +// .partial_cmp(&other.distance) +// .unwrap_or(std::cmp::Ordering::Equal) +// } +// } + +pub mod ord { + use super::Neighbor; + + pub fn fast_distance(x: &Neighbor, y: &Neighbor) -> std::cmp::Ordering { + x.distance() + .partial_cmp(&y.distance()) + .unwrap_or(std::cmp::Ordering::Equal) } -} -impl Eq for Neighbor where VectorIdType: Eq {} + pub fn fast_distance_total(x: &Neighbor, y: &Neighbor) -> std::cmp::Ordering + where + I: Ord, + { + fast_distance(x, y).then(x.id().cmp(y.id())) + } -/// PERF SENSITIVE: does not do well with comparing item with self. -/// Not doing so, allows for a 1% gain. So use it with care. -impl Ord for Neighbor -where - VectorIdType: Eq + Debug, -{ - fn cmp(&self, other: &Self) -> Ordering { - debug_assert!( - self.id.ne(&other.id), - "Neighbor id should not be equal: {:?}, {:?}", - self.id, - other.id - ); - self.distance - .partial_cmp(&other.distance) - .unwrap_or(std::cmp::Ordering::Equal) + pub fn reverse(f: F) -> impl Fn(&Neighbor, &Neighbor) -> std::cmp::Ordering + where + F: Fn(&Neighbor, &Neighbor) -> std::cmp::Ordering, + { + move |x: &Neighbor, y: &Neighbor| f(x, y).reverse() } } -/// PERF SENSITIVE: does not do well with comparing item with self. -/// Not doing so, allows for a 1% gain. So use it with care. -impl PartialOrd for Neighbor -where - VectorIdType: Eq + Debug, -{ - #[inline] - fn lt(&self, other: &Self) -> bool { - debug_assert!( - self.id.ne(&other.id), - "Neighbor id should not be equal: {:?}, {:?}", - self.id, - other.id - ); - self.distance < other.distance +pub mod cmp { + use super::Neighbor; + + pub fn ids(x: &Neighbor, y: &Neighbor) -> bool + where + I: PartialEq, + { + x.id() == y.id() } - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) + pub fn equal(x: &Neighbor, y: &Neighbor) -> bool + where + I: PartialEq, + { + ids(x, y) && x.distance() == y.distance() } } +// /// PERF SENSITIVE: does not do well with comparing item with self. +// /// Not doing so, allows for a 1% gain. So use it with care. +// impl PartialOrd for Neighbor +// where +// VectorIdType: Eq + Debug, +// { +// #[inline] +// fn lt(&self, other: &Self) -> bool { +// debug_assert!( +// self.id.ne(&other.id), +// "Neighbor id should not be equal: {:?}, {:?}", +// self.id, +// other.id +// ); +// self.distance < other.distance +// } +// +// fn partial_cmp(&self, other: &Self) -> Option { +// Some(self.cmp(other)) +// } +// } + /// 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)] @@ -225,177 +272,177 @@ where } } -#[cfg(test)] -mod neighbor_test { - use super::*; - - #[test] - fn eq_lt_works() { - let n1 = Neighbor::new(1, 1.1); - let n2 = Neighbor::new(2, 2.0); - let n3 = Neighbor::new(1, 1.1); - - assert!(n1 != n2); - assert!(n1 < n2); - assert!(n1 == n3); - } - - #[cfg(debug_assertions)] - #[test] - #[should_panic] - fn cmp_same_id_panics() { - let n1 = Neighbor::new(1, 1.1); - let n2 = Neighbor::new(1, 1.1); - - // This should panic - since the ids are the same. - let _: bool = n1 < n2; - } - - #[test] - fn gt_works() { - let n1 = Neighbor::new(1, 1.1); - let n2 = Neighbor::new(2, 2.0); - - let test = n2 > n1; - assert!(test); - } - - #[test] - fn le_works() { - let n1 = Neighbor::new(1, 1.1); - let n2 = Neighbor::new(2, 2.0); - - let test = n1 <= n2; - assert!(test); - } - - #[test] - fn cmp_works() { - let n1 = Neighbor::new(1, 1.1); - let n2 = Neighbor::new(2, 2.0); - let n3 = Neighbor::new(3, 1.1); - - assert_eq!(n1.cmp(&n2), Ordering::Less); - assert_eq!(n2.cmp(&n1), Ordering::Greater); - assert_eq!(n1.cmp(&n3), Ordering::Equal); - } - - #[test] - fn test_search_output_buffer() { - const MAX_LENGTH: usize = 5; - - // Helps with typing. - fn f(i: usize) -> Neighbor { - Neighbor::new(i as u32, i as f32) - } - - // All `push`. - { - let mut buffer = [Neighbor::::default(); MAX_LENGTH]; - let mut inserter = BackInserter::new(&mut buffer); - - assert_eq!(inserter.capacity(), MAX_LENGTH); - assert_eq!(inserter.size_hint(), Some(MAX_LENGTH)); - assert_eq!(inserter.current_len(), 0); - - assert!(inserter.push(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_eq!(inserter.current_len(), 2); - assert_eq!(inserter.size_hint(), Some(MAX_LENGTH - 2)); - - assert!(inserter.push(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_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_eq!(inserter.current_len(), 5); - assert_eq!(inserter.size_hint(), Some(0)); - - assert!(inserter.push(6, 6.0).is_full()); - assert_eq!(inserter.current_len(), 5); - assert_eq!(inserter.size_hint(), Some(0)); - - assert_eq!(&buffer, &[f(1), f(2), f(3), f(4), f(5)]); - } - - // All `iterator`. - { - let mut buffer = [Neighbor::::default(); MAX_LENGTH]; - let mut inserter = BackInserter::new(&mut buffer); - assert_eq!(inserter.capacity(), MAX_LENGTH); - 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)]); - 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()); - - let set = inserter.extend([(10, 10.0), (20, 20.0)]); - assert_eq!(set, 0, "no more items can be added"); - - assert_eq!(&buffer, &[f(1), f(2), f(3), f(4), f(5)]); - } - - // Mixture - { - let mut buffer = [Neighbor::::default(); MAX_LENGTH]; - let mut inserter = BackInserter::new(&mut buffer); - - assert!(inserter.push(1, 1.0).is_available()); - - let set = inserter.extend([(2, 2.0), (3, 3.0)]); - 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_eq!(inserter.current_len(), 4); - assert_eq!(inserter.size_hint(), Some(1)); - - let set = inserter.extend([(5, 5.0), (6, 6.0)]); - assert_eq!( - set, 1, - "there should only be room for one more item in the buffer" - ); - assert_eq!(inserter.current_len(), 5); - assert_eq!(inserter.size_hint(), Some(0)); - - assert_eq!(&buffer, &[f(1), f(2), f(3), f(4), f(5)]); - } - } - - #[test] - fn test_vec_neighbor_search_output_buffer() { - use crate::graph::search_output_buffer::SearchOutputBuffer; - - let mut buf: Vec> = Vec::new(); - assert_eq!(SearchOutputBuffer::::size_hint(&buf), None); - 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_eq!(SearchOutputBuffer::::current_len(&buf), 2); - assert_eq!(buf[0], Neighbor::new(1, 0.5)); - assert_eq!(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)]); - assert_eq!(count, 3); - assert_eq!(SearchOutputBuffer::::current_len(&buf), 5); - assert_eq!(buf[4], Neighbor::new(5, 2.5)); - } -} +// #[cfg(test)] +// mod neighbor_test { +// use super::*; +// +// #[test] +// fn eq_lt_works() { +// let n1 = Neighbor::new(1, 1.1); +// let n2 = Neighbor::new(2, 2.0); +// let n3 = Neighbor::new(1, 1.1); +// +// assert!(n1 != n2); +// assert!(n1 < n2); +// assert!(n1 == n3); +// } +// +// #[cfg(debug_assertions)] +// #[test] +// #[should_panic] +// fn cmp_same_id_panics() { +// let n1 = Neighbor::new(1, 1.1); +// let n2 = Neighbor::new(1, 1.1); +// +// // This should panic - since the ids are the same. +// let _: bool = n1 < n2; +// } +// +// #[test] +// fn gt_works() { +// let n1 = Neighbor::new(1, 1.1); +// let n2 = Neighbor::new(2, 2.0); +// +// let test = n2 > n1; +// assert!(test); +// } +// +// #[test] +// fn le_works() { +// let n1 = Neighbor::new(1, 1.1); +// let n2 = Neighbor::new(2, 2.0); +// +// let test = n1 <= n2; +// assert!(test); +// } +// +// #[test] +// fn cmp_works() { +// let n1 = Neighbor::new(1, 1.1); +// let n2 = Neighbor::new(2, 2.0); +// let n3 = Neighbor::new(3, 1.1); +// +// assert_eq!(n1.cmp(&n2), Ordering::Less); +// assert_eq!(n2.cmp(&n1), Ordering::Greater); +// assert_eq!(n1.cmp(&n3), Ordering::Equal); +// } +// +// #[test] +// fn test_search_output_buffer() { +// const MAX_LENGTH: usize = 5; +// +// // Helps with typing. +// fn f(i: usize) -> Neighbor { +// Neighbor::new(i as u32, i as f32) +// } +// +// // All `push`. +// { +// let mut buffer = [Neighbor::::default(); MAX_LENGTH]; +// let mut inserter = BackInserter::new(&mut buffer); +// +// assert_eq!(inserter.capacity(), MAX_LENGTH); +// assert_eq!(inserter.size_hint(), Some(MAX_LENGTH)); +// assert_eq!(inserter.current_len(), 0); +// +// assert!(inserter.push(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_eq!(inserter.current_len(), 2); +// assert_eq!(inserter.size_hint(), Some(MAX_LENGTH - 2)); +// +// assert!(inserter.push(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_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_eq!(inserter.current_len(), 5); +// assert_eq!(inserter.size_hint(), Some(0)); +// +// assert!(inserter.push(6, 6.0).is_full()); +// assert_eq!(inserter.current_len(), 5); +// assert_eq!(inserter.size_hint(), Some(0)); +// +// assert_eq!(&buffer, &[f(1), f(2), f(3), f(4), f(5)]); +// } +// +// // All `iterator`. +// { +// let mut buffer = [Neighbor::::default(); MAX_LENGTH]; +// let mut inserter = BackInserter::new(&mut buffer); +// assert_eq!(inserter.capacity(), MAX_LENGTH); +// 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)]); +// 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()); +// +// let set = inserter.extend([(10, 10.0), (20, 20.0)]); +// assert_eq!(set, 0, "no more items can be added"); +// +// assert_eq!(&buffer, &[f(1), f(2), f(3), f(4), f(5)]); +// } +// +// // Mixture +// { +// let mut buffer = [Neighbor::::default(); MAX_LENGTH]; +// let mut inserter = BackInserter::new(&mut buffer); +// +// assert!(inserter.push(1, 1.0).is_available()); +// +// let set = inserter.extend([(2, 2.0), (3, 3.0)]); +// 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_eq!(inserter.current_len(), 4); +// assert_eq!(inserter.size_hint(), Some(1)); +// +// let set = inserter.extend([(5, 5.0), (6, 6.0)]); +// assert_eq!( +// set, 1, +// "there should only be room for one more item in the buffer" +// ); +// assert_eq!(inserter.current_len(), 5); +// assert_eq!(inserter.size_hint(), Some(0)); +// +// assert_eq!(&buffer, &[f(1), f(2), f(3), f(4), f(5)]); +// } +// } +// +// #[test] +// fn test_vec_neighbor_search_output_buffer() { +// use crate::graph::search_output_buffer::SearchOutputBuffer; +// +// let mut buf: Vec> = Vec::new(); +// assert_eq!(SearchOutputBuffer::::size_hint(&buf), None); +// 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_eq!(SearchOutputBuffer::::current_len(&buf), 2); +// assert_eq!(buf[0], Neighbor::new(1, 0.5)); +// assert_eq!(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)]); +// assert_eq!(count, 3); +// assert_eq!(SearchOutputBuffer::::current_len(&buf), 5); +// assert_eq!(buf[4], Neighbor::new(5, 2.5)); +// } +// } diff --git a/diskann/src/neighbor/queue.rs b/diskann/src/neighbor/queue.rs index 7c1047f4a..f4991574e 100644 --- a/diskann/src/neighbor/queue.rs +++ b/diskann/src/neighbor/queue.rs @@ -139,7 +139,9 @@ impl NeighborPriorityQueue { if self.size == self.capacity { self.reserve(1.max(self.capacity >> 1)); // 1.5x capacity } - } else if self.size == self.capacity && self.get_unchecked(self.size - 1) < nbr { + } else if self.size == self.capacity + && self.get_unchecked(self.size - 1).distance() < nbr.distance() + { return; } From 03966ef2b6f5b6d7a27bec15e025b588ae4bd25e Mon Sep 17 00:00:00 2001 From: Mark Hildebrand Date: Fri, 24 Jul 2026 14:50:07 -0700 Subject: [PATCH 2/2] Cleanup. --- diskann/src/flat/test/harness.rs | 2 +- .../src/graph/internal/sorted_neighbors.rs | 69 ++- diskann/src/neighbor/mod.rs | 586 ++++++++++-------- diskann/src/test/cmp.rs | 2 +- 4 files changed, 349 insertions(+), 310 deletions(-) diff --git a/diskann/src/flat/test/harness.rs b/diskann/src/flat/test/harness.rs index 6eaaa40ae..30a300d65 100644 --- a/diskann/src/flat/test/harness.rs +++ b/diskann/src/flat/test/harness.rs @@ -8,7 +8,7 @@ //! Use [`KnnOracleRun::run`] to drive `knn_search` under a chosen [`OracleProcessor`] //! and pair the result with the oracle's expected post-processed output. -use std::{cmp::Ordering, convert::Infallible, num::NonZeroUsize}; +use std::{convert::Infallible, num::NonZeroUsize}; use diskann_vector::{PreprocessedDistanceFunction, distance::Metric}; diff --git a/diskann/src/graph/internal/sorted_neighbors.rs b/diskann/src/graph/internal/sorted_neighbors.rs index c40677115..8993ee05a 100644 --- a/diskann/src/graph/internal/sorted_neighbors.rs +++ b/diskann/src/graph/internal/sorted_neighbors.rs @@ -61,6 +61,7 @@ where #[cfg(test)] mod tests { use super::*; + use crate::test::cmp::assert_eq_verbose; use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom}; @@ -73,38 +74,38 @@ mod tests { } } - // #[test] - // fn test_sorted_neighbors() { - // let reference = [ - // Neighbor::new(1, 0.1), - // Neighbor::new(2, 0.2), - // Neighbor::new(3, 0.3), - // Neighbor::new(4, 0.4), - // Neighbor::new(5, 0.5), - // Neighbor::new(6, 0.6), - // Neighbor::new(7, 0.7), - // Neighbor::new(8, 0.8), - // Neighbor::new(9, 0.9), - // Neighbor::new(10, 1.0), - // ]; - - // let mut rng = StdRng::seed_from_u64(0xd6152fb91c744f54); - - // let ntrials = 10; - // for max in 0..reference.len() + 2 { - // for _ in 0..ntrials { - // let mut shuffled = reference.to_vec(); - // shuffled.shuffle(&mut rng); - - // let sorted = SortedNeighbors::new(&mut shuffled, max); - - // let expected_len = reference.len().min(max); - // assert_eq!(sorted.len(), expected_len); - // assert_eq!(&sorted[..expected_len], &reference[..expected_len],); - - // // Changes are visible on the taken vector. - // assert_eq!(shuffled.len(), expected_len) - // } - // } - // } + #[test] + fn test_sorted_neighbors() { + let reference = [ + Neighbor::new(1, 0.1), + Neighbor::new(2, 0.2), + Neighbor::new(3, 0.3), + Neighbor::new(4, 0.4), + Neighbor::new(5, 0.5), + Neighbor::new(6, 0.6), + Neighbor::new(7, 0.7), + Neighbor::new(8, 0.8), + Neighbor::new(9, 0.9), + Neighbor::new(10, 1.0), + ]; + + let mut rng = StdRng::seed_from_u64(0xd6152fb91c744f54); + + let ntrials = 10; + for max in 0..reference.len() + 2 { + for _ in 0..ntrials { + let mut shuffled = reference.to_vec(); + shuffled.shuffle(&mut rng); + + let sorted = SortedNeighbors::new(&mut shuffled, max); + + let expected_len = reference.len().min(max); + assert_eq!(sorted.len(), expected_len); + assert_eq_verbose!(&sorted[..expected_len], &reference[..expected_len]); + + // Changes are visible on the taken vector. + assert_eq!(shuffled.len(), expected_len) + } + } + } } diff --git a/diskann/src/neighbor/mod.rs b/diskann/src/neighbor/mod.rs index ee4f8d742..d99a8efcb 100644 --- a/diskann/src/neighbor/mod.rs +++ b/diskann/src/neighbor/mod.rs @@ -23,91 +23,153 @@ pub use diverse_priority_queue::{ // Neighbor // ////////////// -#[derive(Debug, Clone, Copy)] +/// A pairing of opaque ID with a distance. +/// +/// For all intents and purposes, this is a simple aggregate with minimal additional semantics. +/// +/// # Sorting +/// +/// An exceedingly common algorithmic operation is to sort [`Neighbor`]s according to some +/// protocol. Usually, this is by distance (from smallest to largest) ignoring the ID entirely, +/// but certain situations call for different orderings. This raises two problems: +/// +/// 1. Accurately specifying the protocol. +/// 2. Dealing with the +/// [ordering semantics](https://doc.rust-lang.org/std/cmp/trait.Ord.html#examples-of-incorrect-ord-implementations) +/// of floating point numbers. +/// +/// To that end, the functions in the [`ord`] submodule should be used in combination with +/// the standard library's sorting methods that accept explicit comparison functions like +/// [`std::slice::sort_by`]. +/// +/// ```rust +/// use diskann::neighbor::{self, Neighbor}; +/// +/// let mut neighbors = [ +/// Neighbor::new("a", 10.0), +/// Neighbor::new("b", 5.0), +/// Neighbor::new("c", 7.0), +/// Neighbor::new("d", 5.0), +/// ]; +/// +/// // Sort by increasing distance. +/// neighbors.sort_by(neighbor::ord::fast_distance); +/// assert_eq!( +/// neighbors.map(Neighbor::as_tuple), +/// [("b", 5.0), ("d", 5.0), ("c", 7.0), ("a", 10.0)] +/// ); +/// +/// // Sort reverse distance + IDs. +/// neighbors.sort_by(neighbor::ord::reverse(neighbor::ord::fast_distance_total)); +/// assert_eq!( +/// neighbors.map(Neighbor::as_tuple), +/// [("a", 10.0), ("c", 7.0), ("d", 5.0), ("b", 5.0)] +/// ); +/// ``` +#[derive(Debug, Default, Clone, Copy)] pub struct Neighbor { - /// The id of the node id: I, - - /// The distance from the query node to current node distance: f32, } impl Neighbor { - /// Create the neighbor node and it has not been visited + /// Create a [`Neighbor`] with `id` and `distance`. #[inline] pub fn new(id: I, distance: f32) -> Self { Self { id, distance } } - /// Return the contents of `self` as a tuple. + /// Return the ID and distance in `self` as a tuple. #[inline] pub fn as_tuple(self) -> (I, f32) { (self.id, self.distance) } + /// Return the distance. #[inline] pub fn distance(&self) -> f32 { self.distance } + /// Return the ID. #[inline] pub fn id(&self) -> &I { &self.id } } -impl Default for Neighbor +#[cfg(test)] +impl crate::test::cmp::VerboseEq for Neighbor where - I: Default, + I: crate::test::cmp::VerboseEq, { - fn default() -> Self { - Self { - id: I::default(), - distance: 0.0, + #[inline(never)] + #[track_caller] + fn verbose_eq(&self, other: &Self) -> crate::ANNResult<()> { + if let Err(err) = (self.id).verbose_eq(&other.id) { + return Err(err.context(crate::test::cmp::Field("id"))); + } + + if let Err(err) = (self.distance).verbose_eq(&other.distance) { + return Err(err.context(crate::test::cmp::Field("distance"))); } + + Ok(()) } } -// impl PartialEq for Neighbor -// where -// VectorIdType: Eq, -// { -// #[inline] -// fn eq(&self, other: &Self) -> bool { -// self.id == other.id -// } -// } -// -// impl Eq for Neighbor where VectorIdType: Eq {} - -// /// PERF SENSITIVE: does not do well with comparing item with self. -// /// Not doing so, allows for a 1% gain. So use it with care. -// impl Ord for Neighbor -// where -// VectorIdType: Eq + Debug, -// { -// fn cmp(&self, other: &Self) -> Ordering { -// debug_assert!( -// self.id.ne(&other.id), -// "Neighbor id should not be equal: {:?}, {:?}", -// self.id, -// other.id -// ); -// self.distance -// .partial_cmp(&other.distance) -// .unwrap_or(std::cmp::Ordering::Equal) -// } -// } - pub mod ord { + //! Methods for ordering [`Neighbor`]s. + use super::Neighbor; + /// Return the ordering between `x` and `y` according just to distance. + /// + /// This is a fast, semi-approximate method whose behavior is unspecified when either + /// distance is [`f32::NAN`]. + /// + /// ```rust + /// use diskann::neighbor::{Neighbor, ord::fast_distance}; + /// + /// let x = Neighbor::new(10, 5.0); + /// let y = Neighbor::new(11, 4.0); + /// + /// assert!(fast_distance(&x, &y).is_gt()); + /// assert!(fast_distance(&y, &x).is_lt()); + /// assert!(fast_distance(&x, &x).is_eq()); + /// + /// let z = Neighbor::new(12, f32::NAN); + /// + /// // The following line can return any `Ordering`. + /// // neighbor::ord::fast_distance(&z, &z); + /// ``` pub fn fast_distance(x: &Neighbor, y: &Neighbor) -> std::cmp::Ordering { x.distance() .partial_cmp(&y.distance()) .unwrap_or(std::cmp::Ordering::Equal) } + /// Return the ordering between `x` and `y` according to first distance then ID. + /// + /// This is a fast, semi-approximate method whose behavior is unspecified when either + /// distance is [`f32::NAN`]. + /// + /// ```rust + /// use diskann::neighbor::{Neighbor, ord::fast_distance_total}; + /// + /// let x = Neighbor::new(10, 5.0); + /// let y = Neighbor::new(11, 4.0); + /// let z = Neighbor::new(12, 4.0); + /// + /// // Different distance - ID doesn't matter. + /// assert!(fast_distance_total(&x, &y).is_gt()); + /// assert!(fast_distance_total(&y, &x).is_lt()); + /// assert!(fast_distance_total(&x, &x).is_eq()); + /// + /// // Same distance then compares by ID. + /// assert!(fast_distance_total(&y, &z).is_lt()); + /// assert!(fast_distance_total(&z, &y).is_gt()); + /// ``` pub fn fast_distance_total(x: &Neighbor, y: &Neighbor) -> std::cmp::Ordering where I: Ord, @@ -115,6 +177,26 @@ pub mod ord { fast_distance(x, y).then(x.id().cmp(y.id())) } + /// A combinator for comparisons that reverses the ordering. + /// + /// This can be useful in situations where higher distances need to be ordered first. + /// + /// ```rust + /// use diskann::neighbor::{Neighbor, ord::{reverse, fast_distance}}; + /// + /// let mut neighbors = [ + /// Neighbor::new(1, 1.0), + /// Neighbor::new(2, 2.0), + /// Neighbor::new(3, 3.0), + /// ]; + /// + /// neighbors.sort_by(reverse(fast_distance)); + /// + /// assert_eq!( + /// neighbors.map(Neighbor::as_tuple), + /// [(3, 3.0), (2, 2.0), (1, 1.0)] + /// ); + /// ``` pub fn reverse(f: F) -> impl Fn(&Neighbor, &Neighbor) -> std::cmp::Ordering where F: Fn(&Neighbor, &Neighbor) -> std::cmp::Ordering, @@ -123,61 +205,15 @@ pub mod ord { } } -pub mod cmp { - use super::Neighbor; - - pub fn ids(x: &Neighbor, y: &Neighbor) -> bool - where - I: PartialEq, - { - x.id() == y.id() - } - - pub fn equal(x: &Neighbor, y: &Neighbor) -> bool - where - I: PartialEq, - { - ids(x, y) && x.distance() == y.distance() - } -} - -// /// PERF SENSITIVE: does not do well with comparing item with self. -// /// Not doing so, allows for a 1% gain. So use it with care. -// impl PartialOrd for Neighbor -// where -// VectorIdType: Eq + Debug, -// { -// #[inline] -// fn lt(&self, other: &Self) -> bool { -// debug_assert!( -// self.id.ne(&other.id), -// "Neighbor id should not be equal: {:?}, {:?}", -// self.id, -// other.id -// ); -// self.distance < other.distance -// } -// -// fn partial_cmp(&self, other: &Self) -> Option { -// Some(self.cmp(other)) -// } -// } - /// 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> -where - I: Eq, -{ +pub struct BackInserter<'a, I> { buffer: &'a mut [Neighbor], position: usize, } -impl<'a, I> BackInserter<'a, I> -where - I: Eq, -{ +impl<'a, I> BackInserter<'a, I> { /// Construct a new [`BackInserter`] around the provided slice. /// /// The buffer will have a capacity equal to the length of `buffer`. @@ -188,16 +224,13 @@ where } } - /// Return the overall capacity of the buffer buffer. + /// Return the overall capacity of the buffer. pub fn capacity(&self) -> usize { self.buffer.len() } } -impl SearchOutputBuffer for BackInserter<'_, I> -where - I: Eq, -{ +impl SearchOutputBuffer for BackInserter<'_, I> { fn size_hint(&self) -> Option { // We maintain the invariant that `self.position <= self.buffer.len()`, so this // subtraction should not underflow. @@ -242,10 +275,7 @@ where } } -impl SearchOutputBuffer for Vec> -where - I: Eq, -{ +impl SearchOutputBuffer for Vec> { fn size_hint(&self) -> Option { None } @@ -272,177 +302,185 @@ where } } -// #[cfg(test)] -// mod neighbor_test { -// use super::*; -// -// #[test] -// fn eq_lt_works() { -// let n1 = Neighbor::new(1, 1.1); -// let n2 = Neighbor::new(2, 2.0); -// let n3 = Neighbor::new(1, 1.1); -// -// assert!(n1 != n2); -// assert!(n1 < n2); -// assert!(n1 == n3); -// } -// -// #[cfg(debug_assertions)] -// #[test] -// #[should_panic] -// fn cmp_same_id_panics() { -// let n1 = Neighbor::new(1, 1.1); -// let n2 = Neighbor::new(1, 1.1); -// -// // This should panic - since the ids are the same. -// let _: bool = n1 < n2; -// } -// -// #[test] -// fn gt_works() { -// let n1 = Neighbor::new(1, 1.1); -// let n2 = Neighbor::new(2, 2.0); -// -// let test = n2 > n1; -// assert!(test); -// } -// -// #[test] -// fn le_works() { -// let n1 = Neighbor::new(1, 1.1); -// let n2 = Neighbor::new(2, 2.0); -// -// let test = n1 <= n2; -// assert!(test); -// } -// -// #[test] -// fn cmp_works() { -// let n1 = Neighbor::new(1, 1.1); -// let n2 = Neighbor::new(2, 2.0); -// let n3 = Neighbor::new(3, 1.1); -// -// assert_eq!(n1.cmp(&n2), Ordering::Less); -// assert_eq!(n2.cmp(&n1), Ordering::Greater); -// assert_eq!(n1.cmp(&n3), Ordering::Equal); -// } -// -// #[test] -// fn test_search_output_buffer() { -// const MAX_LENGTH: usize = 5; -// -// // Helps with typing. -// fn f(i: usize) -> Neighbor { -// Neighbor::new(i as u32, i as f32) -// } -// -// // All `push`. -// { -// let mut buffer = [Neighbor::::default(); MAX_LENGTH]; -// let mut inserter = BackInserter::new(&mut buffer); -// -// assert_eq!(inserter.capacity(), MAX_LENGTH); -// assert_eq!(inserter.size_hint(), Some(MAX_LENGTH)); -// assert_eq!(inserter.current_len(), 0); -// -// assert!(inserter.push(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_eq!(inserter.current_len(), 2); -// assert_eq!(inserter.size_hint(), Some(MAX_LENGTH - 2)); -// -// assert!(inserter.push(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_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_eq!(inserter.current_len(), 5); -// assert_eq!(inserter.size_hint(), Some(0)); -// -// assert!(inserter.push(6, 6.0).is_full()); -// assert_eq!(inserter.current_len(), 5); -// assert_eq!(inserter.size_hint(), Some(0)); -// -// assert_eq!(&buffer, &[f(1), f(2), f(3), f(4), f(5)]); -// } -// -// // All `iterator`. -// { -// let mut buffer = [Neighbor::::default(); MAX_LENGTH]; -// let mut inserter = BackInserter::new(&mut buffer); -// assert_eq!(inserter.capacity(), MAX_LENGTH); -// 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)]); -// 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()); -// -// let set = inserter.extend([(10, 10.0), (20, 20.0)]); -// assert_eq!(set, 0, "no more items can be added"); -// -// assert_eq!(&buffer, &[f(1), f(2), f(3), f(4), f(5)]); -// } -// -// // Mixture -// { -// let mut buffer = [Neighbor::::default(); MAX_LENGTH]; -// let mut inserter = BackInserter::new(&mut buffer); -// -// assert!(inserter.push(1, 1.0).is_available()); -// -// let set = inserter.extend([(2, 2.0), (3, 3.0)]); -// 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_eq!(inserter.current_len(), 4); -// assert_eq!(inserter.size_hint(), Some(1)); -// -// let set = inserter.extend([(5, 5.0), (6, 6.0)]); -// assert_eq!( -// set, 1, -// "there should only be room for one more item in the buffer" -// ); -// assert_eq!(inserter.current_len(), 5); -// assert_eq!(inserter.size_hint(), Some(0)); -// -// assert_eq!(&buffer, &[f(1), f(2), f(3), f(4), f(5)]); -// } -// } -// -// #[test] -// fn test_vec_neighbor_search_output_buffer() { -// use crate::graph::search_output_buffer::SearchOutputBuffer; -// -// let mut buf: Vec> = Vec::new(); -// assert_eq!(SearchOutputBuffer::::size_hint(&buf), None); -// 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_eq!(SearchOutputBuffer::::current_len(&buf), 2); -// assert_eq!(buf[0], Neighbor::new(1, 0.5)); -// assert_eq!(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)]); -// assert_eq!(count, 3); -// assert_eq!(SearchOutputBuffer::::current_len(&buf), 5); -// assert_eq!(buf[4], Neighbor::new(5, 2.5)); -// } -// } +#[cfg(test)] +mod neighbor_test { + use super::*; + + use crate::test::cmp::assert_eq_verbose; + + #[test] + fn fast_distance() { + let n1 = Neighbor::new(1, 1.0); + let n2 = Neighbor::new(2, 2.0); + + assert!(ord::fast_distance(&n1, &n2).is_lt()); + assert!(ord::fast_distance(&n2, &n1).is_gt()); + assert!(ord::reverse(ord::fast_distance)(&n1, &n2).is_gt()); + assert!(ord::reverse(ord::fast_distance)(&n2, &n1).is_lt()); + + assert!(ord::fast_distance(&n1, &n1).is_eq()); + assert!(ord::fast_distance(&n2, &n2).is_eq()); + assert!(ord::reverse(ord::fast_distance)(&n1, &n1).is_eq()); + assert!(ord::reverse(ord::fast_distance)(&n2, &n2).is_eq()); + + // The following tests the behavior of NAN. + // + // This **must not** be taken as a guarantee of stability for this behavior. + let nan = Neighbor::new(3, 3.0); + + assert!(ord::fast_distance(&n1, &nan).is_lt()); + assert!(ord::fast_distance(&nan, &n1).is_gt()); + assert!(ord::fast_distance(&nan, &nan).is_eq()); + + assert!(ord::reverse(ord::fast_distance)(&n1, &nan).is_gt()); + assert!(ord::reverse(ord::fast_distance)(&nan, &n1).is_lt()); + assert!(ord::reverse(ord::fast_distance)(&nan, &nan).is_eq()); + } + + #[test] + fn fast_distance_total() { + let n1 = Neighbor::new(1, 1.0); + let n2 = Neighbor::new(2, 2.0); + let n3 = Neighbor::new(3, 2.0); + let n4 = Neighbor::new(4, 3.0); + + assert!(ord::fast_distance_total(&n1, &n1).is_eq()); + assert!(ord::fast_distance_total(&n1, &n2).is_lt()); + assert!(ord::fast_distance_total(&n1, &n3).is_lt()); + assert!(ord::fast_distance_total(&n1, &n4).is_lt()); + + assert!(ord::fast_distance_total(&n2, &n1).is_gt()); + assert!(ord::fast_distance_total(&n2, &n2).is_eq()); + assert!(ord::fast_distance_total(&n2, &n3).is_lt()); + assert!(ord::fast_distance_total(&n2, &n4).is_lt()); + + assert!(ord::fast_distance_total(&n3, &n1).is_gt()); + assert!(ord::fast_distance_total(&n3, &n2).is_gt()); + assert!(ord::fast_distance_total(&n3, &n3).is_eq()); + assert!(ord::fast_distance_total(&n3, &n4).is_lt()); + + assert!(ord::fast_distance_total(&n4, &n1).is_gt()); + assert!(ord::fast_distance_total(&n4, &n2).is_gt()); + assert!(ord::fast_distance_total(&n4, &n3).is_gt()); + assert!(ord::fast_distance_total(&n4, &n4).is_eq()); + } + + #[test] + fn test_search_output_buffer() { + const MAX_LENGTH: usize = 5; + + // Helps with typing. + fn f(i: usize) -> Neighbor { + Neighbor::new(i as u32, i as f32) + } + + // All `push`. + { + let mut buffer = [Neighbor::::default(); MAX_LENGTH]; + let mut inserter = BackInserter::new(&mut buffer); + + assert_eq!(inserter.capacity(), MAX_LENGTH); + assert_eq!(inserter.size_hint(), Some(MAX_LENGTH)); + assert_eq!(inserter.current_len(), 0); + + assert!(inserter.push(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_eq!(inserter.current_len(), 2); + assert_eq!(inserter.size_hint(), Some(MAX_LENGTH - 2)); + + assert!(inserter.push(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_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_eq!(inserter.current_len(), 5); + assert_eq!(inserter.size_hint(), Some(0)); + + assert!(inserter.push(6, 6.0).is_full()); + assert_eq!(inserter.current_len(), 5); + assert_eq!(inserter.size_hint(), Some(0)); + + assert_eq_verbose!(buffer, [f(1), f(2), f(3), f(4), f(5)]); + } + + // All `iterator`. + { + let mut buffer = [Neighbor::::default(); MAX_LENGTH]; + let mut inserter = BackInserter::new(&mut buffer); + assert_eq!(inserter.capacity(), MAX_LENGTH); + 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)]); + 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()); + + let set = inserter.extend([(10, 10.0), (20, 20.0)]); + assert_eq!(set, 0, "no more items can be added"); + + assert_eq_verbose!(buffer, [f(1), f(2), f(3), f(4), f(5)]); + } + + // Mixture + { + let mut buffer = [Neighbor::::default(); MAX_LENGTH]; + let mut inserter = BackInserter::new(&mut buffer); + + assert!(inserter.push(1, 1.0).is_available()); + + let set = inserter.extend([(2, 2.0), (3, 3.0)]); + 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_eq!(inserter.current_len(), 4); + assert_eq!(inserter.size_hint(), Some(1)); + + let set = inserter.extend([(5, 5.0), (6, 6.0)]); + assert_eq!( + set, 1, + "there should only be room for one more item in the buffer" + ); + assert_eq!(inserter.current_len(), 5); + assert_eq!(inserter.size_hint(), Some(0)); + + assert_eq_verbose!(buffer, [f(1), f(2), f(3), f(4), f(5)]); + } + } + + #[test] + fn test_vec_neighbor_search_output_buffer() { + use crate::graph::search_output_buffer::SearchOutputBuffer; + + let mut buf: Vec> = Vec::new(); + assert_eq!(SearchOutputBuffer::::size_hint(&buf), None); + 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_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)]); + assert_eq!(count, 3); + assert_eq!(SearchOutputBuffer::::current_len(&buf), 5); + assert_eq_verbose!(buf[4], Neighbor::new(5, 2.5)); + } +} diff --git a/diskann/src/test/cmp.rs b/diskann/src/test/cmp.rs index e3f9dd488..abcbd2f1d 100644 --- a/diskann/src/test/cmp.rs +++ b/diskann/src/test/cmp.rs @@ -47,7 +47,7 @@ pub(crate) trait VerboseEq { /// error if the lists differ. macro_rules! verbose_eq { ($struct:path { $($fields:ident),+ $(,)? }) => { - impl $crate::test::cmp::VerboseEq for $struct { + impl $crate::test::cmp::VerboseEq for $struct { #[inline(never)] #[track_caller] fn verbose_eq(&self, other: &Self) -> $crate::ANNResult<()> {