Skip to content

Commit 1347d4b

Browse files
carlosdelestcarlosdelest
committed
Adds profiling, including a small refactoring of the QueryProfiler interface
1 parent 81384f2 commit 1347d4b

File tree

8 files changed

+123
-65
lines changed

8 files changed

+123
-65
lines changed

‎server/src/main/java/org/elasticsearch/search/profile/query/QueryProfiler.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,8 @@ public QueryProfiler() {
3939
super(new InternalQueryProfileTree());
4040
}
4141

42-
public void setVectorOpsCount(long vectorOpsCount) {
43-
this.vectorOpsCount = vectorOpsCount;
42+
public void addVectorOpsCount(long vectorOpsCount) {
43+
this.vectorOpsCount += vectorOpsCount;
4444
}
4545

4646
public long getVectorOpsCount() {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,6 @@ protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
4040

4141
@Override
4242
public void profile(QueryProfiler queryProfiler) {
43-
queryProfiler.setVectorOpsCount(vectorOpsCount);
43+
queryProfiler.addVectorOpsCount(vectorOpsCount);
4444
}
4545
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,6 @@ protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
4040

4141
@Override
4242
public void profile(QueryProfiler queryProfiler) {
43-
queryProfiler.setVectorOpsCount(vectorOpsCount);
43+
queryProfiler.addVectorOpsCount(vectorOpsCount);
4444
}
4545
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
3333

3434
@Override
3535
public void profile(QueryProfiler queryProfiler) {
36-
queryProfiler.setVectorOpsCount(vectorOpsCount);
36+
queryProfiler.addVectorOpsCount(vectorOpsCount);
3737
}
3838

3939
public Integer kParam() {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
3333

3434
@Override
3535
public void profile(QueryProfiler queryProfiler) {
36-
queryProfiler.setVectorOpsCount(vectorOpsCount);
36+
queryProfiler.addVectorOpsCount(vectorOpsCount);
3737
}
3838

3939
public Integer kParam() {

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -72,11 +72,6 @@ public RescoreKnnVectorQuery(
7272
public Query rewrite(IndexSearcher searcher) throws IOException {
7373
assert byteTarget == null ^ floatTarget == null : "Either byteTarget or floatTarget must be set";
7474

75-
Query rewritten = super.rewrite(searcher);
76-
if (rewritten != this) {
77-
return rewritten;
78-
}
79-
8075
final DoubleValuesSource valueSource;
8176
if (byteTarget != null) {
8277
valueSource = new VectorSimilarityByteValueSource(fieldName, byteTarget, vectorSimilarityFunction);
@@ -115,7 +110,10 @@ public Integer k() {
115110

116111
@Override
117112
public void profile(QueryProfiler queryProfiler) {
118-
queryProfiler.setVectorOpsCount(vectorOpsCount);
113+
if (innerQuery instanceof ProfilingQuery profilingQuery) {
114+
profilingQuery.profile(queryProfiler);
115+
}
116+
queryProfiler.addVectorOpsCount(vectorOpsCount);
119117
}
120118

121119
@Override

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.apache.lucene.search.ScorerSupplier;
2222
import org.apache.lucene.search.Weight;
2323
import org.elasticsearch.common.lucene.search.function.MinScoreScorer;
24+
import org.elasticsearch.search.profile.query.QueryProfiler;
2425

2526
import java.io.IOException;
2627
import java.util.Objects;
@@ -30,7 +31,7 @@
3031
/**
3132
* This query provides a simple post-filter for the provided Query. The query is assumed to be a Knn(Float|Byte)VectorQuery.
3233
*/
33-
public class VectorSimilarityQuery extends Query {
34+
public class VectorSimilarityQuery extends Query implements ProfilingQuery {
3435
private final float similarity;
3536
private final float docScore;
3637
private final Query innerKnnQuery;
@@ -78,6 +79,13 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo
7879
return new MinScoreWeight(innerWeight, docScore, similarity, this, boost);
7980
}
8081

82+
@Override
83+
public void profile(QueryProfiler queryProfiler) {
84+
if (innerKnnQuery instanceof ProfilingQuery profilingQuery) {
85+
profilingQuery.profile(queryProfiler);
86+
}
87+
}
88+
8189
@Override
8290
public String toString(String field) {
8391
return "VectorSimilarityQuery["

‎server/src/test/java/org/elasticsearch/search/vectors/RescoreKnnVectorQueryTests.java‎

Lines changed: 104 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,13 @@
2525
import org.apache.lucene.index.VectorSimilarityFunction;
2626
import org.apache.lucene.search.IndexSearcher;
2727
import org.apache.lucene.search.MatchAllDocsQuery;
28+
import org.apache.lucene.search.Query;
29+
import org.apache.lucene.search.QueryVisitor;
30+
import org.apache.lucene.search.ScoreMode;
2831
import org.apache.lucene.search.TopDocs;
32+
import org.apache.lucene.search.Weight;
2933
import org.apache.lucene.store.Directory;
34+
import org.elasticsearch.search.profile.query.QueryProfiler;
3035
import org.elasticsearch.test.ESTestCase;
3136

3237
import java.io.IOException;
@@ -41,6 +46,7 @@
4146

4247
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
4348
import static org.hamcrest.Matchers.equalTo;
49+
import static org.hamcrest.Matchers.greaterThan;
4450

4551
public class RescoreKnnVectorQueryTests extends ESTestCase {
4652

@@ -51,7 +57,8 @@ public class RescoreKnnVectorQueryTests extends ESTestCase {
5157

5258
public RescoreKnnVectorQueryTests(VectorProvider vectorProvider, boolean useK) {
5359
this.vectorProvider = vectorProvider;
54-
this.numDocs = randomIntBetween(10, 100);;
60+
this.numDocs = randomIntBetween(10, 100);
61+
;
5562
this.k = useK ? randomIntBetween(1, numDocs - 1) : null;
5663
}
5764

@@ -71,7 +78,11 @@ public void testRescoreDocs() throws Exception {
7178
// Use a RescoreKnnVectorQuery with a match all query, to ensure we get scoring of 1 from the inner query
7279
// and thus we're rescoring the top k docs.
7380
VectorData queryVector = vectorProvider.randomVector(numDims);
74-
RescoreKnnVectorQuery rescoreKnnVectorQuery = vectorProvider.createRescoreQuery(queryVector, adjustedK);
81+
RescoreKnnVectorQuery rescoreKnnVectorQuery = vectorProvider.createRescoreQuery(
82+
queryVector,
83+
adjustedK,
84+
new MatchAllDocsQuery()
85+
);
7586

7687
IndexSearcher searcher = newSearcher(reader, true, false);
7788
TopDocs docs = searcher.search(rescoreKnnVectorQuery, numDocs);
@@ -115,10 +126,90 @@ public void testRescoreDocs() throws Exception {
115126
}
116127
}
117128

129+
public void testProfiling() throws Exception {
130+
int numDims = randomIntBetween(5, 100);
131+
132+
try (Directory d = newDirectory()) {
133+
addRandomDocuments(numDocs, d, numDims, vectorProvider);
134+
135+
try (IndexReader reader = DirectoryReader.open(d)) {
136+
VectorData queryVector = vectorProvider.randomVector(numDims);
137+
138+
checkProfiling(queryVector, reader, new MatchAllDocsQuery());
139+
checkProfiling(queryVector, reader, new MockProfilingQuery(randomIntBetween(1, 100)));
140+
}
141+
}
142+
}
143+
144+
private void checkProfiling(VectorData queryVector, IndexReader reader, Query innerQuery) throws IOException {
145+
RescoreKnnVectorQuery rescoreKnnVectorQuery = vectorProvider.createRescoreQuery(queryVector, k, innerQuery);
146+
IndexSearcher searcher = newSearcher(reader, true, false);
147+
searcher.search(rescoreKnnVectorQuery, numDocs);
148+
149+
QueryProfiler queryProfiler = new QueryProfiler();
150+
rescoreKnnVectorQuery.profile(queryProfiler);
151+
152+
long expectedVectorOpsCount = 0;
153+
if (k != null) {
154+
expectedVectorOpsCount += k;
155+
}
156+
if (innerQuery instanceof ProfilingQuery profilingQuery) {
157+
QueryProfiler anotherProfiler = new QueryProfiler();
158+
profilingQuery.profile(anotherProfiler);
159+
assertThat(anotherProfiler.getVectorOpsCount(), greaterThan(0L));
160+
expectedVectorOpsCount += anotherProfiler.getVectorOpsCount();
161+
}
162+
163+
assertThat(queryProfiler.getVectorOpsCount(), equalTo(expectedVectorOpsCount));
164+
}
165+
166+
/**
167+
* A mock query that is used to test profiling
168+
*/
169+
private static class MockProfilingQuery extends Query implements ProfilingQuery {
170+
171+
private final long vectorOpsCount;
172+
173+
private MockProfilingQuery(long vectorOpsCount) {
174+
this.vectorOpsCount = vectorOpsCount;
175+
}
176+
177+
@Override
178+
public String toString(String field) {
179+
return "";
180+
}
181+
182+
@Override
183+
public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
184+
return new MatchAllDocsQuery().createWeight(searcher, scoreMode, boost);
185+
}
186+
187+
@Override
188+
public void visit(QueryVisitor visitor) {}
189+
190+
@Override
191+
public boolean equals(Object obj) {
192+
return obj instanceof MockProfilingQuery;
193+
}
194+
195+
@Override
196+
public int hashCode() {
197+
return 0;
198+
}
199+
200+
@Override
201+
public void profile(QueryProfiler queryProfiler) {
202+
queryProfiler.addVectorOpsCount(vectorOpsCount);
203+
}
204+
}
205+
206+
/**
207+
* Vector operations depend on the type of vector field used. This interface abstracts the operations needed to perform the tests
208+
*/
118209
private interface VectorProvider {
119210
VectorData randomVector(int numDimensions);
120211

121-
RescoreKnnVectorQuery createRescoreQuery(VectorData queryVector, Integer k);
212+
RescoreKnnVectorQuery createRescoreQuery(VectorData queryVector, Integer k, Query innerQuery);
122213

123214
KnnVectorValues vectorValues(LeafReader leafReader) throws IOException;
124215

@@ -140,14 +231,8 @@ public VectorData randomVector(int numDimensions) {
140231
}
141232

142233
@Override
143-
public RescoreKnnVectorQuery createRescoreQuery(VectorData queryVector, Integer k) {
144-
return new RescoreKnnVectorQuery(
145-
FIELD_NAME,
146-
queryVector.floatVector(),
147-
VectorSimilarityFunction.COSINE,
148-
k,
149-
new MatchAllDocsQuery()
150-
);
234+
public RescoreKnnVectorQuery createRescoreQuery(VectorData queryVector, Integer k, Query innerQuery) {
235+
return new RescoreKnnVectorQuery(FIELD_NAME, queryVector.floatVector(), VectorSimilarityFunction.COSINE, k, innerQuery);
151236
}
152237

153238
@Override
@@ -163,7 +248,7 @@ public void addVectorField(Document document, VectorData vector) {
163248

164249
@Override
165250
public VectorData dataVectorForDoc(KnnVectorValues vectorValues, int docId) throws IOException {
166-
return VectorData.fromFloats(((FloatVectorValues)vectorValues).vectorValue(docId));
251+
return VectorData.fromFloats(((FloatVectorValues) vectorValues).vectorValue(docId));
167252
}
168253

169254
@Override
@@ -183,14 +268,8 @@ public VectorData randomVector(int numDimensions) {
183268
}
184269

185270
@Override
186-
public RescoreKnnVectorQuery createRescoreQuery(VectorData queryVector, Integer k) {
187-
return new RescoreKnnVectorQuery(
188-
FIELD_NAME,
189-
queryVector.byteVector(),
190-
VectorSimilarityFunction.COSINE,
191-
k,
192-
new MatchAllDocsQuery()
193-
);
271+
public RescoreKnnVectorQuery createRescoreQuery(VectorData queryVector, Integer k, Query innerQuery) {
272+
return new RescoreKnnVectorQuery(FIELD_NAME, queryVector.byteVector(), VectorSimilarityFunction.COSINE, k, innerQuery);
194273
}
195274

196275
@Override
@@ -206,7 +285,7 @@ public void addVectorField(Document document, VectorData vector) {
206285

207286
@Override
208287
public VectorData dataVectorForDoc(KnnVectorValues vectorValues, int docId) throws IOException {
209-
return VectorData.fromBytes(((ByteVectorValues)vectorValues).vectorValue(docId));
288+
return VectorData.fromBytes(((ByteVectorValues) vectorValues).vectorValue(docId));
210289
}
211290

212291
@Override
@@ -230,39 +309,12 @@ private static void addRandomDocuments(int numDocs, Directory d, int numDims, Ve
230309

231310
@ParametersFactory
232311
public static Iterable<Object[]> parameters() {
233-
234312
List<Object[]> params = new ArrayList<>();
235-
params.add(new Object[] {new FloatVectorProvider(), true});
236-
params.add(new Object[] {new FloatVectorProvider(), false});
237-
params.add(new Object[] {new ByteVectorProvider(), true});
238-
params.add(new Object[] {new ByteVectorProvider(), false});
313+
params.add(new Object[] { new FloatVectorProvider(), true });
314+
params.add(new Object[] { new FloatVectorProvider(), false });
315+
params.add(new Object[] { new ByteVectorProvider(), true });
316+
params.add(new Object[] { new ByteVectorProvider(), false });
239317

240318
return params;
241319
}
242-
243-
// public void testProfiling() throws Exception {
244-
// int numDocs = randomIntBetween(10, 100);
245-
// int numDims = randomIntBetween(5, 100);
246-
//
247-
// try (Directory d = newDirectory()) {
248-
// addRandomDocuments(numDocs, d, numDims, vectorProvider);
249-
//
250-
// try (IndexReader reader = DirectoryReader.open(d)) {
251-
// float[] queryVector = randomVector(numDims);
252-
//
253-
// RescoreKnnVectorQuery rescoreKnnVectorQuery = new RescoreKnnVectorQuery(
254-
// FIELD_NAME,
255-
// queryVector,
256-
// VectorSimilarityFunction.COSINE,
257-
// randomIntBetween(5, numDocs - 1),
258-
// new MatchAllDocsQuery()
259-
// );
260-
//
261-
// IndexSearcher searcher = newSearcher(reader, true, false);
262-
// QueryProfiler queryProfiler = new QueryProfiler();
263-
// rescoreKnnVectorQuery.profile(queryProfiler);
264-
// }
265-
// }
266-
// }
267-
268320
}

0 commit comments

Comments
 (0)