Skip to content

Commit ad8e490

Browse files
benwtrentbenwtrent
authored
DiskBBQ track global competitiveness of scores between segments (#133005)
This adds global minimum competitive score tracking to disk bbq queries. After every posting list, the collector will sync its competitive score with the global tracker. This allows for other segment collectors to see competitive scores within their own search thread. Right now, this doesn't do much, however, its required if we are to do any early exiting or skipping logic for posting list scoring. When testing this, there is no recall changes, and minimal performance gains (we skip collecting more docs than usual, but right now we still score them...which is the dominating cost).
1 parent 4855b4a commit ad8e490

12 files changed

+313
-37
lines changed

‎server/src/main/java/org/elasticsearch/index/codec/vectors/IVFVectorsReader.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,9 @@ public final void search(String field, float[] target, KnnCollector knnCollector
269269
// is enough?
270270
expectedDocs += scorer.resetPostingsScorer(offsetAndLength.offset());
271271
actualDocs += scorer.visit(knnCollector);
272+
if (knnCollector.getSearchStrategy() != null) {
273+
knnCollector.getSearchStrategy().nextVectorsBlock();
274+
}
272275
}
273276
if (acceptDocs != null) {
274277
float unfilteredRatioVisited = (float) expectedDocs / numVectors;
@@ -278,6 +281,9 @@ public final void search(String field, float[] target, KnnCollector knnCollector
278281
CentroidOffsetAndLength offsetAndLength = centroidPrefetchingIterator.nextPostingListOffsetAndLength();
279282
scorer.resetPostingsScorer(offsetAndLength.offset());
280283
actualDocs += scorer.visit(knnCollector);
284+
if (knnCollector.getSearchStrategy() != null) {
285+
knnCollector.getSearchStrategy().nextVectorsBlock();
286+
}
281287
}
282288
}
283289
}

