Skip to content

Add a tool for computing groundtruth for a runbook#1263

Open
magdalendobson wants to merge 11 commits into
mainfrom
users/magdalen/runbook_gt_tool
Open

Add a tool for computing groundtruth for a runbook#1263
magdalendobson wants to merge 11 commits into
mainfrom
users/magdalen/runbook_gt_tool

Conversation

@magdalendobson

@magdalendobson magdalendobson commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

The groundtruth computation tool in the deprecated C++ repository seems to have some bugs, so I wrote a new one in Rust using our Runbook abstractions. This requires diskann-tools to take a dependency on diskann-benchmark-core, but this is acceptable.

Once this PR lands I will update the streaming groundtruth tool in big-ann-benchmarks to use the Rust code instead.

@codecov-commenter

codecov-commenter commented Jul 17, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 1.06383% with 186 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.30%. Comparing base (6caca35) to head (89cae80).

Files with missing lines Patch % Lines
...ann-tools/src/bin/compute_streaming_groundtruth.rs 0.00% 186 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1263      +/-   ##
==========================================
- Coverage   91.50%   90.30%   -1.21%     
==========================================
  Files         498      499       +1     
  Lines       95522    95708     +186     
==========================================
- Hits        87407    86427     -980     
- Misses       8115     9281    +1166     
Flag Coverage Δ
miri 90.30% <1.06%> (-1.21%) ⬇️
unittests 89.96% <1.06%> (-1.51%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
...ark-core/src/streaming/executors/bigann/runbook.rs 94.11% <100.00%> (ø)
diskann-tools/src/utils/ground_truth.rs 42.28% <100.00%> (ø)
...ann-tools/src/bin/compute_streaming_groundtruth.rs 0.00% <0.00%> (ø)

... and 39 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@magdalendobson
magdalendobson marked this pull request as ready for review July 21, 2026 15:06
@magdalendobson
magdalendobson requested review from a team and Copilot July 21, 2026 15:06

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR adds a Rust-based tool for computing streaming groundtruth for BigANN-style runbooks, and refactors shared runbook parsing/execution abstractions so they can be used from diskann-tools without depending on diskann-benchmark-core.

Changes:

  • Moved/introduced reusable streaming + BigANN runbook parsing/execution APIs under diskann-utils::streaming (and re-exported them from diskann-benchmark-core).
  • Added a new diskann-tools binary (compute_streaming_groundtruth) that simulates runbook stages and writes step<stage>.gt<k> files.
  • Added file-based UX fixtures for runbook parsing/validation errors.

Reviewed changes

Copilot reviewed 19 out of 73 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
Cargo.lock Updates lockfile for new/shifted dependencies (e.g., anyhow, serde_yaml via diskann-utils).
diskann-benchmark-core/Cargo.toml Switches bigann feature to depend on diskann-utils/bigann instead of serde_yaml directly.
diskann-benchmark-core/src/lib.rs Removes benchmark-core-local UX testing helpers now moved to diskann-utils.
diskann-benchmark-core/src/streaming/api.rs Re-exports streaming traits/types from diskann-utils::streaming.
diskann-benchmark-core/src/streaming/executors/bigann/mod.rs Re-exports BigANN runbook types from diskann-utils; keeps benchmark-specific adaptors.
diskann-benchmark-core/src/streaming/executors/bigann/validate.rs Removes benchmark-core-local Validate stream implementation (now in diskann-utils).
diskann-benchmark/src/index/bftree/full_precision_streaming.rs Refactors type aliases/return type for streaming construction to use new BigANN exports.
diskann-benchmark/src/index/bftree/spherical_streaming.rs Refactors type aliases/return type for streaming construction to use new BigANN exports.
diskann-tools/Cargo.toml Enables diskann-utils bigann feature so tools can parse/run BigANN runbooks.
diskann-tools/src/bin/compute_streaming_groundtruth.rs New CLI tool to compute groundtruth for each search stage of a runbook.
diskann-tools/src/utils/ground_truth.rs Makes write_ground_truth public for reuse by the new streaming groundtruth tool.
diskann-utils/Cargo.toml Adds anyhow dependency and optional serde_yaml for bigann; adds test-only UX deps.
diskann-utils/src/lib.rs Exposes new streaming module and adds test-only UX helper module for fixtures.
diskann-utils/src/streaming/api.rs New shared streaming traits (Stream, Executor, etc.) moved out of benchmark-core.
diskann-utils/src/streaming/executors/bigann/mod.rs New shared BigANN runbook executor module under diskann-utils.
diskann-utils/src/streaming/executors/bigann/parsing.rs Shared YAML parsing + UX tests for BigANN runbooks.
diskann-utils/src/streaming/executors/bigann/runbook.rs Shared RunBook executor and Stage model.
diskann-utils/src/streaming/executors/bigann/validate.rs Shared runbook validation stream for active tag-set correctness.
diskann-utils/src/streaming/executors/mod.rs Adds bigann executor module behind the bigann feature.
diskann-utils/src/streaming/mod.rs New streaming module entry point and re-exports.
diskann-utils/tests/bigann-ux/delete-invalid-range/dataset.txt UX fixture: dataset name input for delete invalid range case.
diskann-utils/tests/bigann-ux/delete-invalid-range/expected.txt UX fixture: expected error output for delete invalid range case.
diskann-utils/tests/bigann-ux/delete-invalid-range/runbook.yaml UX fixture: runbook triggering delete invalid range validation.
diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/dataset.txt UX fixture: dataset name input for invalid dataset key type case.
diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/expected.txt UX fixture: expected error output for invalid dataset key type case.
diskann-utils/tests/bigann-ux/dataset-key-not-integer-or-string/runbook.yaml UX fixture: runbook with non-scalar dataset key.
diskann-utils/tests/bigann-ux/dataset-not-found/dataset.txt UX fixture: dataset name input for dataset-not-found case.
diskann-utils/tests/bigann-ux/dataset-not-found/expected.txt UX fixture: expected error output for dataset-not-found case.
diskann-utils/tests/bigann-ux/dataset-not-found/runbook.yaml UX fixture: runbook containing alternative dataset keys.
diskann-utils/tests/bigann-ux/dataset-value-not-mapping/dataset.txt UX fixture: dataset name input for dataset-value-not-mapping case.
diskann-utils/tests/bigann-ux/dataset-value-not-mapping/expected.txt UX fixture: expected error output for dataset-value-not-mapping case.
diskann-utils/tests/bigann-ux/dataset-value-not-mapping/runbook.yaml UX fixture: runbook with a non-mapping stage payload.
diskann-utils/tests/bigann-ux/insert-invalid-range/dataset.txt UX fixture: dataset name input for insert invalid range case.
diskann-utils/tests/bigann-ux/insert-invalid-range/expected.txt UX fixture: expected error output for insert invalid range case.
diskann-utils/tests/bigann-ux/insert-invalid-range/runbook.yaml UX fixture: runbook triggering insert invalid range validation.
diskann-utils/tests/bigann-ux/insert-start-equals-end/dataset.txt UX fixture: dataset name input for insert start==end case.
diskann-utils/tests/bigann-ux/insert-start-equals-end/expected.txt UX fixture: expected error output for insert start==end case.
diskann-utils/tests/bigann-ux/insert-start-equals-end/runbook.yaml UX fixture: runbook triggering insert start==end validation.
diskann-utils/tests/bigann-ux/missing-max-pts/dataset.txt UX fixture: dataset name input for missing max_pts case.
diskann-utils/tests/bigann-ux/missing-max-pts/expected.txt UX fixture: expected error output for missing max_pts case.
diskann-utils/tests/bigann-ux/missing-max-pts/runbook.yaml UX fixture: runbook missing required max_pts key.
diskann-utils/tests/bigann-ux/non-integer-max-pts/dataset.txt UX fixture: dataset name input for non-integer max_pts case.
diskann-utils/tests/bigann-ux/non-integer-max-pts/expected.txt UX fixture: expected error output for non-integer max_pts case.
diskann-utils/tests/bigann-ux/non-integer-max-pts/runbook.yaml UX fixture: runbook with non-integer max_pts value.
diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/dataset.txt UX fixture: dataset name input for replace ids_start==ids_end case.
diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/expected.txt UX fixture: expected error output for replace ids_start==ids_end case.
diskann-utils/tests/bigann-ux/replace-ids-start-equals-end/runbook.yaml UX fixture: runbook triggering replace ids_start==ids_end validation.
diskann-utils/tests/bigann-ux/replace-invalid-ids-range/dataset.txt UX fixture: dataset name input for replace invalid ids range case.
diskann-utils/tests/bigann-ux/replace-invalid-ids-range/expected.txt UX fixture: expected error output for replace invalid ids range case.
diskann-utils/tests/bigann-ux/replace-invalid-ids-range/runbook.yaml UX fixture: runbook triggering replace invalid ids range validation.
diskann-utils/tests/bigann-ux/replace-invalid-tags-range/dataset.txt UX fixture: dataset name input for replace invalid tags range case.
diskann-utils/tests/bigann-ux/replace-invalid-tags-range/expected.txt UX fixture: expected error output for replace invalid tags range case.
diskann-utils/tests/bigann-ux/replace-invalid-tags-range/runbook.yaml UX fixture: runbook triggering replace invalid tags range validation.
diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/dataset.txt UX fixture: dataset name input for replace tags_start==tags_end case.
diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/expected.txt UX fixture: expected error output for replace tags_start==tags_end case.
diskann-utils/tests/bigann-ux/replace-tags-start-equals-end/runbook.yaml UX fixture: runbook triggering replace tags_start==tags_end validation.
diskann-utils/tests/bigann-ux/runbook-does-not-exist/dataset.txt UX fixture: dataset name input for missing runbook file case.
diskann-utils/tests/bigann-ux/runbook-does-not-exist/expected.txt UX fixture: expected error output for missing runbook file case.
diskann-utils/tests/bigann-ux/runbook-not-yaml-map/dataset.txt UX fixture: dataset name input for invalid YAML top-level type case.
diskann-utils/tests/bigann-ux/runbook-not-yaml-map/expected.txt UX fixture: expected error output for invalid YAML top-level type case.
diskann-utils/tests/bigann-ux/runbook-not-yaml-map/runbook.yaml UX fixture: runbook that is not a YAML mapping.
diskann-utils/tests/bigann-ux/stage-key-not-integer/dataset.txt UX fixture: dataset name input for non-integer stage key case.
diskann-utils/tests/bigann-ux/stage-key-not-integer/expected.txt UX fixture: expected error output for non-integer stage key case.
diskann-utils/tests/bigann-ux/stage-key-not-integer/runbook.yaml UX fixture: runbook with float stage key (e.g. 1.1).
diskann-utils/tests/bigann-ux/stage-key-not-number/dataset.txt UX fixture: dataset name input for non-number stage key case.
diskann-utils/tests/bigann-ux/stage-key-not-number/expected.txt UX fixture: expected error output for non-number stage key case.
diskann-utils/tests/bigann-ux/stage-key-not-number/runbook.yaml UX fixture: runbook with string stage key.
diskann-utils/tests/bigann-ux/unrecognized-dataset-key/dataset.txt UX fixture: dataset name input for unknown dataset key case.
diskann-utils/tests/bigann-ux/unrecognized-dataset-key/expected.txt UX fixture: expected error output for unknown dataset key case.
diskann-utils/tests/bigann-ux/unrecognized-dataset-key/runbook.yaml UX fixture: runbook with unrecognized dataset-level key.
diskann-utils/tests/bigann-ux/unrecognized-operation/dataset.txt UX fixture: dataset name input for unrecognized operation case.
diskann-utils/tests/bigann-ux/unrecognized-operation/expected.txt UX fixture: expected error output for unrecognized operation case.
diskann-utils/tests/bigann-ux/unrecognized-operation/runbook.yaml UX fixture: runbook with unrecognized operation value.
Comments suppressed due to low confidence (1)

diskann-utils/src/streaming/executors/bigann/parsing.rs:52

  • Spelling typo in the context message: "precessing" should be "processing". This is user-facing via error contexts and makes UX tests/logs harder to read.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +60 to +62
let n_base = dataset.nrows();
let n_queries = queries.nrows();
let recall_at = args.recall_at as usize;
Comment thread diskann-tools/src/bin/compute_streaming_groundtruth.rs Outdated
Comment on lines +173 to +179
for (offset, eid) in dataset_offsets.clone().zip(ids.clone()) {
if offset < n_base {
active[offset] = true;
ext_id[offset] = eid as u32;
ext_to_offset.insert(eid as u32, offset);
}
}
Copilot AI review requested due to automatic review settings July 21, 2026 16:25

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 67 out of 69 changed files in this pull request and generated 9 comments.

Comment on lines +110 to +119
Value::Number(stage) => match stage.as_i64() {
Some(stage) => {
stages.push(
handle_stage(stage as usize, value)
.with_context(|| format!("processing stage {}", stage))?,
);
Ok(())
}
None => anyhow::bail!("Stage \"{}\" must be an integer", stage),
},
Comment on lines +221 to +225
impl<'a> TryAs<'a, usize> for Value {
fn try_as(&'a self) -> Result<usize, &'static str> {
self.as_i64().map(|i| i as usize).ok_or("usize")
}
}
Comment on lines +146 to +150
let mut validator = validate::Validate::new();
this.run_with(&mut validator, |_| Ok(()))?;

this.max_points = this.max_points.max(validator.max_active());
this.max_tag = validator.max_tag();
Ok(this)
};

inner().context("trying to parse an \"replace\"")

Caused by:
0: processing stage 1
1: trying to parse an "replace"

Caused by:
0: processing stage 1
1: trying to parse an "replace"

Caused by:
0: processing stage 1
1: trying to parse an "replace"

Caused by:
0: processing stage 1
1: trying to parse an "replace"
Comment on lines +48 to +57
raw.stages.sort_by_key(|s| s.index);

let context = |index: usize| {
format!(
"precessing stage {} of dataset \"{}\" in file \"{}\"",
index,
dataset,
path.display()
)
};
Copilot AI review requested due to automatic review settings July 21, 2026 16:33

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 6 out of 7 changed files in this pull request and generated 4 comments.

Comments suppressed due to low confidence (4)

diskann-tools/src/bin/compute_streaming_groundtruth.rs:63

  • The tool doesn’t validate that the base and query matrices have the same dimensionality. If they differ, the distance function may panic (it is specialized with Some(dataset.ncols())) or produce incorrect results. Fail fast with a clear error message before running the runbook.
    let n_base = dataset.nrows();
    let n_queries = queries.nrows();
    let recall_at = args.recall_at as usize;

diskann-tools/src/bin/compute_streaming_groundtruth.rs:137

  • Insert stages silently ignore out-of-bounds offsets and truncate offsets/IDs to u32 via as. Both lead to silently corrupted groundtruth. This should fail fast if the runbook references offsets outside the base dataset or if an external id cannot fit in the u32 groundtruth format.
                for id in dataset_offsets_and_ids.clone() {
                    if id < n_base {
                        active[id] = true;
                        ext_id[id] = id as u32;
                        ext_to_offset.insert(id as u32, id);
                    }
                }

diskann-tools/src/bin/compute_streaming_groundtruth.rs:179

  • Replace stages can silently produce a wrong active set: (1) dataset_offsets and ids ranges aren’t validated to have equal length, so zip may drop unmatched elements; (2) out-of-bounds offsets are silently skipped; (3) usize -> u32 casts can truncate; (4) missing IDs are silently ignored. Validate range lengths, bounds, and ID conversions, and error on inconsistencies.
                // Remove old vectors by external ID.
                for eid in ids.clone() {
                    if let Some(&offset) = ext_to_offset.get(&(eid as u32)) {
                        active[offset] = false;
                        ext_to_offset.remove(&(eid as u32));
                    }
                }
                // Insert new vectors: they inherit the external IDs from `ids`.
                for (offset, eid) in dataset_offsets.clone().zip(ids.clone()) {
                    if offset < n_base {
                        active[offset] = true;
                        ext_id[offset] = eid as u32;
                        ext_to_offset.insert(eid as u32, offset);
                    }
                }

diskann-benchmark-core/src/streaming/executors/bigann/runbook.rs:456

  • Most runbook/ScanDirectory unit tests were removed; only test_runbook remains. This leaves ScanDirectory::find_groundtruth and YAML loading/groundtruth resolution behavior without unit test coverage in this module.
                matches,
            )),
        }
    }
}

