krickert opened a new pull request, #16357:
URL: https://github.com/apache/lucene/pull/16357
# 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. Initial
measurements are promising, but we are still testing, and we're open to any
suggestions.
Here's an example we're testing 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`, package
`org.apache.lucene.sandbox.search.knn`. Zero changes to `lucene/core`. Nothing
touches `HnswGraphSearcher`; everything works through the existing
`minCompetitiveSimilarity()` hook and the `KnnCollector.Decorator` pattern.
- `GlobalKnnFloor` tracks the k-th best similarity across all participating
searches. It is monotonic, updated in batches, and has an `advertise()` method
so a bound computed somewhere else (another JVM, another machine) can be pushed
in from an outside thread.
- `FloorAwareKnnCollector` is a `KnnCollector.Decorator` that folds the
global floor into `minCompetitiveSimilarity()`. 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.
- `SharedFloorKnnCollectorManager` wires the above into the existing
collector manager seam. It reports `isOptimistic() == true`, so it composes
with the optimistic multi-segment strategy rather than competing with it.
- `BlockingFloatHeap` is resurrected from git history (it was removed with
`MultiLeafKnnCollector`). 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 we think it works
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$, `LAMBDA` in `AbstractKnnVectorQuery`):
```math
\mathrm{perShardK}(k, p) = kp + \lambda\sqrt{kp(1-p)}
```
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:
```math
R(k, s) = \frac{ks}{k + \lambda\sqrt{k(s-1)}} = \frac{s}{1 +
\lambda\sqrt{(s-1)/k}}
```
$R \to s$ as $k \to \infty$. The large-k case that hurts today is exactly
where the math is most favorable. At $s = 16$:
| k | perShardK | total collected (vs k*s today) | reduction |
|---|---|---|---|
| 100 | 45 | 720 vs 1,600 | 2.2x |
| 1,000 | 185 | 2,960 vs 16,000 | 5.4x |
| 10,000 | 1,012 | 16,197 vs 160,000 | 9.9x |
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:
```math
\text{wasted fraction} = \frac{s-1}{s} = \frac{15}{16} \approx 94\% \quad
\text{at } s = 16
```
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."
What we have measured so far, on a real corpus (indices took days to build,
so the matrix is still small): a 30% drop in vector comparisons at k=1000 on a
single fast machine, at equal recall. That is the best single point. Other
operating points show less, some show low single digits, because the optimistic
multi-segment search already captures a lot of the same savings inside one JVM.
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. We ran an earlier version of this in
that setup, on a distributed cluster of raspberry pi machines, and saw a
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 were not 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
- No public mutable state bolted onto core collectors. The floor is a
separate object in sandbox; core is untouched.
- The recall-collapse problem from the earlier branch is fixed structurally.
The old code compared best-found against the floor and terminated; this code
only narrows `minCompetitiveSimilarity()`, gated behind a local fill
requirement, so a segment that is still climbing cannot be starved.
- Determinism: the floor is monotonic and batched. Sequential execution is
deterministic; parallel execution varies only in how early pruning engages,
never in whether a collected result is competitive.
- Benchmarked with luceneutil against the current optimistic baseline, not
against a handicapped one.
## 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:sandbox` check 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 queries, and 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 one good number
At k=1000, single JVM, floor sharing on: roughly 30% fewer vector
comparisons than the stock optimistic baseline at equal recall. One operating
point, not a curve, and other points show much less because optimistic search
already avoids most redundant work inside a single process. The reason to keep
going is the gap between 30% measured and the ~94% ceiling from the math above:
the floor engages late in a single fast JVM because every leaf fills at nearly
the same speed. 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,
which is exactly the region the 30% run 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:
```math
\mathrm{floor}(t) \le s^* \quad \text{for all } t, \qquad \mathrm{floor}\
\text{monotone non-decreasing}
```
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
