diff --git a/diskann-providers/src/index/diskann_async.rs b/diskann-providers/src/index/diskann_async.rs index f22fb889f..e8ac3d526 100644 --- a/diskann-providers/src/index/diskann_async.rs +++ b/diskann-providers/src/index/diskann_async.rs @@ -182,7 +182,7 @@ pub(crate) mod tests { use diskann_utils::{test_data_root, views::Matrix}; use diskann_vector::{ DistanceFunction, PureDistanceFunction, - distance::{Metric, SquaredL2}, + distance::{CosineNormalized, Metric, SquaredL2}, }; use rand::{distr::Distribution, rngs::StdRng, seq::SliceRandom}; use rstest::rstest; @@ -1378,29 +1378,36 @@ pub(crate) mod tests { } const SIFTSMALL: &str = "/sift/siftsmall_learn_256pts.fbin"; + const SIFTSMALL_NORMALIZED: &str = "/sift/siftsmall_learn_256pts_normalized.fbin"; #[rstest] #[tokio::test] async fn test_sift_build_and_search( #[values(FullPrecision, Hybrid::new(None))] build_strategy: S, #[values(1, 10)] batchsize: usize, + #[values( + (Metric::L2, SIFTSMALL), + (Metric::CosineNormalized, SIFTSMALL_NORMALIZED), + )] + metric_and_file: (Metric, &str), ) where S: for<'a> InsertStrategy<'a, TestProvider, &'a [f32]> + MultiInsertStrategy> + Clone, { + let (metric, file) = metric_and_file; let ctx = &DefaultContext; let parameters = InitParams { l_build: 64, max_degree: 16, - metric: Metric::L2, + metric, batchsize: NonZeroUsize::new(batchsize).unwrap(), }; let (index, data) = init_from_file( build_strategy.clone(), parameters, - SIFTSMALL, + file, 8, StartPointStrategy::RandomSamples { nsamples: ONE, @@ -1433,7 +1440,11 @@ pub(crate) mod tests { // // Because this dataset is small, we can expect exact equality. for (q, query) in data.row_iter().enumerate() { - let gt = groundtruth(data.as_view(), query, |a, b| SquaredL2::evaluate(a, b)); + let gt = groundtruth(data.as_view(), query, |a, b| match metric { + Metric::L2 => SquaredL2::evaluate(a, b), + Metric::CosineNormalized => CosineNormalized::evaluate(a, b), + _ => unreachable!(), + }); { let mut result_output_buffer = search_output_buffer::IdDistance::new(&mut ids, &mut distances); @@ -2061,8 +2072,11 @@ pub(crate) mod tests { /// PQ only Build & Search /// ////////////////////////////// + #[rstest] + #[case(Metric::L2, SIFTSMALL)] + #[case(Metric::CosineNormalized, SIFTSMALL_NORMALIZED)] #[tokio::test] - async fn test_sift_pq_only_build_and_search() { + async fn test_sift_pq_only_build_and_search(#[case] metric: Metric, #[case] file: &str) { let ctx = &DefaultContext; let create_fn = |data: Arc>, start_points: &[f32]| { let pq_table = train_pq( @@ -2074,8 +2088,7 @@ pub(crate) mod tests { .unwrap(); let (config, parameters) = - simplified_builder(64, 16, Metric::L2, data.ncols(), data.nrows(), no_modify) - .unwrap(); + simplified_builder(64, 16, metric, data.ncols(), data.nrows(), no_modify).unwrap(); let index = Arc::new(new_quant_only_index(config, parameters, pq_table, NoDeletes).unwrap()); @@ -2086,7 +2099,7 @@ pub(crate) mod tests { index }; let (index, data) = - init_and_build_index_from_file(SIFTSMALL, create_fn, build_using_single_insert).await; + init_and_build_index_from_file(file, create_fn, build_using_single_insert).await; let neighbor_accessor = &mut index.provider().neighbors(); // There should be one more reachable node than points in the dataset to account for @@ -2131,7 +2144,17 @@ pub(crate) mod tests { .await .unwrap(); - assert_top_k_exactly_match(q, >, &ids, &distances, top_k); + if metric == Metric::CosineNormalized { + let expected: Vec<_> = gt + .iter() + .rev() + .take(top_k) + .map(|neighbor| *neighbor.id()) + .collect(); + assert_eq!(expected, ids, "failed on query {q}"); + } else { + assert_top_k_exactly_match(q, >, &ids, &distances, top_k); + } } } diff --git a/diskann-providers/src/model/graph/provider/async_/inmem/product.rs b/diskann-providers/src/model/graph/provider/async_/inmem/product.rs index 5b2559e56..8b308ddef 100644 --- a/diskann-providers/src/model/graph/provider/async_/inmem/product.rs +++ b/diskann-providers/src/model/graph/provider/async_/inmem/product.rs @@ -42,6 +42,20 @@ use crate::model::{ /// The default quant provider. pub type DefaultQuant = FastMemoryQuantVectorProviderAsync; +fn quant_pruning_distance_computer( + quant: &DefaultQuant, +) -> (pq::distance::DistanceComputer<'_>, Metric) { + let metric = match quant.metric() { + // Match the squared-L2 approximation used by PQ graph-search queries. + Metric::CosineNormalized => Metric::L2, + metric => metric, + }; + ( + pq::distance::DistanceComputer::new(&quant.pq_chunk_table, metric), + metric, + ) +} + impl CreateVectorStore for FixedChunkPQTable { type Target = DefaultQuant; fn create( @@ -410,10 +424,10 @@ where ) -> Result, Self::PruneAccessorError> { let full = &provider.base_vectors; let quant = &provider.aux_vectors; - let metric = quant.metric(); + let (quant_distance_computer, metric) = quant_pruning_distance_computer(quant); let distance = distances::pq::HybridComputer::new( - quant.distance_computer(), + quant_distance_computer, T::distance(metric, Some(full.dim())), ); @@ -579,10 +593,12 @@ where _context: &'a Ctx, _capacity: usize, ) -> Result, Self::PruneAccessorError> { + let quant = &provider.aux_vectors; + let (distance_computer, _) = quant_pruning_distance_computer(quant); let accessor = PruneAccessor { - provider: &provider.aux_vectors, + provider: quant, neighbors: provider.neighbors(), - distance: provider.aux_vectors.distance_computer(), + distance: distance_computer, }; Ok(accessor) } @@ -647,3 +663,75 @@ where .into_ann_result() } } + +#[cfg(test)] +mod tests { + use diskann::{graph::glue::PruneStrategy, provider::DefaultContext, utils::VectorRepr}; + use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction, distance::Metric}; + + use super::{DefaultQuant, quant_pruning_distance_computer}; + use crate::model::{ + graph::provider::async_::{ + FastMemoryQuantVectorProviderAsync, + common::{NoDeletes, NoStore, Quantized}, + distances::pq::{Hybrid, HybridComputer}, + inmem::{DefaultProvider, DefaultProviderParameters}, + }, + pq::FixedChunkPQTable, + }; + + fn test_table() -> FixedChunkPQTable { + FixedChunkPQTable::new( + 4, + vec![1.0, 0.0, 0.0, 1.0, 2.0, 0.0, 0.0, 2.0].into(), + vec![0, 2, 4].into(), + ) + .unwrap() + } + + #[test] + fn cosine_normalized_query_and_hybrid_pruning_use_squared_l2() { + let provider = + FastMemoryQuantVectorProviderAsync::new(Metric::CosineNormalized, 2, test_table()); + let full0 = [1u8, 0, 0, 2]; + let full1 = [2u8, 0, 0, 1]; + let code0 = [0u8, 1]; + let code1 = [1u8, 0]; + + let query = provider.query_computer(&full0).unwrap(); + assert_eq!(query.evaluate_similarity(&code1), 2.0); + + let (quant_distance_computer, metric) = quant_pruning_distance_computer(&provider); + let computer = + HybridComputer::::new(quant_distance_computer, u8::distance(metric, Some(4))); + for (left, right) in [ + (Hybrid::Full(&full0[..]), Hybrid::Full(&full1[..])), + (Hybrid::Full(&full0[..]), Hybrid::Quant(&code1[..])), + (Hybrid::Quant(&code0[..]), Hybrid::Full(&full1[..])), + (Hybrid::Quant(&code0[..]), Hybrid::Quant(&code1[..])), + ] { + assert_eq!(computer.evaluate_similarity(left, right), 2.0); + } + } + + #[test] + fn cosine_normalized_quantized_pruning_uses_squared_l2() { + let provider: DefaultProvider = DefaultProvider::new_empty( + DefaultProviderParameters::simple(2, 4, Metric::CosineNormalized, 1), + NoStore, + test_table(), + NoDeletes, + ) + .unwrap(); + let accessor = Quantized + .prune_accessor(&provider, &DefaultContext, 2) + .unwrap(); + + assert_eq!( + accessor + .distance + .evaluate_similarity(&[0u8, 1][..], &[1u8, 0][..]), + 2.0, + ); + } +}