///////////
// Tests //
///////////

Comment on lines +193 to +200
let timer = Instant::now();
let n_active = active.iter().filter(|&&b| b).count();

tracing::info!(
"Stage {}: computing top-{} groundtruth for {} active vectors against {} queries",
stage_idx, recall_at, n_active, n_queries,
);

Comment on lines +241 to +252
if !under_k.is_empty() {
tracing::warn!(
"Stage {}: {} / {} queries have fewer than {} results (active set = {}). \
Query indices: {:?}",
stage_idx,
under_k.len(),
n_queries,
recall_at,
n_active,
under_k,
);
}
Comment on lines +147 to +152
for eid in ids.clone() {
if let Some(&offset) = ext_to_offset.get(&(eid as u32)) {
active[offset] = false;
ext_to_offset.remove(&(eid as u32));
}
}
Comment thread diskann-tools/Cargo.toml
Copilot AI review requested due to automatic review settings July 21, 2026 16:40

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 6 out of 7 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (2)

diskann-tools/src/bin/compute_streaming_groundtruth.rs:252

  • The code warns when a query has fewer than recall_at neighbors, but still calls write_ground_truth. write_ground_truth assumes every query has exactly number_of_neighbors entries and will panic when serializing (slice length mismatch) if n_active < recall_at.

