Support filtered search in IVF-RaBitQ - #2373
Draft
jamxia155 wants to merge 9 commits into
Draft
Conversation
Support bitset filters across all IVF-RaBitQ search modes and kernel paths. Use infinite distances as the sole invalid-result sentinel, preserving all finite vector IDs, including UINT32_MAX, and map invalid results to the public out-of-bounds record. Add L2 and inner-product coverage across search modes, bit widths, moderate and 1% pass rates, and under-filled result sets. Validate bitset coverage and skip distance computation for filtered samples.
Extends generate_groundtruth with optional bitset prefiltering so that exact-NN ground truth can be computed over a filtered subset of the dataset. New function create_bitset_filter(n_samples, filter_reject_rate) returns a packed uint32 numpy array (LSB-first, matching cuVS bitset layout) where bit i is set iff vector i passes the filter. Uses a modulo-1000 bucket scheme: vector i passes when i % 1000 >= round(filter_reject_rate * 1000), giving a reject rate within 0.1% of the requested value. calc_truth gains an optional bitset parameter: - GPU path: slices the packed words for each batch, transfers to device, and passes as a cuVS bitset prefilter to brute_force.search - CPU path: unpacks the same words to a boolean accept_mask and masks rejected distances to inf before argpartition Two new mutually exclusive CLI arguments: - --bitset <path>: load a pre-saved .npy uint32 bitset from file - --filter_reject_rate <float>: generate a bitset in memory via create_bitset_filter (cherry picked from commit 09281c1)
The --filter_reject_rate path previously generated a bitset in memory, used it to compute ground truth, then discarded it. The benchmark side (cuvs-bench C++) needs to run searches against the *exact same* bitset that was used to filter the GT, otherwise recall is computed against a stale GT. Persisting the generated bitset is the simplest way to keep producer and consumer in sync. The on-disk format is also switched from .npy to the cuVS .bin format used by the rest of the script (base / query / GT files). .npy is a Python-only format and cannot be read by the C++ blob loader. A single .bin format gives one round-trip across the Python producer, the Python --bitset loader, and the C++ benchmark consumer. With the bitset now persisted, create_bitset_filter no longer needs to be deterministic for reproducibility. Switch it to true Bernoulli sampling via numpy's default_rng to match the C++ generate_bernoulli helper in cpp/bench/ann/src/common/dataset.hpp -- one less subtle divergence between the two sides, and avoids the regular-pattern artifact of the old modulo-1000 scheme. - --filter_reject_rate now writes the generated bitset to <output>/groundtruth.filter.bin via write_bin, with shape (ceil(N/32), 1) -- matching the C++ blob<uint32_t> layout used by cuvs-bench's dataset loader. - --bitset loader switched from np.load to memmap_bin_file with dtype=np.uint32; the (ceil(N/32), 1) memmap is raveled to 1-D for use in calc_truth. - create_bitset_filter switched to numpy Bernoulli sampling (rng.random(n) >= reject_rate) for parity with the C++ side. (cherry picked from commit deabe08)
Loads a precomputed bitset from disk (cuVS .bin format) and feeds it through the existing filter plumbing in the IVF-Flat / CAGRA / brute- force wrappers. Pairs with the Python generate_groundtruth output at <gt_dir>/groundtruth.filter.bin so the benchmark searches against the exact bitset that was used to compute the prefiltered ground truth. - conf.hpp: new optional 'filter_bitset_file' field in dataset_conf; rejected at parse time if 'filtering_rate' is also set. - dataset.hpp: dataset<> ctor accepts the path and loads via the same blob<uint32_t>(file) machinery used for base/query sets. When this path is used, the ground_truth_map is told the GT is pre-filtered; it then counts bitset rejections during GT-map construction and logs a warning if any occur (signals a bitset/GT mismatch). - benchmark.hpp: forwards dataset_conf.filter_bitset_file at the dataset<T> construction site. (cherry picked from commit e2d014a)
The two prefilter examples in the --help epilog were written against an older CLI that took the dataset path as --dataset. It has been a positional argument for some time, and the surrounding examples were corrected earlier in this series, leaving these two as the only stale ones. Copying either of them would fail with an unrecognized-argument error. Also drop the mid-word "base.\<newline>fbin" line continuation so all six examples wrap the same way.
Two related defects in the prefiltered ground truth path. Batch re-basing was only correct for 32-aligned offsets. calc_truth builds an index per batch holding rows [i, i + n_batch), so the prefilter handed to it must be re-based to that batch. Slicing the bitset by word starts at global bit (i / 32) * 32, which equals i only when i is a multiple of 32, making correctness depend on the batch size constant. Replace the word slicing with slice_bitset(), which unpacks, realigns by i % 32, and repacks with a zero-padded tail, so any offset works. The GPU and CPU branches now share one implementation instead of duplicating the arithmetic. Neither side checked that the bitset covers the dataset. A short file was silently accepted, since NumPy clamps the out-of-range slice, and the uncovered rows were read out of bounds. The C++ consumer has the same hole, where the overrun is past the end of an mmap. Both sides now check the covered row count up front and fail with the actual numbers. The check is word-granular, since the file records only a word count. A bitset short by fewer than 32 rows needs the same number of words and is indistinguishable from a correct one; those rows read padding zeros, so they are rejected rather than read out of bounds. The C++ half is restructured so both filter branches derive the row count from one place. It stays behind the has_value() guard, so the base_set_size() header read is still skipped when no filter is configured.
If fewer vectors pass the filter than the requested k, no query can have k valid neighbors and the ground truth is padded with filtered-out ids at infinite distance. That was written out silently. Count the accepted vectors up front and fail before doing any search work. Only the global count matters, since results are merged across batches. Padding bits in the final word are excluded so a bitset with stray padding cannot inflate the count. The count uses a byte-wise popcount table rather than unpackbits, which would allocate eight bytes per bit. Also add --filter_seed, so a generated bitset can be reproduced without the saved file. It is deliberately not named --seed: query generation has its own seed and is unaffected.
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
jamxia155
force-pushed
the
ivf-rabitq-filtering
branch
from
July 29, 2026 02:36
94e9ba5 to
a90521b
Compare
jamxia155
added
feature request
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds support for filtered search in IVF-RaBitQ. Current PR is for bitset filter only.
Also added filtering-related features: