Skip to content

Commit 5b8d499

Browse files
thecoopthecoop
committed
Create new formats for direct IO access
1 parent 9b02a48 commit 5b8d499

File tree

11 files changed

+110
-36
lines changed

11 files changed

+110
-36
lines changed

‎qa/vector/src/main/java/org/elasticsearch/test/knn/KnnIndexTester.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ static Codec createCodec(CmdLineArgs args) {
112112
if (args.indexType() == IndexType.FLAT) {
113113
format = new ES818BinaryQuantizedVectorsFormat();
114114
} else {
115-
format = new ES818HnswBinaryQuantizedVectorsFormat(args.hnswM(), args.hnswEfConstruction(), 1, false, null);
115+
format = new ES818HnswBinaryQuantizedVectorsFormat(args.hnswM(), args.hnswEfConstruction(), 1, null);
116116
}
117117
} else if (args.quantizeBits() < 32) {
118118
if (args.indexType() == IndexType.FLAT) {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -460,6 +460,8 @@
460460
org.elasticsearch.index.codec.vectors.es816.ES816HnswBinaryQuantizedVectorsFormat,
461461
org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat,
462462
org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat,
463+
org.elasticsearch.index.codec.vectors.es818.DirectIOES818BinaryQuantizedVectorsFormat,
464+
org.elasticsearch.index.codec.vectors.es818.DirectIOES818HnswBinaryQuantizedVectorsFormat,
463465
org.elasticsearch.index.codec.vectors.IVFVectorsFormat;
464466

465467
provides org.apache.lucene.codecs.Codec

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
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.codec.vectors.es818;
11+
12+
import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil;
13+
14+
/**
15+
* A variant of
16+
*/
17+
public class DirectIOES818BinaryQuantizedVectorsFormat extends ES818BinaryQuantizedVectorsFormat {
18+
19+
public static final String NAME = "DirectIOES818BinaryQuantizedVectorsFormat";
20+
21+
public DirectIOES818BinaryQuantizedVectorsFormat() {
22+
super(NAME, new DirectIOLucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer()));
23+
}
24+
}

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
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.codec.vectors.es818;
11+
12+
import java.util.concurrent.ExecutorService;
13+
14+
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
15+
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
16+
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_NUM_MERGE_WORKER;
17+
18+
public class DirectIOES818HnswBinaryQuantizedVectorsFormat extends ES818HnswBinaryQuantizedVectorsFormat {
19+
20+
public static final String NAME = "DirectIOES818HnswBinaryQuantizedVectorsFormat";
21+
22+
/** Constructs a format using default graph construction parameters */
23+
public DirectIOES818HnswBinaryQuantizedVectorsFormat() {
24+
this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, DEFAULT_NUM_MERGE_WORKER, null);
25+
}
26+
27+
/**
28+
* Constructs a format using the given graph construction parameters.
29+
*
30+
* @param maxConn the maximum number of connections to a node in the HNSW graph
31+
* @param beamWidth the size of the queue maintained during graph construction.
32+
*/
33+
public DirectIOES818HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth) {
34+
this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null);
35+
}
36+
37+
/**
38+
* Constructs a format using the given graph construction parameters and scalar quantization.
39+
*
40+
* @param maxConn the maximum number of connections to a node in the HNSW graph
41+
* @param beamWidth the size of the queue maintained during graph construction.
42+
* @param numMergeWorkers number of workers (threads) that will be used when doing merge. If
43+
* larger than 1, a non-null {@link ExecutorService} must be passed as mergeExec
44+
* @param mergeExec the {@link ExecutorService} that will be used by ALL vector writers that are
45+
* generated by this format to do the merge
46+
*/
47+
public DirectIOES818HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth, int numMergeWorkers, ExecutorService mergeExec) {
48+
super(NAME, maxConn, beamWidth, numMergeWorkers, new DirectIOES818BinaryQuantizedVectorsFormat(), mergeExec);
49+
}
50+
}

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

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -98,24 +98,20 @@ public class ES818BinaryQuantizedVectorsFormat extends FlatVectorsFormat {
9898
static final String VECTOR_DATA_EXTENSION = "veb";
9999
static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
100100

101-
private static final ES818BinaryFlatVectorsScorer scorer = new ES818BinaryFlatVectorsScorer(
101+
static final ES818BinaryFlatVectorsScorer scorer = new ES818BinaryFlatVectorsScorer(
102102
FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
103103
);
104104

105105
private final FlatVectorsFormat rawVectorFormat;
106106

107107
/** Creates a new instance with the default number of vectors per cluster. */
108108
public ES818BinaryQuantizedVectorsFormat() {
109-
this(false);
109+
this(NAME, new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer()));
110110
}
111111

112-
/** Creates a new instance with the default number of vectors per cluster,
113-
* and whether direct IO should be used to access raw vectors. */
114-
public ES818BinaryQuantizedVectorsFormat(boolean useDirectIO) {
115-
super(NAME);
116-
rawVectorFormat = useDirectIO
117-
? new DirectIOLucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer())
118-
: new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
112+
ES818BinaryQuantizedVectorsFormat(String name, FlatVectorsFormat rawVectorFormat) {
113+
super(name);
114+
this.rawVectorFormat = rawVectorFormat;
119115
}
120116

121117
@Override
@@ -135,6 +131,6 @@ public int getMaxDimensions(String fieldName) {
135131

136132
@Override
137133
public String toString() {
138-
return "ES818BinaryQuantizedVectorsFormat(name=" + NAME + ", flatVectorScorer=" + scorer + ")";
134+
return getName() + "(name=" + getName() + ", flatVectorScorer=" + scorer + ")";
139135
}
140136
}

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

Lines changed: 15 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ public class ES818HnswBinaryQuantizedVectorsFormat extends KnnVectorsFormat {
6969

7070
/** Constructs a format using default graph construction parameters */
7171
public ES818HnswBinaryQuantizedVectorsFormat() {
72-
this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, DEFAULT_NUM_MERGE_WORKER, false, null);
72+
this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, DEFAULT_NUM_MERGE_WORKER, null);
7373
}
7474

