Skip to content

Commit 0d51a86

Browse files
mayya-sharipovamayya-sharipova
committed
And int8_hswn option
1 parent 7e0cc3d commit 0d51a86

File tree

10 files changed

+229
-29
lines changed

10 files changed

+229
-29
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -482,7 +482,7 @@
482482
exports org.elasticsearch.plugins.internal.rewriter to org.elasticsearch.inference;
483483
exports org.elasticsearch.lucene.util.automaton;
484484
exports org.elasticsearch.index.codec.perfield;
485-
exports org.elasticsearch.index.codec.vectors to org.elasticsearch.test.knn;
485+
exports org.elasticsearch.index.codec.vectors to org.elasticsearch.test.knn, org.elasticsearch.gpu;
486486
exports org.elasticsearch.index.codec.vectors.es818 to org.elasticsearch.test.knn;
487487
exports org.elasticsearch.inference.telemetry;
488488
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2082,6 +2082,18 @@ boolean isFlat() {
20822082
return false;
20832083
}
20842084

2085+
public int m() {
2086+
return m;
2087+
}
2088+
2089+
public int efConstruction() {
2090+
return efConstruction;
2091+
}
2092+
2093+
public Float confidenceInterval() {
2094+
return confidenceInterval;
2095+
}
2096+
20852097
@Override
20862098
public String toString() {
20872099
return "{type="

‎x-pack/plugin/gpu/src/main/java/module-info.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
* 2.0.
66
*/
77

8-
/** Provides GPU-accelerated support for vector search. */
8+
/** Provides GPU-accelerated support for vector indexing. */
99
module org.elasticsearch.gpu {
1010
requires org.elasticsearch.logging;
1111
requires org.apache.lucene.core;
@@ -16,6 +16,9 @@
1616

1717
exports org.elasticsearch.xpack.gpu.codec;
1818

19-
provides org.apache.lucene.codecs.KnnVectorsFormat with org.elasticsearch.xpack.gpu.codec.ESGpuHnswVectorsFormat;
2019
provides org.elasticsearch.features.FeatureSpecification with org.elasticsearch.xpack.gpu.GPUFeatures;
20+
provides org.apache.lucene.codecs.KnnVectorsFormat
21+
with
22+
org.elasticsearch.xpack.gpu.codec.ESGpuHnswVectorsFormat,
23+
org.elasticsearch.xpack.gpu.codec.ESGpuHnswSQVectorsFormat;
2124
}

‎x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/GPUPlugin.java‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.elasticsearch.index.mapper.vectors.VectorsFormatProvider;
1515
import org.elasticsearch.plugins.MapperPlugin;
1616
import org.elasticsearch.plugins.Plugin;
17+
import org.elasticsearch.xpack.gpu.codec.ESGpuHnswSQVectorsFormat;
1718
import org.elasticsearch.xpack.gpu.codec.ESGpuHnswVectorsFormat;
1819

1920
public class GPUPlugin extends Plugin implements MapperPlugin {
@@ -28,10 +29,7 @@ public VectorsFormatProvider getVectorsFormatProvider() {
2829
if (gpuMode == IndexSettings.GpuMode.TRUE) {
2930
if (vectorIndexTypeSupported(indexOptions.getType()) == false) {
3031
throw new IllegalArgumentException(
31-
"[index.vectors.indexing.use_gpu] was set to [true], but GPU vector indexing is only supported "
32-
+ "for [hnsw] index_options.type, got: ["
33-
+ indexOptions.getType()
34-
+ "]"
32+
"[index.vectors.indexing.use_gpu] doesn't support [index_options.type] of [" + indexOptions.getType() + "]."
3533
);
3634
}
3735
if (GPUSupport.isSupported(true) == false) {
@@ -52,7 +50,7 @@ && vectorIndexTypeSupported(indexOptions.getType())
5250
}
5351

5452
private boolean vectorIndexTypeSupported(DenseVectorFieldMapper.VectorIndexType type) {
55-
return type == DenseVectorFieldMapper.VectorIndexType.HNSW;
53+
return type == DenseVectorFieldMapper.VectorIndexType.HNSW || type == DenseVectorFieldMapper.VectorIndexType.INT8_HNSW;
5654
}
5755

5856
private static KnnVectorsFormat getVectorsFormat(DenseVectorFieldMapper.DenseVectorIndexOptions indexOptions) {
@@ -63,6 +61,19 @@ private static KnnVectorsFormat getVectorsFormat(DenseVectorFieldMapper.DenseVec
6361
efConstruction = ESGpuHnswVectorsFormat.DEFAULT_BEAM_WIDTH; // default value for GPU graph construction is 128
6462
}
6563
return new ESGpuHnswVectorsFormat(hnswIndexOptions.m(), efConstruction);
64+
} else if (indexOptions.getType() == DenseVectorFieldMapper.VectorIndexType.INT8_HNSW) {
65+
DenseVectorFieldMapper.Int8HnswIndexOptions int8HnswIndexOptions = (DenseVectorFieldMapper.Int8HnswIndexOptions) indexOptions;
66+
int efConstruction = int8HnswIndexOptions.efConstruction();
67+
if (efConstruction == HnswGraphBuilder.DEFAULT_BEAM_WIDTH) {
68+
efConstruction = ESGpuHnswVectorsFormat.DEFAULT_BEAM_WIDTH; // default value for GPU graph construction is 128
69+
}
70+
return new ESGpuHnswSQVectorsFormat(
71+
int8HnswIndexOptions.m(),
72+
efConstruction,
73+
int8HnswIndexOptions.confidenceInterval(),
74+
7,
75+
false
76+
);
6677
} else {
6778
throw new IllegalArgumentException(
6879
"GPU vector indexing is not supported on this vector type: [" + indexOptions.getType() + "]"

‎x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswSQVectorsFormat.java‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
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; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.gpu.codec;
9+
10+
import org.apache.lucene.codecs.KnnVectorsFormat;
11+
import org.apache.lucene.codecs.KnnVectorsReader;
12+
import org.apache.lucene.codecs.KnnVectorsWriter;
13+
import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
14+
import org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader;
15+
import org.apache.lucene.index.SegmentReadState;
16+
import org.apache.lucene.index.SegmentWriteState;
17+
import org.elasticsearch.index.codec.vectors.ES814ScalarQuantizedVectorsFormat;
18+
19+
import java.io.IOException;
20+
21+
import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT;
22+
import static org.elasticsearch.xpack.gpu.codec.ESGpuHnswVectorsFormat.DEFAULT_BEAM_WIDTH;
23+
import static org.elasticsearch.xpack.gpu.codec.ESGpuHnswVectorsFormat.DEFAULT_MAX_CONN;
24+
25+
/**
26+
* Codec format for GPU-accelerated scalar quantized HNSW vector indexes.
27+
* HNSW graph is built on GPU, while scalar quantization and search is performed on CPU.
28+
*/
29+
public class ESGpuHnswSQVectorsFormat extends KnnVectorsFormat {
30+
public static final String NAME = "ESGPUHnswScalarQuantizedVectorsFormat";
31+
static final int MAXIMUM_MAX_CONN = 512;
32+
static final int MAXIMUM_BEAM_WIDTH = 3200;
33+
private final int maxConn;
34+
private final int beamWidth;
35+
36+
/** The format for storing, reading, merging vectors on disk */
37+
private final FlatVectorsFormat flatVectorsFormat;
38+
final CuVSResourceManager cuVSResourceManager;
39+
40+
public ESGpuHnswSQVectorsFormat() {
41+
this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, null, 7, false);
42+
}
43+
44+
public ESGpuHnswSQVectorsFormat(int maxConn, int beamWidth, Float confidenceInterval, int bits, boolean compress) {
45+
super(NAME);
46+
this.cuVSResourceManager = CuVSResourceManager.pooling();
47+
if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) {
48+
throw new IllegalArgumentException(
49+
"maxConn must be positive and less than or equal to " + MAXIMUM_MAX_CONN + "; maxConn=" + maxConn
50+
);
51+
}
52+
if (beamWidth <= 0 || beamWidth > MAXIMUM_BEAM_WIDTH) {
53+
throw new IllegalArgumentException(
54+
"beamWidth must be positive and less than or equal to " + MAXIMUM_BEAM_WIDTH + "; beamWidth=" + beamWidth
55+
);
56+
}
57+
this.maxConn = maxConn;
58+
this.beamWidth = beamWidth;
59+
this.flatVectorsFormat = new ES814ScalarQuantizedVectorsFormat(confidenceInterval, bits, compress);
60+
}
61+
62+
@Override
63+
public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
64+
return new ESGpuHnswVectorsWriter(cuVSResourceManager, state, maxConn, beamWidth, flatVectorsFormat.fieldsWriter(state));
65+
}
66+
67+
@Override
68+
public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
69+
return new Lucene99HnswVectorsReader(state, flatVectorsFormat.fieldsReader(state));
70+
}
71+
72+
@Override
73+
public int getMaxDimensions(String fieldName) {
74+
return MAX_DIMS_COUNT;
75+
}
76+
77+
@Override
78+
public String toString() {
79+
return NAME
80+
+ "(name="
81+
+ NAME
82+
+ ", maxConn="
83+
+ maxConn
84+
+ ", beamWidth="
85+
+ beamWidth
86+
+ ", flatVectorFormat="
87+
+ flatVectorsFormat
88+
+ ")";
89+
}
90+
}

‎x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsFormat.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,15 @@ public int getMaxDimensions(String fieldName) {
8181

8282
@Override
8383
public String toString() {
84-
return NAME + "(maxConn=" + maxConn + ", beamWidth=" + beamWidth + ", flatVectorFormat=" + flatVectorsFormat.getName() + ")";
84+
return NAME
85+
+ "(name="
86+
+ NAME
87+
+ ", maxConn="
88+
+ maxConn
89+
+ ", beamWidth="
90+
+ beamWidth
91+
+ ", flatVectorFormat="
92+
+ flatVectorsFormat.getName()
93+
+ ")";
8594
}
8695
}

‎x-pack/plugin/gpu/src/main/resources/META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11

22
org.elasticsearch.xpack.gpu.codec.ESGpuHnswVectorsFormat
3+
org.elasticsearch.xpack.gpu.codec.ESGpuHnswSQVectorsFormat

‎x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswSQVectorsFormatTests.java‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
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; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
package org.elasticsearch.xpack.gpu.codec;
8+
9+
import org.apache.lucene.codecs.Codec;
10+
import org.apache.lucene.index.VectorEncoding;
11+
import org.apache.lucene.index.VectorSimilarityFunction;
12+
import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase;
13+
import org.apache.lucene.tests.util.TestUtil;
14+
import org.elasticsearch.common.logging.LogConfigurator;
15+
import org.elasticsearch.xpack.gpu.GPUSupport;
16+
import org.junit.BeforeClass;
17+
18+
public class ESGpuHnswSQVectorsFormatTests extends BaseKnnVectorsFormatTestCase {
19+
20+
static {
21+
LogConfigurator.loadLog4jPlugins();
22+
LogConfigurator.configureESLogging(); // native access requires logging to be initialized
23+
}
24+
25+
@BeforeClass
26+
public static void beforeClass() {
27+
assumeTrue("cuvs not supported", GPUSupport.isSupported(false));
28+
}
29+
30+
static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new ESGpuHnswSQVectorsFormat());
31+
32+
@Override
33+
protected Codec getCodec() {
34+
return codec;
35+
}
36+
37+
@Override
38+
protected VectorSimilarityFunction randomSimilarity() {
39+
return VectorSimilarityFunction.values()[random().nextInt(VectorSimilarityFunction.values().length)];
40+
}
41+
42+
@Override
43+
protected VectorEncoding randomVectorEncoding() {
44+
return VectorEncoding.FLOAT32;
45+
}
46+
47+
@Override
48+
public void testRandomBytes() {
49+
// No bytes support
50+
}
51+
52+
@Override
53+
public void testSortedIndexBytes() {
54+
// No bytes support
55+
}
56+
57+
@Override
58+
public void testByteVectorScorerIteration() {
59+
// No bytes support
60+
}
61+
62+
@Override
63+
public void testEmptyByteVectorData() {
64+
// No bytes support
65+
}
66+
67+
@Override
68+
public void testMergingWithDifferentByteKnnFields() {
69+
// No bytes support
70+
}
71+
72+
@Override
73+
public void testMismatchedFields() {
74+
// No bytes support
75+
}
76+
}

‎x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ESGpuHnswVectorsFormatTests.java‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,6 @@
77
package org.elasticsearch.xpack.gpu.codec;
88

99
import org.apache.lucene.codecs.Codec;
10-
import org.apache.lucene.codecs.FilterCodec;
11-
import org.apache.lucene.codecs.KnnVectorsFormat;
1210
import org.apache.lucene.index.VectorEncoding;
1311
import org.apache.lucene.index.VectorSimilarityFunction;
1412
import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase;
@@ -76,15 +74,4 @@ public void testMismatchedFields() throws Exception {
7674
// No bytes support
7775
}
7876

79-
public void testToString() {
80-
FilterCodec customCodec = new FilterCodec("foo", Codec.getDefault()) {
81-
@Override
82-
public KnnVectorsFormat knnVectorsFormat() {
83-
return new ESGpuHnswVectorsFormat();
84-
}
85-
};
86-
String expectedPattern = "ESGpuHnswVectorsFormat(maxConn=16, beamWidth=128, flatVectorFormat=Lucene99FlatVectorsFormat)";
87-
assertEquals(expectedPattern, customCodec.knnVectorsFormat().toString());
88-
}
89-
9077
}

