Skip to content

Commit 0cfb0b3

Browse files
carlosdelestcarlosdelest
committed
Fix DenseVectorFieldIndexTypeUpdateIT test
1 parent 8005a7d commit 0cfb0b3

File tree

2 files changed

+30
-104
lines changed

2 files changed

+30
-104
lines changed

‎muted-tests.yml‎

Lines changed: 0 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -530,60 +530,6 @@ tests:
530530
- class: org.elasticsearch.xpack.test.rest.XPackRestIT
531531
method: test {p0=transform/transforms_crud/Test transform where source query is invalid}
532532
issue: https://github.com/elastic/elasticsearch/issues/132111
533-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
534-
method: testDenseVectorMappingUpdate {initialType=bbq_flat updateType=bbq_hnsw}
535-
issue: https://github.com/elastic/elasticsearch/issues/132112
536-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
537-
method: testDenseVectorMappingUpdate {initialType=hnsw updateType=int4_hnsw}
538-
issue: https://github.com/elastic/elasticsearch/issues/132113
539-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
540-
method: testDenseVectorMappingUpdate {initialType=hnsw updateType=bbq_hnsw}
541-
issue: https://github.com/elastic/elasticsearch/issues/132115
542-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
543-
method: testDenseVectorMappingUpdate {initialType=int8_flat updateType=int8_hnsw}
544-
issue: https://github.com/elastic/elasticsearch/issues/132116
545-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
546-
method: testDenseVectorMappingUpdate {initialType=int4_flat updateType=int8_hnsw}
547-
issue: https://github.com/elastic/elasticsearch/issues/132117
548-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
549-
method: testDenseVectorMappingUpdate {initialType=flat updateType=hnsw}
550-
issue: https://github.com/elastic/elasticsearch/issues/132119
551-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
552-
method: testDenseVectorMappingUpdate {initialType=flat updateType=bbq_flat}
553-
issue: https://github.com/elastic/elasticsearch/issues/132120
554-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
555-
method: testDenseVectorMappingUpdate {initialType=int8_flat updateType=bbq_disk}
556-
issue: https://github.com/elastic/elasticsearch/issues/132122
557-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
558-
method: testDenseVectorMappingUpdate {initialType=int8_flat updateType=int4_hnsw}
559-
issue: https://github.com/elastic/elasticsearch/issues/132123
560-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
561-
method: testDenseVectorMappingUpdate {initialType=int8_flat updateType=hnsw}
562-
issue: https://github.com/elastic/elasticsearch/issues/132124
563-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
564-
method: testDenseVectorMappingUpdate {initialType=int8_flat updateType=int4_flat}
565-
issue: https://github.com/elastic/elasticsearch/issues/132125
566-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
567-
method: testDenseVectorMappingUpdate {initialType=flat updateType=int8_hnsw}
568-
issue: https://github.com/elastic/elasticsearch/issues/132126
569-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
570-
method: testDenseVectorMappingUpdate {initialType=bbq_disk updateType=bbq_disk}
571-
issue: https://github.com/elastic/elasticsearch/issues/132127
572-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
573-
method: testDenseVectorMappingUpdate {initialType=flat updateType=int8_flat}
574-
issue: https://github.com/elastic/elasticsearch/issues/132129
575-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
576-
method: "testDenseVectorMappingUpdate {initialType=flat updateType=bbq_disk #2}"
577-
issue: https://github.com/elastic/elasticsearch/issues/132130
578-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
579-
method: testDenseVectorMappingUpdate {initialType=int4_flat updateType=bbq_flat}
580-
issue: https://github.com/elastic/elasticsearch/issues/132132
581-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
582-
method: testDenseVectorMappingUpdate {initialType=int8_hnsw updateType=int4_hnsw}
583-
issue: https://github.com/elastic/elasticsearch/issues/132133
584-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
585-
method: testDenseVectorMappingUpdate {initialType=bbq_flat updateType=bbq_disk}
586-
issue: https://github.com/elastic/elasticsearch/issues/132134
587533
- class: org.elasticsearch.upgrades.RunningSnapshotIT
588534
method: testRunningSnapshotCompleteAfterUpgrade {upgradedNodes=1}
589535
issue: https://github.com/elastic/elasticsearch/issues/132135
@@ -593,33 +539,9 @@ tests:
593539
- class: org.elasticsearch.upgrades.RunningSnapshotIT
594540
method: testRunningSnapshotCompleteAfterUpgrade {upgradedNodes=3}
595541
issue: https://github.com/elastic/elasticsearch/issues/132137
596-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
597-
method: testDenseVectorMappingUpdate {initialType=int4_flat updateType=int4_hnsw}
598-
issue: https://github.com/elastic/elasticsearch/issues/132140
599-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
600-
method: testDenseVectorMappingUpdate {initialType=int8_flat updateType=bbq_hnsw}
601-
issue: https://github.com/elastic/elasticsearch/issues/132141
602542
- class: org.elasticsearch.index.engine.MergeWithLowDiskSpaceIT
603543
method: testRelocationWhileForceMerging
604544
issue: https://github.com/elastic/elasticsearch/issues/131789
605-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
606-
method: testDenseVectorMappingUpdate {initialType=flat updateType=int4_hnsw}
607-
issue: https://github.com/elastic/elasticsearch/issues/132149
608-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
609-
method: testDenseVectorMappingUpdate {initialType=int4_flat updateType=hnsw}
610-
issue: https://github.com/elastic/elasticsearch/issues/132150
611-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
612-
method: testDenseVectorMappingUpdate {initialType=int8_flat updateType=bbq_flat}
613-
issue: https://github.com/elastic/elasticsearch/issues/132151
614-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
615-
method: "testDenseVectorMappingUpdate {initialType=bbq_hnsw updateType=bbq_disk #2}"
616-
issue: https://github.com/elastic/elasticsearch/issues/132152
617-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
618-
method: testDenseVectorMappingUpdate {initialType=hnsw updateType=int8_hnsw}
619-
issue: https://github.com/elastic/elasticsearch/issues/132164
620-
- class: org.elasticsearch.index.mapper.vectors.DenseVectorFieldIndexTypeUpdateIT
621-
method: testDenseVectorMappingUpdate {initialType=int4_hnsw updateType=bbq_disk}
622-
issue: https://github.com/elastic/elasticsearch/issues/132165
623545

