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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
/*
* Copyright (c) Microsoft Corporation.
* Licensed under the MIT license.
*/
//! # Benchmark Tools for Search Operations.
//!
//! Implementations integrate with search infrastructure via the [`Search`] trait, which
//! defines a loose collection of input parameters and output results. An implementation of
//! [`Search`] is sufficient to use the [`search`] function, which performs a parallelized
//! search with aggregated results.
//!
//! When benchmarking search operations, it can be useful to run the batch search multiple
//! times to obtain a distribution of performance results. This is facilitated by [`search_all`]
//! which automatically runs the batch search for a configurable number of iterations and
//! returns an aggregated result. Custom result aggregation is supported via the [`Aggregate`]
//! trait.
//!
//! ## Customization Points
//!
//! The search API provides several opportunities for customization while still using parts
//! of the predefined infrastructure. The list below summarizes these customization points
//! from the lowest level (with the most control and least reuse) to the highest level.
//!
//! * [`Search`]: The exact mechanics of query search are implementation defined. If
//! additional output metrics are desired, these can be captured in [`Search::Output`],
//! which can be just about any arbitrary type.
//!
//! * [`Aggregate`]: When using [`search_all`], custom aggregators are allowed independently
//! from the [`Search`] implementation. This provides some flexibility in result collection
//! while reusing a specific [`Search`] implementation.
//!
//! On the other hand, if [`Search`] is defined with custom output, then a new aggregator
//! will likely be required.
//!
//! ## Example
//!
//! The example below shows a toy implementation of [`Search`].
//!
//! ```rust
//! use std::{sync::Arc, num::NonZeroUsize};
//! use diskann_benchmark_core::search;
//!
//! /// A simple example implementation of the `Search` trait.
//! #[derive(Debug)]
//! struct Example {
//! /// Implementations are expected to store queries internally with the `Search`
//! /// infrastructure requesting the IDs to search over.
//! num_queries: usize,
//! }
//!
//! /// Example `Search::Parameters`.
//! #[derive(Debug, Clone)]
//! struct Parameters {
//! /// The number of IDs to return per-search. This is meant solely for example purposes.
//! num_ids: usize,
//! }
//!
//! /// Example `Search::Output`.
//! #[derive(Debug, PartialEq)]
//! struct Output {
//! /// The input for the search. This is meant solely for example purposes.
//! index: usize,
//! }
//!
//! impl Output {
//! fn new(index: usize) -> Self {
//! Self { index }
//! }
//! }
//!
//! impl search::Search for Example {
//! type Id = usize;
//! type Parameters = Parameters;
//! type Output = Output;
//!
//! /// Return the number of queries contained in `self`. The search infrastructure will
//! /// generate search requests for all indices in `0..self.num_queries()`.
//! fn num_queries(&self) -> usize {
//! self.num_queries
//! }
//!
//! fn id_count(&self, parameters: &Self::Parameters) -> search::IdCount {
//! search::IdCount::Dynamic(NonZeroUsize::new(parameters.num_ids))
//! }
//!
//! async fn search<O>(
//! &self,
//! parameters: &Self::Parameters,
//! buffer: &mut O,
//! index: usize,
//! ) -> diskann::ANNResult<Self::Output>
//! where
//! O: diskann::graph::SearchOutputBuffer<Self::Id> + Send
//! {
//! use diskann::graph::SearchOutputBuffer;
//!
//! // Fill the buffer with `index`.
//! buffer.extend((0..parameters.num_ids).map(|_| (index, 0.0f32)));
//! Ok(Output::new(index))
//! }
//! }
//!
//! // Run Search
//! let runtime = diskann_benchmark_core::tokio::runtime(1).unwrap();
//!
//! // Search over `Example` that contains 4 queries.
//! let results = search::search(
//! Arc::new(Example { num_queries: 4 }),
//! Parameters { num_ids: 3 },
//! NonZeroUsize::new(1).unwrap(),
//! &runtime
//! ).unwrap();
//!
//! // Number of results is equal to the number of queries.
//! assert_eq!(results.len(), 4);
//! assert_eq!(&*results.output(), &[0, 1, 2, 3].map(Output::new));
//! ```
//!
//! # Built-in Runners
//!
//! ## Graph Index
//!
//! * [`graph::search::Knn`]: K-nearest neighbors search for [`diskann::graph::DiskANNIndex`].
//! * [`graph::search::Range`]: Range search for [`diskann::graph::DiskANNIndex`].
//! * [`graph::MultiHop`]: Multi-hop filtered search for [`diskann::graph::DiskANNIndex`].
pub
pub use ResultIds;
pub use ;