Significantly speed up bitmap computation

[magdalendobson]

Introduction

Bitmap computation in diskann-label-filter is unacceptably slow. Currently, with a 1 million size slice of yfcc and a 10k query set, computing the query bitmaps takes 43.10 seconds. With just a 100K slice of the caselaw dataset and a 10k query set, computing the bitmaps takes 6.03 seconds. This was making it hard to run experiments on filtered search algorithms for the full sizes of these datasets.

Speeding up the bitmap computation is conceptually simple. Instead of iterating over every base label for every query filter, we compute an inverted index for each label type, which maps the label value to the documents with the same value. Then, at query time, we query the inverted index for the relevant label values, and compose the resulting sets as necessary to find the documents satisfying the entire filter expression. At a high level, that is what this PR does.

Lower level details

The overall workflow of the main function, compute_query_bitmaps, is as follows:

1. Check whether the query expression contains any ASTExpr::Not clauses. If so, default to the existing slow path. This is because we don't store the document universe for each label, and thus can't compute the complement of an arbitrary bitset.
2. Otherwise, move to the fast path.
3. Flatten the base labels so that nested values map to a single string (e.g. the JSON string {"car": {"color":"red", "make":Mazda"}} would be transformed to {"car.color":red, "car.make":"Mazda}), and re-organize as a hash map of labels to values.
4. For each label, compute either an inverted index (strings and bools) or an B-tree (ints and floats) depending on its type.
5. At query time, use either the inverted index or the B-tree to produce a bitset for each CompareOp in the clause, and then compose them with AND and OR as needed to produce the final bitset.

We also add a utility to diskann-label-filter for computing the specificity of a set of query filters with respect to a base set, outputting some statistics on it, and optionally outputting the individual specificity values to a file for further processing.

Inverted Index

The inverted index maps each label value, converted to a string, to a bitset containing the doc ids corresponding to that value.

B-Tree

For simplicity, the B-tree implementation converts integers to floats before inserting so that we don't have to deal with two different types of B-tree. The performance of this piece of code isn't sensitive enough that it makes sense to differentiate, but this could be changed in the future.

The B-tree maps collections of ids to vectors instead of bitsets, because concatenating vectors is much cheaper than extending bitsets, and potentially many vectors would be concatenated during a range query.

Timings

Returning to the earlier discussion of timings, for the 1 million size slice of yfcc and a 10k query set, computing the query bitmaps now takes .6 seconds. For the 100K slice of the caselaw dataset and a 10k query set, computing the bitmaps now takes 1.728 seconds.

[Copilot]

Pull request overview

This PR targets a major performance improvement in diskann-label-filter by introducing a fast-path for computing per-query bitmaps using precomputed per-field accelerators (inverted-index style maps for equality and a numeric BTree for range queries), while falling back to the existing evaluator when NOT is present. It also adds an example utility for computing “specificity” statistics over query filters.

Changes:

• Add utils::compute_bitmap::compute_query_bitmaps implementing an accelerated bitmap computation path (with a NOT-guarded slow fallback).
• Export the new bitmap API from diskann-label-filter and add an example (compute_specificities) to compute stats/output.
• Minor doc comment updates in flattening utilities and dependency updates for the new module.

Reviewed changes

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

Show a summary per file

File	Description
diskann-label-filter/src/utils/flatten_utils.rs	Updates doc examples for configurable flattening (one example is currently inconsistent with behavior).
diskann-label-filter/src/utils/compute_bitmap.rs	New accelerated bitmap computation implementation plus unit tests.
diskann-label-filter/src/lib.rs	Exposes the new module and re-exports compute_query_bitmaps.
diskann-label-filter/examples/compute_specificities.rs	New example for computing/saving specificity stats from computed bitmaps.
diskann-label-filter/Cargo.toml	Adds dependencies needed by the new bitmap computation module.
Cargo.lock	Locks new transitive deps (bit-set, rayon) for this crate.

---

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

[magdalendobson]

Performance improvement of bitmap calculation is really needed to unblock serious work on the filtered algorithms, so I have a vested interest in seeing something like this merged ASAP.

That said, there is a lot of work needed to fully productionize it. As examples:

1. The conversion of i64 to f64 is an imprecise conversion. Similarly u64 would have to be treated separately from i64, which is currently isn't in AttributeValue.

With the move to diskann-tools, is this still something you would like to see handled?

2. I'm not convinced the handling of "." as a separator is robust: what happens if a user string contains a "."? For example, currently {"a.b": 1} and {"a": {"b": 1}} get lowered the same, creating ambiguity.

Doesn't our existing code already explicitly treat those two expressions as the same? E.g.

DiskANN/diskann-label-filter/src/parser/query_parser.rs

Line 50 in e3139b4

/// Helper to get a nested value from a label using dot notation (e.g., "specs.cpu")

3. Probably want to use RoaringBitSet instead of BitSet.

After the move to diskann-tools, since the existing drivers use usize for vector ids, I don't think it makes sense to move to RoaringBitSet.

4. This adds anyhow as a low-level library error type, which is not a great fit.

Resolved by the move to diskann-tools.

5. This also unconditionally adds rayon as a dependency of diskann-label-filter and doesn't provide a caller with a clear ability to opt-out.

Resolved by the move to diskann-tools.

6. Copy-paste in eval_query_using_accelerators that could be factored out.

Resolved in latest edits.

7. Many of the helper structs made public that don't necessarily need to be.
   This is probably fine for the datasets we are testing on where everything is "nice", but I have concerns about the overall correctness in the presence of corner cases.

In latest edits I made the helper functions and structs private.

8. Minor, but the PR description says R-tree when the implementation uses a B-tree.

Resolved.

To unblock algorithmic work, though, what if we do the following:

1. Put this behind an "experimental" feature flag and in an "experimental" module to clearly indicate that things can go wrong.
2. Gate the new dependencies on this features to avoid unconditionally adding dependencies (in particularly, anyhow and rayon).
3. Enable this feature in benchmarks to unblock filtered work.

[hildebrandmw]

Performance improvement of bitmap calculation is really needed to unblock serious work on the filtered algorithms, so I have a vested interest in seeing something like this merged ASAP.
That said, there is a lot of work needed to fully productionize it. As examples:

1. The conversion of i64 to f64 is an imprecise conversion. Similarly u64 would have to be treated separately from i64, which is currently isn't in AttributeValue.

With the move to diskann-tools, is this still something you would like to see handled?

2. I'm not convinced the handling of "." as a separator is robust: what happens if a user string contains a "."? For example, currently {"a.b": 1} and {"a": {"b": 1}} get lowered the same, creating ambiguity.

Doesn't our existing code already explicitly treat those two expressions as the same? E.g.

DiskANN/diskann-label-filter/src/parser/query_parser.rs

Line 50 in e3139b4

/// Helper to get a nested value from a label using dot notation (e.g., "specs.cpu")

3. Probably want to use RoaringBitSet instead of BitSet.

After the move to diskann-tools, since the existing drivers use usize for vector ids, I don't think it makes sense to move to RoaringBitSet.

4. This adds anyhow as a low-level library error type, which is not a great fit.

Resolved by the move to diskann-tools.

5. This also unconditionally adds rayon as a dependency of diskann-label-filter and doesn't provide a caller with a clear ability to opt-out.

Resolved by the move to diskann-tools.

6. Copy-paste in eval_query_using_accelerators that could be factored out.

Resolved in latest edits.

7. Many of the helper structs made public that don't necessarily need to be.
   This is probably fine for the datasets we are testing on where everything is "nice", but I have concerns about the overall correctness in the presence of corner cases.

In latest edits I made the helper functions and structs private.

8. Minor, but the PR description says R-tree when the implementation uses a B-tree.

Resolved.

To unblock algorithmic work, though, what if we do the following:

1. Put this behind an "experimental" feature flag and in an "experimental" module to clearly indicate that things can go wrong.
2. Gate the new dependencies on this features to avoid unconditionally adding dependencies (in particularly, anyhow and rayon).
3. Enable this feature in benchmarks to unblock filtered work.

Thanks Magdalen, most of this gets resolved by moving it to diskann-tools, which is a better fit.

[harsha-simhadri]

posted some questions inline. thanks

[magdalendobson]

posted some questions inline. thanks

I believe these are all now resolved, would you mind approving if you agree?