Handle the n_active < recall_at case by padding each per-query queue with sentinel neighbors (unique ids, distance = f32::INFINITY) before writing, and avoid logging a potentially huge under_k index list.

                // Warn about queries that got fewer than K results (active set smaller than K).
                let under_k: Vec<usize> = results
                    .iter()
                    .enumerate()
                    .filter_map(|(qi, pq)| {
                        if pq.size() < recall_at {
                            Some(qi)
                        } else {
                            None
                        }
                    })
                    .collect();
                if !under_k.is_empty() {
                    tracing::warn!(
                        "Stage {}: {} / {} queries have fewer than {} results (active set = {}). \
                         Query indices: {:?}",
                        stage_idx,
                        under_k.len(),
                        n_queries,
                        recall_at,
                        n_active,
                        under_k,
                    );
                }

diskann-tools/src/bin/compute_streaming_groundtruth.rs:63

  • The tool never validates that the base dataset and query matrices have the same dimension. If they differ, distance evaluation can panic or silently compute garbage.

Add an explicit dimension check after loading the matrices and return a clear error before proceeding.

    let n_base = dataset.nrows();
    let n_queries = queries.nrows();
    let recall_at = args.recall_at as usize;

Comment thread diskann-tools/src/bin/compute_streaming_groundtruth.rs Outdated
Copilot AI review requested due to automatic review settings July 21, 2026 20:09

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 6 out of 7 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (3)