‎server/src/main/java/org/elasticsearch/index/codec/vectors/cluster/NeighborQueue.java

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
/**
2626
* Copied from and modified from Apache Lucene.
2727
*/
28-
class NeighborQueue {
28+
public class NeighborQueue {
2929

3030
private enum Order {
3131
MIN_HEAP {
@@ -49,7 +49,7 @@ long apply(long v) {
4949
private final LongHeap heap;
5050
private final Order order;
5151

52-
NeighborQueue(int initialSize, boolean maxHeap) {
52+
public NeighborQueue(int initialSize, boolean maxHeap) {
5353
this.heap = new LongHeap(initialSize);
5454
this.order = maxHeap ? Order.MAX_HEAP : Order.MIN_HEAP;
5555
}
@@ -61,6 +61,10 @@ public int size() {
6161
return heap.size();
6262
}
6363

64+
public long peek() {
65+
return heap.top();
66+
}
67+
6468
/**
6569
* Adds a new graph arc, extending the storage as needed.
6670
*
@@ -84,31 +88,43 @@ public boolean insertWithOverflow(int newNode, float newScore) {
8488
return heap.insertWithOverflow(encode(newNode, newScore));
8589
}
8690

91+
public boolean insertWithOverflow(long encoded) {
92+
return heap.insertWithOverflow(encoded);
93+
}
94+
8795
/**
8896
* Encodes the node ID and its similarity score as long, preserving the Lucene tie-breaking rule
8997
* that when two scores are equal, the smaller node ID must win.
9098
* @param node the node ID
9199
* @param score the node score
92100
* @return the encoded score, node ID
93101
*/
94-
private long encode(int node, float score) {
95-
return order.apply((((long) NumericUtils.floatToSortableInt(score)) << 32) | (0xFFFFFFFFL & ~node));
102+
public long encode(int node, float score) {
103+
return order.apply(encodeRaw(node, score));
96104
}
97105

98106
/** Returns the top element's node id. */
99-
int topNode() {
107+
public int topNode() {
100108
return decodeNodeId(heap.top());
101109
}
102110

111+
public static long encodeRaw(int node, float score) {
112+
return (((long) NumericUtils.floatToSortableInt(score)) << 32) | (0xFFFFFFFFL & ~node);
113+
}
114+
115+
public static float decodeScoreRaw(long heapValue) {
116+
return NumericUtils.sortableIntToFloat((int) (heapValue >> 32));
117+
}
118+
103119
/**
104120
* Returns the top element's node score. For the min heap this is the minimum score. For the max
105121
* heap this is the maximum score.
106122
*/
107-
float topScore() {
123+
public float topScore() {
108124
return decodeScore(heap.top());
109125
}
110126

111-
private float decodeScore(long heapValue) {
127+
public float decodeScore(long heapValue) {
112128
return NumericUtils.sortableIntToFloat((int) (order.apply(heapValue) >> 32));
113129
}
114130

‎server/src/main/java/org/elasticsearch/search/vectors/AbstractIVFKnnVectorQuery.java

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import org.apache.lucene.search.FieldExistsQuery;
2222
import org.apache.lucene.search.FilteredDocIdSetIterator;
2323
import org.apache.lucene.search.IndexSearcher;
24-
import org.apache.lucene.search.KnnCollector;
2524
import org.apache.lucene.search.MatchNoDocsQuery;
2625
import org.apache.lucene.search.Query;
2726
import org.apache.lucene.search.QueryVisitor;
@@ -31,7 +30,6 @@
3130
import org.apache.lucene.search.TaskExecutor;
3231
import org.apache.lucene.search.TopDocs;
3332
import org.apache.lucene.search.TopDocsCollector;
34-
import org.apache.lucene.search.TopKnnCollector;
3533
import org.apache.lucene.search.Weight;
3634
import org.apache.lucene.search.knn.KnnCollectorManager;
3735
import org.apache.lucene.search.knn.KnnSearchStrategy;
@@ -45,6 +43,9 @@
4543
import java.util.List;
4644
import java.util.Objects;
4745
import java.util.concurrent.Callable;
46+
import java.util.concurrent.atomic.LongAccumulator;
47+
48+
import static org.elasticsearch.search.vectors.AbstractMaxScoreKnnCollector.LEAST_COMPETITIVE;
4849

4950
abstract class AbstractIVFKnnVectorQuery extends Query implements QueryProfilerProvider {
5051

@@ -120,7 +121,7 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
120121
// we need to ensure we are getting at least 2*k results to ensure we cover overspill duplicates
121122
// TODO move the logic for automatically adjusting percentages to the query, so we can only pass
122123
// 2k to the collector.
123-
KnnCollectorManager knnCollectorManager = getKnnCollectorManager(Math.round(2f * k), indexSearcher);
124+
IVFCollectorManager knnCollectorManager = getKnnCollectorManager(Math.round(2f * k), indexSearcher);
124125
TaskExecutor taskExecutor = indexSearcher.getTaskExecutor();
125126
List<LeafReaderContext> leafReaderContexts = reader.leaves();
126127

@@ -160,7 +161,7 @@ public Query rewrite(IndexSearcher indexSearcher) throws IOException {
160161
return new KnnScoreDocQuery(topK.scoreDocs, reader);
161162
}
162163

163-
private TopDocs searchLeaf(LeafReaderContext ctx, Weight filterWeight, KnnCollectorManager knnCollectorManager, float visitRatio)
164+
private TopDocs searchLeaf(LeafReaderContext ctx, Weight filterWeight, IVFCollectorManager knnCollectorManager, float visitRatio)
164165
throws IOException {
165166
TopDocs results = getLeafResults(ctx, filterWeight, knnCollectorManager, visitRatio);
166167
IntHashSet dedup = new IntHashSet(results.scoreDocs.length * 4 / 3);
@@ -182,7 +183,7 @@ private TopDocs searchLeaf(LeafReaderContext ctx, Weight filterWeight, KnnCollec
182183
return new TopDocs(results.totalHits, deduplicatedScoreDocs);
183184
}
184185

185-
TopDocs getLeafResults(LeafReaderContext ctx, Weight filterWeight, KnnCollectorManager knnCollectorManager, float visitRatio)
186+
TopDocs getLeafResults(LeafReaderContext ctx, Weight filterWeight, IVFCollectorManager knnCollectorManager, float visitRatio)
186187
throws IOException {
187188
final LeafReader reader = ctx.reader();
188189
final Bits liveDocs = reader.getLiveDocs();
@@ -205,12 +206,12 @@ abstract TopDocs approximateSearch(
205206
LeafReaderContext context,
206207
Bits acceptDocs,
207208
int visitedLimit,
208-
KnnCollectorManager knnCollectorManager,
209+
IVFCollectorManager knnCollectorManager,
209210
float visitRatio
210211
) throws IOException;
211212

212-
protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
213-
return new IVFCollectorManager(k);
213+
protected IVFCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
214+
return new IVFCollectorManager(k, searcher);
214215
}
215216

216217
@Override
@@ -236,14 +237,17 @@ protected boolean match(int doc) {
236237

237238
static class IVFCollectorManager implements KnnCollectorManager {
238239
private final int k;
240+
final LongAccumulator longAccumulator;
239241

240-
IVFCollectorManager(int k) {
242+
IVFCollectorManager(int k, IndexSearcher searcher) {
241243
this.k = k;
244+
longAccumulator = searcher.getIndexReader().leaves().size() > 1 ? new LongAccumulator(Long::max, LEAST_COMPETITIVE) : null;
242245
}
243246

244247
@Override
245-
public KnnCollector newCollector(int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context) throws IOException {
246-
return new TopKnnCollector(k, visitedLimit, searchStrategy);
248+
public AbstractMaxScoreKnnCollector newCollector(int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
249+
throws IOException {
250+
return new MaxScoreTopKnnCollector(k, visitedLimit, searchStrategy);
247251
}
248252
}
249253
}

‎server/src/main/java/org/elasticsearch/search/vectors/AbstractMaxScoreKnnCollector.java

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.search.vectors;
11+
12+
import org.apache.lucene.search.AbstractKnnCollector;
13+
import org.apache.lucene.search.knn.KnnSearchStrategy;
14+
import org.elasticsearch.index.codec.vectors.cluster.NeighborQueue;
15+
16+
/**
17+
* Abstract class for collectors that maintain a maximum score for KNN search.
18+
* It extends the {@link AbstractKnnCollector} and provides methods to manage
19+
* the minimum competitive document score, useful for tracking competitive scores
20+
* over multiple leaves.
21+
*/
22+
abstract class AbstractMaxScoreKnnCollector extends AbstractKnnCollector {
23+
public static final long LEAST_COMPETITIVE = NeighborQueue.encodeRaw(Integer.MAX_VALUE, Float.NEGATIVE_INFINITY);
24+
25+
protected AbstractMaxScoreKnnCollector(int k, long visitLimit, KnnSearchStrategy searchStrategy) {
26+
super(k, visitLimit, searchStrategy);
27+
}
28+
29+
/**
30+
* Returns the minimum competitive document score.
31+
* This is used to determine the global competitiveness of documents in the search.
32+
* This may be a competitive score even if the collector hasn't collected k results yet.
33+
*
34+
* @return the minimum competitive document score
35+
*/
36+
public abstract long getMinCompetitiveDocScore();
37+
38+
/**
39+
* Updates the minimum competitive document score.
40+
*
41+
* @param minCompetitiveDocScore the new minimum competitive document score to set
42+
*/
43+
abstract void updateMinCompetitiveDocScore(long minCompetitiveDocScore);
44+
}

‎server/src/main/java/org/elasticsearch/search/vectors/DiversifiedIVFKnnCollectorManager.java

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,25 +10,26 @@
1010
package org.elasticsearch.search.vectors;
1111

1212
import org.apache.lucene.index.LeafReaderContext;
13-
import org.apache.lucene.search.KnnCollector;
13+
import org.apache.lucene.search.IndexSearcher;
1414
import org.apache.lucene.search.join.BitSetProducer;
15-
import org.apache.lucene.search.knn.KnnCollectorManager;
1615
import org.apache.lucene.search.knn.KnnSearchStrategy;
1716
import org.apache.lucene.util.BitSet;
1817

1918
import java.io.IOException;
2019

21-
public class DiversifiedIVFKnnCollectorManager implements KnnCollectorManager {
20+
public class DiversifiedIVFKnnCollectorManager extends AbstractIVFKnnVectorQuery.IVFCollectorManager {
2221
private final int k;
2322
private final BitSetProducer parentsFilter;
2423

25-
DiversifiedIVFKnnCollectorManager(int k, BitSetProducer parentsFilter) {
24+
DiversifiedIVFKnnCollectorManager(int k, IndexSearcher searcher, BitSetProducer parentsFilter) {
25+
super(k, searcher);
2626
this.k = k;
2727
this.parentsFilter = parentsFilter;
2828
}
2929

3030
@Override
31-
public KnnCollector newCollector(int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context) throws IOException {
31+
public AbstractMaxScoreKnnCollector newCollector(int visitedLimit, KnnSearchStrategy searchStrategy, LeafReaderContext context)
32+
throws IOException {
3233
BitSet parentBitSet = parentsFilter.getBitSet(context);
3334
if (parentBitSet == null) {
3435
return null;

‎server/src/main/java/org/elasticsearch/search/vectors/DiversifyingChildrenIVFKnnFloatVectorQuery.java

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import org.apache.lucene.search.IndexSearcher;
1313
import org.apache.lucene.search.Query;
1414
import org.apache.lucene.search.join.BitSetProducer;
15-
import org.apache.lucene.search.knn.KnnCollectorManager;
1615

1716
import java.util.Objects;
1817

@@ -45,8 +44,8 @@ public DiversifyingChildrenIVFKnnFloatVectorQuery(
4544
}
4645

4746
@Override
48-
protected KnnCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
49-
return new DiversifiedIVFKnnCollectorManager(k, parentsFilter);
47+
protected IVFCollectorManager getKnnCollectorManager(int k, IndexSearcher searcher) {
48+
return new DiversifiedIVFKnnCollectorManager(k, searcher, parentsFilter);
5049
}
5150

5251
@Override

‎server/src/main/java/org/elasticsearch/search/vectors/DiversifyingNearestChildrenKnnCollector.java

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,22 +20,24 @@
2020
package org.elasticsearch.search.vectors;
2121

2222
import org.apache.lucene.internal.hppc.IntIntHashMap;
23-
import org.apache.lucene.search.AbstractKnnCollector;
2423
import org.apache.lucene.search.ScoreDoc;
2524
import org.apache.lucene.search.TopDocs;
2625
import org.apache.lucene.search.TotalHits;
2726
import org.apache.lucene.search.knn.KnnSearchStrategy;
2827
import org.apache.lucene.util.ArrayUtil;
2928
import org.apache.lucene.util.BitSet;
29+
import org.elasticsearch.index.codec.vectors.cluster.NeighborQueue;
3030

3131
/**
3232
* This collects the nearest children vectors. Diversifying the results over the provided parent
3333
* filter. This means the nearest children vectors are returned, but only one per parent
3434
*/
35-
class DiversifyingNearestChildrenKnnCollector extends AbstractKnnCollector {
35+
class DiversifyingNearestChildrenKnnCollector extends AbstractMaxScoreKnnCollector {
3636

3737
private final BitSet parentBitSet;
3838
private final NodeIdCachingHeap heap;
39+
private long minCompetitiveDocScore = LEAST_COMPETITIVE;
40+
private float minCompetitiveScore = Float.NEGATIVE_INFINITY;
3941

4042
/**
4143
* Create a new object for joining nearest child kNN documents with a parent bitset
@@ -72,7 +74,7 @@ public boolean collect(int docId, float nodeScore) {
7274

7375
@Override
7476
public float minCompetitiveSimilarity() {
75-
return heap.size >= k() ? heap.topScore() : Float.NEGATIVE_INFINITY;
77+
return heap.size < k() ? Float.NEGATIVE_INFINITY : Math.max(minCompetitiveScore, heap.topScore());
7678
}
7779

7880
@Override
@@ -101,6 +103,20 @@ public int numCollected() {
101103
return heap.size();
102104
}
103105

106+
@Override
107+
public long getMinCompetitiveDocScore() {
108+
return heap.size() > 0
109+
? Math.max(NeighborQueue.encodeRaw(heap.topNode(), heap.topScore()), minCompetitiveDocScore)
110+
: minCompetitiveDocScore;
111+
}
112+
113+
@Override
114+
void updateMinCompetitiveDocScore(long minCompetitiveDocScore) {
115+
long queueMinCompetitiveDocScore = heap.size() > 0 ? NeighborQueue.encodeRaw(heap.topNode(), heap.topScore()) : LEAST_COMPETITIVE;
116+
this.minCompetitiveDocScore = Math.max(this.minCompetitiveDocScore, Math.max(queueMinCompetitiveDocScore, minCompetitiveDocScore));
117+
this.minCompetitiveScore = NeighborQueue.decodeScoreRaw(this.minCompetitiveDocScore);
118+
}
119+
104120
/**
105121
* This is a minimum binary heap, inspired by {@link org.apache.lucene.util.LongHeap}. But instead
106122
* of encoding and using `long` values. Node ids and scores are kept separate. Additionally, this
@@ -134,10 +150,16 @@ private static class NodeIdCachingHeap {
134150
}
135151

136152
public final int topNode() {
153+
if (size == 0) {
154+
return Integer.MAX_VALUE;
155+
}
137156
return heapNodes[1].child;
138157
}
139158

140159
public final float topScore() {
160+
if (size == 0) {
161+
return Float.NEGATIVE_INFINITY;
162+
}
141163
return heapNodes[1].score;
142164
}
143165

‎server/src/main/java/org/elasticsearch/search/vectors/IVFKnnFloatVectorQuery.java

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,8 @@
1111
import org.apache.lucene.index.FloatVectorValues;
1212
import org.apache.lucene.index.LeafReader;
1313
import org.apache.lucene.index.LeafReaderContext;
14-
import org.apache.lucene.search.KnnCollector;
1514
import org.apache.lucene.search.Query;
1615
import org.apache.lucene.search.TopDocs;
17-
import org.apache.lucene.search.knn.KnnCollectorManager;
18-
import org.apache.lucene.search.knn.KnnSearchStrategy;
1916
import org.apache.lucene.util.Bits;
2017

2118
import java.io.IOException;
@@ -78,7 +75,7 @@ protected TopDocs approximateSearch(
7875
LeafReaderContext context,
7976
Bits acceptDocs,
8077
int visitedLimit,
81-
KnnCollectorManager knnCollectorManager,
78+
IVFCollectorManager knnCollectorManager,
8279
float visitRatio
8380
) throws IOException {
8481
LeafReader reader = context.reader();
@@ -90,11 +87,12 @@ protected TopDocs approximateSearch(
9087
if (floatVectorValues.size() == 0) {
9188
return NO_RESULTS;
9289
}
93-
KnnSearchStrategy strategy = new IVFKnnSearchStrategy(visitRatio);
94-
KnnCollector knnCollector = knnCollectorManager.newCollector(visitedLimit, strategy, context);
90+
IVFKnnSearchStrategy strategy = new IVFKnnSearchStrategy(visitRatio, knnCollectorManager.longAccumulator);
91+
AbstractMaxScoreKnnCollector knnCollector = knnCollectorManager.newCollector(visitedLimit, strategy, context);
9592
if (knnCollector == null) {
9693
return NO_RESULTS;
9794
}
95+
strategy.setCollector(knnCollector);
9896
reader.searchNearestVectors(field, query, knnCollector, acceptDocs);
9997
TopDocs results = knnCollector.topDocs();
10098
return results != null ? results : NO_RESULTS;

0 commit comments

Comments
 (0)