Skip to content

Commit a73950d

Browse files
benwtrentbenwtrent
authored
Using the next disk bbq fmt for snapshot (#137917)
We want to be able to do end-to-end benchmarking for the new format changes. So, let's enable it by default for snapshots while we are working on it.
1 parent 57e9b32 commit a73950d

File tree

6 files changed

+156
-7
lines changed

6 files changed

+156
-7
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat;
3939
import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat;
4040
import org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat;
41+
import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
4142
import org.elasticsearch.logging.Level;
4243
import org.elasticsearch.logging.LogManager;
4344
import org.elasticsearch.logging.Logger;
@@ -131,6 +132,7 @@ static Codec createCodec(CmdLineArgs args) {
131132
encoding,
132133
args.ivfClusterSize(),
133134
ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER,
135+
DenseVectorFieldMapper.ElementType.FLOAT,
134136
args.onDiskRescore()
135137
);
136138
} else if (args.indexType() == IndexType.GPU_HNSW) {

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

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
import org.elasticsearch.index.codec.vectors.DirectIOCapableFlatVectorsFormat;
1919
import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
2020
import org.elasticsearch.index.codec.vectors.es93.DirectIOCapableLucene99FlatVectorsFormat;
21+
import org.elasticsearch.index.codec.vectors.es93.ES93BFloat16FlatVectorsFormat;
22+
import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
2123
import org.elasticsearch.simdvec.ESVectorUtil;
2224

2325
import java.io.IOException;
@@ -51,12 +53,17 @@ public class ESNextDiskBBQVectorsFormat extends KnnVectorsFormat {
5153
public static final int VERSION_START = 1;
5254
public static final int VERSION_CURRENT = VERSION_START;
5355

54-
private static final DirectIOCapableFlatVectorsFormat rawVectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
56+
private static final DirectIOCapableFlatVectorsFormat float32VectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
57+
FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
58+
);
59+
private static final DirectIOCapableFlatVectorsFormat bfloat16VectorFormat = new ES93BFloat16FlatVectorsFormat(
5560
FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
5661
);
5762
private static final Map<String, DirectIOCapableFlatVectorsFormat> supportedFormats = Map.of(
58-
rawVectorFormat.getName(),
59-
rawVectorFormat
63+
float32VectorFormat.getName(),
64+
float32VectorFormat,
65+
bfloat16VectorFormat.getName(),
66+
bfloat16VectorFormat
6067
);
6168

6269
public static final int DEFAULT_VECTORS_PER_CLUSTER = 384;
@@ -205,19 +212,21 @@ public static QuantEncoding fromId(int id) {
205212
private final int vectorPerCluster;
206213
private final int centroidsPerParentCluster;
207214
private final boolean useDirectIO;
215+
private final DirectIOCapableFlatVectorsFormat rawVectorFormat;
208216

209217
public ESNextDiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster) {
210218
this(QuantEncoding.ONE_BIT_4BIT_QUERY, vectorPerCluster, centroidsPerParentCluster);
211219
}
212220

213221
public ESNextDiskBBQVectorsFormat(QuantEncoding quantEncoding, int vectorPerCluster, int centroidsPerParentCluster) {
214-
this(quantEncoding, vectorPerCluster, centroidsPerParentCluster, false);
222+
this(quantEncoding, vectorPerCluster, centroidsPerParentCluster, DenseVectorFieldMapper.ElementType.FLOAT, false);
215223
}
216224

217225
public ESNextDiskBBQVectorsFormat(
218226
QuantEncoding quantEncoding,
219227
int vectorPerCluster,
220228
int centroidsPerParentCluster,
229+
DenseVectorFieldMapper.ElementType elementType,
221230
boolean useDirectIO
222231
) {
223232
super(NAME);
@@ -244,6 +253,11 @@ public ESNextDiskBBQVectorsFormat(
244253
this.vectorPerCluster = vectorPerCluster;
245254
this.centroidsPerParentCluster = centroidsPerParentCluster;
246255
this.quantEncoding = quantEncoding;
256+
this.rawVectorFormat = switch (elementType) {
257+
case FLOAT -> float32VectorFormat;
258+
case BFLOAT16 -> bfloat16VectorFormat;
259+
default -> throw new IllegalArgumentException("Unsupported element type " + elementType);
260+
};
247261
this.useDirectIO = useDirectIO;
248262
}
249263

‎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
@@ -40,6 +40,7 @@
4040
import org.apache.lucene.util.BitUtil;
4141
import org.apache.lucene.util.BytesRef;
4242
import org.apache.lucene.util.VectorUtil;
43+
import org.elasticsearch.Build;
4344
import org.elasticsearch.common.ParsingException;
4445
import org.elasticsearch.common.settings.Setting;
4546
import org.elasticsearch.common.xcontent.support.XContentMapValues;
@@ -54,6 +55,7 @@
5455
import org.elasticsearch.index.codec.vectors.ES815BitFlatVectorFormat;
5556
import org.elasticsearch.index.codec.vectors.ES815HnswBitVectorsFormat;
5657
import org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat;
58+
import org.elasticsearch.index.codec.vectors.diskbbq.next.ESNextDiskBBQVectorsFormat;
5759
import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat;
5860
import org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat;
5961
import org.elasticsearch.index.codec.vectors.es93.ES93BinaryQuantizedVectorsFormat;
@@ -2377,6 +2379,16 @@ static class BBQIVFIndexOptions extends QuantizedIndexOptions {
23772379

23782380
@Override
23792381
KnnVectorsFormat getVectorsFormat(ElementType elementType) {
2382+
assert elementType == ElementType.FLOAT || elementType == ElementType.BFLOAT16;
2383+
if (Build.current().isSnapshot()) {
2384+
return new ESNextDiskBBQVectorsFormat(
2385+
ESNextDiskBBQVectorsFormat.QuantEncoding.ONE_BIT_4BIT_QUERY,
2386+
clusterSize,
2387+
ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER,
2388+
elementType,
2389+
onDiskRescore
2390+
);
2391+
}
23802392
return new ES920DiskBBQVectorsFormat(
23812393
clusterSize,
23822394
ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER,

‎server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQBFloat16VectorsFormatTests.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,15 +49,15 @@ public void setUp() throws Exception {
4949
format = new ES920DiskBBQVectorsFormat(
5050
random().nextInt(2 * MIN_VECTORS_PER_CLUSTER, ES920DiskBBQVectorsFormat.MAX_VECTORS_PER_CLUSTER),
5151
random().nextInt(8, ES920DiskBBQVectorsFormat.MAX_CENTROIDS_PER_PARENT_CLUSTER),
52-
DenseVectorFieldMapper.ElementType.FLOAT,
52+
DenseVectorFieldMapper.ElementType.BFLOAT16,
5353
random().nextBoolean()
5454
);
5555
} else {
5656
// run with low numbers to force many clusters with parents
5757
format = new ES920DiskBBQVectorsFormat(
5858
random().nextInt(MIN_VECTORS_PER_CLUSTER, 2 * MIN_VECTORS_PER_CLUSTER),
5959
random().nextInt(MIN_CENTROIDS_PER_PARENT_CLUSTER, 8),
60-
DenseVectorFieldMapper.ElementType.FLOAT,
60+
DenseVectorFieldMapper.ElementType.BFLOAT16,
6161
random().nextBoolean()
6262
);
6363
}

‎server/src/test/java/org/elasticsearch/index/codec/vectors/diskbbq/next/ESNextDiskBBQBFloat16VectorsFormatTests.java‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
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.diskbbq.next;
11+
12+
import com.carrotsearch.randomizedtesting.generators.RandomPicks;
13+
14+
import org.apache.lucene.codecs.Codec;
15+
import org.apache.lucene.codecs.KnnVectorsFormat;
16+
import org.apache.lucene.codecs.KnnVectorsReader;
17+
import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat;
18+
import org.apache.lucene.index.CodecReader;
19+
import org.apache.lucene.index.LeafReader;
20+
import org.apache.lucene.index.VectorSimilarityFunction;
21+
import org.apache.lucene.tests.util.TestUtil;
22+
import org.elasticsearch.common.logging.LogConfigurator;
23+
import org.elasticsearch.index.codec.vectors.BaseBFloat16KnnVectorsFormatTestCase;
24+
import org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat;
25+
import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
26+
import org.junit.AssumptionViolatedException;
27+
import org.junit.Before;
28+
29+
import java.io.IOException;
30+
import java.util.List;
31+
32+
import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MIN_CENTROIDS_PER_PARENT_CLUSTER;
33+
import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.MIN_VECTORS_PER_CLUSTER;
34+
import static org.hamcrest.Matchers.anEmptyMap;
35+
import static org.hamcrest.Matchers.equalTo;
36+
37+
public class ESNextDiskBBQBFloat16VectorsFormatTests extends BaseBFloat16KnnVectorsFormatTestCase {
38+
39+
static {
40+
LogConfigurator.loadLog4jPlugins();
41+
LogConfigurator.configureESLogging(); // native access requires logging to be initialized
42+
}
43+
44+
private KnnVectorsFormat format;
45+
46+
@Before
47+
@Override
48+
public void setUp() throws Exception {
49+
ESNextDiskBBQVectorsFormat.QuantEncoding encoding = ESNextDiskBBQVectorsFormat.QuantEncoding.values()[random().nextInt(
50+
ESNextDiskBBQVectorsFormat.QuantEncoding.values().length
51+
)];
52+
if (rarely()) {
53+
format = new ESNextDiskBBQVectorsFormat(
54+
encoding,
55+
random().nextInt(2 * MIN_VECTORS_PER_CLUSTER, ES920DiskBBQVectorsFormat.MAX_VECTORS_PER_CLUSTER),
56+
random().nextInt(8, ES920DiskBBQVectorsFormat.MAX_CENTROIDS_PER_PARENT_CLUSTER),
57+
DenseVectorFieldMapper.ElementType.BFLOAT16,
58+
random().nextBoolean()
59+
);
60+
} else {
61+
// run with low numbers to force many clusters with parents
62+
format = new ESNextDiskBBQVectorsFormat(
63+
encoding,
64+
random().nextInt(MIN_VECTORS_PER_CLUSTER, 2 * MIN_VECTORS_PER_CLUSTER),
65+
random().nextInt(MIN_CENTROIDS_PER_PARENT_CLUSTER, 8),
66+
DenseVectorFieldMapper.ElementType.BFLOAT16,
67+
random().nextBoolean()
68+
);
69+
}
70+
super.setUp();
71+
}
72+
73+
@Override
74+
protected Codec getCodec() {
75+
return TestUtil.alwaysKnnVectorsFormat(format);
76+
}
77+
78+
@Override
79+
protected VectorSimilarityFunction randomSimilarity() {
80+
return RandomPicks.randomFrom(
81+
random(),
82+
List.of(
83+
VectorSimilarityFunction.DOT_PRODUCT,
84+
VectorSimilarityFunction.EUCLIDEAN,
85+
VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT
86+
)
87+
);
88+
}
89+
90+
@Override
91+
public void testSearchWithVisitedLimit() {
92+
throw new AssumptionViolatedException("ivf doesn't enforce visitation limit");
93+
}
94+
95+
@Override
96+
public void testAdvance() throws Exception {
97+
// TODO re-enable with hierarchical IVF, clustering as it is is flaky
98+
}
99+
100+
@Override
101+
protected void assertOffHeapByteSize(LeafReader r, String fieldName) throws IOException {
102+
var fieldInfo = r.getFieldInfos().fieldInfo(fieldName);
103+
104+
if (r instanceof CodecReader codecReader) {
105+
KnnVectorsReader knnVectorsReader = codecReader.getVectorReader();
106+
if (knnVectorsReader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) {
107+
knnVectorsReader = fieldsReader.getFieldReader(fieldName);
108+
}
109+
var offHeap = knnVectorsReader.getOffHeapByteSize(fieldInfo);
110+
long totalByteSize = offHeap.values().stream().mapToLong(Long::longValue).sum();
111+
// IVF doesn't report stats at the moment
112+
assertThat(offHeap, anEmptyMap());
113+
assertThat(totalByteSize, equalTo(0L));
114+
} else {
115+
throw new AssertionError("unexpected:" + r.getClass());
116+
}
117+
}
118+
}

‎server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapperTests.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.apache.lucene.util.BytesRef;
2525
import org.apache.lucene.util.IOConsumer;
2626
import org.apache.lucene.util.VectorUtil;
27+
import org.elasticsearch.Build;
2728
import org.elasticsearch.common.bytes.BytesReference;
2829
import org.elasticsearch.common.util.BigArrays;
2930
import org.elasticsearch.common.xcontent.XContentHelper;
@@ -2114,7 +2115,9 @@ public void testKnnBBQIVFVectorsFormat() throws IOException {
21142115
assertThat(codec, instanceOf(LegacyPerFieldMapperCodec.class));
21152116
knnVectorsFormat = ((LegacyPerFieldMapperCodec) codec).getKnnVectorsFormatForField("field");
21162117
}
2117-
String expectedString = "ES920DiskBBQVectorsFormat(vectorPerCluster=384)";
2118+
String expectedString = Build.current().isSnapshot()
2119+
? "ESNextDiskBBQVectorsFormat(vectorPerCluster=384)"
2120+
: "ES920DiskBBQVectorsFormat(vectorPerCluster=384)";
21182121
assertEquals(expectedString, knnVectorsFormat.toString());
21192122
}
21202123

0 commit comments

Comments
 (0)