Skip to content

Commit cf67fc8

Browse files
thecoopthecoop
committed
Add Direct IO support to DiskBBQ
1 parent 38c05ea commit cf67fc8

File tree

13 files changed

+383
-49
lines changed

13 files changed

+383
-49
lines changed

‎rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/46_knn_search_bbq_ivf.yml‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,98 @@ setup:
306306
- match: { hits.hits.0._score: $rescore_score0 }
307307
- match: { hits.hits.1._score: $rescore_score1 }
308308
- match: { hits.hits.2._score: $rescore_score2 }
309+
310+
---
311+
"Test index configured rescore vector with on-disk rescore":
312+
- requires:
313+
cluster_features: [ "mapper.vectors.diskbbq_on_disk_rescoring" ]
314+
reason: Needs on_disk_rescoring feature for DiskBBQ
315+
- skip:
316+
features: "headers"
317+
- do:
318+
indices.create:
319+
index: bbq_on_disk_rescore_ivf
320+
body:
321+
settings:
322+
index:
323+
number_of_shards: 1
324+
mappings:
325+
properties:
326+
vector:
327+
type: dense_vector
328+
dims: 64
329+
index: true
330+
similarity: max_inner_product
331+
index_options:
332+
type: bbq_disk
333+
on_disk_rescore: true
334+
rescore_vector:
335+
oversample: 1.5
336+
337+
- do:
338+
bulk:
339+
index: bbq_on_disk_rescore_ivf
340+
refresh: true
341+
body: |
342+
{ "index": {"_id": "1"}}
343+
{ "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] }
344+
{ "index": {"_id": "2"}}
345+
{ "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] }
346+
{ "index": {"_id": "3"}}
347+
{ "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] }
348+
349+
- do:
350+
headers:
351+
Content-Type: application/json
352+
search:
353+
rest_total_hits_as_int: true
354+
index: bbq_on_disk_rescore_ivf
355+
body:
356+
knn:
357+
field: vector
358+
query_vector: [ 0.128, 0.067, -0.08 , 0.395, -0.11 , -0.259, 0.473, -0.393,
359+
0.292, 0.571, -0.491, 0.444, -0.288, 0.198, -0.343, 0.015,
360+
0.232, 0.088, 0.228, 0.151, -0.136, 0.236, -0.273, -0.259,
361+
-0.217, 0.359, -0.207, 0.352, -0.142, 0.192, -0.061, -0.17 ,
362+
-0.343, 0.189, -0.221, 0.32 , -0.301, -0.1 , 0.005, 0.232,
363+
-0.344, 0.136, 0.252, 0.157, -0.13 , -0.244, 0.193, -0.034,
364+
-0.12 , -0.193, -0.102, 0.252, -0.185, -0.167, -0.575, 0.582,
365+
-0.426, 0.983, 0.212, 0.204, 0.03 , -0.276, -0.425, -0.158 ]
366+
k: 3
367+
num_candidates: 3
368+
369+
- match: { hits.total: 3 }
370+
- set: { hits.hits.0._score: rescore_score0 }
371+
- set: { hits.hits.1._score: rescore_score1 }
372+
- set: { hits.hits.2._score: rescore_score2 }
373+
374+
- do:
375+
headers:
376+
Content-Type: application/json
377+
search:
378+
rest_total_hits_as_int: true
379+
index: bbq_on_disk_rescore_ivf
380+
body:
381+
query:
382+
script_score:
383+
query: { match_all: { } }
384+
script:
385+
source: "double similarity = dotProduct(params.query_vector, 'vector'); return similarity < 0 ? 1 / (1 + -1 * similarity) : similarity + 1"
386+
params:
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+
396+
# Compare scores as hit IDs may change depending on how things are distributed
397+
- match: { hits.total: 3 }
398+
- match: { hits.hits.0._score: $rescore_score0 }
399+
- match: { hits.hits.1._score: $rescore_score1 }
400+
- match: { hits.hits.2._score: $rescore_score2 }
309401
---
310402
"Test index configured rescore vector updateable and settable to 0":
311403
- do:

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

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

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

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,11 @@
1313
import org.apache.lucene.codecs.KnnVectorsReader;
1414
import org.apache.lucene.codecs.KnnVectorsWriter;
1515
import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil;
16-
import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
17-
import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
1816
import org.apache.lucene.index.SegmentReadState;
1917
import org.apache.lucene.index.SegmentWriteState;
18+
import org.elasticsearch.index.codec.vectors.DirectIOCapableFlatVectorsFormat;
2019
import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
20+
import org.elasticsearch.index.codec.vectors.es93.DirectIOCapableLucene99FlatVectorsFormat;
2121

2222
import java.io.IOException;
2323
import java.util.Map;
@@ -55,12 +55,16 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
5555
static final String IVF_META_EXTENSION = "mivf";
5656

5757
public static final int VERSION_START = 0;
58-
public static final int VERSION_CURRENT = VERSION_START;
58+
public static final int VERSION_DIRECT_IO = 1;
59+
public static final int VERSION_CURRENT = VERSION_DIRECT_IO;
5960

60-
private static final Lucene99FlatVectorsFormat rawVectorFormat = new Lucene99FlatVectorsFormat(
61+
private static final DirectIOCapableFlatVectorsFormat rawVectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
6162
FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
6263
);
63-
private static final Map<String, FlatVectorsFormat> supportedFormats = Map.of(rawVectorFormat.getName(), rawVectorFormat);
64+
private static final Map<String, DirectIOCapableFlatVectorsFormat> supportedFormats = Map.of(
65+
rawVectorFormat.getName(),
66+
rawVectorFormat
67+
);
6468

6569
// This dynamically sets the cluster probe based on the `k` requested and the number of clusters.
6670
// useful when searching with 'efSearch' type parameters instead of requiring a specific ratio.
@@ -74,8 +78,13 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
7478

7579
private final int vectorPerCluster;
7680
private final int centroidsPerParentCluster;
81+
private final boolean useDirectIO;
7782

7883
public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster) {
84+
this(vectorPerCluster, centroidsPerParentCluster, false);
85+
}
86+
87+
public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster, boolean useDirectIO) {
7988
super(NAME);
8089
if (vectorPerCluster < MIN_VECTORS_PER_CLUSTER || vectorPerCluster > MAX_VECTORS_PER_CLUSTER) {
8190
throw new IllegalArgumentException(
@@ -99,6 +108,7 @@ public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentClu
99108
}
100109
this.vectorPerCluster = vectorPerCluster;
101110
this.centroidsPerParentCluster = centroidsPerParentCluster;
111+
this.useDirectIO = useDirectIO;
102112
}
103113

104114
/** Constructs a format using the given graph construction parameters and scalar quantization. */
@@ -109,8 +119,9 @@ public ES920DiskBBQVectorsFormat() {
109119
@Override
110120
public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
111121
return new ES920DiskBBQVectorsWriter(
112-
rawVectorFormat.getName(),
113122
state,
123+
rawVectorFormat.getName(),
124+
useDirectIO,
114125
rawVectorFormat.fieldsWriter(state),
115126
vectorPerCluster,
116127
centroidsPerParentCluster
@@ -119,10 +130,10 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
119130

120131
@Override
121132
public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
122-
return new ES920DiskBBQVectorsReader(state, f -> {
133+
return new ES920DiskBBQVectorsReader(state, (f, dio) -> {
123134
var format = supportedFormats.get(f);
124135
if (format == null) return null;
125-
return format.fieldsReader(state);
136+
return format.fieldsReader(state, dio);
126137
});
127138
}
128139

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,12 @@
99

1010
package org.elasticsearch.index.codec.vectors.diskbbq;
1111

12-
import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
1312
import org.apache.lucene.index.FieldInfo;
1413
import org.apache.lucene.index.SegmentReadState;
1514
import org.apache.lucene.index.VectorSimilarityFunction;
1615
import org.apache.lucene.search.KnnCollector;
1716
import org.apache.lucene.store.IndexInput;
1817
import org.apache.lucene.util.Bits;
19-
import org.apache.lucene.util.IOFunction;
2018
import org.apache.lucene.util.VectorUtil;
2119
import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
2220
import org.elasticsearch.index.codec.vectors.cluster.NeighborQueue;
@@ -40,7 +38,7 @@
4038
*/
4139
public class ES920DiskBBQVectorsReader extends IVFVectorsReader {
4240

43-
public ES920DiskBBQVectorsReader(SegmentReadState state, IOFunction<String, FlatVectorsReader> getFormatReader) throws IOException {
41+
ES920DiskBBQVectorsReader(SegmentReadState state, GetFormatReader getFormatReader) throws IOException {
4442
super(state, getFormatReader);
4543
}
4644

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,13 +52,14 @@ public class ES920DiskBBQVectorsWriter extends IVFVectorsWriter {
5252
private final int centroidsPerParentCluster;
5353

5454
public ES920DiskBBQVectorsWriter(
55-
String rawVectorFormatName,
5655
SegmentWriteState state,
56+
String rawVectorFormatName,
57+
boolean useDirectIOReads,
5758
FlatVectorsWriter rawVectorDelegate,
5859
int vectorPerCluster,
5960
int centroidsPerParentCluster
6061
) throws IOException {
61-
super(state, rawVectorFormatName, rawVectorDelegate);
62+
super(state, rawVectorFormatName, useDirectIOReads, rawVectorDelegate);
6263
this.vectorPerCluster = vectorPerCluster;
6364
this.centroidsPerParentCluster = centroidsPerParentCluster;
6465
}

0 commit comments

Comments
 (0)