624546
# Examples:
625547
#

‎server/src/internalClusterTest/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldIndexTypeUpdateIT.java‎

Lines changed: 30 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,14 @@
1212
import com.carrotsearch.randomizedtesting.annotations.Name;
1313
import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
1414

15+
import org.elasticsearch.action.DocWriteResponse;
1516
import org.elasticsearch.action.admin.indices.mapping.get.GetFieldMappingsResponse;
1617
import org.elasticsearch.action.index.IndexRequest;
1718
import org.elasticsearch.cluster.metadata.IndexMetadata;
1819
import org.elasticsearch.common.settings.Settings;
20+
import org.elasticsearch.rest.RestStatus;
21+
import org.elasticsearch.search.vectors.KnnVectorQueryBuilder;
22+
import org.elasticsearch.search.vectors.VectorData;
1923
import org.elasticsearch.test.ESIntegTestCase;
2024
import org.elasticsearch.xcontent.XContentBuilder;
2125
import org.elasticsearch.xcontent.XContentFactory;
@@ -25,6 +29,7 @@
2529
import java.util.Collection;
2630
import java.util.List;
2731
import java.util.Map;
32+
import java.util.concurrent.ExecutionException;
2833

2934
import static org.elasticsearch.index.IndexSettings.INDEX_MAPPER_SOURCE_MODE_SETTING;
3035
import static org.elasticsearch.index.mapper.SourceFieldMapper.Mode.SYNTHETIC;
@@ -75,7 +80,7 @@ public static Collection<Object[]> params() {
7580
@SuppressWarnings("unchecked")
7681
public void testDenseVectorMappingUpdate() throws Exception {
7782
dimensions = randomIntBetween(1, 10) * 64;
78-
var client = client().admin().indices();
83+
var indicesClient = client().admin().indices();
7984

8085
Settings.Builder settingsBuilder = Settings.builder()
8186
.put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
@@ -85,7 +90,7 @@ public void testDenseVectorMappingUpdate() throws Exception {
8590
}
8691

8792
// Create index with initial mapping
88-
var createRequest = client.prepareCreate(INDEX_NAME)
93+
var createRequest = indicesClient.prepareCreate(INDEX_NAME)
8994
.setSettings(Settings.builder().put("index.number_of_shards", randomIntBetween(1, 5)))
9095
.setMapping(updateMapping(dimensions, initialType))
9196
.setSettings(settingsBuilder.build());
@@ -97,16 +102,16 @@ public void testDenseVectorMappingUpdate() throws Exception {
97102
indexDoc(i);
98103
}
99104

100-
client.prepareFlush(INDEX_NAME).get();
101-
client.prepareRefresh(INDEX_NAME).get();
105+
indicesClient.prepareFlush(INDEX_NAME).get();
106+
indicesClient.prepareRefresh(INDEX_NAME).get();
102107

103108
// Update mapping to new type
104-
var putMappingRequest = client.preparePutMapping(INDEX_NAME).setSource(updateMapping(dimensions, updateType)).request();
105-
assertAcked(client.putMapping(putMappingRequest));
109+
var putMappingRequest = indicesClient.preparePutMapping(INDEX_NAME).setSource(updateMapping(dimensions, updateType)).request();
110+
assertAcked(indicesClient.putMapping(putMappingRequest));
106111

107112
// Validate mapping
108-
GetFieldMappingsResponse fieldMapping = client.getFieldMappings(
109-
client.prepareGetFieldMappings(INDEX_NAME).setFields(VECTOR_FIELD).request()
113+
GetFieldMappingsResponse fieldMapping = indicesClient.getFieldMappings(
114+
indicesClient.prepareGetFieldMappings(INDEX_NAME).setFields(VECTOR_FIELD).request()
110115
).get();
111116
var fieldMappingMetadata = fieldMapping.fieldMappings(INDEX_NAME, VECTOR_FIELD);
112117
var fieldMap = (Map<String, Object>) fieldMappingMetadata.sourceAsMap().get(VECTOR_FIELD);
@@ -119,28 +124,26 @@ public void testDenseVectorMappingUpdate() throws Exception {
119124
indexDoc(i);
120125
}
121126

122-
client.prepareFlush(INDEX_NAME).get();
123-
client.prepareRefresh(INDEX_NAME).get();
127+
indicesClient.prepareFlush(INDEX_NAME).get();
128+
indicesClient.prepareRefresh(INDEX_NAME).get();
124129

125130
// Search to ensure all documents are present
126131
int expectedDocs = docsBefore + docsAfter;
127-
assertNoFailuresAndResponse(client().prepareSearch(INDEX_NAME).setSize(expectedDocs + 10), response -> {
132+
133+
// Count query
134+
assertNoFailuresAndResponse(client().prepareSearch(INDEX_NAME).setSize(0).setTrackTotalHits(true).setSize(expectedDocs), response -> {
128135
assertHitCount(response, expectedDocs);
129136
});
130-
}
131137

132-
private XContentBuilder initialMapping(int dimensions, String type) throws IOException {
133-
XContentBuilder builder = XContentFactory.jsonBuilder();
134-
builder.startObject();
135-
{
136-
builder.startObject("mappings");
137-
{
138-
createFieldMapping(dimensions, type, builder);
139-
}
140-
builder.endObject();
138+
// KNN query
139+
float[] queryVector = new float[dimensions];
140+
for (int i = 0; i < queryVector.length; i++) {
141+
queryVector[i] = randomFloatBetween(-1, 1, true);
141142
}
142-
builder.endObject();
143-
return builder;
143+
KnnVectorQueryBuilder queryBuilder = new KnnVectorQueryBuilder(VECTOR_FIELD, queryVector, null, null, null, null);
144+
assertNoFailuresAndResponse(client().prepareSearch(INDEX_NAME).setQuery(queryBuilder).setTrackTotalHits(true).setSize(expectedDocs), response -> {
145+
assertHitCount(response, expectedDocs);
146+
});
144147
}
145148

146149
private XContentBuilder updateMapping(int dimensions, String type) throws IOException {
@@ -168,9 +171,10 @@ private static void createFieldMapping(int dimensions, String type, XContentBuil
168171
builder.endObject();
169172
}
170173

171-
private void indexDoc(int id) throws IOException {
174+
private void indexDoc(int id) throws ExecutionException, InterruptedException {
172175
Float[] vector = randomArray(dimensions, dimensions, Float[]::new, () -> randomFloatBetween(-1, 1, true));
173-
IndexRequest req = prepareIndex(INDEX_NAME).setSource(VECTOR_FIELD, vector).setId(Integer.toString(id)).request();
174-
client().index(req);
176+
IndexRequest indexRequest = prepareIndex(INDEX_NAME).setSource(VECTOR_FIELD, vector).setId(Integer.toString(id)).request();
177+
DocWriteResponse indexResponse = client().index(indexRequest).get();
178+
assertEquals(RestStatus.CREATED, indexResponse.status());
175179
}
176180
}

0 commit comments

Comments
 (0)