Add shared-floor kNN collection to the sandbox module#16357
Conversation
Adds GlobalKnnFloor, FloorAwareKnnCollector, SharedFloorKnnCollectorManager, and a resurrected BlockingFloatHeap under org.apache.lucene.sandbox.search.knn. Opt-in, zero changes to lucene/core, composes with optimistic multi-segment search via the existing minCompetitiveSimilarity hook.
|
@vigyasharma @benwtrent still not giving up on a concurrent streaming search. I've been testing it for a few months on my spare time. Thanks for the guidance so far - any additional feedback is welcome and encouraged. No rush - but I think sandbox is a much better home for this until I can show more solid numbers. Keeping it draft form until I run a full sweep of tests - which takes awhile. |
|
Cross-shard optimistic kNN - Phase 1 measurements @ k=10000
Summary: on uniform data the floor is cheaper than the static quota but doesn't match its recall - an open, non-dominant tradeoff. On skewed data the static quota's savings come from missing 3 in 4 true results; the floor is the only arm that is both cheap and correct. Progress on the recall/visits tradeoff raised: Rebuilt the mechanism as a sandbox-only module (no Lucene core changes): To measure it accurately, a distributed lucene setup in one JVM was set up where every shard returns full Setup: ~962 GB for the full 247M-vector index used in T1 (16 shards × ~61GB each; the joined single-index copy is a separate 962GB copy of the same data - 1.88TB combined on disk if you count both). The shards are set up Not yet run: shard-count scaling (S=4/8/16) and the k=100/1000 rows - full sweep plan is in the same folder ( |
…ish-once Four fixes from review, each with tests: - GlobalKnnFloor.offer read the heap-fullness flag after offering, so a concurrent publisher filling the heap in that window could trick a thread into publishing a partial-heap top that overshoots the true k-th best (over-pruning). The flag is now read before the offer. - BlockingFloatHeap.poll checked emptiness outside the lock, letting two concurrent polls of a one-element heap both pass and corrupt the heap. - FloorAwareKnnCollector's synchronization trigger was an exact bit-mask match over the visited count, evaluated only inside collect(); boundaries falling between two collects were skipped entirely, starving the floor in the late phase of the search where few candidates beat the bar. The schedule is now a visited-count threshold, which also lifts the power-of-two restriction on syncInterval. - SharedFloorKnnCollectorManager gains globalShare and a perShardGate helper mirroring the optimistic strategy's quota formula: a shard of a larger corpus now gates at its expected contribution to the merged top-k instead of at full k (its own index looks like the whole corpus to it, so it could never derive this locally). Full-k collectors over multi-segment indexes likewise gate at each leaf's pro-rata share. The manager also enforces publish-once-per-leaf: the optimistic second pass re-collects the same documents, and republishing them would put duplicate scores in the floor's heap, which can lift the floor above the true merged cutoff. Also documents the greediness/gateK interaction: a low greediness with a below-queue gate sizes the clamp at (1-g)*gateK, which can silently reduce the adaptive floor to a fixed per-shard quota.
What protects recall is the clamp's absolute width, (1 - greediness) * gateK, not the greediness fraction: the same 0.9 that leaves a 600-slot clamp at gateK 6000 collapses to the bare minimum at gateK 184, which benchmarked at double-digit recall loss. greedinessForClamp(gateK, clampSlots) translates a required width into the greediness that produces it, and returns 0 (floor neutralized) when the width cannot fit inside the gate. Javadoc on both classes now points width-first configuration at the helper.
1 participant
Sandbox: shared-floor kNN collection for multi-shard search
The problem
Distributed HNSW search is ineffective on distributed platforms today, as large K requires each shard to calculate the same value K and a parent node will join the results and discard 1/(number of shards) of calculated results.
The goal of this PR is to open up the ability to test ways to bring this number down which should naturally help latency in responses. Although initial measurements suggest around a 30% reduction, testing is underway to produce a more rigorous data driven approach.
Here's an example of what is being tested now:
A 16 shard HNSW cluster at k=10000 computes 160,000 candidates, ships them to a coordinator, and throws away 150,000 of them. That is 94% of the work wasted, and it gets worse as k and shard count grow. Every shard searches as if it were alone, because Lucene has no way to tell a running search "the rest of the cluster already has better results than what you're chasing."
This PR is a follow-up to #15676. I'm closing that one in favor of this approach, which took the feedback there seriously: no changes to core, no new interfaces on the search path, and honest numbers instead of a muzzled baseline.
What this adds
Four classes in
lucene/sandbox, packageorg.apache.lucene.sandbox.search.knn. Zero changes tolucene/core. Nothing touchesHnswGraphSearcher; everything works through the existingminCompetitiveSimilarity()hook and theKnnCollector.Decoratorpattern.GlobalKnnFloortracks the k-th best similarity across all participating searches. It is monotonic, updated in batches, and has anadvertise()method so a bound computed somewhere else (another JVM, another machine) can be pushed in from an outside thread.FloorAwareKnnCollectoris aKnnCollector.Decoratorthat folds the global floor intominCompetitiveSimilarity(). It has an ascent gate so a search can't be pruned before it has collected enough locally to know the floor applies to it, and a greediness clamp that bounds how much exploration the floor can take away.SharedFloorKnnCollectorManagerwires the above into the existing collector manager seam. It reportsisOptimistic() == true, so it composes with the optimistic multi-segment strategy rather than competing with it.BlockingFloatHeapis resurrected from git history (it was removed withMultiLeafKnnCollector). The floor needs a bounded concurrent min-heap and this one already existed and was tested.Everything is opt-in. Stock searches are untouched. All tuning values (greediness, sync interval, exploration slots, activation threshold, gate size) are constructor parameters with documented defaults, and there is an activation threshold so small-k searches never engage the floor at all.
Why it should work
Lucene's optimistic path already has the right idea for segments. If shards are random samples of the corpus (hash and round-robin sharding both give you this), the number of global top-k hits living in a shard with share$p$ is $\mathrm{Binomial}(k, p)$ , and the optimistic quota formula sizes each participant accordingly ($\lambda = 16$ ,
LAMBDAinAbstractKnnVectorQuery):For$s$ equal shards the total collected telescopes to $k + \lambda\sqrt{k(s-1)}$ , so the fan-in reduction versus today's "every shard returns k" contract is:
The visit-cost side follows from where HNSW spends its time. Per-shard visits decompose into a descent term proportional to$\ln N$ (which no correct scheme can skip; a shard cannot know it is non-competitive before reaching its neighborhood) and a term proportional to $\mathrm{ef}$ . On the 247M index the $\mathrm{ef}$ term is about 99% of visits, so quota reduction transfers nearly linearly into visit reduction.
The floor's job is the remainder. Stock behavior makes each shard refine until its frontier drops below its own local k-th best, which under random sharding sits near the corpus$(k \cdot s)$ -th best. So each shard certifies $k$ results of which only $\sim k/s$ survive the merge. The refinement spent on the other $k(s-1)/s$ is the waste:
A converged floor lifts every shard's stopping bar from rank-$(k \cdot s)$ territory to rank- $k$ territory, making that work skippable. How much you actually skip depends on how early the floor rises, which is a timing property of the deployment. That is the main thing left to measure.
Why sandbox and not core
Because the numbers say "promising" and not "proven."
Measured so far, on a real corpus (indices took days to build, so the matrix is still small) — two different baselines, kept separate. Against a single JVM's stock optimistic search, which already avoids most redundant per-segment work, the floor saves only low single digits at equal recall. Against the distributed contract — every shard returns full k and a coordinator merges, simulated as 16 disjoint single-segment indexes searched concurrently in one JVM — the floor traced −24% visits at k=1000 within 0.003 recall of the full-k baseline, and the reduction grows with k (see the measurement comment below for the k=10,000 sweep). The first number is why this can't justify touching core; the second is the regime the mechanism exists for.
The regime this is actually built for is the one that is hardest to benchmark honestly: separate JVMs per shard, where each shard would otherwise run to completion and ship its full k. An earlier version was distributed on a cluster of Raspberry Pi machines, where we saw consistent savings of around 50%. That cluster existed to see how the floor behaves on live running nodes with heavy latency, not to be a realistic field test, and the measurements weren't made with luceneutil, so I don't lean on the number. Mentioning it because it did add confidence that the direction is right: when nodes are slow and staggered, the floor arrives while shards are still working, and that is when it earns its keep. Reproducing that result properly is the main open work, and it's why I want this in sandbox where the API can still move while that measurement happens.
So the ask is: sandbox now, with the understanding that graduation to core requires a much larger body of evidence. If the distributed numbers don't hold up, it stays in sandbox or comes back out.
Prior feedback, addressed
minCompetitiveSimilarity(), gated behind a local fill requirement, so a segment that is still climbing cannot be starved.Testing
Unit tests cover the floor, the collector (gate behavior, clamp sizing, parameter validation), and the heap. Integration tests assert recall parity within 0.05 of stock search when external bounds are advertised. The full
:lucene:sandboxcheck passes (compile, test, ecj lint, javadoc, licenses).Benchmark methodology and the math behind the floor (why it cannot admit a false positive past the gate) are in the design doc on the branch.
Work-in-progress so far
The data
247M Cohere embeddings, 1024-dim, dot product. Sharded round-robin into 16 shards (globalId = rowInShard * 16 + shardId), each shard force-merged to a single segment, then joined with addIndexes under NoMergePolicy so leaf ord equals shard id. That gives a single index whose 16 leaves are exact stand-ins for 16 remote shards, which lets luceneutil measure the cross-shard behavior without introducing network timing as a variable. Ground truth is exact NN over the full corpus (10,000 held-out query vectors available; the runs so far use samples of 30–100 queries, which at k=10,000 means each recall figure averages 300k–1M neighbor judgments). Recall is always measured on the merged global result, never per shard.
Building and force-merging the shards took several days, which is why the measurement matrix is still thin. More points are running.
The numbers so far
At k=1000, 16 simulated shards sharing one floor: −24% vector comparisons versus the every-shard-returns-full-k baseline, within 0.003 recall. At k=10,000 the same comparison reaches −88 to −91% at a recall cost of 0.02–0.04 (full tables in the measurement comment below). Against a single JVM's stock optimistic baseline the saving is only low single digits, because optimistic search already avoids most redundant work inside one process. The reason to keep going is the gap between what's measured and the ~94% ceiling from the math above: in a single fast JVM every leaf fills at nearly the same speed, so the floor engages late. In a real cluster with heterogeneous shards and a scout seeding the floor early, the floor should be live during ascent instead of after it — exactly the region these runs never got to exercise.
Why pruning on the floor cannot cost recall
Let$s^*$ be the k-th best similarity over the final merged result. The floor at any time $t$ is the k-th best over the hits observed so far, and hits observed so far are a subset of all hits that will ever exist. A k-th best over a subset can only be less than or equal to the k-th best over the superset, so:
Any candidate whose reachable score is below$\mathrm{floor}(t)$ is below $s^*$ and can never enter the final top-k. Pruning it is loss-free. This holds for externally advertised bounds too, as long as the sender only advertises scores of real hits, which is the contract on
advertise().The one caveat is the same one stock HNSW already lives with: the frontier top is treated as the reachable bound, and a path through a low-scoring bridge node can reach a better cluster. Raising the bar from the local k-th best toward$s^*$ prunes in exactly the band where bridges can be lost. That is what the greediness clamp bounds: the effective bar never exceeds the $((1-g) \cdot \mathrm{gateK})$ -th best similarity the leaf has seen, so a $(1-g)$ -deep exploration frontier always survives. $g = 0$ is bit-identical to stock, $g = 1$ is the hard-stop failure mode from #15676. The dial has hard endpoints and the default sits at 0.5.