diskann-tools/src/bin/compute_streaming_groundtruth.rs:134

  • The Insert stage silently ignores dataset offsets >= n_base and truncates external IDs with as u32. That can produce incorrect groundtruth (missing active vectors) or corrupt IDs if the runbook refers to out-of-bounds offsets / large tags. It also assumes a dataset offset cannot be activated twice; if the runbook reuses offsets, this implementation will silently overwrite state and can later generate duplicate neighbor IDs (undefined behavior for NeighborPriorityQueue).
                for id in dataset_offsets_and_ids.clone() {
                    if id < n_base {
                        active[id] = true;
                        ext_id[id] = id as u32;
                        ext_to_offset.insert(id as u32, id);
                    }

diskann-tools/src/bin/compute_streaming_groundtruth.rs:149

  • The Delete stage currently ignores IDs that aren't present in ext_to_offset (and also truncates IDs via as u32). Given RunBook validation requires delete ranges to be valid, silently skipping a delete will desync the active set and yield wrong groundtruth; better to fail fast with a clear error.
                for eid in ids.clone() {
                    if let Some(&offset) = ext_to_offset.get(&(eid as u32)) {
                        active[offset] = false;
                        ext_to_offset.remove(&(eid as u32));
                    }

diskann-tools/src/bin/compute_streaming_groundtruth.rs:177

  • Replace assumes the dataset_offsets and ids ranges are the same length (it zips them), silently ignores out-of-bounds dataset offsets, and truncates external IDs via as u32. If the YAML has mismatched ranges or large IDs, this will silently drop replacements and produce incorrect groundtruth. Also, removing a tag that isn't present currently gets ignored, which can hide state-tracking bugs.
                // Remove old vectors by external ID.
                for eid in ids.clone() {
                    if let Some(&offset) = ext_to_offset.get(&(eid as u32)) {
                        active[offset] = false;
                        ext_to_offset.remove(&(eid as u32));
                    }
                }
                // Insert new vectors: they inherit the external IDs from `ids`.
                for (offset, eid) in dataset_offsets.clone().zip(ids.clone()) {
                    if offset < n_base {
                        active[offset] = true;
                        ext_id[offset] = eid as u32;
                        ext_to_offset.insert(eid as u32, offset);
                    }
                }


let distance_fn = V::distance(args.distance_function, Some(dataset.ncols()));

for (stage_idx, stage) in runbook.stages().iter().enumerate() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can this be re-written as a Stream instead of iterating over Stages?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, done

Copilot AI review requested due to automatic review settings July 24, 2026 16:10

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

Copilot AI review requested due to automatic review settings July 24, 2026 16:11

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

Copilot AI review requested due to automatic review settings July 24, 2026 16:18

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

}
impl FindGroundtruth for AllowMissing {
fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result<PathBuf> {
Ok(self.dir.join(format!("step{}.{}", stage, self.suffix)))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a small detail, but we can expose the name mangling for ground-truth as part of the ScanDirectory struct in diskann-benchmark-core - just as a single source of truth for the expected file name.

queries: &'a Matrix<V>,
distance_fn: V::Distance,
n_base: usize,
n_queries: usize,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are n_queries and n_base tracked here? Can't they be fetched from queries and dataset respectively?

recall_at: usize,
active: Vec<bool>,
ext_id: Vec<u32>,
ext_to_offset: HashMap<u32, usize>,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unless I'm missing something, isn't ext_to_offset sufficient? It can track whether or not an ID is active by its presence or absence. And you can iterate over (k, v) pairs to get all the necessary distances.

})
.collect();

let under_k: Vec<usize> = results

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wouldn't it be a lot easier to check if the number of active IDs was ever less than K? I think that's the only reason this could happen - and the error message would be much clearer. This is also something that can be easily added to the validation step for RunBook and could be queried through an API. That way, we'd get the error almost immediately instead of potentially deep into a computation.

/// Returns a reference to the stages in this runbook.
#[cfg(test)]
pub(super) fn stages(&self) -> &[Stage] {
pub fn stages(&self) -> &[Stage] {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think for now this can be left as #[cfg(test)].

DataType::Uint8 => run::<u8>(&args),
DataType::Int8 => run::<i8>(&args),
}
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IF you want a small compile-time win, we can do something like this:

trait Distance {
    fn distance(&self, query: usize, data: usize) -> anyhow::Result<f32>;
}

struct Foo<T: VectorRepr> {
    data: Matrix<T>,
    queries: Matrix<T>,
    distance: T::Distance,
}

impl<T: VectorRepr> Distance for Foo<T> {
    fn distance(&self, query: usize, data: usize) -> anyhow::Result<f32> {
        // Should return a better error type for out-of-bounds
        distance.evaluate_similarity(self.queries.row[query], self.data.row[data])
    }
}

The rest of the code can take &dyn Distance and that way all the streaming/computation loops are only compiled once. Alternaticely, the distance could take a &HashMap<u32, usize> and a &mut NeighborPriorityQueue and do all the distance computations behind the trait object boundary.

Just a thought to move us towards faster compiling patterns.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants