diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt index 72fcd7427cdb..ee713a89d58f 100644 --- a/lucene/CHANGES.txt +++ b/lucene/CHANGES.txt @@ -103,6 +103,13 @@ API Changes New Features --------------------- +* GITHUB#16357: Add shared-floor kNN collection to the sandbox module: GlobalKnnFloor, + FloorAwareKnnCollector and SharedFloorKnnCollectorManager let concurrent searchers of one query + (segments in one process, or shards across processes via GlobalKnnFloor#advertise) prune graph + exploration against a shared lower bound of the query's final merged top-k cutoff. This remains + an approximate-search optimization: deployments must measure its recall/visit trade-off. Opt-in, + no changes to core search. (Kristian Rickert) + * GITHUB#16241: Add HangulCompositionCharFilter to analysis-nori for opt-in composition of modern Hangul conjoining jamo before KoreanTokenizer. (Mingi Jeong) diff --git a/lucene/sandbox/src/java/module-info.java b/lucene/sandbox/src/java/module-info.java index ee9be3227de2..8d593a2adf04 100644 --- a/lucene/sandbox/src/java/module-info.java +++ b/lucene/sandbox/src/java/module-info.java @@ -28,6 +28,7 @@ exports org.apache.lucene.sandbox.document; exports org.apache.lucene.sandbox.queries; exports org.apache.lucene.sandbox.search; + exports org.apache.lucene.sandbox.search.knn; exports org.apache.lucene.sandbox.index; exports org.apache.lucene.sandbox.facet; exports org.apache.lucene.sandbox.facet.recorders; diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/BlockingFloatHeap.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/BlockingFloatHeap.java new file mode 100644 index 000000000000..4c2246d51ed6 --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/BlockingFloatHeap.java @@ -0,0 +1,200 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.lucene.sandbox.search.knn; + +import java.util.concurrent.locks.ReentrantLock; + +/** + * A blocking bounded min heap that stores floats. The top element is the lowest value of the heap. + * + *

