diskann_benchmark_core/
recall.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::{
7    collections::{BTreeMap, HashSet},
8    hash::Hash,
9};
10
11use diskann_utils::{
12    strided::StridedView,
13    views::{Matrix, MatrixView},
14};
15use thiserror::Error;
16
17#[derive(Debug, Clone)]
18#[non_exhaustive]
19pub struct RecallMetrics {
20    /// The `k` value for `k-recall-at-n`.
21    pub recall_k: usize,
22    /// The `n` value for `k-recall-at-n`.
23    pub recall_n: usize,
24    /// The number of queries.
25    pub num_queries: usize,
26    /// The average recall across queries with non-empty groundtruth.
27    /// Queries with zero groundtruth results are excluded from the average.
28    pub average: f64,
29}
30
31#[derive(Debug, Error)]
32pub enum ComputeRecallError {
33    #[error("results matrix has {0} rows but ground truth has {1}")]
34    RowsMismatch(usize, usize),
35    #[error("distances matrix has {0} rows but ground truth has {1}")]
36    DistanceRowsMismatch(usize, usize),
37    #[error("recall k value {0} must be less than or equal to recall n {1}")]
38    RecallKAndNError(usize, usize),
39    #[error(
40        "number of groundtruth values per query {0} must be at least the specified recall n {1}"
41    )]
42    NotEnoughGroundTruth(usize, usize),
43    #[error("number of groundtruth distances {0} does not match groundtruth entries {1}")]
44    GroundTruthDistanceMismatch(usize, usize),
45}
46
47/// An abstraction over data-structures such as vector-of-vectors.
48///
49/// This is used in recall calculations such as [`knn`] and [`average_precision`] and
50/// is purposely `dyn` compatible to reduce compilation overhead.
51///
52/// Implementations should ensure that if [`Self::nrows`] returns a value `N` that `row(i)`
53/// returns a slice for all `i` in `0..N`. Accesses outside of that range are allowed to
54/// panic.
55///
56/// The implementation [`Self::ncols`] is optional and can be implemented if the length of
57/// all inner vectors is known and identical for all inner vectors, enabling faster error
58/// paths during recall calculation.
59///
60/// If [`Self::ncols`] returns `Some(K)`, then the length of each slice returned from
61/// [`Self::row`] should have a length equal to `K`. Note that unsafe code may **not** rely
62/// on this behavior.
63///
64/// If [`Self::ncols`] returns `None` then no assumption can be made about the length of
65/// the slices yielded from [`Self::row`].
66pub trait Rows<T> {
67    /// Return the number of subslices in `Self`.
68    fn nrows(&self) -> usize;
69
70    /// Return the `i`th subslice contained in self.
71    fn row(&self, i: usize) -> &[T];
72
73    /// Return `Some(K)` if all subslices are known to have length `K`. Otherwise, return
74    /// `None`.
75    ///
76    /// # Provided Implementation
77    ///
78    /// The provided implementation returns `None`.
79    fn ncols(&self) -> Option<usize> {
80        None
81    }
82}
83
84impl<T> Rows<T> for Matrix<T> {
85    fn nrows(&self) -> usize {
86        Matrix::<T>::nrows(self)
87    }
88    fn row(&self, i: usize) -> &[T] {
89        Matrix::<T>::row(self, i)
90    }
91    fn ncols(&self) -> Option<usize> {
92        Some(Matrix::<T>::ncols(self))
93    }
94}
95
96impl<T> Rows<T> for MatrixView<'_, T> {
97    fn nrows(&self) -> usize {
98        MatrixView::<'_, T>::nrows(self)
99    }
100    fn row(&self, i: usize) -> &[T] {
101        MatrixView::<'_, T>::row(self, i)
102    }
103    fn ncols(&self) -> Option<usize> {
104        Some(MatrixView::<'_, T>::ncols(self))
105    }
106}
107
108impl<T> Rows<T> for Vec<Vec<T>> {
109    fn nrows(&self) -> usize {
110        self.len()
111    }
112    fn row(&self, i: usize) -> &[T] {
113        &self[i]
114    }
115}
116
117/// Aggregate trait for required behavior when computing recall and average precision.
118pub trait RecallCompatible: Eq + Hash + Clone + std::fmt::Debug {}
119
120impl<T> RecallCompatible for T where T: Eq + Hash + Clone + std::fmt::Debug {}
121
122// Enum representing whether the ground truth has fixed size
123// (every row should have the same number, >= recall_k entries)
124// or flexible size (row may have any number of entries, including none)
125#[derive(Copy, Clone, Debug)]
126pub enum GroundTruthMode {
127    Fixed,
128    Flexible,
129}
130
131/// Compute the K-nearest-neighbors recall value "K-recall-at-N".
132///
133/// For each entry in `groundtruth` and `results`, this computes the `recall_k` number of
134/// elements of `groundtruth` that are present in the first `recall_n` entries of `results`.
135///
136/// If `groundtruth_distances` is provided, then it will be used to allow ties when matching
137/// the last values of each entry of `results`. Values will be counted towards the recall if
138/// they have the same distance as the last ordered candidate.
139///
140/// Note that an error will NOT be given if an entry in `results`
141/// has fewer than `recall_n` candidates.
142///
143/// If `ground_truth_mode` is `GroundTruthMode::Fixed`, then an
144/// error will be given if any entry in `groundtruth` has fewer
145/// than `recall_k` candidates.
146pub fn knn<T>(
147    groundtruth: &dyn Rows<T>,
148    groundtruth_distances: Option<StridedView<'_, f32>>,
149    results: &dyn Rows<T>,
150    recall_k: usize,
151    recall_n: usize,
152    ground_truth_mode: GroundTruthMode,
153) -> Result<RecallMetrics, ComputeRecallError>
154where
155    T: RecallCompatible,
156{
157    if recall_k > recall_n {
158        return Err(ComputeRecallError::RecallKAndNError(recall_k, recall_n));
159    }
160
161    let nrows = results.nrows();
162    if nrows != groundtruth.nrows() {
163        return Err(ComputeRecallError::RowsMismatch(nrows, groundtruth.nrows()));
164    }
165
166    if let GroundTruthMode::Fixed = ground_truth_mode {
167        // Validate that all rows in `groundtruth` have at least `recall_k` entries.
168        for i in 0..nrows {
169            let gt_row = groundtruth.row(i);
170            if gt_row.len() < recall_k {
171                return Err(ComputeRecallError::NotEnoughGroundTruth(
172                    gt_row.len(),
173                    recall_k,
174                ));
175            }
176        }
177    }
178
179    // If `groundtruth_distances` are present, validate that there are
180    // enough rows and that each row has the same
181    // number of entries as the corresponding row in `groundtruth`.
182    if let Some(distances) = groundtruth_distances {
183        if nrows != distances.nrows() {
184            return Err(ComputeRecallError::DistanceRowsMismatch(
185                distances.nrows(),
186                nrows,
187            ));
188        }
189
190        for i in 0..nrows {
191            let gt_row = groundtruth.row(i);
192            let distances_row = distances.row(i);
193            if gt_row.len() != distances_row.len() {
194                return Err(ComputeRecallError::GroundTruthDistanceMismatch(
195                    distances_row.len(),
196                    gt_row.len(),
197                ));
198            }
199        }
200    }
201
202    // The actual recall computation for groundtruth.
203    //
204    // To avoid floating-point accumulation error we do not sum per-query `f64`
205    // ratios. Each ratio `r / this_recall_k` is generally not exactly
206    // representable (e.g. `99.0 / 100.0`), and summing thousands of them
207    // compounds the rounding, producing values like `0.9999899999999999` where
208    // the closest `f64` to the true answer is `0.99999`. Instead we accumulate
209    // integer hit counts, grouped by their (per-row) denominator
210    // `this_recall_k`, and divide once per distinct denominator at the end. This
211    // is algebraically identical to the mean of per-query recalls; the result is
212    // still rounded to `f64`, but whenever all scored rows share a denominator
213    // (the common fixed-size-groundtruth case) the rounding is limited to that
214    // single division rather than accumulating across every query.
215    //
216    // A `BTreeMap` (rather than `HashMap`) keeps the final summation order
217    // deterministic across runs, so the reported recall is bit-reproducible.
218    let mut hits_by_k: BTreeMap<usize, u64> = BTreeMap::new();
219    let mut num_scored_queries: usize = 0;
220    let mut this_groundtruth = HashSet::new();
221    let mut this_results = HashSet::new();
222
223    for i in 0..results.nrows() {
224        let result = results.row(i);
225
226        let gt_row = groundtruth.row(i);
227        // `groundtruth` does not have to be fixed-size,
228        // so we compute `recall_k` for this row based on its gt length
229        let this_recall_k = gt_row.len().min(recall_k);
230
231        if this_recall_k == 0 {
232            continue;
233        }
234
235        // Populate the groundtruth using the top-k
236        this_groundtruth.clear();
237        this_groundtruth.extend(gt_row.iter().take(this_recall_k).cloned());
238
239        // If we have distances, then continue to append distances as long as the distance
240        // value is constant
241        if let Some(distances) = groundtruth_distances {
242            let distances_row = distances.row(i);
243
244            // we've already checked that `results` and `distances` have at lesat
245            // `recall_k >= this_recall_k` entries, so it's safe to access `distances_row[this_recall_k - 1]`
246            let last_distance = distances_row[this_recall_k - 1];
247            for (d, g) in distances_row.iter().zip(gt_row.iter()).skip(this_recall_k) {
248                if *d == last_distance {
249                    this_groundtruth.insert(g.clone());
250                } else {
251                    break;
252                }
253            }
254        }
255
256        this_results.clear();
257        this_results.extend(result.iter().take(recall_n).cloned());
258
259        // Count the overlap
260        let r = this_groundtruth
261            .iter()
262            .filter(|i| this_results.contains(i))
263            .count()
264            .min(this_recall_k);
265
266        *hits_by_k.entry(this_recall_k).or_insert(0) += r as u64;
267        num_scored_queries += 1;
268    }
269
270    // Compute the average recall as the mean of the per-query recalls. Grouping
271    // by denominator lets each group be reduced with an integer numerator,
272    // limiting rounding to at most one division per distinct `this_recall_k`.
273    //
274    // `num_scored_queries` is the *global* count of scored queries, not a
275    // per-denominator count: the mean of per-query recalls divides the grand
276    // total by the number of queries, so each query contributes `1 / N`
277    // regardless of its `this_recall_k`. Concretely,
278    // `(1/N) * sum_i (r_i / k_i) == sum_d (hits_by_k[d] / (d * N))`.
279    //
280    // The denominator is formed as a `u128` product so `k * num_scored_queries`
281    // cannot overflow on very large benchmarks; the single division to `f64` is
282    // the only rounding step.
283    let average = if num_scored_queries == 0 {
284        0.0
285    } else {
286        hits_by_k
287            .into_iter()
288            .map(|(k, hits)| {
289                let denom = (k as u128) * (num_scored_queries as u128);
290                (hits as f64) / (denom as f64)
291            })
292            .sum()
293    };
294
295    Ok(RecallMetrics {
296        recall_k,
297        recall_n,
298        num_queries: nrows,
299        average,
300    })
301}
302
303#[derive(Debug, Clone)]
304#[non_exhaustive]
305pub struct AveragePrecisionMetrics {
306    /// The number of queries.
307    pub num_queries: usize,
308    /// The average precision
309    pub average_precision: f64,
310}
311
312#[derive(Debug, Error)]
313pub enum AveragePrecisionError {
314    #[error("results has {0} elements but ground truth has {1}")]
315    EntriesMismatch(usize, usize),
316}
317
318/// Compute average precision of a range search result
319pub fn average_precision<T>(
320    results: &dyn Rows<T>,
321    groundtruth: &dyn Rows<T>,
322) -> Result<AveragePrecisionMetrics, AveragePrecisionError>
323where
324    T: RecallCompatible,
325{
326    let nrows = results.nrows();
327    let groundtruth_nrows = groundtruth.nrows();
328    if nrows != groundtruth_nrows {
329        return Err(AveragePrecisionError::EntriesMismatch(
330            nrows,
331            groundtruth_nrows,
332        ));
333    }
334
335    // The actual recall computation.
336    let mut num_gt_results = 0;
337    let mut num_reported_results = 0;
338
339    let mut scratch = HashSet::new();
340    let nrows = results.nrows();
341
342    for i in 0..nrows {
343        let result = results.row(i);
344        let gt = groundtruth.row(i);
345
346        scratch.clear();
347        scratch.extend(result.iter().cloned());
348        num_reported_results += gt.iter().filter(|i| scratch.contains(i)).count();
349        num_gt_results += gt.len();
350    }
351
352    // Perform post-processing.
353    let average_precision = (num_reported_results as f64) / (num_gt_results as f64);
354
355    Ok(AveragePrecisionMetrics {
356        average_precision,
357        num_queries: nrows,
358    })
359}
360
361///////////
362// Tests //
363///////////
364
365#[cfg(test)]
366mod tests {
367    use diskann_utils::views::{self, Matrix};
368
369    use super::*;
370
371    fn test_rows_inner(rows: &dyn Rows<usize>, ncols: Option<usize>) {
372        assert_eq!(rows.ncols(), ncols);
373        assert_eq!(rows.nrows(), 3);
374        assert_eq!(rows.row(0), &[0, 1, 2, 3]);
375        assert_eq!(rows.row(1), &[4, 5, 6, 7]);
376        assert_eq!(rows.row(2), &[8, 9, 10, 11]);
377    }
378
379    #[test]
380    fn test_rows() {
381        let mut i = 0usize;
382        let mat = Matrix::new(
383            views::Init(|| {
384                let v = i;
385                i += 1;
386                v
387            }),
388            3,
389            4,
390        );
391
392        test_rows_inner(&mat, Some(4));
393        test_rows_inner(&(mat.as_view()), Some(4));
394
395        let vecs = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]];
396        test_rows_inner(&vecs, None);
397    }
398
399    struct ExpectedRecall {
400        recall_k: usize,
401        recall_n: usize,
402        // Recall for each component.
403        components: Vec<usize>,
404    }
405
406    impl ExpectedRecall {
407        fn new(recall_k: usize, recall_n: usize, components: Vec<usize>) -> Self {
408            assert!(recall_k <= recall_n);
409            components.iter().for_each(|x| {
410                assert!(*x <= recall_k);
411            });
412            Self {
413                recall_k,
414                recall_n,
415                components,
416            }
417        }
418
419        fn compute_recall(&self) -> f64 {
420            (self.components.iter().sum::<usize>() as f64)
421                / ((self.components.len() * self.recall_k) as f64)
422        }
423    }
424
425    #[test]
426    fn test_happy_path() {
427        let groundtruth = Matrix::try_from(
428            vec![
429                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, // row 0
430                5, 6, 7, 8, 9, 10, 11, 12, 13, 14, // row 1
431                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, // row 2
432                0, 1, 2, 3, 4, 5, 6, 7, 8, 9, // row 3
433            ]
434            .into(),
435            4,
436            10,
437        )
438        .unwrap();
439
440        let distances = Matrix::try_from(
441            vec![
442                0.0, 1.0, 2.0, 3.0, 3.0, 3.0, 3.0, 4.0, 5.0, 6.0, // row 0
443                2.0, 3.0, 3.0, 3.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, // row 1
444                0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, // row 2
445                0.0, 1.0, 2.0, 3.0, 3.0, 3.0, 3.0, 4.0, 5.0, 6.0, // row 3
446            ]
447            .into(),
448            4,
449            10,
450        )
451        .unwrap();
452
453        // Shift row 0 by one and row 1 by two.
454        let our_results = Matrix::try_from(
455            vec![
456                100, 0, 1, 2, 5, 6, // row 0
457                100, 101, 7, 8, 9, 10, // row 1
458                0, 1, 2, 3, 4, 5, // row 2
459                0, 1, 2, 3, 4, 5, // row 3
460            ]
461            .into(),
462            4,
463            6,
464        )
465        .unwrap();
466
467        //---------//
468        // No Ties //
469        //---------//
470        let expected_no_ties = vec![
471            // Equal `k` and `n`
472            ExpectedRecall::new(1, 1, vec![0, 0, 1, 1]),
473            ExpectedRecall::new(2, 2, vec![1, 0, 2, 2]),
474            ExpectedRecall::new(3, 3, vec![2, 1, 3, 3]),
475            ExpectedRecall::new(4, 4, vec![3, 2, 4, 4]),
476            ExpectedRecall::new(5, 5, vec![3, 3, 5, 5]),
477            ExpectedRecall::new(6, 6, vec![4, 4, 6, 6]),
478            // Unequal `k` and `n`.
479            ExpectedRecall::new(1, 2, vec![1, 0, 1, 1]),
480            ExpectedRecall::new(1, 3, vec![1, 0, 1, 1]),
481            ExpectedRecall::new(2, 3, vec![2, 0, 2, 2]),
482            ExpectedRecall::new(3, 5, vec![3, 1, 3, 3]),
483        ];
484        let epsilon = 1e-6; // Define a small tolerance
485
486        for (i, expected) in expected_no_ties.iter().enumerate() {
487            assert_eq!(expected.components.len(), our_results.nrows());
488            let recall = knn(
489                &groundtruth,
490                None,
491                &our_results,
492                expected.recall_k,
493                expected.recall_n,
494                GroundTruthMode::Fixed,
495            )
496            .unwrap();
497
498            let left = recall.average;
499            let right = expected.compute_recall();
500            assert!(
501                (left - right).abs() < epsilon,
502                "left = {}, right = {} on input {}",
503                left,
504                right,
505                i
506            );
507
508            assert_eq!(recall.num_queries, our_results.nrows());
509            assert_eq!(recall.recall_k, expected.recall_k);
510            assert_eq!(recall.recall_n, expected.recall_n);
511        }
512
513        //-----------//
514        // With Ties //
515        //-----------//
516        let expected_with_ties = vec![
517            // Equal `k` and `n`
518            ExpectedRecall::new(1, 1, vec![0, 0, 1, 1]),
519            ExpectedRecall::new(2, 2, vec![1, 0, 2, 2]),
520            ExpectedRecall::new(3, 3, vec![2, 1, 3, 3]),
521            ExpectedRecall::new(4, 4, vec![3, 2, 4, 4]),
522            ExpectedRecall::new(5, 5, vec![4, 3, 5, 5]), // tie-breaker kicks in
523            ExpectedRecall::new(6, 6, vec![5, 4, 6, 6]), // tie-breaker kicks in
524            // Unequal `k` and `n`.
525            ExpectedRecall::new(1, 2, vec![1, 0, 1, 1]),
526            ExpectedRecall::new(1, 3, vec![1, 0, 1, 1]),
527            ExpectedRecall::new(2, 3, vec![2, 1, 2, 2]),
528            ExpectedRecall::new(4, 5, vec![4, 3, 4, 4]),
529        ];
530
531        for (i, expected) in expected_with_ties.iter().enumerate() {
532            assert_eq!(expected.components.len(), our_results.nrows());
533            let recall = knn(
534                &groundtruth,
535                Some(distances.as_view().into()),
536                &our_results,
537                expected.recall_k,
538                expected.recall_n,
539                GroundTruthMode::Fixed,
540            )
541            .unwrap();
542
543            let left = recall.average;
544            let right = expected.compute_recall();
545            assert!(
546                (left - right).abs() < epsilon,
547                "left = {}, right = {} on input {}",
548                left,
549                right,
550                i
551            );
552
553            assert_eq!(recall.num_queries, our_results.nrows());
554            assert_eq!(recall.recall_k, expected.recall_k);
555            assert_eq!(recall.recall_n, expected.recall_n);
556        }
557    }
558
559    #[test]
560    fn test_error_recall_k_and_n() {
561        let groundtruth = Matrix::<u32>::new(0, 10, 10);
562        let results = Matrix::<u32>::new(0, 10, 10);
563        let err = knn(&groundtruth, None, &results, 11, 10, GroundTruthMode::Fixed).unwrap_err();
564        assert!(matches!(err, ComputeRecallError::RecallKAndNError(..)));
565    }
566
567    #[test]
568    fn test_error_rows_mismatch() {
569        let groundtruth = Matrix::<u32>::new(0, 11, 10);
570        let results = Matrix::<u32>::new(0, 10, 10);
571        let err = knn(&groundtruth, None, &results, 10, 10, GroundTruthMode::Fixed).unwrap_err();
572        assert!(matches!(err, ComputeRecallError::RowsMismatch(..)));
573        let err_allow_insufficient_results =
574            knn(&groundtruth, None, &results, 10, 10, GroundTruthMode::Fixed).unwrap_err();
575        assert!(matches!(
576            err_allow_insufficient_results,
577            ComputeRecallError::RowsMismatch(..)
578        ));
579    }
580
581    #[test]
582    fn test_error_not_enough_groundtruth() {
583        let groundtruth = Matrix::<u32>::new(0, 10, 5);
584        let results = Matrix::<u32>::new(0, 10, 10);
585        let err = knn(&groundtruth, None, &results, 10, 10, GroundTruthMode::Fixed).unwrap_err();
586        assert!(matches!(err, ComputeRecallError::NotEnoughGroundTruth(..)));
587        let err_allow_insufficient_results =
588            knn(&groundtruth, None, &results, 10, 10, GroundTruthMode::Fixed).unwrap_err();
589        assert!(matches!(
590            err_allow_insufficient_results,
591            ComputeRecallError::NotEnoughGroundTruth(..)
592        ));
593    }
594
595    #[test]
596    fn test_dynamic_groundtruth_valid() {
597        let groundtruth: Vec<_> = (0..10).map(|_| vec![0u32; 5]).collect();
598        let results = Matrix::<u32>::new(0, 10, 10);
599        // Should succeed: each row uses this_recall_k = min(5, 10) = 5
600        // Should succeed in Flexible mode, but fail in Fixed mode
601        let recall_flexible = knn(
602            &groundtruth,
603            None,
604            &results,
605            10,
606            10,
607            GroundTruthMode::Flexible,
608        )
609        .unwrap();
610        assert_eq!(recall_flexible.num_queries, 10);
611        // Should fail in Fixed mode
612        let err = knn(&groundtruth, None, &results, 10, 10, GroundTruthMode::Fixed).unwrap_err();
613        assert!(matches!(err, ComputeRecallError::NotEnoughGroundTruth(..)));
614        assert_eq!(recall_flexible.num_queries, 10);
615    }
616
617    #[test]
618    fn test_dynamic_groundtruth_full_match() {
619        let gt_row: Vec<u32> = (1..=5).collect();
620        let groundtruth: Vec<_> = (0..10).map(|_| gt_row.clone()).collect();
621        let mut results = Matrix::<u32>::new(0, 10, 10);
622        for i in 0..10 {
623            for (j, v) in (1u32..=10).enumerate() {
624                results[(i, j)] = v;
625            }
626        }
627        let recall = knn(
628            &groundtruth,
629            None,
630            &results,
631            10,
632            10,
633            GroundTruthMode::Flexible,
634        )
635        .unwrap();
636        assert!((recall.average - 1.0).abs() < 1e-10);
637    }
638
639    #[test]
640    fn test_dynamic_groundtruth_partial_match() {
641        // groundtruth: [1, 2, 3, 4, 5]; results contain [1, 2, 3, 6, 7, 8, 9, 10, 11, 12]
642        let gt_row: Vec<u32> = (1..=5).collect();
643        let groundtruth: Vec<_> = (0..10).map(|_| gt_row.clone()).collect();
644        let mut results = Matrix::<u32>::new(0, 10, 10);
645        let res_row: Vec<u32> = vec![1, 2, 3, 6, 7, 8, 9, 10, 11, 12];
646        for i in 0..10 {
647            for (j, &v) in res_row.iter().enumerate() {
648                results[(i, j)] = v;
649            }
650        }
651        let recall = knn(
652            &groundtruth,
653            None,
654            &results,
655            10,
656            10,
657            GroundTruthMode::Flexible,
658        )
659        .unwrap();
660        assert!((recall.average - 0.6).abs() < 1e-10);
661    }
662
663    #[test]
664    fn test_dynamic_groundtruth_mixed_zero_nonzero() {
665        let mut groundtruth: Vec<Vec<u32>> = Vec::new();
666        // First 5 rows: non-empty groundtruth
667        for _ in 0..5 {
668            groundtruth.push((1..=5).collect());
669        }
670        // Last 5 rows: empty groundtruth
671        for _ in 0..5 {
672            groundtruth.push(vec![]);
673        }
674
675        let mut results = Matrix::<u32>::new(0, 10, 10);
676        for i in 0..10 {
677            for (j, v) in (1u32..=10).enumerate() {
678                results[(i, j)] = v;
679            }
680        }
681
682        let recall = knn(
683            &groundtruth,
684            None,
685            &results,
686            10,
687            10,
688            GroundTruthMode::Flexible,
689        )
690        .unwrap();
691        assert_eq!(recall.num_queries, 10);
692        assert!((recall.average - 1.0).abs() < 1e-10);
693    }
694
695    #[test]
696    fn test_dynamic_groundtruth_all_zero() {
697        let groundtruth: Vec<Vec<u32>> = (0..10).map(|_| vec![]).collect();
698        let results = Matrix::<u32>::new(0, 10, 10);
699
700        let recall = knn(
701            &groundtruth,
702            None,
703            &results,
704            10,
705            10,
706            GroundTruthMode::Flexible,
707        )
708        .unwrap();
709        assert_eq!(recall.num_queries, 10);
710        assert_eq!(recall.average, 0.0);
711        assert!(!recall.average.is_nan());
712        assert!(!recall.average.is_infinite());
713    }
714
715    // Regression test for the reported "0.9999899999999999" recall (issue #1171).
716    //
717    // With a fixed denominator, the average of the per-query recalls must be
718    // computed without floating-point accumulation error: a perfect run reports
719    // exactly `1.0`, and a run missing exactly 5 of 500_000 neighbours reports
720    // exactly `0.99999` rather than `0.9999899999999999`.
721    #[test]
722    fn test_fixed_denominator_no_precision_loss() {
723        let k = 100usize;
724        let num_queries = 5000usize;
725
726        // Groundtruth: every row is `0..k`.
727        let gt_row: Vec<u32> = (0..k as u32).collect();
728        let groundtruth: Vec<_> = (0..num_queries).map(|_| gt_row.clone()).collect();
729
730        // Perfect results: every query recovers all `k` neighbours.
731        let perfect: Vec<_> = (0..num_queries).map(|_| gt_row.clone()).collect();
732        let recall = knn(&groundtruth, None, &perfect, k, k, GroundTruthMode::Fixed).unwrap();
733        assert_eq!(recall.average, 1.0);
734
735        // Miss exactly one neighbour on 5 queries (5 misses out of 500_000).
736        let mut near_perfect = perfect.clone();
737        for row in near_perfect.iter_mut().take(5) {
738            // Replace the last groundtruth id with a non-matching one.
739            *row.last_mut().unwrap() = u32::MAX;
740        }
741        let recall = knn(
742            &groundtruth,
743            None,
744            &near_perfect,
745            k,
746            k,
747            GroundTruthMode::Fixed,
748        )
749        .unwrap();
750        assert_eq!(recall.average, 0.99999);
751    }
752
753    #[test]
754    fn test_error_distance_rows_mismatch() {
755        let groundtruth = Matrix::<u32>::new(0, 10, 10);
756        let distances = Matrix::<f32>::new(0.0, 9, 10);
757        let results = Matrix::<u32>::new(0, 10, 10);
758        let err = knn(
759            &groundtruth,
760            Some(distances.as_view().into()),
761            &results,
762            10,
763            10,
764            GroundTruthMode::Fixed,
765        )
766        .unwrap_err();
767        assert!(matches!(err, ComputeRecallError::DistanceRowsMismatch(..)));
768    }
769
770    #[test]
771    fn test_error_distance_cols_mismatch() {
772        let groundtruth = Matrix::<u32>::new(0, 10, 10);
773        let distances = Matrix::<f32>::new(0.0, 10, 9);
774        let results = Matrix::<u32>::new(0, 10, 10);
775        let err = knn(
776            &groundtruth,
777            Some(distances.as_view().into()),
778            &results,
779            10,
780            10,
781            GroundTruthMode::Fixed,
782        )
783        .unwrap_err();
784        assert!(matches!(
785            err,
786            ComputeRecallError::GroundTruthDistanceMismatch(..)
787        ));
788    }
789
790    #[test]
791    fn test_error_distance_cols_mismatch_variable_size_groundtruth() {
792        // groundtruth: 2 rows, first row has 3 elements, second row has 2
793        let groundtruth: Vec<Vec<u32>> = vec![vec![1, 2, 3], vec![4, 5]];
794        // distances: first row has 2 elements (should be 3), second row has 2 (matches)
795        let distances: Vec<Vec<f32>> = vec![vec![0.1, 0.2], vec![0.3, 0.4]];
796        let distances = Matrix::try_from(
797            distances.into_iter().flatten().collect::<Vec<_>>().into(),
798            2,
799            2,
800        )
801        .unwrap();
802        // results: 2 rows, each with 3 elements
803        let results: Vec<Vec<u32>> = vec![vec![1, 2, 3], vec![4, 5, 6]];
804        let err = knn(
805            &groundtruth,
806            Some(distances.as_view().into()),
807            &results,
808            3,
809            3,
810            GroundTruthMode::Flexible,
811        )
812        .unwrap_err();
813
814        println!("{err}");
815
816        assert!(matches!(
817            err,
818            ComputeRecallError::GroundTruthDistanceMismatch(2, 3)
819        ));
820    }
821}
diskann_benchmark_core/recall.rs

diskann_benchmark_core/
recall.rs
