Skip to content

Commit d854513

Browse files
thecoopthecoop
committed
Create a generic quantized vectors format that supports direct IO
1 parent 38c05ea commit d854513

File tree

20 files changed

+715
-183
lines changed

20 files changed

+715
-183
lines changed

‎rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,68 @@ setup:
338338
- match: { hits.hits.1._score: $rescore_score1 }
339339
- match: { hits.hits.2._score: $rescore_score2 }
340340
---
341+
"Test index configured rescore vector with on-disk rescoring":
342+
- requires:
343+
cluster_features: ["mapper.vectors.bbq_hnsw_on_disk_rescoring"]
344+
reason: Needs on_disk_rescoring feature
345+
- skip:
346+
features: "headers"
347+
- do:
348+
indices.create:
349+
index: bbq_on_disk_rescore_hnsw
350+
body:
351+
settings:
352+
index:
353+
number_of_shards: 1
354+
mappings:
355+
properties:
356+
vector:
357+
type: dense_vector
358+
dims: 64
359+
index: true
360+
similarity: max_inner_product
361+
index_options:
362+
type: bbq_hnsw
363+
on_disk_rescore: true
364+
rescore_vector:
365+
oversample: 1.5
366+
367+
- do:
368+
bulk:
369+
index: bbq_on_disk_rescore_hnsw
370+
refresh: true
371+
body: |
372+
{ "index": {"_id": "1"}}
373+
{ "vector": [0.077, 0.32 , -0.205, 0.63 , 0.032, 0.201, 0.167, -0.313, 0.176, 0.531, -0.375, 0.334, -0.046, 0.078, -0.349, 0.272, 0.307, -0.083, 0.504, 0.255, -0.404, 0.289, -0.226, -0.132, -0.216, 0.49 , 0.039, 0.507, -0.307, 0.107, 0.09 , -0.265, -0.285, 0.336, -0.272, 0.369, -0.282, 0.086, -0.132, 0.475, -0.224, 0.203, 0.439, 0.064, 0.246, -0.396, 0.297, 0.242, -0.028, 0.321, -0.022, -0.009, -0.001 , 0.031, -0.533, 0.45, -0.683, 1.331, 0.194, -0.157, -0.1 , -0.279, -0.098, -0.176] }
374+
{ "index": {"_id": "2"}}
375+
{ "vector": [0.196, 0.514, 0.039, 0.555, -0.042, 0.242, 0.463, -0.348, -0.08 , 0.442, -0.067, -0.05 , -0.001, 0.298, -0.377, 0.048, 0.307, 0.159, 0.278, 0.119, -0.057, 0.333, -0.289, -0.438, -0.014, 0.361, -0.169, 0.292, -0.229, 0.123, 0.031, -0.138, -0.139, 0.315, -0.216, 0.322, -0.445, -0.059, 0.071, 0.429, -0.602, -0.142, 0.11 , 0.192, 0.259, -0.241, 0.181, -0.166, 0.082, 0.107, -0.05 , 0.155, 0.011, 0.161, -0.486, 0.569, -0.489, 0.901, 0.208, 0.011, -0.209, -0.153, -0.27 , -0.013] }
376+
{ "index": {"_id": "3"}}
377+
{ "vector": [0.196, 0.514, 0.039, 0.555, -0.042, 0.242, 0.463, -0.348, -0.08 , 0.442, -0.067, -0.05 , -0.001, 0.298, -0.377, 0.048, 0.307, 0.159, 0.278, 0.119, -0.057, 0.333, -0.289, -0.438, -0.014, 0.361, -0.169, 0.292, -0.229, 0.123, 0.031, -0.138, -0.139, 0.315, -0.216, 0.322, -0.445, -0.059, 0.071, 0.429, -0.602, -0.142, 0.11 , 0.192, 0.259, -0.241, 0.181, -0.166, 0.082, 0.107, -0.05 , 0.155, 0.011, 0.161, -0.486, 0.569, -0.489, 0.901, 0.208, 0.011, -0.209, -0.153, -0.27 , -0.013] }
378+
- do:
379+
headers:
380+
Content-Type: application/json
381+
search:
382+
rest_total_hits_as_int: true
383+
index: bbq_on_disk_rescore_hnsw
384+
body:
385+
knn:
386+
field: vector
387+
query_vector: [0.128, 0.067, -0.08 , 0.395, -0.11 , -0.259, 0.473, -0.393,
388+
0.292, 0.571, -0.491, 0.444, -0.288, 0.198, -0.343, 0.015,
389+
0.232, 0.088, 0.228, 0.151, -0.136, 0.236, -0.273, -0.259,
390+
-0.217, 0.359, -0.207, 0.352, -0.142, 0.192, -0.061, -0.17 ,
391+
-0.343, 0.189, -0.221, 0.32 , -0.301, -0.1 , 0.005, 0.232,
392+
-0.344, 0.136, 0.252, 0.157, -0.13 , -0.244, 0.193, -0.034,
393+
-0.12 , -0.193, -0.102, 0.252, -0.185, -0.167, -0.575, 0.582,
394+
-0.426, 0.983, 0.212, 0.204, 0.03 , -0.276, -0.425, -0.158]
395+
k: 3
396+
num_candidates: 3
397+
398+
- match: { hits.total: 3 }
399+
- set: { hits.hits.0._score: rescore_score0 }
400+
- set: { hits.hits.1._score: rescore_score1 }
401+
- set: { hits.hits.2._score: rescore_score2 }
402+
---
341403
"Test index configured rescore vector updateable and settable to 0":
342404
- requires:
343405
cluster_features: ["mapper.dense_vector.rescore_zero_vector"]

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

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.index.store;
11+
12+
import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
13+
14+
import org.apache.logging.log4j.Level;
15+
import org.apache.lucene.misc.store.DirectIODirectory;
16+
import org.apache.lucene.store.Directory;
17+
import org.apache.lucene.store.FSDirectory;
18+
import org.apache.lucene.store.IOContext;
19+
import org.apache.lucene.store.IndexOutput;
20+
import org.apache.lucene.tests.util.LuceneTestCase;
21+
import org.elasticsearch.common.settings.Settings;
22+
import org.elasticsearch.core.Strings;
23+
import org.elasticsearch.plugins.Plugin;
24+
import org.elasticsearch.search.vectors.KnnSearchBuilder;
25+
import org.elasticsearch.search.vectors.VectorData;
26+
import org.elasticsearch.test.ESIntegTestCase;
27+
import org.elasticsearch.test.ESTestCase;
28+
import org.elasticsearch.test.InternalSettingsPlugin;
29+
import org.elasticsearch.test.MockLog;
30+
import org.elasticsearch.test.junit.annotations.TestLogging;
31+
import org.junit.BeforeClass;
32+
33+
import java.io.IOException;
34+
import java.nio.file.Path;
35+
import java.util.Collection;
36+
import java.util.List;
37+
import java.util.Map;
38+
import java.util.OptionalLong;
39+
import java.util.stream.IntStream;
40+
41+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
42+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
43+
import static org.hamcrest.Matchers.equalTo;
44+
import static org.hamcrest.Matchers.is;
45+
46+
@LuceneTestCase.SuppressCodecs("*") // only use our own codecs
47+
@ESTestCase.WithoutEntitlements // requires entitlement delegation ES-10920
48+
public class DirectIOIT extends ESIntegTestCase {
49+
50+
private static boolean SUPPORTED;
51+
52+
@BeforeClass
53+
public static void checkSupported() {
54+
Path path = createTempDir("directIOProbe");
55+
try (Directory dir = open(path); IndexOutput out = dir.createOutput("out", IOContext.DEFAULT)) {
56+
out.writeString("test");
57+
SUPPORTED = true;
58+
} catch (IOException e) {
59+
SUPPORTED = false;
60+
}
61+
}
62+
63+
static DirectIODirectory open(Path path) throws IOException {
64+
return new DirectIODirectory(FSDirectory.open(path)) {
65+
@Override
66+
protected boolean useDirectIO(String name, IOContext context, OptionalLong fileLength) {
67+
return true;
68+
}
69+
};
70+
}
71+
72+
private final String type;
73+
74+
@ParametersFactory
75+
public static Iterable<Object[]> parameters() {
76+
return List.<Object[]>of(new Object[] { "bbq_hnsw" });
77+
}
78+
79+
public DirectIOIT(String type) {
80+
this.type = type;
81+
}
82+
83+
@Override
84+
protected Collection<Class<? extends Plugin>> nodePlugins() {
85+
return List.of(InternalSettingsPlugin.class);
86+
}
87+
88+
private String indexVectors(boolean directIO) {
89+
String indexName = "test-vectors-" + directIO;
90+
assertAcked(
91+
prepareCreate(indexName).setSettings(Settings.builder().put(InternalSettingsPlugin.USE_COMPOUND_FILE.getKey(), false))
92+
.setMapping(Strings.format("""
93+
{
94+
"properties": {
95+
"fooVector": {
96+
"type": "dense_vector",
97+
"dims": 64,
98+
"element_type": "float",
99+
"index": true,
100+
"similarity": "l2_norm",
101+
"index_options": {
102+
"type": "%s",
103+
"on_disk_rescore": %s
104+
}
105+
}
106+
}
107+
}
108+
""", type, directIO))
109+
);
110+
ensureGreen(indexName);
111+
112+
for (int i = 0; i < 1000; i++) {
113+
indexDoc(indexName, Integer.toString(i), "fooVector", IntStream.range(0, 64).mapToDouble(d -> randomFloat()).toArray());
114+
}
115+
refresh();
116+
assertBBQIndexType(indexName, type); // test assertion to ensure that the correct index type is being used
117+
return indexName;
118+
}
119+
120+
@SuppressWarnings("unchecked")
121+
static void assertBBQIndexType(String indexName, String type) {
122+
var response = indicesAdmin().prepareGetFieldMappings(indexName).setFields("fooVector").get();
123+
var map = (Map<String, Object>) response.fieldMappings(indexName, "fooVector").sourceAsMap().get("fooVector");
124+
assertThat((String) ((Map<String, Object>) map.get("index_options")).get("type"), is(equalTo(type)));
125+
}
126+
127+
@TestLogging(value = "org.elasticsearch.index.store.FsDirectoryFactory:DEBUG", reason = "to capture trace logging for direct IO")
128+
public void testDirectIOUsed() {
129+
try (MockLog mockLog = MockLog.capture(FsDirectoryFactory.class)) {
130+
// we're just looking for some evidence direct IO is used (or not)
131+
MockLog.LoggingExpectation expectation = SUPPORTED
132+
? new MockLog.PatternSeenEventExpectation(
133+
"Direct IO used",
134+
FsDirectoryFactory.class.getCanonicalName(),
135+
Level.DEBUG,
136+
"Opening .*\\.vec with direct IO"
137+
)
138+
: new MockLog.PatternSeenEventExpectation(
139+
"Direct IO not used",
140+
FsDirectoryFactory.class.getCanonicalName(),
141+
Level.DEBUG,
142+
"Could not open .*\\.vec with direct IO"
143+
);
144+
mockLog.addExpectation(expectation);
145+
146+
String indexName = indexVectors(true);
147+
148+
// do a search
149+
var knn = List.of(new KnnSearchBuilder("fooVector", new VectorData(null, new byte[64]), 10, 20, 10f, null, null));
150+
assertHitCount(prepareSearch(indexName).setKnnSearch(knn), 10);
151+
mockLog.assertAllExpectationsMatched();
152+
}
153+
}
154+
155+
@TestLogging(value = "org.elasticsearch.index.store.FsDirectoryFactory:DEBUG", reason = "to capture trace logging for direct IO")
156+
public void testDirectIONotUsed() {
157+
try (MockLog mockLog = MockLog.capture(FsDirectoryFactory.class)) {
158+
// nothing about direct IO should be logged at all
159+
MockLog.LoggingExpectation expectation = SUPPORTED
160+
? new MockLog.PatternNotSeenEventExpectation(
161+
"Direct IO used",
162+
FsDirectoryFactory.class.getCanonicalName(),
163+
Level.DEBUG,
164+
"Opening .*\\.vec with direct IO"
165+
)
166+
: new MockLog.PatternNotSeenEventExpectation(
167+
"Direct IO not used",
168+
FsDirectoryFactory.class.getCanonicalName(),
169+
Level.DEBUG,
170+
"Could not open .*\\.vec with direct IO"
171+
);
172+
mockLog.addExpectation(expectation);
173+
174+
String indexName = indexVectors(false);
175+
176+
// do a search
177+
var knn = List.of(new KnnSearchBuilder("fooVector", new VectorData(null, new byte[64]), 10, 20, 10f, null, null));
178+
assertHitCount(prepareSearch(indexName).setKnnSearch(knn), 10);
179+
mockLog.assertAllExpectationsMatched();
180+
}
181+
}
182+
183+
@Override
184+
protected boolean addMockFSIndexStore() {
185+
return false; // we require to always use the "real" hybrid directory
186+
}
187+
}

