feat(index): use ACORN-1 traversal for prefiltered HNSW search#7927
Conversation
When a prefilter keeps 10% or more of a partition, search the graph computing distances only for mask-passing nodes, bridging filtered-out neighbors with a two-hop expansion (ACORN-1) and seeding the frontier with a sample of mask-passing nodes so the search cannot strand far from the passing set. The previous traversal scored every visited node and only gated the result heap, wasting most distance computations under selective filters. Sparse filters keep the exact flat fallback and unfiltered search is unchanged.
|
ACTION NEEDED The PR title and description are used as the merge commit message. Please update your PR title and description to match the specification. For details on the error please inspect the "PR Title Check" action. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: QUIET Plan: Pro Plus Run ID: 📒 Files selected for processing (9)
📝 WalkthroughWalkthroughAdds ACORN-1 masked HNSW traversal, filtered IVF routing between flat, basic, and ACORN search, approximate-mode wiring, correctness tests, and synthetic/SIFT benchmarks comparing masked search strategies. ChangesACORN masked search
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Scanner
participant IvfSubIndex
participant HNSW
participant beam_search_acorn
Scanner->>IvfSubIndex: apply ApproxMode::Fast and prefilter
IvfSubIndex->>IvfSubIndex: build mask and choose search path
IvfSubIndex->>HNSW: search_acorn with mask
HNSW->>beam_search_acorn: traverse masked graph
beam_search_acorn-->>HNSW: return ordered results
HNSW-->>IvfSubIndex: return filtered neighbors
IvfSubIndex-->>Scanner: return scan results
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
cb507d2 to
10f81ac
Compare
|
Great work, will take a review today |
There was a problem hiding this comment.
Thanks for exploring this approach. Avoiding distance calculations for filtered-out vertices looks promising. The inline cardinality case is the main concern to address before making this traversal the default for dense filters.
A couple of possible directions may be worth considering:
- One option could be to keep filtered vertices as unscored structural waypoints, with adaptive bridge expansion bounded by an adjacency budget. If the projected traversal still produces fewer than
kin-mask/in-range candidates, it could fall back to the legacy traversal or exact scan. - Another option could be to initially keep this heuristic behind
ApproxMode::Fastor an explicit opt-in, whileNormalpreserves the existing full-topology traversal. That would provide a per-query rollback while fixed-recall coverage is added for defaultef, deletion-only/all-pass masks, and loaded FLAT/PQ/SQ indices.
In either design, it may also be useful for a non-empty mask that passes every node to shortcut to the existing unfiltered path.
Happy to take another look once this has been updated.
| visited.insert(neighbor); | ||
| passing.push(neighbor); | ||
| } | ||
| } else if !expanded.contains(neighbor) { |
There was a problem hiding this comment.
This traversal can return fewer than k results even when at least k matching nodes exist. Two consecutive filtered vertices disconnect passing components from this one-hop bridge, and the only global restarts are capped at 16. On a connected graph with 20 passing nodes, k=20, and the default ef=30, the existing traversal returns 20 while this traversal returns 16. The mask selectivity is 20/58 (34.5%), so the production dispatcher selects this path; increasing ef cannot recover an unreachable component.
Reproducer run at 10f81ac259
I replaced the empty test module with the following regression test:
#[cfg(test)]
mod tests {
use super::*;
struct ChainGraph {
neighbors: Vec<Vec<u32>>,
}
impl BorrowingGraph for ChainGraph {
fn len(&self) -> usize {
self.neighbors.len()
}
fn neighbors(&self, key: u32) -> &[u32] {
&self.neighbors[key as usize]
}
}
struct ZeroDistance;
impl DistCalculator for ZeroDistance {
fn distance(&self, _id: u32) -> f32 {
0.0
}
fn distance_all(&self, _k_hint: usize) -> Vec<f32> {
Vec::new()
}
}
#[test]
fn acorn_can_return_fewer_than_all_matching_nodes() {
const PASSING_COUNT: usize = 20;
const FAILING_COUNT: usize = (PASSING_COUNT - 1) * 2;
let mut neighbors = vec![Vec::new(); PASSING_COUNT + FAILING_COUNT];
for index in 0..PASSING_COUNT - 1 {
let left = index as u32;
let first_failing = (PASSING_COUNT + index * 2) as u32;
let second_failing = first_failing + 1;
let right = left + 1;
neighbors[left as usize].push(first_failing);
neighbors[first_failing as usize].extend([left, second_failing]);
neighbors[second_failing as usize].extend([first_failing, right]);
neighbors[right as usize].push(second_failing);
}
let graph = ChainGraph { neighbors };
let params = HnswQueryParams {
ef: 30,
lower_bound: None,
upper_bound: None,
dist_q_c: 0.0,
};
let entry = OrderedNode::new(0, 0.0.into());
let mut mask_generator = VisitedGenerator::new(graph.len());
let mut mask = mask_generator.generate(graph.len());
for id in 0..PASSING_COUNT as u32 {
mask.insert(id);
}
let mut acorn_visited_generator = VisitedGenerator::new(graph.len());
let mut acorn_expanded_generator = VisitedGenerator::new(graph.len());
let acorn_results = beam_search_acorn(
&graph,
&entry,
¶ms,
&ZeroDistance,
&mask,
None,
&mut acorn_visited_generator.generate(graph.len()),
&mut acorn_expanded_generator.generate(graph.len()),
);
let mut basic_visited_generator = VisitedGenerator::new(graph.len());
let basic_results = beam_search_borrowed(
&graph,
&entry,
¶ms,
&ZeroDistance,
Some(&mask),
None,
&mut basic_visited_generator.generate(graph.len()),
);
assert_eq!(basic_results.len(), PASSING_COUNT);
assert_eq!(acorn_results.len(), PASSING_COUNT);
}
}$ cargo test -p lance-index acorn_can_return_fewer_than_all_matching_nodes -- --nocapture
assertion `left == right` failed
left: 16
right: 20|
Thanks for the review Changes in the latest commit:
Also re-ran on GIST1M (1M x 960, L2) to check the story holds at higher dimensionality. Clustered filters, latency / recall@10:
Two places it loses, uniform random masks at 2% drop recall (0.775 vs 0.975), and at 50% the waypoint bookkeeping makes it slower than the current traversal. |
… chains Review follow-ups: masked chains deeper than one node are crossed via unscored waypoints under a 4 * ef budget (reviewer's chain graph adopted as a regression test), the traversal is opt-in via the existing ApproxMode::Fast so the default path is unchanged, all-pass masks shortcut to the unfiltered path, and under-delivery falls back to the existing traversal. Coverage extended to default ef, deletion-style masks, and FLAT/PQ/SQ under both modes.
Xuanwo
left a comment
There was a problem hiding this comment.
Thank you for working on this!
Addresses #2628.
Problem
With
prefilter=true, the HNSW sub-index runs the full beam search and only gates the result heap by the mask. Every visited node gets a distance computation whether or not it can be returned, so under selective filters most of the work is wasted. The worst case is a filter whose matches cluster far from the query's region of the graph: on SIFT1M that costs 9-240 ms per query (table below).Change
Adds
beam_search_acorn, a filtered traversal in the style of ACORN-1 (Patel et al. 2024), used by the HNSW sub-index for prefiltered queries when the query opts in withApproxMode::Fast.Normal(the default) andAccuratekeep the existing traversal, so default behavior is unchanged.min(ef, matching)results, waypoints are expanded under a budget (4 * ef), so passing components behind masked chains deeper than one node are still reached. This addresses the review counterexample, whose regression test is included.The traversal is generic over the vector storage, so
Fastcovers IVF_HNSW_FLAT, IVF_HNSW_SQ, and IVF_HNSW_PQ, on both the in-memory and disk-loaded graph backends. Deletion masks compose unchanged since the bitset is built from the existingPreFilter.Benchmarks
SIFT1M (1M x 128, L2), single graph, default build params, k=10, median of 20 official queries, Apple M-series. Unfiltered control: 0.30 ms at 1.000 recall@10 against the official ground truth.
currentis today's traversal at ef=100. Cells are latency / recall@10.Filter matches a cluster unrelated to the query (the pathological case today):
At the same ef the traversal is 10-250x faster and gives up at most 4 points of recall. Since ef is a per-query parameter, spending some of that headroom (ef=400) matches the current traversal's recall on every row while staying 4-100x faster.
Filter matches the query's own region: 0.09-0.27 ms at recall equal to the current traversal (which costs 0.34-0.60 ms). Uniform random masks: 0.85-1.45 ms at 0.945-1.00 vs 0.84-6.49 ms at ~1.00, converging at 25%+ selectivity.
GIST1M (1M x 960, L2, unfiltered ceiling 0.820 at ef=100) confirms the story at higher dimensionality, where the flat scan stops being competitive even at 2% selectivity (10 ms). Clustered filters at 2%: 2.4 ms / 0.985 vs 72.8 ms / 0.985 for the current traversal, and
Fastat ef=400 reaches 0.990-1.000 at 5-7.6 ms.Where it loses, so it is on the record: uniform random masks at 2% selectivity drop recall (0.775 vs 0.975 on GIST1M), and at 50% random the waypoint bookkeeping makes it slower than the current traversal (15.3 vs 4.1 ms). Random masks spread the passing set too thin for seeds plus bridges. Clustered filters (categories, languages, tenants) are the win case, and the opt-in leaves that judgment with the caller.
Reproduce with
examples/acorn_bench_sift.rs(SIFT_DIR=... cargo run --release -p lance-index --example acorn_bench_sift) andexamples/acorn_bench.rs(synthetic, no download). Happy to drop or relocate the examples if you'd rather not carry them.Tests
test_acorn_reaches_across_masked_chains).test_acorn_filtered_search: results are mask-only, built and loaded graphs return identical results, recall floors at default ef and on deletion-style (nearly all-pass) masks.test_subindex_prefilter_dispatch: dense masks in both modes, all-pass shortcut equals unfiltered, sparse masks take the exact flat scan.test_ann_prefilterextended to IVF_HNSW_FLAT / PQ / SQ xApproxMode::Normal/Fast(32 cases).Notes
ApproxModequery setting, no new API surface, default path untouched. If the recall coverage looks sufficient down the line, makingFastthe dense-prefilter default can be revisited separately.