diskann-benchmark 0.51.0

DiskANN is a fast approximate nearest neighbor search library for high dimensional data
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

use std::{num::NonZeroUsize, sync::Arc};

use diskann_benchmark_core::{self as benchmark_core, search as core_search};

use crate::{backend::index::result::RangeSearchResults, inputs::graph_index::GraphRangeSearch};

#[derive(Debug, Clone, Copy)]
pub(crate) struct RangeSearchSteps<'a> {
    pub(crate) reps: NonZeroUsize,
    pub(crate) num_tasks: &'a [NonZeroUsize],
    pub(crate) runs: &'a [GraphRangeSearch],
}

impl<'a> RangeSearchSteps<'a> {
    pub(crate) fn new(
        reps: NonZeroUsize,
        num_tasks: &'a [NonZeroUsize],
        runs: &'a [GraphRangeSearch],
    ) -> Self {
        Self {
            reps,
            num_tasks,
            runs,
        }
    }
}

type Run = core_search::Run<diskann::graph::search::Range>;

pub(crate) trait Range<I> {
    fn search_all(
        &self,
        parameters: Vec<Run>,
        groundtruth: &dyn benchmark_core::recall::Rows<I>,
    ) -> anyhow::Result<Vec<RangeSearchResults>>;
}

pub(crate) fn run<I>(
    runner: &dyn Range<I>,
    groundtruth: &dyn benchmark_core::recall::Rows<I>,
    steps: RangeSearchSteps<'_>,
) -> anyhow::Result<Vec<RangeSearchResults>> {
    let mut all = Vec::new();

    for threads in steps.num_tasks.iter() {
        for run in steps.runs.iter() {
            let setup = core_search::Setup {
                threads: *threads,
                tasks: *threads,
                reps: steps.reps,
            };

            let parameters: Vec<_> = run
                .construct_params()?
                .into_iter()
                .map(|range_search_params| {
                    core_search::Run::new(range_search_params, setup.clone())
                })
                .collect();

            all.extend(runner.search_all(parameters, groundtruth)?);
        }
    }

    Ok(all)
}

///////////
// Impls //
///////////

impl<DP, T, S> Range<DP::InternalId> for Arc<core_search::graph::Range<DP, T, S>>
where
    DP: diskann::provider::DataProvider,
    core_search::graph::Range<DP, T, S>: core_search::Search<
        Id = DP::InternalId,
        Parameters = diskann::graph::search::Range,
        Output = core_search::graph::range::Metrics,
    >,
{
    fn search_all(
        &self,
        parameters: Vec<core_search::Run<diskann::graph::search::Range>>,
        groundtruth: &dyn benchmark_core::recall::Rows<DP::InternalId>,
    ) -> anyhow::Result<Vec<RangeSearchResults>> {
        let results = core_search::search_all(
            self.clone(),
            parameters.into_iter(),
            core_search::graph::range::Aggregator::new(groundtruth),
        )?;

        Ok(results.into_iter().map(RangeSearchResults::new).collect())
    }
}