‎x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/GPUDenseVectorFieldMapperTests.java‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -39,31 +39,42 @@ protected Collection<Plugin> getPlugins() {
3939
return Collections.singletonList(plugin);
4040
}
4141

42-
public void testKnnESGPUHnswVectorsFormat() throws IOException {
42+
public void testESGPUHnswVectorsFormat() throws IOException {
43+
KnnVectorsFormat knnVectorsFormat = getKnnVectorsFormat("hnsw");
44+
String expectedStr = "ESGpuHnswVectorsFormat(name=ESGpuHnswVectorsFormat, "
45+
+ "maxConn=16, beamWidth=128, flatVectorFormat=Lucene99FlatVectorsFormat)";
46+
assertEquals(expectedStr, knnVectorsFormat.toString());
47+
}
48+
49+
public void testESGpuHnswScalarQuantizedVectorsFormat() throws IOException {
50+
KnnVectorsFormat knnVectorsFormat = getKnnVectorsFormat("int8_hnsw");
51+
String expectedStr = "ESGPUHnswScalarQuantizedVectorsFormat(name=ESGPUHnswScalarQuantizedVectorsFormat, "
52+
+ "maxConn=16, beamWidth=128, flatVectorFormat=ES814ScalarQuantizedVectorsFormat";
53+
assertTrue(knnVectorsFormat.toString().startsWith(expectedStr));
54+
}
55+
56+
private KnnVectorsFormat getKnnVectorsFormat(String indexOptionsType) throws IOException {
4357
final int dims = randomIntBetween(128, 4096);
4458
MapperService mapperService = createMapperService(fieldMapping(b -> {
4559
b.field("type", "dense_vector");
4660
b.field("dims", dims);
4761
b.field("index", true);
4862
b.field("similarity", "dot_product");
4963
b.startObject("index_options");
50-
b.field("type", "hnsw");
64+
b.field("type", indexOptionsType);
5165
b.endObject();
5266
}));
5367
CodecService codecService = new CodecService(mapperService, BigArrays.NON_RECYCLING_INSTANCE);
5468
Codec codec = codecService.codec("default");
55-
KnnVectorsFormat knnVectorsFormat;
5669
if (CodecService.ZSTD_STORED_FIELDS_FEATURE_FLAG) {
5770
assertThat(codec, instanceOf(PerFieldMapperCodec.class));
58-
knnVectorsFormat = ((PerFieldMapperCodec) codec).getKnnVectorsFormatForField("field");
71+
return ((PerFieldMapperCodec) codec).getKnnVectorsFormatForField("field");
5972
} else {
6073
if (codec instanceof CodecService.DeduplicateFieldInfosCodec deduplicateFieldInfosCodec) {
6174
codec = deduplicateFieldInfosCodec.delegate();
6275
}
6376
assertThat(codec, instanceOf(LegacyPerFieldMapperCodec.class));
64-
knnVectorsFormat = ((LegacyPerFieldMapperCodec) codec).getKnnVectorsFormatForField("field");
77+
return ((LegacyPerFieldMapperCodec) codec).getKnnVectorsFormatForField("field");
6578
}
66-
String expectedString = "ESGpuHnswVectorsFormat(maxConn=16, beamWidth=128, flatVectorFormat=Lucene99FlatVectorsFormat)";
67-
assertEquals(expectedString, knnVectorsFormat.toString());
6879
}
6980
}

0 commit comments

Comments
 (0)