Skip to content

Commit 3508148

Browse files
benwtrentbenwtrent
authored
Adding tests for bit fmt usage (#139217)
This adds some rolling upgrade tests for the formats to ensure things are supported.
1 parent 33b3ff7 commit 3508148

File tree

1 file changed

+147
-0
lines changed

1 file changed

+147
-0
lines changed

‎qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/upgrades/VectorSearchIT.java‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,13 +37,160 @@ public VectorSearchIT(@Name("upgradedNodes") int upgradedNodes) {
3737
private static final String BBQ_INDEX_NAME = "bbq_vector_index";
3838
private static final String FLAT_QUANTIZED_INDEX_NAME = "flat_quantized_vector_index";
3939
private static final String FLAT_BBQ_INDEX_NAME = "flat_bbq_vector_index";
40+
private static final String HNSW_BIT_INDEX_NAME = "hnsw_bit_vector_index";
41+
private static final String FLAT_BIT_INDEX_NAME = "flat_bit_vector_index";
4042

4143
// TODO: replace these with proper test features
4244
private static final String FLOAT_VECTOR_SEARCH_TEST_FEATURE = "gte_v8.4.0";
4345
private static final String BYTE_VECTOR_SEARCH_TEST_FEATURE = "gte_v8.6.0";
4446
private static final String QUANTIZED_VECTOR_SEARCH_TEST_FEATURE = "gte_v8.12.1";
4547
private static final String FLAT_QUANTIZED_VECTOR_SEARCH_TEST_FEATURE = "gte_v8.13.0";
4648
private static final String BBQ_VECTOR_SEARCH_TEST_FEATURE = "gte_v8.18.0";
49+
private static final String BIT_VECTOR_SEARCH_TEST_FEATURE = "gte_v8.15.0";
50+
51+
public void testBitVectors() throws Exception {
52+
assumeTrue("Bit vector search is not supported on this version", oldClusterHasFeature(BIT_VECTOR_SEARCH_TEST_FEATURE));
53+
if (isOldCluster()) {
54+
String mapping = """
55+
{
56+
"properties": {
57+
"vector": {
58+
"type": "dense_vector",
59+
"dims": 24,
60+
"element_type": "bit",
61+
"index": true,
62+
"similarity": "l2_norm",
63+
"index_options": {
64+
"type": "hnsw",
65+
"ef_construction": 100,
66+
"m": 16
67+
}
68+
}
69+
}
70+
}
71+
""";
72+
// create index and index 10 random floating point vectors
73+
createIndex(HNSW_BIT_INDEX_NAME, Settings.EMPTY, mapping);
74+
indexVectors(HNSW_BIT_INDEX_NAME);
75+
// force merge the index
76+
client().performRequest(new Request("POST", "/" + HNSW_BIT_INDEX_NAME + "/_forcemerge?max_num_segments=1"));
77+
78+
mapping = """
79+
{
80+
"properties": {
81+
"vector": {
82+
"type": "dense_vector",
83+
"dims": 24,
84+
"element_type": "bit",
85+
"index": true,
86+
"similarity": "l2_norm",
87+
"index_options": {
88+
"type": "flat"
89+
}
90+
}
91+
}
92+
}
93+
""";
94+
// create index and index 10 random floating point vectors
95+
createIndex(FLAT_BIT_INDEX_NAME, Settings.EMPTY, mapping);
96+
indexVectors(FLAT_BIT_INDEX_NAME);
97+
// force merge the index
98+
client().performRequest(new Request("POST", "/" + FLAT_BIT_INDEX_NAME + "/_forcemerge?max_num_segments=1"));
99+
}
100+
101+
// search hnsw index first
102+
{
103+
Request searchRequest = new Request("POST", "/" + HNSW_BIT_INDEX_NAME + "/_search");
104+
searchRequest.setJsonEntity("""
105+
{
106+
"query": {
107+
"script_score": {
108+
"query": {
109+
"exists": {
110+
"field": "vector"
111+
}
112+
},
113+
"script": {
114+
"source": "hamming(params.query, 'vector')",
115+
"params": {
116+
"query": [4, 5, 6]
117+
}
118+
}
119+
}
120+
}
121+
}
122+
""");
123+
Map<String, Object> response = search(searchRequest);
124+
assertThat(extractValue(response, "hits.total.value"), equalTo(7));
125+
List<Map<String, Object>> hits = extractValue(response, "hits.hits");
126+
assertThat(hits.get(0).get("_id"), equalTo("3"));
127+
assertThat((double) hits.get(0).get("_score"), closeTo(8.0, 0.0001));
128+
129+
// search with knn
130+
searchRequest = new Request("POST", "/" + HNSW_BIT_INDEX_NAME + "/_search");
131+
searchRequest.setJsonEntity("""
132+
{
133+
"knn": {
134+
"field": "vector",
135+
"query_vector": [4, 5, 6],
136+
"k": 2,
137+
"num_candidates": 5
138+
}
139+
}
140+
""");
141+
response = search(searchRequest);
142+
assertThat(extractValue(response, "hits.total.value"), equalTo(2));
143+
hits = extractValue(response, "hits.hits");
144+
assertThat(hits.get(0).get("_id"), equalTo("1"));
145+
assertThat((double) hits.get(0).get("_score"), closeTo(0.8333333, 0.0001));
146+
}
147+
// search the flat
148+
{
149+
Request searchRequest = new Request("POST", "/" + FLAT_BIT_INDEX_NAME + "/_search");
150+
searchRequest.setJsonEntity("""
151+
{
152+
"query": {
153+
"script_score": {
154+
"query": {
155+
"exists": {
156+
"field": "vector"
157+
}
158+
},
159+
"script": {
160+
"source": "hamming(params.query, 'vector')",
161+
"params": {
162+
"query": [4, 5, 6]
163+
}
164+
}
165+
}
166+
}
167+
}
168+
""");
169+
Map<String, Object> response = search(searchRequest);
170+
assertThat(extractValue(response, "hits.total.value"), equalTo(7));
171+
List<Map<String, Object>> hits = extractValue(response, "hits.hits");
172+
assertThat(hits.get(0).get("_id"), equalTo("3"));
173+
assertThat((double) hits.get(0).get("_score"), closeTo(8.0, 0.0001));
174+
175+
// search with knn
176+
searchRequest = new Request("POST", "/" + FLAT_BIT_INDEX_NAME + "/_search");
177+
searchRequest.setJsonEntity("""
178+
{
179+
"knn": {
180+
"field": "vector",
181+
"query_vector": [4, 5, 6],
182+
"k": 2,
183+
"num_candidates": 5
184+
}
185+
}
186+
""");
187+
response = search(searchRequest);
188+
assertThat(extractValue(response, "hits.total.value"), equalTo(2));
189+
hits = extractValue(response, "hits.hits");
190+
assertThat(hits.get(0).get("_id"), equalTo("1"));
191+
assertThat((double) hits.get(0).get("_score"), closeTo(0.833333, 0.0001));
192+
}
193+
}
47194

48195
public void testScriptByteVectorSearch() throws Exception {
49196
assumeTrue("byte vector search is not supported on this version", oldClusterHasFeature(BYTE_VECTOR_SEARCH_TEST_FEATURE));

0 commit comments

Comments
 (0)