diff --git a/lucene/core/src/java/org/apache/lucene/search/knn/MultiLeafKnnCollector.java b/lucene/core/src/java/org/apache/lucene/search/knn/MultiLeafKnnCollector.java index 1ca979d6794..051cd9ed633 100644 --- a/lucene/core/src/java/org/apache/lucene/search/knn/MultiLeafKnnCollector.java +++ b/lucene/core/src/java/org/apache/lucene/search/knn/MultiLeafKnnCollector.java @@ -33,17 +33,17 @@ public final class MultiLeafKnnCollector implements KnnCollector { // greediness of globally non-competitive search: (0,1] private static final float DEFAULT_GREEDINESS = 0.9f; + private static final int DEFAULT_INTERVAL = 0xff; // the global queue of the highest similarities collected so far across all segments private final BlockingFloatHeap globalSimilarityQueue; // the local queue of the highest similarities if we are not competitive globally // the size of this queue is defined by greediness private final FloatHeap nonCompetitiveQueue; - private final float greediness; // the queue of the local similarities to periodically update with the global queue private final FloatHeap updatesQueue; private final float[] updatesScratch; // interval to synchronize the local and global queues, as a number of visited vectors - private final int interval = 0xff; // 255 + private final int interval; private boolean kResultsCollected = false; private float cachedGlobalMinSim = Float.NEGATIVE_INFINITY; private final AbstractKnnCollector subCollector; @@ -58,7 +58,32 @@ public final class MultiLeafKnnCollector implements KnnCollector { */ public MultiLeafKnnCollector( int k, BlockingFloatHeap globalSimilarityQueue, AbstractKnnCollector subCollector) { - this.greediness = DEFAULT_GREEDINESS; + this(k, DEFAULT_GREEDINESS, DEFAULT_INTERVAL, globalSimilarityQueue, subCollector); + } + + /** + * Create a new MultiLeafKnnCollector. + * + * @param k the number of neighbors to collect + * @param greediness the greediness of the global search + * @param interval (by number of collected values) the interval to synchronize the local and + * global queues + * @param globalSimilarityQueue the global queue of the highest similarities collected so far + * @param subCollector the local collector + */ + public MultiLeafKnnCollector( + int k, + float greediness, + int interval, + BlockingFloatHeap globalSimilarityQueue, + AbstractKnnCollector subCollector) { + if (greediness < 0 || greediness > 1) { + throw new IllegalArgumentException("greediness must be in [0,1]"); + } + if (interval <= 0) { + throw new IllegalArgumentException("interval must be positive"); + } + this.interval = interval; this.subCollector = subCollector; this.globalSimilarityQueue = globalSimilarityQueue; this.nonCompetitiveQueue = new FloatHeap(Math.max(1, Math.round((1 - greediness) * k)));