Beta filter for disk search (impl, RFC #1101)#1121
Conversation
There was a problem hiding this comment.
Pull request overview
Implements RFC #1101 by introducing a typed SearchPlan<'_> (with hierarchical GraphMode) to replace the (vector_filter, is_flat_search) pair on disk search, and adds the new disk-side BetaFilter traversal bias (PQ-distance scaling by β for predicate matches) while keeping post-filter and rerank semantics intact.
Changes:
- Replaces the disk search API’s filter/flat-search parameter pair with
SearchPlan+GraphMode, and projects this into an internalFilterModeused by the disk search strategy. - Adds disk-path
GraphMode::BetaFilterbehavior by scaling PQ distances for predicate matches during traversal. - Updates tools/benchmarks and benchmark JSON inputs to construct
SearchPlanand migrateis_flat_search→search_mode.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| diskann-tools/src/utils/search_disk_index.rs | Constructs SearchPlan at the CLI boundary and passes it into disk search. |
| diskann-disk/src/search/provider/disk_provider.rs | Plumbs SearchPlan/FilterMode through disk search, applies β bias in pq_distances, and updates tests. |
| diskann-disk/src/search/mod.rs | Exposes the new search::filter_parameter module. |
| diskann-disk/src/search/filter_parameter.rs | New module defining SearchPlan, GraphMode, β validation, and internal FilterMode. |
| diskann-disk/src/lib.rs | Removes build-time filter_parameter re-export now that filtering is search-scoped. |
| diskann-disk/src/build/mod.rs | Drops filter_parameter re-export from build module. |
| diskann-disk/src/build/configuration/mod.rs | Removes filter_parameter submodule from build configuration. |
| diskann-disk/src/build/configuration/filter_parameter.rs | Deletes the old build/configuration filter parameter module. |
| diskann-disk/src/build/builder/core.rs | Updates tests/call sites to use SearchPlan::graph(). |
| diskann-benchmark/src/inputs/disk.rs | Introduces SearchMode and migrates JSON schema from is_flat_search to search_mode. |
| diskann-benchmark/src/backend/disk_index/search.rs | Builds SearchPlan from (search_mode, vector_filters_file) and passes it into disk search. |
| diskann-benchmark/perf_test_inputs/wikipedia-100K-disk-index.json | Migrates input JSON to search_mode. |
| diskann-benchmark/perf_test_inputs/openai-100K-disk-index.json | Migrates input JSON to search_mode. |
| diskann-benchmark/example/disk-index.json | Migrates example JSON to search_mode. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| let mut distance = self.scratch.pq_scratch.aligned_dist_scratch[i]; | ||
| if let FilterMode::BetaFilter { predicate, beta } = self.filter { | ||
| if predicate(*id) { | ||
| distance *= beta; | ||
| } |
| impl<Data, VP> SearchExt for DiskAccessor<'_, Data, VP> | ||
| where | ||
| Data: GraphDataType<VectorIdType = u32>, | ||
| Data: GraphDataType<VectorIdType = u32>, |
| /// Beta-biased beam: matching vectors' PQ distances multiplied by | ||
| /// `beta` ∈ (0, 1] in `pq_distances`. Predicate also post-filters. | ||
| BetaFilter { predicate: Predicate<'a>, beta: f32 }, |
| @@ -71,7 +87,8 @@ pub(crate) struct DiskSearchPhase { | |||
| pub(crate) beam_width: usize, | |||
| pub(crate) search_list: Vec<u32>, | |||
| pub(crate) recall_at: u32, | |||
| pub(crate) is_flat_search: bool, | |||
| #[serde(default)] | |||
| pub(crate) search_mode: SearchMode, | |||
| pub(crate) distance: SimilarityMeasure, | |||
| @@ -8,8 +8,12 @@ use std::{collections::HashSet, sync::atomic::AtomicBool, time::Instant}; | |||
| use diskann::utils::IntoUsize; | |||
There was a problem hiding this comment.
Worth noting this whole module diskann-tools/src/utils/search_disk_index.r looks dead. pub fn search_disk_index / SearchDiskIndexParameters aren't referenced from any binary.
I recommend deleting search_disk_index.rs (and its pub mod / pub use in utils/mod.rs) in a follow-up rather than carry the is_flat_search: bool migration.
|
|
||
| use thiserror::Error; | ||
|
|
||
| /// Closure used to filter vector IDs during disk search. |
There was a problem hiding this comment.
Consider pruning the implementation-detail half of this doc. The invocation-site list (flat_search, pq_distances, RerankAndFilter::post_process) is internal — consumers of Predicate don't need it to use the alias, and naming those sites here couples the public doc to the current internal layout.
| OB: search_output_buffer::SearchOutputBuffer<(u32, Data::AssociatedDataType)> + Send, | ||
| { | ||
| let provider = self.index.provider(); | ||
| let predicate = strategy.filter.post_filter(); |
There was a problem hiding this comment.
For flat_searchcalling FilterMode::post_filter(), flat scan applies the predicate before PQ scoring (inline pre-filter), it's not "post" anything.
I suggest flat_search accepting filter directly, like
async fn flat_search<OB>(
&self,
query: &[Data::VectorDataType],
filter: Option<&(dyn Fn(u32) -> bool + Send + Sync)>,
neighbors_before_reranking: usize,
output: &mut OB,
) -> ANNResult<...>
// dispatch
SearchPlan::FlatScan { filter } => self.runtime.block_on(self.flat_search(
query, filter.as_deref(), l, &mut result_output_buffer,
))?,|
|
||
| /// Disk-local projection of `SearchPlan` used inside the strategy. | ||
| /// | ||
| /// `search_strategy()` is the only site that introspects `GraphMode`'s |
There was a problem hiding this comment.
DiskSearchStrategy is only available in graph search and it can hold mode: &'a GraphMode<'a>, directly, we don't need to introduce a new type to represent similar thing.
| associated_data: &mut [Data::AssociatedDataType], | ||
| vector_filter: &(dyn Fn(&Data::VectorIdType) -> bool + Send + Sync), | ||
| is_flat_search: bool, | ||
| plan: &SearchPlan<'_>, |
There was a problem hiding this comment.
nit: Accept own type? plan: SearchPlan<'_>
Reference Issues/PRs
Implementation of RFC #1101.
What does this implement/fix? Briefly explain your changes.
Replaces
searcher.search()'s(vector_filter: Option<VectorFilter>, is_flat_search: bool)parameter pair with a singleplan: SearchPlan<'_>enum.SearchPlanis hierarchical (FlatScan { filter }/Graph(GraphMode));GraphModevariants areUnfiltered,PostFilter, and the newBetaFilter { predicate, beta }.Internally,
search_strategy()projects the plan into a disk-localFilterModesum type carried byDiskSearchStrategy/DiskAccessor/RerankAndFilter.pq_distancesapplies β via a one-lineif let FilterMode::BetaFilter { .. }gate.The new BetaFilter biases beam traversal toward predicate matches by multiplying their PQ distances by β ∈ (0, 1] in pq_distances; the hard post-filter and full-precision reranking are unchanged.
Also:
filter_parameter.rsfromdiskann-disk/src/build/configuration/todiskann-disk/src/search/(it's a search-time concept).SearchPlan.is_flat_search: bool→search_mode: "graph" | "flat".