Add a tool for computing groundtruth for a runbook#1263
Add a tool for computing groundtruth for a runbook#1263magdalendobson wants to merge 11 commits into
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ 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
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
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 fromdiskann-benchmark-core). - Added a new
diskann-toolsbinary (compute_streaming_groundtruth) that simulates runbook stages and writesstep<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.
| let n_base = dataset.nrows(); | ||
| let n_queries = queries.nrows(); | ||
| let recall_at = args.recall_at as usize; |
| 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); | ||
| } | ||
| } |
| 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), | ||
| }, |
| 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") | ||
| } | ||
| } |
| 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" |
| raw.stages.sort_by_key(|s| s.index); | ||
|
|
||
| let context = |index: usize| { | ||
| format!( | ||
| "precessing stage {} of dataset \"{}\" in file \"{}\"", | ||
| index, | ||
| dataset, | ||
| path.display() | ||
| ) | ||
| }; |
There was a problem hiding this comment.
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
u32viaas. 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 theu32groundtruth 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_offsetsandidsranges aren’t validated to have equal length, sozipmay drop unmatched elements; (2) out-of-bounds offsets are silently skipped; (3)usize -> u32casts 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_runbookremains. This leavesScanDirectory::find_groundtruthand YAML loading/groundtruth resolution behavior without unit test coverage in this module.
matches,
)),
}
}
}
///////////
// Tests //
///////////
| 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, | ||
| ); | ||
|
|
| 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, | ||
| ); | ||
| } |
| 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)); | ||
| } | ||
| } |
There was a problem hiding this comment.
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_atneighbors, but still callswrite_ground_truth.write_ground_truthassumes every query has exactlynumber_of_neighborsentries and will panic when serializing (slice length mismatch) ifn_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;
There was a problem hiding this comment.
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 viaas 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_offsetsandidsranges are the same length (it zips them), silently ignores out-of-bounds dataset offsets, and truncates external IDs viaas 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() { |
There was a problem hiding this comment.
Can this be re-written as a Stream instead of iterating over Stages?
| } | ||
| impl FindGroundtruth for AllowMissing { | ||
| fn find_groundtruth(&mut self, stage: usize) -> anyhow::Result<PathBuf> { | ||
| Ok(self.dir.join(format!("step{}.{}", stage, self.suffix))) |
There was a problem hiding this comment.
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, |
There was a problem hiding this comment.
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>, |
There was a problem hiding this comment.
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 |
There was a problem hiding this comment.
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] { |
There was a problem hiding this comment.
I think for now this can be left as #[cfg(test)].
| DataType::Uint8 => run::<u8>(&args), | ||
| DataType::Int8 => run::<i8>(&args), | ||
| } | ||
| } |
There was a problem hiding this comment.
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.
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-toolsto take a dependency ondiskann-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.