Skip to content

Commit 401fede

Browse files
committed
Fix test
1 parent ab484b5 commit 401fede

File tree

1 file changed

+9
-6
lines changed

1 file changed

+9
-6
lines changed

server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
import java.util.stream.Stream;
4848

4949
import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.OVERSAMPLE_LIMIT;
50+
import static org.elasticsearch.search.SearchService.DEFAULT_SIZE;
5051
import static org.hamcrest.Matchers.containsString;
5152
import static org.hamcrest.Matchers.equalTo;
5253
import static org.hamcrest.Matchers.hasSize;
@@ -175,8 +176,13 @@ protected void doAssertLuceneQuery(KnnVectorQueryBuilder queryBuilder, Query que
175176
assertThat(((VectorSimilarityQuery) query).getSimilarity(), equalTo(queryBuilder.getVectorSimilarity()));
176177
query = ((VectorSimilarityQuery) query).getInnerKnnQuery();
177178
}
179+
Integer k = queryBuilder.k();
180+
if (k == null) {
181+
k = context.requestSize() == null || context.requestSize() < 0 ? DEFAULT_SIZE : context.requestSize();
182+
}
178183
if (queryBuilder.rescoreVectorBuilder() != null && isQuantizedElementType()) {
179184
RescoreKnnVectorQuery rescoreQuery = (RescoreKnnVectorQuery) query;
185+
assertEquals(k.intValue(), (rescoreQuery.k()));
180186
query = rescoreQuery.innerQuery();
181187
}
182188
switch (elementType()) {
@@ -190,14 +196,11 @@ protected void doAssertLuceneQuery(KnnVectorQueryBuilder queryBuilder, Query que
190196
}
191197
BooleanQuery booleanQuery = builder.build();
192198
Query filterQuery = booleanQuery.clauses().isEmpty() ? null : booleanQuery;
193-
// The field should always be resolved to the concrete field
194-
Integer k = queryBuilder.k();
195199
Integer numCands = queryBuilder.numCands();
196200
if (queryBuilder.rescoreVectorBuilder() != null && isQuantizedElementType()) {
197-
Float numCandsFactor = queryBuilder.rescoreVectorBuilder().overesample();
198-
int minCands = k == null ? 1 : k;
199-
numCands = Math.max(minCands, (int) Math.ceil(numCands * numCandsFactor));
200-
numCands = Math.min(numCands, OVERSAMPLE_LIMIT);
201+
Float oversample = queryBuilder.rescoreVectorBuilder().oversample();
202+
k = Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * oversample));
203+
numCands = Math.max(numCands, k);
201204
}
202205

203206
Query knnVectorQueryBuilt = switch (elementType()) {

0 commit comments

Comments
 (0)