‎server/src/main/java/module-info.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -464,6 +464,7 @@
464464
org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat,
465465
org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat,
466466
org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat,
467+
org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat,
467468
org.elasticsearch.index.codec.vectors.es93.ES93HnswBinaryQuantizedVectorsFormat;
468469

469470
provides org.apache.lucene.codecs.Codec

‎server/src/main/java/org/elasticsearch/index/codec/vectors/DirectIOCapableFlatVectorsFormat.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,5 +19,10 @@ protected DirectIOCapableFlatVectorsFormat(String name) {
1919
super(name);
2020
}
2121

22+
@Override
23+
public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
24+
return fieldsReader(state, false);
25+
}
26+
2227
public abstract FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectIO) throws IOException;
2328
}

‎server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93BinaryQuantizedVectorsFormat.java‎

Lines changed: 34 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsWriter;
3333

3434
import java.io.IOException;
35+
import java.util.Map;
3536

3637
/**
3738
* Copied from Lucene, replace with Lucene's implementation sometime after Lucene 10
@@ -86,19 +87,33 @@
8687
* <li>The sparse vector information, if required, mapping vector ordinal to doc ID
8788
* </ul>
8889
*/
89-
public class ES93BinaryQuantizedVectorsFormat extends DirectIOCapableFlatVectorsFormat {
90+
public class ES93BinaryQuantizedVectorsFormat extends ES93GenericFlatVectorsFormat {
9091

9192
public static final String NAME = "ES93BinaryQuantizedVectorsFormat";
9293

93-
private final DirectIOCapableLucene99FlatVectorsFormat rawVectorFormat;
94+
private static final DirectIOCapableFlatVectorsFormat rawVectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
95+
FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
96+
);
97+
98+
private static final Map<String, DirectIOCapableFlatVectorsFormat> supportedFormats = Map.of(
99+
rawVectorFormat.getName(),
100+
rawVectorFormat
101+
);
94102

95103
private static final ES818BinaryFlatVectorsScorer scorer = new ES818BinaryFlatVectorsScorer(
96104
FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
97105
);
98106

107+
private final boolean useDirectIO;
108+
99109
public ES93BinaryQuantizedVectorsFormat() {
100110
super(NAME);
101-
rawVectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
111+
this.useDirectIO = false;
112+
}
113+
114+
public ES93BinaryQuantizedVectorsFormat(boolean useDirectIO) {
115+
super(NAME);
116+
this.useDirectIO = useDirectIO;
102117
}
103118

104119
@Override
@@ -107,17 +122,27 @@ protected FlatVectorsScorer flatVectorsScorer() {
107122
}
108123

109124
@Override
110-
public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
111-
return new ES818BinaryQuantizedVectorsWriter(scorer, rawVectorFormat.fieldsWriter(state), state);
125+
protected boolean useDirectIOReads() {
126+
return useDirectIO;
112127
}
113128

114129
@Override
115-
public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
116-
return new ES818BinaryQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state), scorer);
130+
protected DirectIOCapableFlatVectorsFormat writeFlatVectorsFormat() {
131+
return rawVectorFormat;
117132
}
118133

119134
@Override
120-
public FlatVectorsReader fieldsReader(SegmentReadState state, boolean useDirectIO) throws IOException {
121-
return new ES818BinaryQuantizedVectorsReader(state, rawVectorFormat.fieldsReader(state, useDirectIO), scorer);
135+
protected Map<String, DirectIOCapableFlatVectorsFormat> supportedReadFlatVectorsFormats() {
136+
return supportedFormats;
137+
}
138+
139+
@Override
140+
public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
141+
return new ES818BinaryQuantizedVectorsWriter(scorer, super.fieldsWriter(state), state);
142+
}
143+
144+
@Override
145+
public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
146+
return new ES818BinaryQuantizedVectorsReader(state, super.fieldsReader(state), scorer);
122147
}
123148
}

0 commit comments

Comments
 (0)