7575
/**
@@ -79,18 +79,7 @@ public ES818HnswBinaryQuantizedVectorsFormat() {
7979
* @param beamWidth the size of the queue maintained during graph construction.
8080
*/
8181
public ES818HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth) {
82-
this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, false, null);
83-
}
84-
85-
/**
86-
* Constructs a format using the given graph construction parameters.
87-
*
88-
* @param maxConn the maximum number of connections to a node in the HNSW graph
89-
* @param beamWidth the size of the queue maintained during graph construction.
90-
* @param useDirectIO whether direct IO should be used to access raw vectors
91-
*/
92-
public ES818HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth, boolean useDirectIO) {
93-
this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, useDirectIO, null);
82+
this(maxConn, beamWidth, DEFAULT_NUM_MERGE_WORKER, null);
9483
}
9584

9685
/**
@@ -103,14 +92,19 @@ public ES818HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth, boolean
10392
* @param mergeExec the {@link ExecutorService} that will be used by ALL vector writers that are
10493
* generated by this format to do the merge
10594
*/
106-
public ES818HnswBinaryQuantizedVectorsFormat(
95+
public ES818HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth, int numMergeWorkers, ExecutorService mergeExec) {
96+
this(NAME, maxConn, beamWidth, numMergeWorkers, new ES818BinaryQuantizedVectorsFormat(), mergeExec);
97+
}
98+
99+
ES818HnswBinaryQuantizedVectorsFormat(
100+
String name,
107101
int maxConn,
108102
int beamWidth,
109103
int numMergeWorkers,
110-
boolean useDirectIO,
104+
FlatVectorsFormat flatVectorsFormat,
111105
ExecutorService mergeExec
112106
) {
113-
super(NAME);
107+
super(name);
114108
if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) {
115109
throw new IllegalArgumentException(
116110
"maxConn must be positive and less than or equal to " + MAXIMUM_MAX_CONN + "; maxConn=" + maxConn
@@ -128,7 +122,7 @@ public ES818HnswBinaryQuantizedVectorsFormat(
128122
}
129123
this.numMergeWorkers = numMergeWorkers;
130124

131-
flatVectorsFormat = new ES818BinaryQuantizedVectorsFormat(useDirectIO);
125+
this.flatVectorsFormat = flatVectorsFormat;
132126

133127
if (mergeExec != null) {
134128
this.mergeExec = new TaskExecutor(mergeExec);
@@ -154,7 +148,10 @@ public int getMaxDimensions(String fieldName) {
154148

155149
@Override
156150
public String toString() {
157-
return "ES818HnswBinaryQuantizedVectorsFormat(name=ES818HnswBinaryQuantizedVectorsFormat, maxConn="
151+
return getName()
152+
+ "(name="
153+
+ getName()
154+
+ ", maxConn="
158155
+ maxConn
159156
+ ", beamWidth="
160157
+ beamWidth

‎server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656
import org.elasticsearch.index.codec.vectors.ES815BitFlatVectorFormat;
5757
import org.elasticsearch.index.codec.vectors.ES815HnswBitVectorsFormat;
5858
import org.elasticsearch.index.codec.vectors.IVFVectorsFormat;
59+
import org.elasticsearch.index.codec.vectors.es818.DirectIOES818HnswBinaryQuantizedVectorsFormat;
5960
import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat;
6061
import org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat;
6162
import org.elasticsearch.index.fielddata.FieldDataContext;
@@ -2199,7 +2200,9 @@ public BBQHnswIndexOptions(int m, int efConstruction, RescoreVector rescoreVecto
21992200
@Override
22002201
KnnVectorsFormat getVectorsFormat(ElementType elementType) {
22012202
assert elementType == ElementType.FLOAT;
2202-
return new ES818HnswBinaryQuantizedVectorsFormat(m, efConstruction, disableOffheapCacheRescoring);
2203+
return disableOffheapCacheRescoring
2204+
? new DirectIOES818HnswBinaryQuantizedVectorsFormat(m, efConstruction)
2205+
: new ES818HnswBinaryQuantizedVectorsFormat(m, efConstruction);
22032206
}
22042207

22052208
@Override

‎server/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,6 @@ org.elasticsearch.index.codec.vectors.es816.ES816BinaryQuantizedVectorsFormat
77
org.elasticsearch.index.codec.vectors.es816.ES816HnswBinaryQuantizedVectorsFormat
88
org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat
99
org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat
10+
org.elasticsearch.index.codec.vectors.es818.DirectIOES818BinaryQuantizedVectorsFormat
11+
org.elasticsearch.index.codec.vectors.es818.DirectIOES818HnswBinaryQuantizedVectorsFormat
1012
org.elasticsearch.index.codec.vectors.IVFVectorsFormat

‎server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818DirectIOBinaryQuantizedVectorsFormatTests.java‎ renamed to ‎server/src/test/java/org/elasticsearch/index/codec/vectors/es818/DirectIOES818BinaryQuantizedVectorsFormatTests.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,9 @@
3232
import java.util.Locale;
3333
import java.util.OptionalLong;
3434

35-
public class ES818DirectIOBinaryQuantizedVectorsFormatTests extends ES818BinaryQuantizedVectorsFormatTests {
35+
public class DirectIOES818BinaryQuantizedVectorsFormatTests extends ES818BinaryQuantizedVectorsFormatTests {
3636

37-
static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new ES818BinaryQuantizedVectorsFormat(true));
37+
static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new DirectIOES818BinaryQuantizedVectorsFormat());
3838

3939
@Override
4040
protected Codec getCodec() {

‎server/src/test/java/org/elasticsearch/index/codec/vectors/es818/ES818DirectIOHnswBinaryQuantizedVectorsFormatTests.java‎ renamed to ‎server/src/test/java/org/elasticsearch/index/codec/vectors/es818/DirectIOES818HnswBinaryQuantizedVectorsFormatTests.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,10 +35,10 @@
3535
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH;
3636
import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN;
3737

38-
public class ES818DirectIOHnswBinaryQuantizedVectorsFormatTests extends ES818HnswBinaryQuantizedVectorsFormatTests {
38+
public class DirectIOES818HnswBinaryQuantizedVectorsFormatTests extends ES818HnswBinaryQuantizedVectorsFormatTests {
3939

4040
static final Codec codec = TestUtil.alwaysKnnVectorsFormat(
41-
new ES818HnswBinaryQuantizedVectorsFormat(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, true)
41+
new DirectIOES818HnswBinaryQuantizedVectorsFormat(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH)
4242
);
4343

4444
@Override

0 commit comments

Comments
 (0)