A primitive priority queue that maintains a partial ordering of its elements such that the + * least element can always be found in constant time. Implementation is based on {@link + * org.apache.lucene.util.LongHeap} + * + * @lucene.internal + */ +final class BlockingFloatHeap { + private final int maxSize; + private final float[] heap; + private final ReentrantLock lock; + private int size; + + public BlockingFloatHeap(int maxSize) { + if (maxSize < 1) { + throw new IllegalArgumentException("maxSize must be at least 1, got: " + maxSize); + } + this.maxSize = maxSize; + this.heap = new float[maxSize + 1]; + this.lock = new ReentrantLock(); + this.size = 0; + } + + /** + * Inserts a value into this heap. + * + *

If the number of values would exceed the heap's maxSize, the least value is discarded + * + * @param value the value to add + * @return the new 'top' element in the queue. + */ + public float offer(float value) { + lock.lock(); + try { + if (size < maxSize) { + push(value); + return heap[1]; + } else { + if (value >= heap[1]) { + updateTop(value); + } + return heap[1]; + } + } finally { + lock.unlock(); + } + } + + /** + * Inserts array of values into this heap. + * + *

Values must be sorted in ascending order. + * + * @param values a set of values to insert, must be sorted in ascending order + * @param len number of values from the {@code values} array to insert + * @return the new 'top' element in the queue. + */ + public float offer(float[] values, int len) { + if (len < 0 || len > values.length) { + throw new IllegalArgumentException( + "len must be in [0, values.length=" + values.length + "], got: " + len); + } + lock.lock(); + try { + for (int i = len - 1; i >= 0; i--) { + if (size < maxSize) { + push(values[i]); + } else { + if (values[i] >= heap[1]) { + updateTop(values[i]); + } else { + break; + } + } + } + return heap[1]; + } finally { + lock.unlock(); + } + } + + /** + * Removes and returns the head of the heap + * + * @return the head of the heap, the smallest value + * @throws IllegalStateException if the heap is empty + */ + public float poll() { + lock.lock(); + try { + // The emptiness check must happen under the lock: checking outside lets two concurrent + // polls of a one-element heap both pass, and the second corrupts the heap state. + if (size > 0) { + float result = heap[1]; // save first value + heap[1] = heap[size]; // move last to first + size--; + downHeap(1); // adjust heap + return result; + } else { + throw new IllegalStateException("The heap is empty"); + } + } finally { + lock.unlock(); + } + } + + /** + * Retrieves, but does not remove, the head of this heap. + * + * @return the head of the heap, the smallest value + */ + public float peek() { + lock.lock(); + try { + return heap[1]; + } finally { + lock.unlock(); + } + } + + /** + * Returns the number of elements in this heap. + * + * @return the number of elements in this heap + */ + public int size() { + lock.lock(); + try { + return size; + } finally { + lock.unlock(); + } + } + + private void push(float element) { + size++; + heap[size] = element; + upHeap(size); + } + + private float updateTop(float value) { + heap[1] = value; + downHeap(1); + return heap[1]; + } + + private void downHeap(int i) { + float value = heap[i]; // save top value + int j = i << 1; // find smaller child + int k = j + 1; + if (k <= size && heap[k] < heap[j]) { + j = k; + } + while (j <= size && heap[j] < value) { + heap[i] = heap[j]; // shift up child + i = j; + j = i << 1; + k = j + 1; + if (k <= size && heap[k] < heap[j]) { + j = k; + } + } + heap[i] = value; // install saved value + } + + private void upHeap(int origPos) { + int i = origPos; + float value = heap[i]; // save bottom value + int j = i >>> 1; + while (j > 0 && value < heap[j]) { + heap[i] = heap[j]; // shift parents down + i = j; + j = j >>> 1; + } + heap[i] = value; // install saved value + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/FloorAwareKnnCollector.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/FloorAwareKnnCollector.java new file mode 100644 index 000000000000..1adcd635d151 --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/FloorAwareKnnCollector.java @@ -0,0 +1,399 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.lucene.sandbox.search.knn; + +import org.apache.lucene.search.AbstractKnnCollector; +import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.util.hnsw.FloatHeap; + +/** + * A {@link org.apache.lucene.search.KnnCollector.Decorator} that raises its {@link + * #minCompetitiveSimilarity()} using a {@link GlobalKnnFloor} shared with the other searchers of + * the same query, so that a graph search can reduce exploration of candidates whose score cannot + * enter the query's final, merged top-k, not merely candidates that cannot enter this collector's + * local top-k. Because HNSW may need low-scoring bridge nodes to reach better neighborhoods, this + * is an approximate-search optimization: the exploration clamp controls its recall/visit trade-off. + * + *

The graph searcher already re-reads {@link #minCompetitiveSimilarity()} whenever a collected + * hit reports an improvement, and uses it both to stop the search when the best remaining candidate + * falls below it and to avoid enqueueing hopeless candidates. Folding the shared floor into that + * method is therefore the entire integration: no searcher changes are required, and a collector + * whose floor never rises behaves exactly like its delegate. + * + *

Three guards keep the shared bound from harming recall or becoming a synchronization hot spot: + * + *

+ * + *

The effective bound derived from the floor is one ulp below the floor itself. The floor is the + * similarity of a real collected hit, and when several hits tie at the cutoff the merged result set + * is decided by tie-breaking, not by score; a search must therefore remain willing to find + * candidates exactly at the floor, otherwise a document that would have won the tie-break could be + * abandoned. Keeping the bound strictly below the floor preserves those ties without publishing + * document identities. + * + *

Instances are confined to a single thread, like every {@link KnnCollector}; only the shared + * {@link GlobalKnnFloor} is touched by multiple threads. + * + * @lucene.experimental + */ +public final class FloorAwareKnnCollector extends KnnCollector.Decorator { + + /** + * Default fraction of the search effort that follows the shared floor rather than the local + * frontier; see the class comment for the roles of the two extremes. The default is deliberately + * conservative. Callers who have verified recall on their own data may trade some of it for fewer + * visits by raising this. + */ + public static final float DEFAULT_GREEDINESS = 0.5f; + + /** + * Default minimum number of non-competitive queue slots, whatever the greediness. The protection + * a graph search needs against a tight external bound is an absolute number of below-bound + * candidates it may keep routing through, not a fraction of k: with a purely fractional clamp, a + * small-k search under high greediness is left a clamp one or two candidates wide, and randomized + * testing showed that costing double-digit recall. A useful side effect is that when k does not + * exceed the configured minimum, the clamp is at least as wide as the local queue and the shared + * floor is neutralized entirely: small-k searches behave exactly like stock search no matter what + * has been advertised. + */ + public static final int DEFAULT_MIN_EXPLORATION_SLOTS = 16; + + /** + * Default number of visited vectors between synchronizations with the shared floor. The value + * balances floor freshness (a staler floor prunes less, costing visits but never recall) against + * cross-thread traffic; it is the interval the removed {@code MultiLeafKnnCollector} shipped + * with. + */ + public static final int DEFAULT_SYNC_INTERVAL = 256; + + /** + * The greediness at which the clamp keeps {@code clampSlots} exploration slots against a gate of + * {@code gateK}: {@code 1 - clampSlots / gateK}, or 0 when the requested width is the whole gate + * or more (the floor is then fully neutralized). + * + *

Prefer deriving greediness through this method over choosing a constant. What protects + * recall is the clamp's absolute width, not the greediness fraction: the same {@code + * greediness = 0.9} that leaves a comfortable 600-slot clamp at {@code gateK = 6000} collapses to + * the bare minimum at {@code gateK = 184}, which benchmarked at double-digit recall loss. A + * caller who instead fixes the width it needs, and lets this method translate it, is safe at + * every gate size. + * + * @param gateK the ascent gate the collector will use, at least 1 + * @param clampSlots the exploration width to preserve, at least 1 + */ + public static float greedinessForClamp(int gateK, int clampSlots) { + if (gateK < 1) { + throw new IllegalArgumentException("gateK must be at least 1, got: " + gateK); + } + if (clampSlots < 1) { + throw new IllegalArgumentException("clampSlots must be at least 1, got: " + clampSlots); + } + if (clampSlots >= gateK) { + return 0f; + } + return 1f - clampSlots / (float) gateK; + } + + private final AbstractKnnCollector subCollector; + private final GlobalKnnFloor globalFloor; + + /** + * Number of locally collected results after which the shared floor engages; the ascent gate's + * threshold. See the class comment. + */ + private final int gateK; + + /** Number of visited vectors between synchronizations with the shared floor. */ + private final int syncInterval; + + /** + * Whether this collector publishes its observed similarities into the shared floor. A + * non-publishing collector still reads the floor and prunes against it; it only never feeds it. + * Used when the same documents' scores have already been published by an earlier search of the + * same leaf (the optimistic strategy's second pass), where publishing again would insert + * duplicate scores into the floor's heap and could inflate the floor above the true merged + * cutoff, violating {@link GlobalKnnFloor}'s distinct-document contract. + */ + private final boolean publishToFloor; + + /** + * The best {@code max(minExplorationSlots, (1 - greediness) * k)} similarities seen by this + * collector, competitive or not. Its minimum caps the effective bound, implementing the + * greediness clamp. + */ + private final FloatHeap nonCompetitiveQueue; + + /** + * Similarities observed since the last synchronization, awaiting publication; {@code null} for a + * non-publishing collector. + */ + private final FloatHeap updatesQueue; + + /** + * Scratch used to drain {@link #updatesQueue} in ascending order for batch publication; {@code + * null} for a non-publishing collector. + */ + private final float[] updatesScratch; + + private boolean gateOpened; + + /** + * The visited count at or beyond which the next synchronization with the shared floor happens. + * Deliberately a threshold rather than an exact boundary: {@link #collect} is invoked only for + * candidates that beat the current bound, so visited counts are consulted at irregular strides + * and an equality test would skip every boundary that falls between two collects — precisely in + * the late phase of the search, when few candidates beat the bar and the floor matters most. + */ + private long nextSyncAt; + + private float cachedGlobalFloor = Float.NEGATIVE_INFINITY; + + /** + * Create a collector applying the {@link #DEFAULT_GREEDINESS default greediness}. + * + * @param subCollector the collector gathering this searcher's local results + * @param globalFloor the floor shared by all searchers of this query + */ + public FloorAwareKnnCollector(AbstractKnnCollector subCollector, GlobalKnnFloor globalFloor) { + this(subCollector, globalFloor, DEFAULT_GREEDINESS); + } + + /** + * Create a collector with an explicit greediness and default slot minimum and sync interval. + * + * @param subCollector the collector gathering this searcher's local results + * @param globalFloor the floor shared by all searchers of this query + * @param greediness fraction of the search effort that follows the shared floor, in {@code [0, + * 1]}; see the class comment + */ + public FloorAwareKnnCollector( + AbstractKnnCollector subCollector, GlobalKnnFloor globalFloor, float greediness) { + this( + subCollector, + globalFloor, + greediness, + DEFAULT_MIN_EXPLORATION_SLOTS, + DEFAULT_SYNC_INTERVAL); + } + + /** + * Create a collector whose ascent gate opens when its local queue fills, with an explicit + * greediness, slot minimum and sync interval. + * + * @param subCollector the collector gathering this searcher's local results + * @param globalFloor the floor shared by all searchers of this query + * @param greediness fraction of the search effort that follows the shared floor, in {@code [0, + * 1]}; see the class comment + * @param minExplorationSlots smallest permitted size of the greediness clamp's queue, whatever + * the greediness; must be at least 1. See {@link #DEFAULT_MIN_EXPLORATION_SLOTS} for the role + * this plays in protecting recall. + * @param syncInterval number of visited vectors between synchronizations with the shared floor; + * must be positive. Smaller intervals keep the floor fresher (fewer visits) at the price of + * more cross-thread traffic. + */ + public FloorAwareKnnCollector( + AbstractKnnCollector subCollector, + GlobalKnnFloor globalFloor, + float greediness, + int minExplorationSlots, + int syncInterval) { + this( + subCollector, globalFloor, greediness, minExplorationSlots, syncInterval, subCollector.k()); + } + + /** + * Create a fully configured collector. + * + * @param subCollector the collector gathering this searcher's local results + * @param globalFloor the floor shared by all searchers of this query + * @param greediness fraction of the search effort that follows the shared floor, in {@code [0, + * 1]}; see the class comment + * @param minExplorationSlots smallest permitted size of the greediness clamp's queue, whatever + * the greediness; must be at least 1. See {@link #DEFAULT_MIN_EXPLORATION_SLOTS} for the role + * this plays in protecting recall. + * @param syncInterval number of visited vectors between synchronizations with the shared floor; + * must be positive. Smaller intervals keep the floor fresher (fewer visits) at the price of + * more cross-thread traffic. + * @param gateK the number of locally collected results after which the shared floor engages, in + * {@code [1, subCollector.k()]}. The default, used by the other constructors, is {@code + * subCollector.k()}: the floor engages when the local queue fills. Set it below the queue + * size when this collector's expected share of the merged top-k is smaller than the queue it + * collects into; see the class comment on the ascent gate for when that applies, and the + * greediness clamp caution for the interaction between a below-queue gate and low greediness. + */ + public FloorAwareKnnCollector( + AbstractKnnCollector subCollector, + GlobalKnnFloor globalFloor, + float greediness, + int minExplorationSlots, + int syncInterval, + int gateK) { + this(subCollector, globalFloor, greediness, minExplorationSlots, syncInterval, gateK, true); + } + + /** + * Create a fully configured collector that may be barred from publishing. Package-private: used + * by {@link SharedFloorKnnCollectorManager} to enforce {@link GlobalKnnFloor}'s distinct-document + * contract by publishing each leaf's scores at most once per query; see {@link #publishToFloor}. + */ + FloorAwareKnnCollector( + AbstractKnnCollector subCollector, + GlobalKnnFloor globalFloor, + float greediness, + int minExplorationSlots, + int syncInterval, + int gateK, + boolean publishToFloor) { + super(subCollector); + if (greediness < 0 || greediness > 1 || Float.isNaN(greediness)) { + throw new IllegalArgumentException("greediness must be in [0,1], got: " + greediness); + } + if (minExplorationSlots < 1) { + throw new IllegalArgumentException( + "minExplorationSlots must be at least 1, got: " + minExplorationSlots); + } + if (syncInterval < 1) { + throw new IllegalArgumentException("syncInterval must be positive, got: " + syncInterval); + } + if (gateK < 1 || gateK > subCollector.k()) { + throw new IllegalArgumentException( + "gateK must be in [1, subCollector.k()=" + subCollector.k() + "], got: " + gateK); + } + this.subCollector = subCollector; + this.globalFloor = globalFloor; + this.gateK = gateK; + this.syncInterval = syncInterval; + this.publishToFloor = publishToFloor; + this.nonCompetitiveQueue = + new FloatHeap(Math.max(minExplorationSlots, Math.round((1 - greediness) * gateK))); + this.updatesQueue = publishToFloor ? new FloatHeap(globalFloor.k()) : null; + this.updatesScratch = publishToFloor ? new float[globalFloor.k()] : null; + } + + /** The ascent gate's threshold; package-private, for tests and the manager. */ + int gateK() { + return gateK; + } + + /** Whether this collector feeds the shared floor; package-private, for tests. */ + boolean publishesToFloor() { + return publishToFloor; + } + + @Override + public boolean collect(int docId, float similarity) { + boolean localSimUpdated = subCollector.collect(docId, similarity); + boolean gateJustOpened = gateOpened == false && subCollector.numCollected() >= gateK; + if (gateJustOpened) { + gateOpened = true; + } + if (publishToFloor) { + updatesQueue.offer(similarity); + } + boolean globalSimUpdated = nonCompetitiveQueue.offer(similarity); + + if (gateOpened && (gateJustOpened || visitedCount() >= nextSyncAt)) { + nextSyncAt = visitedCount() + syncInterval; + if (publishToFloor) { + // The shared heap requires ascending input; draining the pending min-heap yields exactly + // that. Scores observed before the gate opened are included in the first batch, so nothing + // seen during the ascent is lost to the shared floor. + int len = updatesQueue.size(); + if (len > 0) { + for (int i = 0; i < len; i++) { + updatesScratch[i] = updatesQueue.poll(); + } + assert updatesQueue.size() == 0; + cachedGlobalFloor = globalFloor.offer(updatesScratch, len); + globalSimUpdated = true; + } + } else { + float refreshed = globalFloor.floor(); + if (refreshed > cachedGlobalFloor) { + cachedGlobalFloor = refreshed; + globalSimUpdated = true; + } + } + } + // Reporting an update whenever the effective bound may have moved (locally or via the shared + // floor) prompts the graph searcher to re-read minCompetitiveSimilarity() and raise its + // termination bar promptly. + return localSimUpdated || globalSimUpdated; + } + + @Override + public float minCompetitiveSimilarity() { + if (gateOpened == false) { + // Ascent gate: until gateK results have been collected, expose only the delegate's bound, + // which is NEGATIVE_INFINITY by the KnnCollector contract. See the class comment. + return subCollector.minCompetitiveSimilarity(); + } + // nextDown keeps the bound strictly below the floor so exact score ties at the cutoff remain + // reachable; nextDown of NEGATIVE_INFINITY is NEGATIVE_INFINITY, so an undefined floor is a + // no-op here. + return Math.max( + subCollector.minCompetitiveSimilarity(), + Math.min(nonCompetitiveQueue.peek(), Math.nextDown(cachedGlobalFloor))); + } + + @Override + public String toString() { + return "FloorAwareKnnCollector[subCollector=" + subCollector + "]"; + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/GlobalKnnFloor.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/GlobalKnnFloor.java new file mode 100644 index 000000000000..12f6562acf9e --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/GlobalKnnFloor.java @@ -0,0 +1,208 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.lucene.sandbox.search.knn; + +import java.util.Objects; +import java.util.concurrent.atomic.LongAccumulator; +import org.apache.lucene.util.NumericUtils; + +/** + * Shared, per-query lower bound on the similarity of the k-th best result of a kNN search whose + * hits are produced by multiple independent searchers (index segments searched concurrently, or + * remote shards in separate JVMs). + * + *

The floor is the k-th highest similarity among all hits observed so far, tracked with a + * bounded min-heap of the k best observed scores. Because the observed hits are always a subset of + * the hits that will ever exist for this query, the k-th best of the observed subset can never + * exceed the k-th best of the complete result set. The floor is therefore, at every moment, a valid + * lower bound on the final top-k similarity cutoff: a candidate that cannot beat the current floor + * can never enter the final merged top-k. This is a score-bound invariant, not a proof of HNSW + * recall: graph navigation may need to traverse a low-scoring bridge node to reach a better region. + * {@link FloorAwareKnnCollector} therefore applies the bound through a configurable exploration + * clamp, whose recall/visit trade-off must be measured for the target corpus. This invariant + * requires that every score offered or advertised belongs to a distinct document; feeding duplicate + * documents (for example, hits for the same document from two replicas) may inflate the floor above + * the true cutoff and callers must deduplicate before publishing. + * + *

The floor is monotonic: it starts at {@link Float#NEGATIVE_INFINITY}, becomes defined only + * once k scores have been observed, and afterwards only rises. Monotonicity is what makes the floor + * safe to share across threads and processes with no coordination beyond a max-reduction: updates + * commute, duplicated or reordered deliveries are harmless, and a stale read can only under-prune + * (costing visits), never over-prune (costing correctness). + * + *

There are two ways to feed the floor: + * + *

+ * + *

An instance carries the state of a single query execution and must not be reused across + * queries: a floor derived from one query's hits is meaningless, and unsafe, as a bound for another + * query. + * + * @lucene.experimental + */ +public final class GlobalKnnFloor { + + private final int k; + private final BlockingFloatHeap heap; + + /** + * The current floor as a sortable-int-encoded float (see {@link + * NumericUtils#floatToSortableInt(float)}), widened to a long. The sortable encoding preserves + * float ordering under integer comparison, so a {@code Long::max} accumulator implements a + * lock-free monotonic maximum over float values. + */ + private final LongAccumulator floorBits; + + /** + * Latched to true once the heap has received k scores. Written by whichever publishing thread + * observes the heap reaching size k; the heap never shrinks, so the transition happens once and + * the flag is stable afterwards. Until the latch is set, the heap's minimum is the minimum of + * fewer than k scores, which is not a valid k-th-best bound and must not feed the floor. + */ + private volatile boolean heapFull; + + /** + * Create a floor for a query collecting {@code k} results. + * + * @param k the number of results the query collects; the floor becomes defined once k distinct + * scores have been observed + */ + public GlobalKnnFloor(int k) { + if (k < 1) { + throw new IllegalArgumentException("k must be at least 1, got: " + k); + } + this.k = k; + this.heap = new BlockingFloatHeap(k); + this.floorBits = new LongAccumulator(Long::max, encode(Float.NEGATIVE_INFINITY)); + } + + /** Return the number of results the query collects, which is also the heap capacity. */ + public int k() { + return k; + } + + /** + * Publish a batch of locally observed similarities and return the resulting floor. + * + *

Each published score must belong to a distinct document (across all calls for this query, + * from all publishers). Scores that are not competitive are cheaply discarded by the bounded + * heap; dropping scores never invalidates the floor, it only makes it less tight. + * + * @param scores similarities to publish, sorted in ascending order + * @param len the number of leading entries of {@code scores} to publish, must be positive + * @return the floor after this batch, {@link Float#NEGATIVE_INFINITY} if fewer than k scores have + * been observed so far + */ + public float offer(float[] scores, int len) { + Objects.requireNonNull(scores, "scores"); + if (len <= 0) { + throw new IllegalArgumentException("len must be positive, got: " + len); + } + if (len > scores.length) { + throw new IllegalArgumentException( + "len must not exceed scores.length: len=" + len + ", scores.length=" + scores.length); + } + if (isSorted(scores, len) == false) { + throw new IllegalArgumentException( + "scores must be sorted in ascending order and must not contain NaN"); + } + // The flag must be read BEFORE the offer. heapTop is only a valid k-th-best bound if the heap + // was already full when our scores went in; reading the flag afterwards would let a concurrent + // publisher fill the heap and set the flag in between, tricking this thread into publishing a + // heapTop computed over fewer than k scores, which can overshoot the true k-th best and cause + // over-pruning. + boolean wasFull = heapFull; + float heapTop = heap.offer(scores, len); + if (wasFull) { + // The heap already held k scores before this call, so the value returned by offer() is the + // minimum of a full heap: the k-th best of everything observed, a valid bound. + floorBits.accumulate(encode(heapTop)); + } else if (heap.size() >= k) { + heapFull = true; + // The heap was not full before this call but is now (filled by this batch, a concurrent + // publisher, or both). heapTop may predate the fill, so re-read the top of the now-full + // heap instead of trusting it. + floorBits.accumulate(encode(heap.peek())); + } + return floor(); + } + + /** + * Raise the floor to an externally computed lower bound of the final k-th best similarity, if it + * exceeds the current floor. Values at or below the current floor are ignored, so duplicated or + * reordered deliveries from remote feeders need no special handling. + * + *

Contract: {@code kthBestLowerBound} must not exceed the k-th best similarity of the query's + * final, merged result set. The k-th best similarity among any k or more real, distinct hits for + * this query (for example, a remote shard's converged top-k) always satisfies this. A value that + * violates the contract makes over-pruning, and therefore missing results, possible. + * + * @param kthBestLowerBound the bound to advertise; must not be NaN + */ + public void advertise(float kthBestLowerBound) { + if (Float.isNaN(kthBestLowerBound)) { + throw new IllegalArgumentException("advertised bound must not be NaN"); + } + floorBits.accumulate(encode(kthBestLowerBound)); + } + + /** + * Return the current floor: a lower bound of the final top-k similarity cutoff, or {@link + * Float#NEGATIVE_INFINITY} while fewer than k scores have been observed and nothing has been + * advertised. The returned value never decreases over the lifetime of this object. + */ + public float floor() { + return decode(floorBits.get()); + } + + private static long encode(float value) { + return NumericUtils.floatToSortableInt(value); + } + + private static float decode(long bits) { + return NumericUtils.sortableIntToFloat((int) bits); + } + + private static boolean isSorted(float[] scores, int len) { + // The last element needs an explicit NaN check: NaN compares false against everything, so a + // trailing NaN never trips the pairwise comparison below, and Arrays.sort places NaN exactly + // there. + if (Float.isNaN(scores[len - 1])) { + return false; + } + for (int i = 1; i < len; i++) { + if (Float.isNaN(scores[i - 1]) || scores[i - 1] > scores[i]) { + return false; + } + } + return true; + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/SharedFloorKnnCollectorManager.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/SharedFloorKnnCollectorManager.java new file mode 100644 index 000000000000..47c46b57b38c --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/SharedFloorKnnCollectorManager.java @@ -0,0 +1,382 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.lucene.sandbox.search.knn; + +import java.io.IOException; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.search.TopKnnCollector; +import org.apache.lucene.search.knn.KnnCollectorManager; +import org.apache.lucene.search.knn.KnnSearchStrategy; + +/** + * A {@link KnnCollectorManager} whose collectors share a {@link GlobalKnnFloor}, so that every + * segment of a query prunes against a lower bound of the query's final merged top-k cutoff rather + * than only against its own local top-k. + * + *

This manager reports {@link #isOptimistic()} and provides {@link #newOptimisticCollector}, + * which means the kNN query machinery runs its regular pro-rata collection strategy on top of it: + * each segment initially collects only its statistically expected share of the global top-k, and + * segments whose results prove globally competitive are searched again. The shared floor + * complements that strategy rather than replacing it. During the first pass it lets segments stop + * refining results that the merged cutoff has already outrun, and because the floor object is + * shared across both passes, segments searched again start with the bound the first pass + * established instead of rediscovering it. + * + *

The floor may also be fed from outside this process through {@link + * GlobalKnnFloor#advertise(float)}, allowing several searchers of disjoint indexes, each holding a + * manager wired to the same query, to bound each other's work. Coordinating who advertises what, + * and deduplicating documents that may appear in more than one index, is the caller's + * responsibility; see {@link GlobalKnnFloor} for the exact contract. + * + *

When the index this manager searches is one shard of a larger corpus, pass the shard's {@code + * globalShare} (the fraction of the corpus it holds). The manager then opens each collector's + * ascent gate at the shard's statistically expected contribution to the merged top-k, computed by + * {@link #perShardGate(int, double)}, instead of at the full local queue size. Without this, a + * shard's local search cannot know it is one of many — its own index looks like the whole corpus — + * so it pays the full cost of collecting k results before the shared floor is allowed to end its + * search, which forfeits most of the cross-shard saving. With the gate at the expected share, a + * shard pays only the share's fill cost, and then continues past it exactly as long as the shared + * floor says its results remain globally competitive: hot shards run long, cold shards stop early. + * + *

Each leaf's scores are published into the shared floor at most once per query. The optimistic + * strategy searches competitive segments a second time, re-collecting the same documents; if the + * second pass republished them, the floor's heap would hold duplicate scores for the same + * documents, and a floor over a multiset with duplicates can exceed the true merged cutoff, + * breaking the bound's safety (see {@link GlobalKnnFloor}'s distinct-document contract). Collectors + * created for a leaf that has already been searched therefore read the floor without feeding it. + * + *

Floor sharing engages only when the query's k reaches an activation threshold; below it this + * manager creates plain, undecorated collectors and the search is exactly stock search. The two + * regimes justify the cutoff. The savings available to a shared bound grow with k: the pro-rata + * strategy sizes each segment's share as {@code k * proportion} plus {@code 16} standard deviations + * of a binomial, so at small k the statistical padding swamps the share and there is little + * redundant work left for a floor to eliminate, while at large k the padding is relatively + * negligible and most per-segment work above the merged cutoff is redundant. The recall risk moves + * the opposite way: the smaller k is, the closer any valid floor sits to scores an incomplete graph + * search routes through. Spending nothing and risking nothing below the threshold is therefore the + * right default, and deployments that have measured recall on their own data can lower it + * explicitly. + * + *

A manager, like the floor it holds, carries the state of a single query execution: create one + * per query, and never share one across queries. Both collector-creation methods may be called + * concurrently, as segments are searched in parallel. The same execution-scoped manager must also + * be returned for optimistic reentry; callers should construct their query and manager together for + * one search rather than retaining either as reusable query state. + * + *

Filtered kNN queries may fall back to exact search in core when floor pruning intentionally + * returns fewer than the core per-leaf quota. This preserves correctness but can erase the expected + * visit reduction. Until core exposes an intentional-partial-result contract, deployments should + * measure filtered queries separately and disable floor sharing where that fallback is + * unacceptable. + * + * @lucene.experimental + */ +public final class SharedFloorKnnCollectorManager implements KnnCollectorManager { + + /** + * Default smallest k at which floor sharing engages. See the class comment for why small-k + * queries are better served by stock search. + */ + public static final int DEFAULT_FLOOR_ACTIVATION_K = 100; + + /** + * The number of standard deviations of statistical padding {@link #perShardGate(int, double)} + * adds to a shard's expected share of the merged top-k; the same value the optimistic + * multi-segment strategy uses to size per-segment collection ({@code LAMBDA} in {@code + * AbstractKnnVectorQuery}). + */ + private static final int LAMBDA = 16; + + private final int k; + private final GlobalKnnFloor globalFloor; + private final float greediness; + private final int floorActivationK; + private final int minExplorationSlots; + private final int syncInterval; + private final float globalShare; + + /** + * The leaves whose scores have been claimed for publication into the shared floor, so that a + * second search of the same leaf gets a non-publishing collector; see the class comment. Identity + * semantics are correct here: a manager lives for one query execution against one reader, and the + * optimistic strategy passes the same {@link LeafReaderContext} instances to both passes. + */ + private final Set publishedLeaves = ConcurrentHashMap.newKeySet(); + + /** + * Create a manager with its own floor and the {@link FloorAwareKnnCollector#DEFAULT_GREEDINESS + * default greediness}, for queries whose searchers all live in this process. + * + * @param k the number of neighbors the query collects + */ + public SharedFloorKnnCollectorManager(int k) { + this(k, new GlobalKnnFloor(k), FloorAwareKnnCollector.DEFAULT_GREEDINESS); + } + + /** + * Create a manager around an externally provided floor, for queries whose floor is also fed by + * searchers outside this process. + * + * @param k the number of neighbors the query collects + * @param globalFloor the floor shared by all searchers of this query; its {@link + * GlobalKnnFloor#k()} must equal {@code k}, since a floor tracking the wrong result-set size + * is not a valid bound + */ + public SharedFloorKnnCollectorManager(int k, GlobalKnnFloor globalFloor) { + this(k, globalFloor, FloorAwareKnnCollector.DEFAULT_GREEDINESS); + } + + /** + * Create a manager around an externally provided floor with an explicit greediness, applying the + * {@link #DEFAULT_FLOOR_ACTIVATION_K default activation threshold}. + * + * @param k the number of neighbors the query collects + * @param globalFloor the floor shared by all searchers of this query; its {@link + * GlobalKnnFloor#k()} must equal {@code k} + * @param greediness fraction of each segment's search effort that follows the shared floor, in + * {@code [0, 1]}; see {@link FloorAwareKnnCollector} + */ + public SharedFloorKnnCollectorManager(int k, GlobalKnnFloor globalFloor, float greediness) { + this(k, globalFloor, greediness, DEFAULT_FLOOR_ACTIVATION_K); + } + + /** + * Create a manager around an externally provided floor with an explicit greediness and activation + * threshold, applying the default slot minimum and sync interval. + * + * @param k the number of neighbors the query collects + * @param globalFloor the floor shared by all searchers of this query; its {@link + * GlobalKnnFloor#k()} must equal {@code k} + * @param greediness fraction of each segment's search effort that follows the shared floor, in + * {@code [0, 1]}; see {@link FloorAwareKnnCollector} + * @param floorActivationK the smallest k at which floor sharing engages; for smaller k this + * manager creates plain collectors and the search is exactly stock search. See the class + * comment for the reasoning behind the {@link #DEFAULT_FLOOR_ACTIVATION_K default}. + */ + public SharedFloorKnnCollectorManager( + int k, GlobalKnnFloor globalFloor, float greediness, int floorActivationK) { + this( + k, + globalFloor, + greediness, + floorActivationK, + FloorAwareKnnCollector.DEFAULT_MIN_EXPLORATION_SLOTS, + FloorAwareKnnCollector.DEFAULT_SYNC_INTERVAL); + } + + /** + * Create a fully configured manager for an index holding the whole corpus. Equivalent to the + * seven-argument constructor with {@code globalShare = 1}. + * + * @param k the number of neighbors the query collects + * @param globalFloor the floor shared by all searchers of this query; its {@link + * GlobalKnnFloor#k()} must equal {@code k} + * @param greediness fraction of each segment's search effort that follows the shared floor, in + * {@code [0, 1]}; see {@link FloorAwareKnnCollector} + * @param floorActivationK the smallest k at which floor sharing engages; for smaller k this + * manager creates plain collectors and the search is exactly stock search. See the class + * comment for the reasoning behind the {@link #DEFAULT_FLOOR_ACTIVATION_K default}. + * @param minExplorationSlots smallest permitted size of each collector's greediness clamp queue; + * must be at least 1. See {@link FloorAwareKnnCollector#DEFAULT_MIN_EXPLORATION_SLOTS}. + * @param syncInterval number of visited vectors between each collector's synchronizations with + * the shared floor; must be positive. See {@link + * FloorAwareKnnCollector#DEFAULT_SYNC_INTERVAL}. + */ + public SharedFloorKnnCollectorManager( + int k, + GlobalKnnFloor globalFloor, + float greediness, + int floorActivationK, + int minExplorationSlots, + int syncInterval) { + this(k, globalFloor, greediness, floorActivationK, minExplorationSlots, syncInterval, 1f); + } + + /** + * Create a fully configured manager. Every tuning value the mechanism has is a parameter here; + * the shorter constructors exist only to supply defaults. + * + * @param k the number of neighbors the query collects + * @param globalFloor the floor shared by all searchers of this query; its {@link + * GlobalKnnFloor#k()} must equal {@code k} + * @param greediness fraction of each segment's search effort that follows the shared floor, in + * {@code [0, 1]}; see {@link FloorAwareKnnCollector}. When {@code globalShare} is below 1, + * the gate is small and the clamp {@code (1 - greediness) * gateK} with it: derive this value + * from the exploration width the graph needs via {@link + * FloorAwareKnnCollector#greedinessForClamp(int, int)} (with {@code gateK} from {@link + * #perShardGate(int, double)}) rather than choosing a constant. See the caution on {@link + * FloorAwareKnnCollector}'s greediness clamp for what goes wrong with constants. + * @param floorActivationK the smallest k at which floor sharing engages; for smaller k this + * manager creates plain collectors and the search is exactly stock search. See the class + * comment for the reasoning behind the {@link #DEFAULT_FLOOR_ACTIVATION_K default}. + * @param minExplorationSlots smallest permitted size of each collector's greediness clamp queue; + * must be at least 1. See {@link FloorAwareKnnCollector#DEFAULT_MIN_EXPLORATION_SLOTS}. + * @param syncInterval number of visited vectors between each collector's synchronizations with + * the shared floor; must be positive. See {@link + * FloorAwareKnnCollector#DEFAULT_SYNC_INTERVAL}. + * @param globalShare the fraction of the whole corpus held by the index this manager searches, in + * {@code (0, 1]}. Pass a value below 1 when this index is one shard of a sharded corpus whose + * searchers share the floor across processes; each collector's ascent gate then opens at the + * shard's expected contribution to the merged top-k instead of at its full local queue. See + * the class comment. + */ + public SharedFloorKnnCollectorManager( + int k, + GlobalKnnFloor globalFloor, + float greediness, + int floorActivationK, + int minExplorationSlots, + int syncInterval, + float globalShare) { + if (k < 1) { + throw new IllegalArgumentException("k must be at least 1, got: " + k); + } + Objects.requireNonNull(globalFloor, "globalFloor"); + if (globalFloor.k() != k) { + throw new IllegalArgumentException( + "the floor must track the same result-set size as the query: floor k=" + + globalFloor.k() + + ", query k=" + + k); + } + if (greediness < 0 || greediness > 1 || Float.isNaN(greediness)) { + throw new IllegalArgumentException("greediness must be in [0,1], got: " + greediness); + } + if (floorActivationK < 1) { + throw new IllegalArgumentException( + "floorActivationK must be at least 1, got: " + floorActivationK); + } + if (minExplorationSlots < 1) { + throw new IllegalArgumentException( + "minExplorationSlots must be at least 1, got: " + minExplorationSlots); + } + if (syncInterval < 1) { + throw new IllegalArgumentException("syncInterval must be positive, got: " + syncInterval); + } + if (globalShare <= 0 || globalShare > 1 || Float.isNaN(globalShare)) { + throw new IllegalArgumentException("globalShare must be in (0,1], got: " + globalShare); + } + this.k = k; + this.globalFloor = globalFloor; + this.greediness = greediness; + this.floorActivationK = floorActivationK; + this.minExplorationSlots = minExplorationSlots; + this.syncInterval = syncInterval; + this.globalShare = globalShare; + } + + /** + * The number of results a searcher holding fraction {@code share} of the corpus is expected to + * contribute to the merged top-k, with the same statistical padding the optimistic multi-segment + * strategy applies: under random sharding the number of global top-k hits in the shard is + * binomial, and the gate is sized {@code k * share} plus {@value #LAMBDA} standard deviations, + * clamped to {@code [1, k]}. This is where a shard's ascent gate should open: below it the shard + * has not yet collected its statistically due contribution and the shared floor must not end its + * search; above it, continuing is worthwhile exactly as long as the shard's results remain + * globally competitive. + * + * @param k the number of neighbors the query collects, at least 1 + * @param share the fraction of the corpus the searcher holds, in {@code (0, 1]} + * @return the gate size, in {@code [1, k]} + */ + public static int perShardGate(int k, double share) { + if (k < 1) { + throw new IllegalArgumentException("k must be at least 1, got: " + k); + } + if (share <= 0 || share > 1 || Double.isNaN(share)) { + throw new IllegalArgumentException("share must be in (0,1], got: " + share); + } + // Mirrors perLeafTopKCalculation in AbstractKnnVectorQuery so that a gate computed for a + // shard equals the quota the optimistic strategy would give a same-sized segment. + int gate = (int) Math.max(1, k * share + LAMBDA * Math.sqrt(k * share * (1 - share))); + return Math.min(k, gate); + } + + /** Return the floor shared by this manager's collectors, so that callers may feed or read it. */ + public GlobalKnnFloor getGlobalFloor() { + return globalFloor; + } + + @Override + public KnnCollector newCollector( + int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context) + throws IOException { + TopKnnCollector collector = new TopKnnCollector(k, visitedLimit, searchStrategy); + if (k < floorActivationK) { + return collector; + } + return floorAware(collector, context); + } + + @Override + public KnnCollector newOptimisticCollector( + int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context, int perLeafK) + throws IOException { + // The local queue is sized to the segment's pro-rata share; the floor itself always tracks + // the full k best across the query. Activation is decided by the query's k, not the segment's + // share: the policy is about the query. + TopKnnCollector collector = new TopKnnCollector(perLeafK, visitedLimit, searchStrategy); + if (k < floorActivationK) { + return collector; + } + return floorAware(collector, context); + } + + private FloorAwareKnnCollector floorAware(TopKnnCollector collector, LeafReaderContext context) { + // Publish each leaf's scores at most once per query; a second search of the same leaf + // re-collects the same documents, and republishing them would violate the floor's + // distinct-document contract. See the class comment. + boolean publish = context == null || publishedLeaves.add(context); + return new FloorAwareKnnCollector( + collector, + globalFloor, + greediness, + minExplorationSlots, + syncInterval, + gateFor(context, collector.k()), + publish); + } + + /** + * The ascent gate for a collector over {@code context}: the leaf's statistically expected + * contribution to the merged top-k, never more than the local queue it collects into. The leaf's + * share of the whole corpus is its share of this index scaled by the index's {@code globalShare}. + * When that combined share is 1 — a single-segment index holding the entire corpus — the gate is + * the queue size and behavior is exactly the default: the local queue is the query's own top-k, + * and only external advertisements could justify ending the search below it. + */ + private int gateFor(LeafReaderContext context, int queueSize) { + double leafGlobalShare = globalShare; + if (context != null && context.parent != null) { + leafGlobalShare *= context.reader().maxDoc() / (double) context.parent.reader().maxDoc(); + } + if (leafGlobalShare <= 0 || leafGlobalShare >= 1) { + return queueSize; + } + return Math.min(queueSize, perShardGate(k, leafGlobalShare)); + } + + @Override + public boolean isOptimistic() { + return true; + } +} diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/package-info.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/package-info.java new file mode 100644 index 000000000000..4ae3ff79cb8b --- /dev/null +++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/search/knn/package-info.java @@ -0,0 +1,24 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * Shared-floor kNN collection: collectors that prune each searcher's graph exploration against a + * lower bound of the query's final merged top-k cutoff, shared across index segments and, through + * {@link org.apache.lucene.sandbox.search.knn.GlobalKnnFloor#advertise(float)}, across searchers in + * other processes. + */ +package org.apache.lucene.sandbox.search.knn; diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestBlockingFloatHeap.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestBlockingFloatHeap.java new file mode 100644 index 000000000000..d4a1f8770f4c --- /dev/null +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestBlockingFloatHeap.java @@ -0,0 +1,162 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.lucene.sandbox.search.knn; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.randomIntBetween; + +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.SuppressForbidden; + +public class TestBlockingFloatHeap extends LuceneTestCase { + + public void testRejectsInvalidCapacity() { + expectThrows(IllegalArgumentException.class, () -> new BlockingFloatHeap(0)); + expectThrows(IllegalArgumentException.class, () -> new BlockingFloatHeap(-1)); + } + + public void testBasicOperations() { + BlockingFloatHeap heap = new BlockingFloatHeap(3); + heap.offer(2); + heap.offer(4); + heap.offer(1); + heap.offer(3); + assertEquals(3, heap.size()); + assertEquals(2, heap.peek(), 0); + + assertEquals(2, heap.poll(), 0); + assertEquals(3, heap.poll(), 0); + assertEquals(4, heap.poll(), 0); + assertEquals(0, heap.size(), 0); + } + + public void testBasicOperations2() { + int size = atLeast(10); + BlockingFloatHeap heap = new BlockingFloatHeap(size); + double sum = 0, sum2 = 0; + + for (int i = 0; i < size; i++) { + float next = random().nextFloat(100f); + sum += next; + heap.offer(next); + } + + float last = Float.NEGATIVE_INFINITY; + for (long i = 0; i < size; i++) { + float next = heap.poll(); + assertTrue(next >= last); + last = next; + sum2 += last; + } + assertEquals(sum, sum2, 0.01); + } + + @SuppressForbidden(reason = "Thread sleep") + public void testMultipleThreads() throws Exception { + Thread[] threads = new Thread[randomIntBetween(3, 20)]; + final CountDownLatch latch = new CountDownLatch(1); + BlockingFloatHeap globalHeap = new BlockingFloatHeap(1); + + for (int i = 0; i < threads.length; i++) { + threads[i] = + new Thread( + () -> { + try { + latch.await(); + int numIterations = randomIntBetween(10, 100); + float bottomValue = 0; + + while (numIterations-- > 0) { + bottomValue += randomIntBetween(0, 5); + globalHeap.offer(bottomValue); + Thread.sleep(randomIntBetween(0, 50)); + + float globalBottomValue = globalHeap.peek(); + assertTrue(globalBottomValue >= bottomValue); + bottomValue = globalBottomValue; + } + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + threads[i].start(); + } + + latch.countDown(); + for (Thread t : threads) { + t.join(); + } + } + + public void testConcurrentPollsDoNotCorruptTheHeap() throws Exception { + // Twice as many pollers as elements: the surplus polls must fail cleanly with + // IllegalStateException, never pass an emptiness check raced by another poller and corrupt + // the heap's size or return garbage. Every successfully polled value must be one of the + // values actually offered, each exactly once. + int size = 8; + int pollers = 2 * size; + BlockingFloatHeap heap = new BlockingFloatHeap(size); + for (int i = 0; i < size; i++) { + heap.offer(i + 1f); + } + + CountDownLatch startingGun = new CountDownLatch(1); + ConcurrentLinkedQueue polled = new ConcurrentLinkedQueue<>(); + AtomicInteger emptyPolls = new AtomicInteger(); + Thread[] threads = new Thread[pollers]; + for (int i = 0; i < pollers; i++) { + threads[i] = + new Thread( + () -> { + try { + startingGun.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + try { + polled.add(heap.poll()); + } catch ( + @SuppressWarnings("unused") + IllegalStateException e) { + emptyPolls.incrementAndGet(); + } + }); + threads[i].start(); + } + startingGun.countDown(); + for (Thread t : threads) { + t.join(); + } + + assertEquals("every element must be polled exactly once", size, polled.size()); + assertEquals("every surplus poll must fail cleanly", pollers - size, emptyPolls.get()); + assertEquals(0, heap.size()); + float[] values = new float[polled.size()]; + int i = 0; + for (float value : polled) { + values[i++] = value; + } + java.util.Arrays.sort(values); + for (int v = 0; v < size; v++) { + assertEquals("polled values must be exactly the offered values", v + 1f, values[v], 0.0f); + } + } +} diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestFloorAwareKnnCollector.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestFloorAwareKnnCollector.java new file mode 100644 index 000000000000..a899370993f4 --- /dev/null +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestFloorAwareKnnCollector.java @@ -0,0 +1,447 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.lucene.sandbox.search.knn; + +import org.apache.lucene.search.TopKnnCollector; +import org.apache.lucene.tests.util.LuceneTestCase; + +public class TestFloorAwareKnnCollector extends LuceneTestCase { + + public void testInvalidGreediness() { + GlobalKnnFloor floor = new GlobalKnnFloor(4); + TopKnnCollector delegate = new TopKnnCollector(4, Integer.MAX_VALUE); + expectThrows( + IllegalArgumentException.class, () -> new FloorAwareKnnCollector(delegate, floor, -0.1f)); + expectThrows( + IllegalArgumentException.class, () -> new FloorAwareKnnCollector(delegate, floor, 1.1f)); + expectThrows( + IllegalArgumentException.class, + () -> new FloorAwareKnnCollector(delegate, floor, Float.NaN)); + } + + public void testInvalidMinExplorationSlots() { + GlobalKnnFloor floor = new GlobalKnnFloor(4); + TopKnnCollector delegate = new TopKnnCollector(4, Integer.MAX_VALUE); + expectThrows( + IllegalArgumentException.class, + () -> + new FloorAwareKnnCollector( + delegate, floor, 0.5f, 0, FloorAwareKnnCollector.DEFAULT_SYNC_INTERVAL)); + expectThrows( + IllegalArgumentException.class, + () -> + new FloorAwareKnnCollector( + delegate, floor, 0.5f, -1, FloorAwareKnnCollector.DEFAULT_SYNC_INTERVAL)); + } + + public void testInvalidSyncInterval() { + GlobalKnnFloor floor = new GlobalKnnFloor(4); + TopKnnCollector delegate = new TopKnnCollector(4, Integer.MAX_VALUE); + expectThrows( + IllegalArgumentException.class, + () -> + new FloorAwareKnnCollector( + delegate, floor, 0.5f, FloorAwareKnnCollector.DEFAULT_MIN_EXPLORATION_SLOTS, 0)); + expectThrows( + IllegalArgumentException.class, + () -> + new FloorAwareKnnCollector( + delegate, floor, 0.5f, FloorAwareKnnCollector.DEFAULT_MIN_EXPLORATION_SLOTS, -8)); + // Any positive interval is valid: synchronization is scheduled against a visited-count + // threshold, not a bit mask, so no power-of-two restriction applies. + new FloorAwareKnnCollector( + delegate, floor, 0.5f, FloorAwareKnnCollector.DEFAULT_MIN_EXPLORATION_SLOTS, 100); + } + + public void testInvalidGateK() { + GlobalKnnFloor floor = new GlobalKnnFloor(8); + TopKnnCollector delegate = new TopKnnCollector(8, Integer.MAX_VALUE); + expectThrows( + IllegalArgumentException.class, + () -> + new FloorAwareKnnCollector( + delegate, + floor, + 0.5f, + FloorAwareKnnCollector.DEFAULT_MIN_EXPLORATION_SLOTS, + FloorAwareKnnCollector.DEFAULT_SYNC_INTERVAL, + 0)); + expectThrows( + IllegalArgumentException.class, + () -> + new FloorAwareKnnCollector( + delegate, + floor, + 0.5f, + FloorAwareKnnCollector.DEFAULT_MIN_EXPLORATION_SLOTS, + FloorAwareKnnCollector.DEFAULT_SYNC_INTERVAL, + -1)); + // A gate above the queue size could never open: the queue's size saturates at its capacity. + expectThrows( + IllegalArgumentException.class, + () -> + new FloorAwareKnnCollector( + delegate, + floor, + 0.5f, + FloorAwareKnnCollector.DEFAULT_MIN_EXPLORATION_SLOTS, + FloorAwareKnnCollector.DEFAULT_SYNC_INTERVAL, + 9)); + } + + public void testGreedinessForClampRoundTrips() { + // The helper must invert the clamp sizing: constructing a collector with the derived + // greediness yields a clamp of exactly the requested width, at any gate size. This is the + // property that makes width-first configuration safe where a constant greediness is not. + int[][] cases = {{6000, 600}, {184, 64}, {1012, 101}, {117, 58}, {10000, 16}}; + for (int[] c : cases) { + int gateK = c[0]; + int clampSlots = c[1]; + float greediness = FloorAwareKnnCollector.greedinessForClamp(gateK, clampSlots); + assertEquals( + "gateK=" + gateK + " clampSlots=" + clampSlots, + clampSlots, + Math.round((1 - greediness) * gateK)); + } + // A width of the whole gate or more means the floor cannot bind: greediness 0, stock search. + assertEquals(0f, FloorAwareKnnCollector.greedinessForClamp(44, 64), 0.0f); + assertEquals(0f, FloorAwareKnnCollector.greedinessForClamp(64, 64), 0.0f); + + expectThrows( + IllegalArgumentException.class, () -> FloorAwareKnnCollector.greedinessForClamp(0, 16)); + expectThrows( + IllegalArgumentException.class, () -> FloorAwareKnnCollector.greedinessForClamp(100, 0)); + } + + public void testGateBelowQueueSizeOpensAtGateKResults() { + // A share-sized gate: the local queue holds 32 results but the collector's expected share of + // the merged top-k is only 8, so the floor must engage after 8 collected results, long before + // the queue fills. The local k-th best is still undefined at that point, so the exposed bound + // comes entirely from min(clamp, floor). + int queueSize = 32; + int gateK = 8; + GlobalKnnFloor floor = new GlobalKnnFloor(queueSize); + floor.advertise(1000f); + TopKnnCollector delegate = new TopKnnCollector(queueSize, Integer.MAX_VALUE); + // greediness 1 with a slot minimum of 2 keeps the two best scores seen as exploration slots. + FloorAwareKnnCollector collector = + new FloorAwareKnnCollector( + delegate, floor, 1f, 2, FloorAwareKnnCollector.DEFAULT_SYNC_INTERVAL, gateK); + + for (int doc = 0; doc < gateK - 1; doc++) { + collector.incVisitedCount(1); + collector.collect(doc, doc + 1f); + assertEquals( + "the shared floor must be invisible until gateK results are collected", + Float.NEGATIVE_INFINITY, + collector.minCompetitiveSimilarity(), + 0.0f); + } + + collector.incVisitedCount(1); + collector.collect(gateK - 1, (float) gateK); + // Scores are 1..8: the clamp queue keeps {7, 8}, so the bound is min(7, nextDown(1000)) = 7, + // while the delegate's own bound is still NEGATIVE_INFINITY because its queue is not full. + assertEquals(Float.NEGATIVE_INFINITY, delegate.minCompetitiveSimilarity(), 0.0f); + assertEquals( + "once the gate opens, the bound must come from the clamp and the floor", + 7f, + collector.minCompetitiveSimilarity(), + 0.0f); + } + + public void testGreedinessClampIsSizedFromGateK() { + // With a queue of 40 but a gate of 8 at greediness 0.5, the clamp must keep (1 - 0.5) * 8 = 4 + // slots, not (1 - 0.5) * 40 = 20: the clamp protects the share-sized search the gate defines, + // not the queue's capacity. + int queueSize = 40; + int gateK = 8; + GlobalKnnFloor floor = new GlobalKnnFloor(queueSize); + floor.advertise(1000f); + FloorAwareKnnCollector collector = + new FloorAwareKnnCollector( + new TopKnnCollector(queueSize, Integer.MAX_VALUE), + floor, + 0.5f, + 1, + FloorAwareKnnCollector.DEFAULT_SYNC_INTERVAL, + gateK); + + for (int doc = 0; doc < gateK; doc++) { + collector.incVisitedCount(1); + collector.collect(doc, doc + 1f); + } + + // Scores are 1..8 and the clamp keeps the best 4 of them, {5, 6, 7, 8}: the bound must be the + // clamp's minimum, 5. A clamp sized from the queue would hold all 8 scores and expose 1. + assertEquals( + "the clamp must keep (1 - greediness) * gateK slots", + 5f, + collector.minCompetitiveSimilarity(), + 0.0f); + } + + public void testConfigurableSlotMinimumControlsNeutralization() { + // The neutralization property follows the configured minimum, not a baked-in number: with a + // slot minimum of 4 at k=8, the clamp is narrower than the local queue and an advertised + // floor must bind; with a slot minimum of 8 it must not. + int k = 8; + for (int slots : new int[] {4, 8}) { + GlobalKnnFloor floor = new GlobalKnnFloor(k); + floor.advertise(1000f); + TopKnnCollector delegate = new TopKnnCollector(k, Integer.MAX_VALUE); + FloorAwareKnnCollector collector = + new FloorAwareKnnCollector( + delegate, floor, 1f, slots, FloorAwareKnnCollector.DEFAULT_SYNC_INTERVAL); + for (int doc = 0; doc < k; doc++) { + collector.incVisitedCount(1); + collector.collect(doc, doc + 1f); + } + if (slots < k) { + assertTrue( + "a clamp narrower than the local queue must let the floor bind", + collector.minCompetitiveSimilarity() > delegate.minCompetitiveSimilarity()); + } else { + assertEquals( + "a clamp as wide as the local queue must neutralize the floor", + delegate.minCompetitiveSimilarity(), + collector.minCompetitiveSimilarity(), + 0.0f); + } + } + } + + public void testAscentGateIgnoresFloorUntilLocalQueueFills() { + int k = FloorAwareKnnCollector.DEFAULT_MIN_EXPLORATION_SLOTS + 4; + GlobalKnnFloor floor = new GlobalKnnFloor(k); + // A sibling searcher has already converged and established a high floor. A correct collector + // must not expose it before this searcher has escaped its own ascent, otherwise the graph + // search would be terminated at its entry point before finding anything. + floor.advertise(1000f); + FloorAwareKnnCollector collector = + new FloorAwareKnnCollector(new TopKnnCollector(k, Integer.MAX_VALUE), floor); + + for (int doc = 0; doc < k - 1; doc++) { + collector.incVisitedCount(1); + collector.collect(doc, doc + 1f); + assertEquals( + "the shared floor must be invisible while the local queue is filling", + Float.NEGATIVE_INFINITY, + collector.minCompetitiveSimilarity(), + 0.0f); + } + + collector.incVisitedCount(1); + collector.collect(k - 1, (float) k); + assertTrue( + "once the local queue is full, the shared floor must start binding", + collector.minCompetitiveSimilarity() > Float.NEGATIVE_INFINITY); + // With k exceeding the clamp's minimum size, the bound must exceed the local k-th best (1): + // the floor is genuinely binding, not merely echoing the local queue. + assertTrue(collector.minCompetitiveSimilarity() > 1f); + } + + public void testGreedinessClampCapsTheSharedFloor() { + // k and greediness sized so the fractional clamp exceeds the absolute minimum: 40 collected + // scores with greediness 0.5 keep a non-competitive queue of 20 entries, so the effective + // bound may never exceed the 20th best similarity this collector has seen. + int k = 40; + float greediness = 0.5f; + GlobalKnnFloor floor = new GlobalKnnFloor(k); + floor.advertise(1000f); + FloorAwareKnnCollector collector = + new FloorAwareKnnCollector(new TopKnnCollector(k, Integer.MAX_VALUE), floor, greediness); + + for (int doc = 0; doc < k; doc++) { + collector.incVisitedCount(1); + collector.collect(doc, doc + 1f); + } + + // Scores are 1..40: the local k-th best is 1, the 20th best seen is 21, and the floor is + // 1000. The clamp must win. + assertEquals( + "the bound must be capped by the clamp queue's minimum, not jump to the shared floor", + 21f, + collector.minCompetitiveSimilarity(), + 0.0f); + } + + public void testSharedFloorBindsOneUlpBelowItsValue() { + int localK = 20; + // Size the floor for a larger result set so the local scores cannot define it: the only floor + // source in this test is the advertised bound. + GlobalKnnFloor floor = new GlobalKnnFloor(100); + floor.advertise(2.5f); + // greediness 1 collapses the clamp to its absolute minimum of DEFAULT_MIN_EXPLORATION_SLOTS + // entries. + FloorAwareKnnCollector collector = + new FloorAwareKnnCollector(new TopKnnCollector(localK, Integer.MAX_VALUE), floor, 1f); + + // Three scores below the advertised floor, then 17 above it: the local k-th best (1) stays + // below the floor while the clamp queue's minimum (the 16th best seen, 3.1) rises above it, + // so min(clamp, nextDown(floor)) selects the floor term. + collector.incVisitedCount(1); + collector.collect(0, 1f); + collector.incVisitedCount(1); + collector.collect(1, 1.2f); + collector.incVisitedCount(1); + collector.collect(2, 1.4f); + for (int doc = 3; doc < localK; doc++) { + collector.incVisitedCount(1); + collector.collect(doc, 3f + 0.1f * (doc - 2)); + } + + // The bound must sit strictly below the floor: a hit scoring exactly at the floor may still + // win the merged tie-break and must remain findable. + assertEquals(Math.nextDown(2.5f), collector.minCompetitiveSimilarity(), 0.0f); + assertTrue(collector.minCompetitiveSimilarity() < 2.5f); + } + + public void testFloorIsNeutralizedAtSmallK() { + // When k does not exceed the clamp's absolute minimum, the clamp queue is at least as large + // as the local queue, its minimum can never exceed the local k-th best, and the shared floor + // must have no effect at all: even a hostile advertised bound cannot change the bound stock + // search would have used. + int k = FloorAwareKnnCollector.DEFAULT_MIN_EXPLORATION_SLOTS / 2; + GlobalKnnFloor floor = new GlobalKnnFloor(k); + floor.advertise(Float.MAX_VALUE); + TopKnnCollector delegate = new TopKnnCollector(k, Integer.MAX_VALUE); + FloorAwareKnnCollector collector = new FloorAwareKnnCollector(delegate, floor, 1f); + + for (int doc = 0; doc < 3 * k; doc++) { + collector.incVisitedCount(1); + collector.collect(doc, random().nextFloat()); + assertEquals( + "at k <= the clamp's slot minimum the bound must be exactly the delegate's", + delegate.minCompetitiveSimilarity(), + collector.minCompetitiveSimilarity(), + 0.0f); + } + } + + public void testCollectReportsSharedFloorUpdates() { + // k equal to the clamp's absolute minimum makes both local structures the same size, so a + // score rejected by the local queue is also rejected by the clamp queue and cannot report an + // update through either. + int k = FloorAwareKnnCollector.DEFAULT_MIN_EXPLORATION_SLOTS; + GlobalKnnFloor floor = new GlobalKnnFloor(k); + FloorAwareKnnCollector collector = + new FloorAwareKnnCollector(new TopKnnCollector(k, Integer.MAX_VALUE), floor, 0f); + + for (int doc = 0; doc < k; doc++) { + collector.incVisitedCount(1); + assertTrue("a locally accepted hit must report an update", collector.collect(doc, doc + 5f)); + } + + // Both queues hold 5..k+4. A worse score away from a synchronization boundary changes + // nothing and must say so, otherwise the searcher would re-derive its bound for no reason. + collector.incVisitedCount(1); + assertFalse( + "a rejected hit between synchronizations must not report an update", + collector.collect(k, 1f)); + + // Advance a full sync interval past the gate-open synchronization: even a rejected hit must + // report an update there, because the re-read of the shared floor may have moved the + // effective bound. + collector.incVisitedCount(FloorAwareKnnCollector.DEFAULT_SYNC_INTERVAL); + assertTrue( + "a hit past the synchronization threshold must report an update after the floor re-read", + collector.collect(k + 1, 1f)); + } + + public void testSyncHappensOnceIntervalElapsesRegardlessOfAlignment() { + // The synchronization schedule must be a threshold over the visited count, not an exact + // boundary match: collect() is invoked only for candidates that beat the current bound, so + // visited counts arrive at irregular strides, and a boundary that falls between two collects + // must trigger on the next collect rather than being skipped. Skipped boundaries starve the + // floor exactly in the late phase of the search, when few candidates beat the bar. + int k = 4; + GlobalKnnFloor floor = new GlobalKnnFloor(k); + FloorAwareKnnCollector collector = + new FloorAwareKnnCollector( + new TopKnnCollector(k, Integer.MAX_VALUE), floor, 0.5f, 1, 256, k); + + // Fill to the gate: the first batch publishes and defines the floor at the k-th best, 1. + for (int doc = 0; doc < k; doc++) { + collector.incVisitedCount(1); + collector.collect(doc, doc + 1f); + } + assertEquals(1f, floor.floor(), 0.0f); + + // A better score arrives, then the visited count jumps far past the next synchronization + // threshold without ever landing on a multiple of the interval. + collector.incVisitedCount(1); + collector.collect(k, 100f); + collector.incVisitedCount(300); + assertTrue("test setup: land off any interval multiple", collector.visitedCount() % 256 != 0); + collector.incVisitedCount(1); + collector.collect(k + 1, 101f); + + assertTrue( + "a collect past the synchronization threshold must publish the pending scores", + floor.floor() > 1f); + } + + public void testNonPublishingCollectorReadsButNeverFeedsTheFloor() { + int k = 20; + GlobalKnnFloor floor = new GlobalKnnFloor(k); + FloorAwareKnnCollector collector = + new FloorAwareKnnCollector( + new TopKnnCollector(k, Integer.MAX_VALUE), floor, 1f, 2, 256, k, false); + assertFalse(collector.publishesToFloor()); + + // Collect a full queue of scores: a publishing collector would have defined the floor here. + for (int doc = 0; doc < k; doc++) { + collector.incVisitedCount(1); + collector.collect(doc, doc + 1f); + } + assertEquals( + "a non-publishing collector must never feed the shared floor", + Float.NEGATIVE_INFINITY, + floor.floor(), + 0.0f); + + // The floor rises externally; the collector must still read it and prune against it. After + // the final collect the clamp holds the two best scores seen ({20, 30}), so the bound is + // max(local k-th best = 2, min(clamp = 20, nextDown(1000))) = 20 — above anything the local + // queue alone would justify, provable only by a floor that was actually read. + floor.advertise(1000f); + collector.incVisitedCount(FloorAwareKnnCollector.DEFAULT_SYNC_INTERVAL); + collector.collect(k, 30f); + assertEquals( + "a non-publishing collector must still prune against the shared floor", + 20f, + collector.minCompetitiveSimilarity(), + 0.0f); + } + + public void testDelegationOfCollectorPlumbing() { + GlobalKnnFloor floor = new GlobalKnnFloor(3); + TopKnnCollector delegate = new TopKnnCollector(3, 17); + FloorAwareKnnCollector collector = new FloorAwareKnnCollector(delegate, floor); + assertEquals(3, collector.k()); + assertEquals(17, collector.visitLimit()); + collector.incVisitedCount(5); + assertEquals(5, collector.visitedCount()); + assertEquals(5, delegate.visitedCount()); + assertFalse(collector.earlyTerminated()); + collector.incVisitedCount(12); + assertTrue( + "the visit limit must keep terminating through the decorator", collector.earlyTerminated()); + } +} diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestGlobalKnnFloor.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestGlobalKnnFloor.java new file mode 100644 index 000000000000..6519fc02355d --- /dev/null +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestGlobalKnnFloor.java @@ -0,0 +1,294 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.lucene.sandbox.search.knn; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Comparator; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.NamedThreadFactory; + +public class TestGlobalKnnFloor extends LuceneTestCase { + + public void testInvalidK() { + IllegalArgumentException e = + expectThrows(IllegalArgumentException.class, () -> new GlobalKnnFloor(0)); + assertTrue(e.getMessage(), e.getMessage().contains("k must be at least 1")); + expectThrows(IllegalArgumentException.class, () -> new GlobalKnnFloor(-1)); + } + + public void testFloorUndefinedUntilKScoresObserved() { + int k = TestUtil.nextInt(random(), 1, 50); + GlobalKnnFloor floor = new GlobalKnnFloor(k); + assertEquals(Float.NEGATIVE_INFINITY, floor.floor(), 0.0f); + for (int i = 0; i < k - 1; i++) { + float returned = floor.offer(new float[] {random().nextFloat()}, 1); + assertEquals( + "floor must stay undefined while fewer than k scores have been observed", + Float.NEGATIVE_INFINITY, + returned, + 0.0f); + } + float kthArrival = random().nextFloat(); + float returned = floor.offer(new float[] {kthArrival}, 1); + assertTrue( + "floor must be defined once k scores have been observed", + returned > Float.NEGATIVE_INFINITY); + } + + public void testFloorIsKthBestOfEverythingOffered() { + int k = TestUtil.nextInt(random(), 1, 32); + GlobalKnnFloor floor = new GlobalKnnFloor(k); + int totalScores = TestUtil.nextInt(random(), k, 500); + List allScores = new ArrayList<>(totalScores); + int offered = 0; + while (offered < totalScores) { + int batchSize = TestUtil.nextInt(random(), 1, Math.min(16, totalScores - offered)); + float[] batch = new float[batchSize]; + for (int i = 0; i < batchSize; i++) { + batch[i] = random().nextFloat() * 10; + allScores.add(batch[i]); + } + Arrays.sort(batch); + floor.offer(batch, batchSize); + offered += batchSize; + } + allScores.sort(Comparator.reverseOrder()); + float expectedKthBest = allScores.get(k - 1); + assertEquals(expectedKthBest, floor.floor(), 0.0f); + } + + public void testFloorIsMonotonic() { + int k = TestUtil.nextInt(random(), 1, 16); + GlobalKnnFloor floor = new GlobalKnnFloor(k); + float previous = floor.floor(); + for (int batch = 0; batch < 50; batch++) { + int batchSize = TestUtil.nextInt(random(), 1, 8); + float[] scores = new float[batchSize]; + for (int i = 0; i < batchSize; i++) { + scores[i] = (random().nextFloat() - 0.5f) * 100; + } + Arrays.sort(scores); + float current = floor.offer(scores, batchSize); + assertTrue("floor must never decrease", current >= previous); + previous = current; + } + } + + public void testAdvertiseIsIndependentOfHeapFill() { + // An advertised bound is valid on its own authority; it must define the floor even when no + // local scores have been observed yet. + GlobalKnnFloor floor = new GlobalKnnFloor(10); + floor.advertise(0.5f); + assertEquals(0.5f, floor.floor(), 0.0f); + } + + public void testAdvertiseBelowCurrentFloorIsIgnored() { + GlobalKnnFloor floor = new GlobalKnnFloor(10); + floor.advertise(0.5f); + floor.advertise(0.4f); + assertEquals( + "a lower advertised bound must not lower the floor: duplicated or reordered remote " + + "deliveries rely on this", + 0.5f, + floor.floor(), + 0.0f); + floor.advertise(0.6f); + assertEquals(0.6f, floor.floor(), 0.0f); + } + + public void testAdvertiseCombinesWithLocalScores() { + int k = 4; + GlobalKnnFloor floor = new GlobalKnnFloor(k); + float[] scores = new float[] {0.1f, 0.2f, 0.3f, 0.4f}; + floor.offer(scores, scores.length); + assertEquals(0.1f, floor.floor(), 0.0f); + // A remote bound above the local k-th best takes over. + floor.advertise(0.25f); + assertEquals(0.25f, floor.floor(), 0.0f); + // Local scores that push the local k-th best above the advertised bound take over again. + floor.offer(new float[] {0.5f, 0.6f, 0.7f}, 3); + // Heap now holds {0.4, 0.5, 0.6, 0.7}: local k-th best is 0.4. + assertEquals(0.4f, floor.floor(), 0.0f); + } + + public void testAdvertiseRejectsNaN() { + GlobalKnnFloor floor = new GlobalKnnFloor(10); + expectThrows(IllegalArgumentException.class, () -> floor.advertise(Float.NaN)); + } + + public void testOfferRejectsNonPositiveLen() { + GlobalKnnFloor floor = new GlobalKnnFloor(10); + expectThrows(IllegalArgumentException.class, () -> floor.offer(new float[] {1f}, 0)); + expectThrows(IllegalArgumentException.class, () -> floor.offer(new float[] {1f}, -1)); + } + + public void testOfferValidatesItsPublicBatchContract() { + GlobalKnnFloor floor = new GlobalKnnFloor(2); + expectThrows(NullPointerException.class, () -> floor.offer(null, 1)); + expectThrows(IllegalArgumentException.class, () -> floor.offer(new float[] {1f}, 2)); + expectThrows(IllegalArgumentException.class, () -> floor.offer(new float[] {2f, 1f}, 2)); + expectThrows(IllegalArgumentException.class, () -> floor.offer(new float[] {Float.NaN}, 1)); + // A trailing NaN is the layout Arrays.sort produces and never trips a pairwise comparison, + // so it needs its own rejection path. Encoded, NaN sorts above positive infinity: letting one + // into a full heap would publish a floor no real score can beat. + expectThrows(IllegalArgumentException.class, () -> floor.offer(new float[] {1f, Float.NaN}, 2)); + expectThrows( + IllegalArgumentException.class, () -> floor.offer(new float[] {1f, 2f, Float.NaN}, 3)); + } + + public void testConcurrentPublishersConvergeToKthBest() throws Exception { + int k = 64; + int numThreads = TestUtil.nextInt(random(), 2, 8); + int batchesPerThread = 50; + int maxBatchSize = 16; + GlobalKnnFloor floor = new GlobalKnnFloor(k); + + // Pre-generate every thread's scores on the test thread so the expected k-th best can be + // computed independently of scheduling. + float[][][] batches = new float[numThreads][batchesPerThread][]; + List allScores = new ArrayList<>(); + for (int t = 0; t < numThreads; t++) { + for (int b = 0; b < batchesPerThread; b++) { + int batchSize = TestUtil.nextInt(random(), 1, maxBatchSize); + float[] batch = new float[batchSize]; + for (int i = 0; i < batchSize; i++) { + batch[i] = random().nextFloat(); + allScores.add(batch[i]); + } + Arrays.sort(batch); + batches[t][b] = batch; + } + } + allScores.sort(Comparator.reverseOrder()); + float expectedKthBest = allScores.get(k - 1); + + ExecutorService executor = + Executors.newFixedThreadPool(numThreads, new NamedThreadFactory("global-knn-floor-test")); + try { + CountDownLatch startingGun = new CountDownLatch(1); + List> futures = new ArrayList<>(numThreads); + for (int t = 0; t < numThreads; t++) { + final int thread = t; + futures.add( + executor.submit( + () -> { + try { + startingGun.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + float previous = Float.NEGATIVE_INFINITY; + for (float[] batch : batches[thread]) { + float current = floor.offer(batch, batch.length); + assertTrue( + "floor must appear monotonic to every publisher", current >= previous); + previous = current; + } + })); + } + startingGun.countDown(); + for (Future future : futures) { + future.get(); + } + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(30, TimeUnit.SECONDS)); + } + assertEquals( + "after all publishers finish, the floor must be the k-th best of every score offered", + expectedKthBest, + floor.floor(), + 0.0f); + } + + public void testConcurrentFillNeverOvershootsTrueKthBest() throws Exception { + // Regression tripwire for a heap-fill race: a publisher whose batch lands while the heap is + // not yet full must not treat its returned heap top as a valid bound, even if a concurrent + // publisher fills the heap in the window between the offer returning and the fullness flag + // being read. The scenario is adversarial by construction: one thread offers high scores that + // do not fill the heap on their own, the other offers low scores that complete the fill. The + // true k-th best is a low score, so a stale partial-heap top (a high score) published as the + // floor overshoots it — and because the floor is a monotonic maximum, a single overshoot + // survives to the final assertion. The race window is narrow, hence the many short rounds. + int rounds = 500; + int k = 32; + ExecutorService executor = + Executors.newFixedThreadPool(2, new NamedThreadFactory("global-knn-floor-race-test")); + try { + for (int round = 0; round < rounds; round++) { + GlobalKnnFloor floor = new GlobalKnnFloor(k); + float[] highs = new float[k - 1]; + for (int i = 0; i < highs.length; i++) { + highs[i] = 0.9f + i * 0.001f; + } + float[] lows = new float[k]; + for (int i = 0; i < lows.length; i++) { + lows[i] = 0.1f + i * 0.001f; + } + // Top-k of the union: all k-1 highs plus the single best low, whose score is the k-th + // best and the highest floor any correct execution can reach. + float trueKthBest = lows[lows.length - 1]; + + CountDownLatch startingGun = new CountDownLatch(1); + Future highPublisher = + executor.submit( + () -> { + try { + startingGun.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + floor.offer(highs, highs.length); + }); + Future lowPublisher = + executor.submit( + () -> { + try { + startingGun.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + floor.offer(lows, lows.length); + }); + startingGun.countDown(); + highPublisher.get(); + lowPublisher.get(); + + assertEquals( + "round " + round + ": the floor must be exactly the k-th best of everything offered", + trueKthBest, + floor.floor(), + 0.0f); + } + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(30, TimeUnit.SECONDS)); + } + } +} diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestSharedFloorKnnCollectorManager.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestSharedFloorKnnCollectorManager.java new file mode 100644 index 000000000000..8c0e0271f55e --- /dev/null +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestSharedFloorKnnCollectorManager.java @@ -0,0 +1,347 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.lucene.sandbox.search.knn; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.search.knn.KnnCollectorManager; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; + +/** + * Unit tests of {@link SharedFloorKnnCollectorManager}'s collector-creation policy: which gate its + * collectors get, which of them may publish into the shared floor, and how a shard's {@code + * globalShare} shapes both. The end-to-end behavior of the collectors themselves is covered by + * {@link TestSharedFloorKnnSearch}. + */ +public class TestSharedFloorKnnCollectorManager extends LuceneTestCase { + + private static final String FIELD = "vector"; + private static final VectorSimilarityFunction SIMILARITY = VectorSimilarityFunction.EUCLIDEAN; + + public void testPerShardGateMatchesOptimisticQuota() { + // The gate must equal the quota the optimistic multi-segment strategy would give a segment of + // the same proportion (perLeafTopKCalculation in AbstractKnnVectorQuery): k * share plus 16 + // standard deviations of the binomial, truncated. The two spot values below are the gates for + // a 1-of-16 shard at the two k values of the benchmark this parameter exists for. + assertEquals(1012, SharedFloorKnnCollectorManager.perShardGate(10000, 1 / 16.0)); + assertEquals(184, SharedFloorKnnCollectorManager.perShardGate(1000, 1 / 16.0)); + // A searcher holding the whole corpus is due all of k, with no padding to add. + assertEquals(10000, SharedFloorKnnCollectorManager.perShardGate(10000, 1.0)); + // The gate never exceeds k, however much statistical padding the formula would add. + assertEquals(100, SharedFloorKnnCollectorManager.perShardGate(100, 0.99)); + // The gate is at least 1, however small the share. + assertEquals(1, SharedFloorKnnCollectorManager.perShardGate(1, 1e-9)); + + expectThrows( + IllegalArgumentException.class, () -> SharedFloorKnnCollectorManager.perShardGate(0, 0.5)); + expectThrows( + IllegalArgumentException.class, + () -> SharedFloorKnnCollectorManager.perShardGate(100, 0.0)); + expectThrows( + IllegalArgumentException.class, + () -> SharedFloorKnnCollectorManager.perShardGate(100, 1.1)); + expectThrows( + IllegalArgumentException.class, + () -> SharedFloorKnnCollectorManager.perShardGate(100, Double.NaN)); + } + + public void testInvalidGlobalShare() { + int k = 1000; + GlobalKnnFloor floor = new GlobalKnnFloor(k); + for (float share : new float[] {0f, -0.5f, 1.1f, Float.NaN}) { + expectThrows( + IllegalArgumentException.class, + () -> new SharedFloorKnnCollectorManager(k, floor, 0.9f, 1, 16, 256, share)); + } + } + + public void testGlobalShareOpensTheGateAtTheExpectedShare() throws IOException { + // A single-segment index declared to hold 1/16 of the corpus: collectors must open their + // ascent gate at the shard's expected contribution to the merged top-k, not at the full local + // queue — a shard's own index looks like its entire corpus, so without globalShare the gate + // would sit at k and the floor could never end the fill. + int k = 1000; + try (Directory dir = newDirectory()) { + indexInSegments(dir, randomVectors(20, 8), 1); + try (DirectoryReader reader = DirectoryReader.open(dir)) { + LeafReaderContext context = reader.leaves().get(0); + + SharedFloorKnnCollectorManager sharded = + new SharedFloorKnnCollectorManager(k, new GlobalKnnFloor(k), 0.9f, 1, 16, 256, 1f / 16); + FloorAwareKnnCollector collector = + (FloorAwareKnnCollector) sharded.newCollector(Integer.MAX_VALUE, null, context); + assertEquals( + "the gate must open at the shard's expected share of the merged top-k", + SharedFloorKnnCollectorManager.perShardGate(k, 1.0 / 16), + collector.gateK()); + + SharedFloorKnnCollectorManager whole = + new SharedFloorKnnCollectorManager(k, new GlobalKnnFloor(k), 0.9f, 1, 16, 256, 1f); + FloorAwareKnnCollector wholeCollector = + (FloorAwareKnnCollector) whole.newCollector(Integer.MAX_VALUE, null, context); + assertEquals( + "an index holding the whole corpus must gate at the full queue", + k, + wholeCollector.gateK()); + } + } + } + + public void testMultiSegmentGateIsProRataPerLeaf() throws IOException { + // Full-k collectors over a multi-segment index (the optimistic strategy's second pass) must + // gate at the leaf's pro-rata share, not at k: the floor established during the first pass is + // otherwise barred from ending the second pass's re-fill until it has collected all of k. + int k = 1000; + int segments = 4; + try (Directory dir = newDirectory()) { + indexInSegments(dir, randomVectors(40, 8), segments); + try (DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals("test setup: segment count", segments, reader.leaves().size()); + SharedFloorKnnCollectorManager manager = + new SharedFloorKnnCollectorManager(k, new GlobalKnnFloor(k), 0.9f, 1, 16, 256); + int totalDocs = reader.maxDoc(); + for (LeafReaderContext context : reader.leaves()) { + FloorAwareKnnCollector collector = + (FloorAwareKnnCollector) manager.newCollector(Integer.MAX_VALUE, null, context); + double leafShare = context.reader().maxDoc() / (double) totalDocs; + assertEquals( + "leaf " + context.ord + " must gate at its pro-rata share", + SharedFloorKnnCollectorManager.perShardGate(k, leafShare), + collector.gateK()); + assertTrue(collector.gateK() < k); + } + } + } + } + + public void testEachLeafPublishesIntoTheFloorAtMostOnce() throws IOException { + // The optimistic strategy searches a competitive leaf twice, re-collecting the same + // documents. Republishing them would put duplicate scores in the floor's heap, and a k-th + // best over a multiset with duplicates can exceed the true merged cutoff — so the second + // collector for the same leaf must read the floor without feeding it. + int k = 1000; + try (Directory dir = newDirectory()) { + indexInSegments(dir, randomVectors(30, 8), 2); + try (DirectoryReader reader = DirectoryReader.open(dir)) { + SharedFloorKnnCollectorManager manager = + new SharedFloorKnnCollectorManager(k, new GlobalKnnFloor(k), 0.9f, 1, 16, 256); + LeafReaderContext first = reader.leaves().get(0); + LeafReaderContext second = reader.leaves().get(1); + + FloorAwareKnnCollector pass1 = + (FloorAwareKnnCollector) + manager.newOptimisticCollector(Integer.MAX_VALUE, null, first, 100); + assertTrue("the first search of a leaf must publish", pass1.publishesToFloor()); + + FloorAwareKnnCollector pass2 = + (FloorAwareKnnCollector) manager.newCollector(Integer.MAX_VALUE, null, first); + assertFalse("a second search of the same leaf must not publish", pass2.publishesToFloor()); + + FloorAwareKnnCollector otherLeaf = + (FloorAwareKnnCollector) manager.newCollector(Integer.MAX_VALUE, null, second); + assertTrue( + "the first search of a different leaf must publish", otherLeaf.publishesToFloor()); + } + } + } + + /** + * The distributed wiring this manager exists for, in miniature: two disjoint indexes standing in + * for two shards of one corpus, each searched through its own manager with its true {@code + * globalShare}, both wired to one shared floor. The merged result must not lose recall against + * stock search of the same two indexes. + */ + public void testDisjointIndexesSharingOneFloorKeepMergedRecall() throws IOException { + int dim = 16; + int docsPerShard = 600; + int k = 64; + float[][] shardA = randomVectors(docsPerShard, dim); + float[][] shardB = randomVectors(docsPerShard, dim); + + try (Directory dirA = newDirectory(); + Directory dirB = newDirectory()) { + indexInSegments(dirA, shardA, 1); + indexInSegments(dirB, shardB, 1); + try (DirectoryReader readerA = DirectoryReader.open(dirA); + DirectoryReader readerB = DirectoryReader.open(dirB)) { + IndexSearcher searcherA = new IndexSearcher(readerA); + IndexSearcher searcherB = new IndexSearcher(readerB); + + double stockRecallSum = 0; + double flooredRecallSum = 0; + int queries = 10; + for (int i = 0; i < queries; i++) { + float[] query = randomVector(dim); + Set truth = mergedExactTopK(readerA, readerB, query, k); + + TopDocs stockA = searcherA.search(new KnnFloatVectorQuery(FIELD, query, k), k); + TopDocs stockB = searcherB.search(new KnnFloatVectorQuery(FIELD, query, k), k); + stockRecallSum += mergedRecall(stockA, stockB, docsPerShard, truth, k); + + // One floor for the query; each shard's manager knows its share of the whole corpus. + GlobalKnnFloor floor = new GlobalKnnFloor(k); + TopDocs flooredA = + searcherA.search(new SharedFloorShardQuery(FIELD, query, k, floor, 0.5f), k); + TopDocs flooredB = + searcherB.search(new SharedFloorShardQuery(FIELD, query, k, floor, 0.5f), k); + flooredRecallSum += mergedRecall(flooredA, flooredB, docsPerShard, truth, k); + } + double stockRecall = stockRecallSum / queries; + double flooredRecall = flooredRecallSum / queries; + assertTrue( + "cross-shard floor sharing lost recall: stock=" + + stockRecall + + " floored=" + + flooredRecall, + flooredRecall >= stockRecall - 0.05); + } + } + } + + /** + * A {@link KnnFloatVectorQuery} searching one shard of a two-shard corpus: its manager is wired + * to the query's shared floor and declares this shard's share of the whole corpus. A query + * instance carries single-execution state and must not be reused. + */ + private static class SharedFloorShardQuery extends KnnFloatVectorQuery { + private final SharedFloorKnnCollectorManager manager; + + SharedFloorShardQuery(String field, float[] target, int k, GlobalKnnFloor floor, float share) { + super(field, target, k); + this.manager = new SharedFloorKnnCollectorManager(k, floor, 0.9f, 1, 16, 256, share); + } + + @Override + protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) { + return manager; + } + } + + private float[][] randomVectors(int count, int dim) { + float[][] vectors = new float[count][]; + for (int i = 0; i < count; i++) { + vectors[i] = randomVector(dim); + } + return vectors; + } + + private float[] randomVector(int dim) { + float[] vector = new float[dim]; + for (int i = 0; i < dim; i++) { + vector[i] = (float) random().nextGaussian(); + } + return vector; + } + + private void indexInSegments(Directory dir, float[][] vectors, int segments) throws IOException { + try (IndexWriter writer = + new IndexWriter(dir, new IndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) { + int docsPerSegment = (vectors.length + segments - 1) / segments; + for (int i = 0; i < vectors.length; i++) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField(FIELD, vectors[i], SIMILARITY)); + writer.addDocument(doc); + if ((i + 1) % docsPerSegment == 0) { + writer.flush(); + } + } + writer.commit(); + } + } + + private record DocScore(int globalDoc, float score) {} + + /** + * Exact top-k over both shards, in a merged id space where shard A's docs keep their ids and + * shard B's are offset by A's shard size. + */ + private Set mergedExactTopK( + IndexReader readerA, IndexReader readerB, float[] query, int k) throws IOException { + List scored = new ArrayList<>(); + collectExact(readerA, query, 0, scored); + collectExact(readerB, query, readerA.maxDoc(), scored); + scored.sort( + Comparator.comparingDouble(DocScore::score) + .reversed() + .thenComparingInt(DocScore::globalDoc)); + Set topK = new HashSet<>(); + for (int i = 0; i < k && i < scored.size(); i++) { + topK.add(scored.get(i).globalDoc()); + } + return topK; + } + + private void collectExact(IndexReader reader, float[] query, int offset, List out) + throws IOException { + for (LeafReaderContext ctx : reader.leaves()) { + FloatVectorValues values = ctx.reader().getFloatVectorValues(FIELD); + assertNotNull(values); + KnnVectorValues.DocIndexIterator iterator = values.iterator(); + for (int doc = iterator.nextDoc(); + doc != DocIdSetIterator.NO_MORE_DOCS; + doc = iterator.nextDoc()) { + float score = SIMILARITY.compare(query, values.vectorValue(iterator.index())); + out.add(new DocScore(offset + ctx.docBase + doc, score)); + } + } + } + + /** Merge the two shards' results by score, keep the top k, and score recall against truth. */ + private static double mergedRecall( + TopDocs shardA, TopDocs shardB, int shardBOffset, Set truth, int k) { + List merged = new ArrayList<>(); + for (ScoreDoc scoreDoc : shardA.scoreDocs) { + merged.add(new DocScore(scoreDoc.doc, scoreDoc.score)); + } + for (ScoreDoc scoreDoc : shardB.scoreDocs) { + merged.add(new DocScore(shardBOffset + scoreDoc.doc, scoreDoc.score)); + } + merged.sort( + Comparator.comparingDouble(DocScore::score) + .reversed() + .thenComparingInt(DocScore::globalDoc)); + int hits = 0; + for (int i = 0; i < k && i < merged.size(); i++) { + if (truth.contains(merged.get(i).globalDoc())) { + hits++; + } + } + return hits / (double) k; + } +} diff --git a/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestSharedFloorKnnSearch.java b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestSharedFloorKnnSearch.java new file mode 100644 index 000000000000..f5275849fe45 --- /dev/null +++ b/lucene/sandbox/src/test/org/apache/lucene/sandbox/search/knn/TestSharedFloorKnnSearch.java @@ -0,0 +1,538 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.lucene.sandbox.search.knn; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.search.knn.KnnCollectorManager; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.NamedThreadFactory; + +/** + * End-to-end tests of kNN search through {@link SharedFloorKnnCollectorManager}: the shared floor + * must reduce work without changing what stock search would have found. + * + *

Queries here activate the floor at every k (see {@link SharedFloorKnnQuery}), because the + * point is to exercise the floor path; the activation default is covered by its own test. The k + * values exceed {@link FloorAwareKnnCollector#DEFAULT_MIN_EXPLORATION_SLOTS}, since at or below it + * the clamp neutralizes the floor by design and the tests would not be testing anything. + */ +public class TestSharedFloorKnnSearch extends LuceneTestCase { + + private static final String FIELD = "vector"; + private static final VectorSimilarityFunction SIMILARITY = VectorSimilarityFunction.EUCLIDEAN; + + /** + * The floor must not cost recall relative to stock search, regardless of how the index is cut + * into segments. + */ + public void testRecallParityAcrossSegmentCounts() throws IOException { + int dim = 16; + int numDocs = 1200; + int k = 64; + int numQueries = 10; + for (int segments : new int[] {1, 2, 5}) { + float[][] vectors = new float[numDocs][]; + for (int i = 0; i < numDocs; i++) { + vectors[i] = randomVector(dim); + } + try (Directory dir = newDirectory()) { + indexInSegments(dir, vectors, segments); + try (DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals(segments, reader.leaves().size()); + IndexSearcher searcher = new IndexSearcher(reader); + double stockRecallSum = 0; + double flooredRecallSum = 0; + for (int i = 0; i < numQueries; i++) { + float[] query = randomVector(dim); + Set truth = exactTopK(reader, query, k); + stockRecallSum += + recall(searcher.search(new KnnFloatVectorQuery(FIELD, query, k), k), truth, k); + flooredRecallSum += + recall(searcher.search(new SharedFloorKnnQuery(FIELD, query, k), k), truth, k); + } + double stockRecall = stockRecallSum / numQueries; + double flooredRecall = flooredRecallSum / numQueries; + assertTrue( + "shared-floor search lost recall at " + + segments + + " segments: stock=" + + stockRecall + + " floored=" + + flooredRecall, + flooredRecall >= stockRecall - 0.05); + } + } + } + } + + /** + * With no executor, the leaves of a query are searched one after another, and a segment that + * converges early establishes a floor before later segments have collected anything. This is the + * harshest ordering for a shared bound, and the ascent gate is what makes it safe. The scenario + * is made adversarial: the first segment is a tight cluster of mediocre near-duplicates that + * converges quickly to a high local cutoff, while every true neighbor lives in later segments, + * still unsearched when that cutoff is published. + */ + public void testDecoyFirstSegmentDoesNotStarveLaterSegments() throws IOException { + int dim = 16; + int k = 64; + int trueNeighborSegments = 4; + int trueNeighborsPerSegment = 16; + int backgroundPerSegment = 300; + int decoyCount = 400; + + float[] center = randomVector(dim); + float[] decoyDirection = randomUnitVector(dim); + + // Segment 0: the decoy cluster, at moderate distance from the query, internally very dense. + List decoySegment = new ArrayList<>(decoyCount); + for (int i = 0; i < decoyCount; i++) { + decoySegment.add(displaced(center, decoyDirection, 2f, 0.02f)); + } + + // Segments 1..4: a few true nearest neighbors each, hidden among distant background vectors. + List> laterSegments = new ArrayList<>(trueNeighborSegments); + for (int s = 0; s < trueNeighborSegments; s++) { + List segment = new ArrayList<>(trueNeighborsPerSegment + backgroundPerSegment); + for (int i = 0; i < trueNeighborsPerSegment; i++) { + segment.add( + displaced(center, randomUnitVector(dim), 0.9f + random().nextFloat() * 0.2f, 0f)); + } + for (int i = 0; i < backgroundPerSegment; i++) { + segment.add(displaced(center, randomUnitVector(dim), 5f + random().nextFloat(), 0f)); + } + laterSegments.add(segment); + } + + try (Directory dir = newDirectory()) { + try (IndexWriter writer = + new IndexWriter(dir, new IndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) { + addSegment(writer, decoySegment); + for (List segment : laterSegments) { + addSegment(writer, segment); + } + writer.commit(); + } + try (DirectoryReader reader = DirectoryReader.open(dir)) { + assertEquals(1 + trueNeighborSegments, reader.leaves().size()); + IndexSearcher searcher = new IndexSearcher(reader); + int numQueries = 5; + double stockRecallSum = 0; + double flooredRecallSum = 0; + for (int i = 0; i < numQueries; i++) { + float[] query = displaced(center, randomUnitVector(dim), 0.05f, 0f); + Set truth = exactTopK(reader, query, k); + stockRecallSum += + recall(searcher.search(new KnnFloatVectorQuery(FIELD, query, k), k), truth, k); + flooredRecallSum += + recall(searcher.search(new SharedFloorKnnQuery(FIELD, query, k), k), truth, k); + } + double stockRecall = stockRecallSum / numQueries; + double flooredRecall = flooredRecallSum / numQueries; + assertTrue( + "an early-converging decoy segment starved the segments holding the true neighbors: " + + "stock=" + + stockRecall + + " floored=" + + flooredRecall, + flooredRecall >= stockRecall - 0.05); + } + } + } + + /** + * Without an executor, execution is single-threaded and the floor evolves identically on every + * run, so two executions of the same search must return exactly the same documents and scores. + */ + public void testSequentialSearchIsDeterministic() throws IOException { + int dim = 16; + int numDocs = 1000; + int k = 32; + float[][] vectors = new float[numDocs][]; + for (int i = 0; i < numDocs; i++) { + vectors[i] = randomVector(dim); + } + try (Directory dir = newDirectory()) { + indexInSegments(dir, vectors, 4); + try (DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + for (int i = 0; i < 5; i++) { + float[] query = randomVector(dim); + // Fresh query objects: a manager and its floor carry single-execution state. + TopDocs first = searcher.search(new SharedFloorKnnQuery(FIELD, query, k), k); + TopDocs second = searcher.search(new SharedFloorKnnQuery(FIELD, query, k), k); + assertEquals(first.scoreDocs.length, second.scoreDocs.length); + for (int j = 0; j < first.scoreDocs.length; j++) { + assertEquals("doc at rank " + j, first.scoreDocs[j].doc, second.scoreDocs[j].doc); + assertEquals( + "score at rank " + j, first.scoreDocs[j].score, second.scoreDocs[j].score, 0.0f); + } + } + } + } + } + + /** + * Under an executor the leaves are searched concurrently and the floor's evolution depends on + * thread interleaving, which may legitimately vary visit counts; result quality must not suffer. + */ + public void testParallelRecallMatchesSequential() throws Exception { + int dim = 16; + int numDocs = 1200; + int k = 64; + int numQueries = 10; + float[][] vectors = new float[numDocs][]; + for (int i = 0; i < numDocs; i++) { + vectors[i] = randomVector(dim); + } + ExecutorService executor = + Executors.newFixedThreadPool(4, new NamedThreadFactory("shared-floor-knn-test")); + try (Directory dir = newDirectory()) { + indexInSegments(dir, vectors, 5); + try (DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher sequentialSearcher = new IndexSearcher(reader); + IndexSearcher parallelSearcher = new IndexSearcher(reader, executor); + double sequentialRecallSum = 0; + double parallelRecallSum = 0; + for (int i = 0; i < numQueries; i++) { + float[] query = randomVector(dim); + Set truth = exactTopK(reader, query, k); + sequentialRecallSum += + recall( + sequentialSearcher.search(new SharedFloorKnnQuery(FIELD, query, k), k), truth, k); + parallelRecallSum += + recall( + parallelSearcher.search(new SharedFloorKnnQuery(FIELD, query, k), k), truth, k); + } + double sequentialRecall = sequentialRecallSum / numQueries; + double parallelRecall = parallelRecallSum / numQueries; + assertTrue( + "parallel execution lost recall: sequential=" + + sequentialRecall + + " parallel=" + + parallelRecall, + parallelRecall >= sequentialRecall - 0.05); + } + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(30, TimeUnit.SECONDS)); + } + } + + /** + * A bound advertised from outside the searching process with realistic slack, here the query's + * exact similarity at global rank 4k, must not cost recall. Tightness, not validity, is what + * risks recall with an externally advertised bound: a bar close to the final cutoff, imposed + * before the graph search has discovered its neighborhood, can sever the paths through mediocre + * intermediate nodes that graph navigation depends on. An advertiser is therefore expected to + * leave rank headroom (for example, a scout advertising its k'-th best for k' several times k), + * and this test pins down that a bound with such headroom is harmless. + */ + public void testSlackAdvertisedBoundPreservesRecall() throws IOException { + int dim = 16; + int numDocs = 1200; + int k = 64; + int numQueries = 10; + float[][] vectors = new float[numDocs][]; + for (int i = 0; i < numDocs; i++) { + vectors[i] = randomVector(dim); + } + try (Directory dir = newDirectory()) { + indexInSegments(dir, vectors, 4); + try (DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + double unadvertisedRecallSum = 0; + double advertisedRecallSum = 0; + for (int i = 0; i < numQueries; i++) { + float[] query = randomVector(dim); + Set truth = exactTopK(reader, query, k); + unadvertisedRecallSum += + recall(searcher.search(new SharedFloorKnnQuery(FIELD, query, k), k), truth, k); + + SharedFloorKnnQuery advertisedQuery = new SharedFloorKnnQuery(FIELD, query, k); + advertisedQuery + .manager + .getGlobalFloor() + .advertise(exactKthBestScore(reader, query, 4 * k)); + advertisedRecallSum += recall(searcher.search(advertisedQuery, k), truth, k); + } + double unadvertisedRecall = unadvertisedRecallSum / numQueries; + double advertisedRecall = advertisedRecallSum / numQueries; + assertTrue( + "an advertised bound with rank headroom cost recall: unadvertised=" + + unadvertisedRecall + + " advertised=" + + advertisedRecall, + advertisedRecall >= unadvertisedRecall - 0.05); + } + } + } + + /** + * At {@code greediness = 0} the non-competitive queue is at least as large as the local queue, so + * the effective bound can never exceed what the local search would have imposed on itself: the + * shared floor is fully neutralized. Recall must then match stock search even under the tightest + * bound that exists, the query's exact final k-th best similarity, advertised before the search + * starts. This pins down the greediness dial's safe endpoint; the recall cost of tighter settings + * under tight bounds is a measured trade, not a correctness property. + */ + public void testZeroGreedinessNeutralizesTightestBound() throws IOException { + int dim = 16; + int numDocs = 1200; + int k = 64; + int numQueries = 10; + float[][] vectors = new float[numDocs][]; + for (int i = 0; i < numDocs; i++) { + vectors[i] = randomVector(dim); + } + try (Directory dir = newDirectory()) { + indexInSegments(dir, vectors, 4); + try (DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + double stockRecallSum = 0; + double advertisedRecallSum = 0; + for (int i = 0; i < numQueries; i++) { + float[] query = randomVector(dim); + Set truth = exactTopK(reader, query, k); + stockRecallSum += + recall(searcher.search(new KnnFloatVectorQuery(FIELD, query, k), k), truth, k); + + SharedFloorKnnQuery advertisedQuery = new SharedFloorKnnQuery(FIELD, query, k, 0f); + advertisedQuery.manager.getGlobalFloor().advertise(exactKthBestScore(reader, query, k)); + advertisedRecallSum += recall(searcher.search(advertisedQuery, k), truth, k); + } + double stockRecall = stockRecallSum / numQueries; + double advertisedRecall = advertisedRecallSum / numQueries; + assertTrue( + "greediness 0 must neutralize even the tightest advertised bound: stock=" + + stockRecall + + " advertised=" + + advertisedRecall, + advertisedRecall >= stockRecall - 0.05); + } + } + } + + /** + * Below the activation threshold the manager creates plain collectors, so the search must be + * stock search to the last bit: identical documents and scores, even when a hostile (invalid) + * bound has been advertised. This is the policy layer that keeps small-k queries, where a floor + * has little to save and the most recall to lose, entirely out of the mechanism. + */ + public void testBelowActivationThresholdSearchIsExactlyStock() throws IOException { + int dim = 16; + int numDocs = 1000; + int k = 10; + float[][] vectors = new float[numDocs][]; + for (int i = 0; i < numDocs; i++) { + vectors[i] = randomVector(dim); + } + assertTrue( + "this test requires k below the default activation threshold", + k < SharedFloorKnnCollectorManager.DEFAULT_FLOOR_ACTIVATION_K); + try (Directory dir = newDirectory()) { + indexInSegments(dir, vectors, 4); + try (DirectoryReader reader = DirectoryReader.open(dir)) { + IndexSearcher searcher = new IndexSearcher(reader); + for (int i = 0; i < 5; i++) { + float[] query = randomVector(dim); + TopDocs stock = searcher.search(new KnnFloatVectorQuery(FIELD, query, k), k); + + SharedFloorKnnQuery flooredQuery = + new SharedFloorKnnQuery( + FIELD, + query, + k, + FloorAwareKnnCollector.DEFAULT_GREEDINESS, + SharedFloorKnnCollectorManager.DEFAULT_FLOOR_ACTIVATION_K); + // Deliberately invalid: far above any real similarity. Below the activation threshold + // it must not matter, because no collector ever consults the floor. + flooredQuery.manager.getGlobalFloor().advertise(Float.MAX_VALUE); + TopDocs floored = searcher.search(flooredQuery, k); + + assertEquals(stock.scoreDocs.length, floored.scoreDocs.length); + for (int j = 0; j < stock.scoreDocs.length; j++) { + assertEquals("doc at rank " + j, stock.scoreDocs[j].doc, floored.scoreDocs[j].doc); + assertEquals( + "score at rank " + j, stock.scoreDocs[j].score, floored.scoreDocs[j].score, 0.0f); + } + } + } + } + } + + /** + * A {@link KnnFloatVectorQuery} routed through a {@link SharedFloorKnnCollectorManager}. The + * manager is created with the query and returned for both collection passes, so both share one + * floor; consequently a query instance carries single-execution state and must not be reused. + * Unless a threshold is given, the floor activates at every k, because these tests exist to + * exercise the floor path. + */ + private static class SharedFloorKnnQuery extends KnnFloatVectorQuery { + final SharedFloorKnnCollectorManager manager; + + SharedFloorKnnQuery(String field, float[] target, int k) { + this(field, target, k, FloorAwareKnnCollector.DEFAULT_GREEDINESS); + } + + SharedFloorKnnQuery(String field, float[] target, int k, float greediness) { + this(field, target, k, greediness, 1); + } + + SharedFloorKnnQuery( + String field, float[] target, int k, float greediness, int floorActivationK) { + super(field, target, k); + this.manager = + new SharedFloorKnnCollectorManager( + k, new GlobalKnnFloor(k), greediness, floorActivationK); + } + + @Override + protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) { + return manager; + } + } + + private float[] randomVector(int dim) { + float[] vector = new float[dim]; + for (int i = 0; i < dim; i++) { + vector[i] = (float) random().nextGaussian(); + } + return vector; + } + + private float[] randomUnitVector(int dim) { + float[] vector = randomVector(dim); + double norm = 0; + for (float component : vector) { + norm += component * component; + } + norm = Math.sqrt(norm); + for (int i = 0; i < dim; i++) { + vector[i] /= (float) norm; + } + return vector; + } + + /** + * Return {@code center + distance * direction + jitter}, where the jitter is a Gaussian + * perturbation of the given magnitude in each dimension. + */ + private float[] displaced(float[] center, float[] direction, float distance, float jitter) { + float[] vector = new float[center.length]; + for (int i = 0; i < center.length; i++) { + vector[i] = center[i] + distance * direction[i] + jitter * (float) random().nextGaussian(); + } + return vector; + } + + private void indexInSegments(Directory dir, float[][] vectors, int segments) throws IOException { + try (IndexWriter writer = + new IndexWriter(dir, new IndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) { + int docsPerSegment = (vectors.length + segments - 1) / segments; + int written = 0; + while (written < vectors.length) { + List segment = new ArrayList<>(docsPerSegment); + for (int i = 0; i < docsPerSegment && written < vectors.length; i++, written++) { + segment.add(vectors[written]); + } + addSegment(writer, segment); + } + writer.commit(); + } + } + + private void addSegment(IndexWriter writer, List vectors) throws IOException { + for (float[] vector : vectors) { + Document doc = new Document(); + doc.add(new KnnFloatVectorField(FIELD, vector, SIMILARITY)); + writer.addDocument(doc); + } + writer.flush(); + } + + private record DocScore(int doc, float score) {} + + private List exactSearch(IndexReader reader, float[] query) throws IOException { + List scored = new ArrayList<>(); + for (LeafReaderContext ctx : reader.leaves()) { + FloatVectorValues values = ctx.reader().getFloatVectorValues(FIELD); + assertNotNull(values); + KnnVectorValues.DocIndexIterator iterator = values.iterator(); + for (int doc = iterator.nextDoc(); + doc != DocIdSetIterator.NO_MORE_DOCS; + doc = iterator.nextDoc()) { + float score = SIMILARITY.compare(query, values.vectorValue(iterator.index())); + scored.add(new DocScore(ctx.docBase + doc, score)); + } + } + scored.sort( + Comparator.comparingDouble(DocScore::score).reversed().thenComparingInt(DocScore::doc)); + return scored; + } + + private Set exactTopK(IndexReader reader, float[] query, int k) throws IOException { + List scored = exactSearch(reader, query); + Set topK = new HashSet<>(); + for (int i = 0; i < k && i < scored.size(); i++) { + topK.add(scored.get(i).doc()); + } + return topK; + } + + private float exactKthBestScore(IndexReader reader, float[] query, int k) throws IOException { + List scored = exactSearch(reader, query); + assertTrue(scored.size() >= k); + return scored.get(k - 1).score(); + } + + private static double recall(TopDocs topDocs, Set truth, int k) { + int hits = 0; + for (ScoreDoc scoreDoc : topDocs.scoreDocs) { + if (truth.contains(scoreDoc.doc)) { + hits++; + } + } + return hits / (double) k; + } +}