Skip to content

Commit 75ef977

Browse files
authored
Create benchmarks for the flat index (full-precision in-memory implementation) (#1170)
Establishes a baseline benchmark for `FlatIndex::knn_search`, enabling evaluation of new capabilities and innovations in DiskANN's flat-scan implementation. ### Why Add Flat Index Benchmarks if Exhaustive Exists? **Exhaustive benchmarks** are lightweight benchmarks focused on quantizer speed and accuracy: they measure compress + distance performance in isolation, without involving the `FlatIndex` machinery or any data provider. They operate entirely outside the DiskANN provider/index abstraction. **Flat benchmarks** are end-to-end benchmarks that exercise `FlatIndex` with different strategy configurations. They exercise `FlatIndex` through the actual provider and search strategy abstractions, so they capture realistic runtime behavior, and integration overhead that exhaustive benchmarks intentionally skip. This PR enables full-precision in-memory flat search. Future work may add: - Quantized flat search - Quantized flat search with reranking (disk-based or in-memory) - Batched SIMD scan - Integration with diversity search ### Changes - New `diskann-benchmark/src/flat` benchmark module with: - `InMemProvider` wrapping `FastMemoryVectorProviderAsync` (cache-line-aligned vector storage) with identity `DataProvider` mapping - `FlatScanStrategy` and `FlatVisitor` implementing `SearchStrategy`/`DistancesUnordered` - `FlatSearcher` implementing the `Search` trait from `benchmark_core` - Input schema (`diskann-benchmark/src/inputs/flat.rs`) with dataset, distance metric, queries, groundtruth, k, thread counts, and reps - Integration test (`flat_search_integration`) and example input (`diskann-benchmark/example/flat-index.json`) - Performance test input for wikipedia-100K dataset (`diskann-benchmark/perf_test_inputs/wikipedia-100K-flat-index.json`) - Unrelated fix: renamed temp file in `run_integration_test` from `graph-index.json` to `input.json` - the helper is shared by all backends, not just graph-index Follows established patterns from the graph-index and exhaustive backends. ## Output example: ``` cargo run --package diskann-benchmark --release -- run --input-file diskann-benchmark\perf_test_inputs\wikipedia-100K-flat-index.json --output-file ./target/tmp/flat-index-output.json ###################### # Running Job 1 of 1 # ###################### Data: target/tmp\wikipedia_cohere/wikipedia_base.bin.crop_nb_100000 Data Type: float32 Distance: inner_product Queries: target/tmp\wikipedia_cohere/wikipedia_query.bin Groundtruth: target/tmp\wikipedia_cohere/wikipedia-100K K: 100 Threads: 4, 8 Reps: 1 Loading dataset... Loaded 100000 vectors of dimension 768 Queries: 5000, Groundtruth: 5000x100 K, Avg cmps, QPS - mean(max), Avg Latency, p99 Latency, Recall, Threads =================================================================================================================== 100, 100000, 123.0 (123.0), 32430.0us (32430.0us), 35711.0us (35711us), 0.9999899999999999, 4 100, 100000, 250.4 (250.4), 31909.2us (31909.2us), 33323.0us (33323us), 0.9999899999999999, 8 ```
1 parent ee7dbaa commit 75ef977

9 files changed

Lines changed: 831 additions & 3 deletions

File tree

diskann-benchmark-core/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040
//! in which method can fail, the [`anyhow::Error`] type balances generality and fidelity.
4141
4242
mod internal;
43-
pub(crate) mod utils;
43+
pub mod utils;
4444

4545
// Public Utility Modules
4646
pub mod recall;

diskann-benchmark-core/src/utils.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@
55

66
use diskann_benchmark_runner::utils::percentiles::AsF64Lossy;
77

8-
pub(crate) fn average_all<I>(x: I) -> f64
8+
/// Computes the arithmetic mean of an iterator of values, returning `0.0` when empty.
9+
pub fn average_all<I>(x: I) -> f64
910
where
1011
I: IntoIterator<Item: AsF64Lossy>,
1112
{
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
{
2+
"search_directories": [
3+
"test_data/disk_index_search"
4+
],
5+
"jobs": [
6+
{
7+
"type": "flat-search",
8+
"content": {
9+
"data": "disk_index_siftsmall_learn_256pts_data.fbin",
10+
"data_type": "float32",
11+
"distance": "squared_l2",
12+
"search": {
13+
"queries": "disk_index_sample_query_10pts.fbin",
14+
"groundtruth": "disk_index_10pts_idx_uint32_truth_search_res.bin",
15+
"k": 10,
16+
"num_threads": [1],
17+
"reps": 1
18+
}
19+
}
20+
}
21+
]
22+
}
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
{
2+
"search_directories": [
3+
"target/tmp"
4+
],
5+
"jobs": [
6+
{
7+
"type": "flat-search",
8+
"content": {
9+
"data": "wikipedia_cohere/wikipedia_base.bin.crop_nb_100000",
10+
"data_type": "float32",
11+
"distance": "inner_product",
12+
"search": {
13+
"queries": "wikipedia_cohere/wikipedia_query.bin",
14+
"groundtruth": "wikipedia_cohere/wikipedia-100K",
15+
"k": 100,
16+
"num_threads": [4, 8],
17+
"reps": 1
18+
}
19+
}
20+
}
21+
]
22+
}

diskann-benchmark/src/flat/mod.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
/*
2+
* Copyright (c) Microsoft Corporation.
3+
* Licensed under the MIT license.
4+
*/
5+
6+
use diskann_benchmark_runner::Registry;
7+
8+
mod search;
9+
10+
pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> {
11+
search::register_benchmarks(registry)
12+
}

0 commit comments

Comments
 (0)