Skip to content

Commit 634fe3e

Browse files
john-wagsterjohn-wagster
committed
iter
1 parent c13fc1f commit 634fe3e

File tree

34 files changed

+780
-188
lines changed

34 files changed

+780
-188
lines changed

‎modules/percolator/src/internalClusterTest/java/org/elasticsearch/percolator/PercolatorQuerySearchIT.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import java.util.Collections;
4141
import java.util.Map;
4242

43+
import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.IVF_FORMAT;
4344
import static org.elasticsearch.index.query.QueryBuilders.boolQuery;
4445
import static org.elasticsearch.index.query.QueryBuilders.combinedFieldsQuery;
4546
import static org.elasticsearch.index.query.QueryBuilders.constantScoreQuery;
@@ -1359,7 +1360,15 @@ public void testKnnQueryNotSupportedInPercolator() throws IOException {
13591360
""");
13601361
indicesAdmin().prepareCreate("index1").setMapping(mappings).get();
13611362
ensureGreen();
1362-
QueryBuilder knnVectorQueryBuilder = new KnnVectorQueryBuilder("my_vector", new float[] { 1, 1, 1, 1, 1 }, 10, 10, 10f, null, null);
1363+
QueryBuilder knnVectorQueryBuilder = new KnnVectorQueryBuilder(
1364+
"my_vector",
1365+
new float[] { 1, 1, 1, 1, 1 },
1366+
10,
1367+
10,
1368+
IVF_FORMAT.isEnabled() ? 10f : null,
1369+
null,
1370+
null
1371+
);
13631372

13641373
IndexRequestBuilder indexRequestBuilder = prepareIndex("index1").setId("knn_query1")
13651374
.setSource(jsonBuilder().startObject().field("my_query", knnVectorQueryBuilder).endObject());

‎server/src/internalClusterTest/java/org/elasticsearch/index/store/DirectIOIT.java‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import java.util.OptionalLong;
3737
import java.util.stream.IntStream;
3838

39+
import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.IVF_FORMAT;
3940
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
4041
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
4142
import static org.hamcrest.Matchers.equalTo;
@@ -124,7 +125,17 @@ public void testDirectIOUsed() {
124125
indexVectors();
125126

126127
// do a search
127-
var knn = List.of(new KnnSearchBuilder("fooVector", new VectorData(null, new byte[64]), 10, 20, 10f, null, null));
128+
var knn = List.of(
129+
new KnnSearchBuilder(
130+
"fooVector",
131+
new VectorData(null, new byte[64]),
132+
10,
133+
20,
134+
IVF_FORMAT.isEnabled() ? 10f : null,
135+
null,
136+
null
137+
)
138+
);
128139
assertHitCount(prepareSearch("foo-vectors").setKnnSearch(knn), 10);
129140
mockLog.assertAllExpectationsMatched();
130141
}

‎server/src/internalClusterTest/java/org/elasticsearch/search/KnnSearchIT.java‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424

2525
import java.util.List;
2626

27+
import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.IVF_FORMAT;
2728
import static org.hamcrest.Matchers.notNullValue;
2829

2930
@ESIntegTestCase.ClusterScope(minNumDataNodes = 2)
@@ -77,29 +78,45 @@ public void testKnnSearchWithScroll() throws Exception {
7778
// test top level knn search
7879
{
7980
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
80-
sourceBuilder.knnSearch(List.of(new KnnSearchBuilder(VECTOR_FIELD, new float[] { 0, 0 }, k, 100, 10f, null, null)));
81+
sourceBuilder.knnSearch(
82+
List.of(new KnnSearchBuilder(VECTOR_FIELD, new float[] { 0, 0 }, k, 100, IVF_FORMAT.isEnabled() ? 10f : null, null, null))
83+
);
8184
executeScrollSearch(client, sourceBuilder, k);
8285
}
8386
// test top level knn search + another query
8487
{
8588
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
86-
sourceBuilder.knnSearch(List.of(new KnnSearchBuilder(VECTOR_FIELD, new float[] { 0, 0 }, k, 100, 10f, null, null)));
89+
sourceBuilder.knnSearch(
90+
List.of(new KnnSearchBuilder(VECTOR_FIELD, new float[] { 0, 0 }, k, 100, IVF_FORMAT.isEnabled() ? 10f : null, null, null))
91+
);
8792
sourceBuilder.query(QueryBuilders.existsQuery("category").boost(10));
8893
executeScrollSearch(client, sourceBuilder, k + 10);
8994
}
9095

9196
// test knn query
9297
{
9398
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
94-
sourceBuilder.query(new KnnVectorQueryBuilder(VECTOR_FIELD, new float[] { 0, 0 }, k, 100, 10f, null, null));
99+
sourceBuilder.query(
100+
new KnnVectorQueryBuilder(VECTOR_FIELD, new float[] { 0, 0 }, k, 100, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
101+
);
95102
executeScrollSearch(client, sourceBuilder, k * numShards);
96103
}
97104
// test knn query + another query
98105
{
99106
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
100107
sourceBuilder.query(
101108
QueryBuilders.boolQuery()
102-
.should(new KnnVectorQueryBuilder(VECTOR_FIELD, new float[] { 0, 0 }, k, 100, 10f, null, null))
109+
.should(
110+
new KnnVectorQueryBuilder(
111+
VECTOR_FIELD,
112+
new float[] { 0, 0 },
113+
k,
114+
100,
115+
IVF_FORMAT.isEnabled() ? 10f : null,
116+
null,
117+
null
118+
)
119+
)
103120
.should(QueryBuilders.existsQuery("category").boost(10))
104121
);
105122
executeScrollSearch(client, sourceBuilder, k * numShards + 10);

‎server/src/internalClusterTest/java/org/elasticsearch/search/nested/VectorNestedIT.java‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.util.List;
2020
import java.util.Map;
2121

22+
import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.IVF_FORMAT;
2223
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
2324
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
2425
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
@@ -73,7 +74,8 @@ public void testSimpleNested() throws Exception {
7374
assertResponse(
7475
prepareSearch("test").setKnnSearch(
7576
List.of(
76-
new KnnSearchBuilder("nested.vector", new float[] { 1, 1, 1 }, 1, 1, 10f, null, null).innerHit(new InnerHitBuilder())
77+
new KnnSearchBuilder("nested.vector", new float[] { 1, 1, 1 }, 1, 1, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
78+
.innerHit(new InnerHitBuilder())
7779
)
7880
).setAllowPartialSearchResults(false),
7981
response -> assertThat(response.getHits().getHits().length, greaterThan(0))
@@ -155,7 +157,15 @@ private void testNestedWithTwoSegments(boolean flush) {
155157
waitForRelocation(ClusterHealthStatus.GREEN);
156158
refresh();
157159

158-
var knn = new KnnSearchBuilder("nested.vector", new float[] { -0.5f, 90.0f, -10f, 14.8f, -156.0f }, 2, 3, 10f, null, null);
160+
var knn = new KnnSearchBuilder(
161+
"nested.vector",
162+
new float[] { -0.5f, 90.0f, -10f, 14.8f, -156.0f },
163+
2,
164+
3,
165+
IVF_FORMAT.isEnabled() ? 10f : null,
166+
null,
167+
null
168+
);
159169
var request = prepareSearch("test").addFetchField("name").setKnnSearch(List.of(knn));
160170
assertNoFailuresAndResponse(request, response -> {
161171
assertHitCount(response, 2);

‎server/src/internalClusterTest/java/org/elasticsearch/search/profile/dfs/DfsProfilerIT.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import java.util.List;
2828
import java.util.Map;
2929

30+
import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.IVF_FORMAT;
3031
import static org.elasticsearch.search.profile.query.RandomQueryGenerator.randomQueryBuilder;
3132
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
3233
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
@@ -72,7 +73,7 @@ public void testProfileDfs() throws Exception {
7273
new float[] { randomFloat(), randomFloat(), randomFloat() },
7374
randomIntBetween(5, 10),
7475
50,
75-
10f,
76+
IVF_FORMAT.isEnabled() ? 10f : null,
7677
randomBoolean() ? null : new RescoreVectorBuilder(randomFloatBetween(1.0f, 10.0f, false)),
7778
randomBoolean() ? null : randomFloat()
7879
);

‎server/src/internalClusterTest/java/org/elasticsearch/search/query/RescoreKnnVectorQueryIT.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
import java.util.function.Function;
5050
import java.util.stream.Collectors;
5151

52+
import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.IVF_FORMAT;
5253
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
5354
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
5455
import static org.hamcrest.Matchers.equalTo;
@@ -129,7 +130,7 @@ public static TestParams generate() {
129130
randomVector(numDims),
130131
k,
131132
(int) (k * randomFloatBetween(1.0f, 10.0f, true)),
132-
randomBoolean() ? null : randomFloatBetween(0.0f, 100.0f, true),
133+
(IVF_FORMAT.isEnabled() == false && randomBoolean()) ? null : randomFloatBetween(0.0f, 100.0f, true),
133134
new RescoreVectorBuilder(randomFloatBetween(1.0f, 100f, true))
134135
);
135136
}

‎server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverTelemetryIT.java‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import java.io.IOException;
2828
import java.util.List;
2929

30+
import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.IVF_FORMAT;
3031
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
3132
import static org.hamcrest.Matchers.equalTo;
3233

@@ -84,7 +85,9 @@ public void testTelemetryForRetrievers() throws IOException {
8485
// search#1 - this will record 1 entry for "retriever" in `sections`, and 1 for "knn" under `retrievers`
8586
{
8687
performSearch(
87-
new SearchSourceBuilder().retriever(new KnnRetrieverBuilder("vector", new float[] { 1.0f }, null, 10, 15, 10f, null, null))
88+
new SearchSourceBuilder().retriever(
89+
new KnnRetrieverBuilder("vector", new float[] { 1.0f }, null, 10, 15, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
90+
)
8891
);
8992
}
9093

@@ -99,7 +102,9 @@ public void testTelemetryForRetrievers() throws IOException {
99102
{
100103
performSearch(
101104
new SearchSourceBuilder().retriever(
102-
new StandardRetrieverBuilder(new KnnVectorQueryBuilder("vector", new float[] { 1.0f }, 10, 15, 10f, null, null))
105+
new StandardRetrieverBuilder(
106+
new KnnVectorQueryBuilder("vector", new float[] { 1.0f }, 10, 15, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
107+
)
103108
)
104109
);
105110
}
@@ -114,7 +119,9 @@ public void testTelemetryForRetrievers() throws IOException {
114119
// his will record 1 entry for "knn" in `sections`
115120
{
116121
performSearch(
117-
new SearchSourceBuilder().knnSearch(List.of(new KnnSearchBuilder("vector", new float[] { 1.0f }, 10, 15, 10f, null, null)))
122+
new SearchSourceBuilder().knnSearch(
123+
List.of(new KnnSearchBuilder("vector", new float[] { 1.0f }, 10, 15, IVF_FORMAT.isEnabled() ? 10f : null, null, null))
124+
)
118125
);
119126
}
120127

‎server/src/test/java/org/elasticsearch/action/search/DfsQueryPhaseTests.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
import java.util.List;
4949
import java.util.concurrent.atomic.AtomicReference;
5050

51+
import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.IVF_FORMAT;
5152
import static org.hamcrest.Matchers.hasSize;
5253
import static org.hamcrest.Matchers.instanceOf;
5354
import static org.mockito.Mockito.mock;
@@ -353,8 +354,8 @@ public void testRewriteShardSearchRequestWithRank() {
353354
SearchSourceBuilder ssb = new SearchSourceBuilder().query(bm25)
354355
.knnSearch(
355356
List.of(
356-
new KnnSearchBuilder("vector", new float[] { 0.0f }, 10, 100, 10f, null, null),
357-
new KnnSearchBuilder("vector2", new float[] { 0.0f }, 10, 100, 10f, null, null)
357+
new KnnSearchBuilder("vector", new float[] { 0.0f }, 10, 100, IVF_FORMAT.isEnabled() ? 10f : null, null, null),
358+
new KnnSearchBuilder("vector2", new float[] { 0.0f }, 10, 100, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
358359
)
359360
)
360361
.rankBuilder(new TestRankBuilder(100));

‎server/src/test/java/org/elasticsearch/action/search/KnnSearchSingleNodeTests.java‎

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import java.io.IOException;
2727
import java.util.List;
2828

29+
import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.IVF_FORMAT;
2930
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
3031
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
3132
import static org.hamcrest.Matchers.equalTo;
@@ -63,7 +64,8 @@ public void testKnnSearchRemovedVector() throws IOException {
6364
client().prepareUpdate("index", "0").setDoc("vector", (Object) null).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get();
6465

6566
float[] queryVector = randomVector();
66-
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 20, 50, 10f, null, null).boost(5.0f);
67+
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 20, 50, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
68+
.boost(5.0f);
6769
assertResponse(
6870
client().prepareSearch("index")
6971
.setKnnSearch(List.of(knnSearch))
@@ -107,7 +109,9 @@ public void testKnnWithQuery() throws IOException {
107109
indicesAdmin().prepareRefresh("index").get();
108110

109111
float[] queryVector = randomVector();
110-
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 5, 50, 10f, null, null).boost(5.0f).queryName("knn");
112+
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 5, 50, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
113+
.boost(5.0f)
114+
.queryName("knn");
111115
assertResponse(
112116
client().prepareSearch("index")
113117
.setKnnSearch(List.of(knnSearch))
@@ -156,9 +160,8 @@ public void testKnnFilter() throws IOException {
156160
indicesAdmin().prepareRefresh("index").get();
157161

158162
float[] queryVector = randomVector();
159-
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 5, 50, 10f, null, null).addFilterQuery(
160-
QueryBuilders.termsQuery("field", "second")
161-
);
163+
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 5, 50, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
164+
.addFilterQuery(QueryBuilders.termsQuery("field", "second"));
162165
assertResponse(client().prepareSearch("index").setKnnSearch(List.of(knnSearch)).addFetchField("*").setSize(10), response -> {
163166
assertHitCount(response, 5);
164167
assertEquals(5, response.getHits().getHits().length);
@@ -199,9 +202,8 @@ public void testKnnFilterWithRewrite() throws IOException {
199202
indicesAdmin().prepareRefresh("index").get();
200203

201204
float[] queryVector = randomVector();
202-
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 5, 50, 10f, null, null).addFilterQuery(
203-
QueryBuilders.termsLookupQuery("field", new TermsLookup("index", "lookup-doc", "other-field"))
204-
);
205+
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 5, 50, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
206+
.addFilterQuery(QueryBuilders.termsLookupQuery("field", new TermsLookup("index", "lookup-doc", "other-field")));
205207
assertResponse(client().prepareSearch("index").setKnnSearch(List.of(knnSearch)).setSize(10), response -> {
206208
assertHitCount(response, 5);
207209
assertEquals(5, response.getHits().getHits().length);
@@ -246,8 +248,10 @@ public void testMultiKnnClauses() throws IOException {
246248
indicesAdmin().prepareRefresh("index").get();
247249

248250
float[] queryVector = randomVector(20f, 21f);
249-
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 5, 50, 10f, null, null).boost(5.0f);
250-
KnnSearchBuilder knnSearch2 = new KnnSearchBuilder("vector_2", queryVector, 5, 50, 10f, null, null).boost(10.0f);
251+
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 5, 50, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
252+
.boost(5.0f);
253+
KnnSearchBuilder knnSearch2 = new KnnSearchBuilder("vector_2", queryVector, 5, 50, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
254+
.boost(10.0f);
251255
assertResponse(
252256
client().prepareSearch("index")
253257
.setKnnSearch(List.of(knnSearch, knnSearch2))
@@ -308,8 +312,8 @@ public void testMultiKnnClausesSameDoc() throws IOException {
308312

309313
float[] queryVector = randomVector();
310314
// Having the same query vector and same docs should mean our KNN scores are linearly combined if the same doc is matched
311-
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 5, 50, 10f, null, null);
312-
KnnSearchBuilder knnSearch2 = new KnnSearchBuilder("vector_2", queryVector, 5, 50, 10f, null, null);
315+
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 5, 50, IVF_FORMAT.isEnabled() ? 10f : null, null, null);
316+
KnnSearchBuilder knnSearch2 = new KnnSearchBuilder("vector_2", queryVector, 5, 50, IVF_FORMAT.isEnabled() ? 10f : null, null, null);
313317
assertResponse(
314318
client().prepareSearch("index")
315319
.setKnnSearch(List.of(knnSearch))
@@ -383,7 +387,7 @@ public void testKnnFilteredAlias() throws IOException {
383387
indicesAdmin().prepareRefresh("index").get();
384388

385389
float[] queryVector = randomVector();
386-
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 10, 50, 10f, null, null);
390+
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 10, 50, IVF_FORMAT.isEnabled() ? 10f : null, null, null);
387391
final int expectedHitCount = expectedHits;
388392
assertResponse(client().prepareSearch("test-alias").setKnnSearch(List.of(knnSearch)).setSize(10), response -> {
389393
assertHitCount(response, expectedHitCount);
@@ -420,7 +424,7 @@ public void testKnnSearchAction() throws IOException {
420424
float[] queryVector = randomVector();
421425
assertResponse(
422426
client().prepareSearch("index1", "index2")
423-
.setQuery(new KnnVectorQueryBuilder("vector", queryVector, 5, 5, 10f, null, null))
427+
.setQuery(new KnnVectorQueryBuilder("vector", queryVector, 5, 5, IVF_FORMAT.isEnabled() ? 10f : null, null, null))
424428
.setSize(2),
425429
response -> {
426430
// The total hits is num_cands * num_shards, since the query gathers num_cands hits from each shard
@@ -454,7 +458,8 @@ public void testKnnVectorsWith4096Dims() throws IOException {
454458
indicesAdmin().prepareRefresh("index").get();
455459

456460
float[] queryVector = randomVector(4096);
457-
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 3, 50, 10f, null, null).boost(5.0f);
461+
KnnSearchBuilder knnSearch = new KnnSearchBuilder("vector", queryVector, 3, 50, IVF_FORMAT.isEnabled() ? 10f : null, null, null)
462+
.boost(5.0f);
458463
assertResponse(client().prepareSearch("index").setKnnSearch(List.of(knnSearch)).addFetchField("*").setSize(10), response -> {
459464
assertHitCount(response, 3);
460465
assertEquals(3, response.getHits().getHits().length);

0 commit comments

Comments
 (0)