diskann-benchmark 0.55.0

DiskANN3 is a composable library for bringing scalable, accurate and cost-effective vector indexing to multiple databases.
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

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

use std::io::Write;

use diskann_benchmark_core::streaming::{executors::bigann, Executor};
use diskann_benchmark_runner::output::Output;

use crate::index::streaming::{
    managed::{self, Managed},
    stats::StreamStats,
};
use crate::inputs::graph_index::DynamicRunbookParams;

/// Run a streaming benchmark using the given runbook parameters.
///
/// `make_streamer` receives `max_points` from the loaded runbook and returns the
/// constructed streamer. This is shared between full-precision and spherical streaming
/// benchmarks to avoid duplicating the runbook load → run_with → stage banner → summary logic.
pub(super) fn run_streaming<T, F>(
    runbook_params: &DynamicRunbookParams,
    make_streamer: F,
    mut output: &mut dyn Output,
) -> anyhow::Result<Vec<managed::Stats<StreamStats>>>
where
    T: 'static,
    F: FnOnce(usize) -> anyhow::Result<bigann::WithData<T, u32, Managed<T, StreamStats>>>,
{
    let groundtruth_directory = runbook_params
        .resolved_gt_directory
        .as_ref()
        .ok_or_else(|| {
            anyhow::anyhow!("Ground truth directory path was not resolved during validation")
        })?;

    let mut runbook = bigann::RunBook::load(
        &runbook_params.runbook_path,
        &runbook_params.dataset_name,
        &mut bigann::ScanDirectory::new(groundtruth_directory)?,
    )?;

    let mut streamer = make_streamer(runbook.max_points())?;

    let mut results = Vec::new();
    let stages = runbook.len();
    let mut i = 1;

    runbook.run_with(
        &mut streamer,
        |o: managed::Stats<StreamStats>| -> anyhow::Result<()> {
            if o.inner().is_maintain() {
                let message = format!("Ran maintenance before stage {}", i);
                write!(output, "{}", crate::utils::SmallBanner(&message))?;
            } else {
                let message = format!("Finished stage {} of {}: {}", i, stages, o.inner().kind());
                write!(output, "{}", crate::utils::SmallBanner(&message))?;
                i += 1;
            }
            writeln!(output, "{}", o)?;
            results.push(o);
            Ok(())
        },
    )?;

    write!(
        output,
        "{}",
        crate::utils::SmallBanner("End of Run Summary")
    )?;

    writeln!(
        output,
        "{}",
        crate::index::streaming::stats::Summary::new(results.iter().map(|r| r.inner()))
    )?;

    Ok(results)
}