Skip to content

Commit 19902ac

Browse files
thecoopthecoop
committed
Tests!
1 parent 391e2b7 commit 19902ac

File tree

8 files changed

+220
-28
lines changed

8 files changed

+220
-28
lines changed

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

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -91,16 +91,16 @@ public class ES93BinaryQuantizedVectorsFormat extends ES93GenericFlatVectorsForm
9191

9292
public static final String NAME = "ES93BinaryQuantizedVectorsFormat";
9393

94-
private static final DirectIOCapableFlatVectorsFormat rawVectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
94+
private static final DirectIOCapableFlatVectorsFormat float32VectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
9595
FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
9696
);
9797
private static final DirectIOCapableFlatVectorsFormat bfloat16VectorFormat = new ES93BFloat16FlatVectorsFormat(
9898
FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
9999
);
100100

101101
private static final Map<String, DirectIOCapableFlatVectorsFormat> supportedFormats = Map.of(
102-
rawVectorFormat.getName(),
103-
rawVectorFormat,
102+
float32VectorFormat.getName(),
103+
float32VectorFormat,
104104
bfloat16VectorFormat.getName(),
105105
bfloat16VectorFormat
106106
);
@@ -110,15 +110,16 @@ public class ES93BinaryQuantizedVectorsFormat extends ES93GenericFlatVectorsForm
110110
);
111111

112112
private final boolean useDirectIO;
113+
private final DirectIOCapableFlatVectorsFormat rawFormat;
113114

114115
public ES93BinaryQuantizedVectorsFormat() {
115-
super(NAME);
116-
this.useDirectIO = false;
116+
this(false, false);
117117
}
118118

119-
public ES93BinaryQuantizedVectorsFormat(boolean useDirectIO) {
119+
public ES93BinaryQuantizedVectorsFormat(boolean useDirectIO, boolean useBFloat16) {
120120
super(NAME);
121121
this.useDirectIO = useDirectIO;
122+
this.rawFormat = useBFloat16 ? bfloat16VectorFormat : float32VectorFormat;
122123
}
123124

124125
@Override
@@ -133,7 +134,7 @@ protected boolean useDirectIOReads() {
133134

134135
@Override
135136
protected DirectIOCapableFlatVectorsFormat writeFlatVectorsFormat() {
136-
return rawVectorFormat;
137+
return rawFormat;
137138
}
138139

139140
@Override

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,9 +62,9 @@ public ES93HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth) {
6262
* @param beamWidth the size of the queue maintained during graph construction.
6363
* @param useDirectIO whether to use direct IO when reading raw vectors
6464
*/
65-
public ES93HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth, boolean useDirectIO) {
65+
public ES93HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth, boolean useDirectIO, boolean useBFloat16) {
6666
super(NAME, maxConn, beamWidth);
67-
flatVectorsFormat = new ES93BinaryQuantizedVectorsFormat(useDirectIO);
67+
flatVectorsFormat = new ES93BinaryQuantizedVectorsFormat(useDirectIO, useBFloat16);
6868
}
6969

7070
/**
@@ -82,11 +82,12 @@ public ES93HnswBinaryQuantizedVectorsFormat(
8282
int maxConn,
8383
int beamWidth,
8484
boolean useDirectIO,
85+
boolean useBFloat16,
8586
int numMergeWorkers,
8687
ExecutorService mergeExec
8788
) {
8889
super(NAME, maxConn, beamWidth, numMergeWorkers, mergeExec);
89-
flatVectorsFormat = new ES93BinaryQuantizedVectorsFormat(useDirectIO);
90+
flatVectorsFormat = new ES93BinaryQuantizedVectorsFormat(useDirectIO, useBFloat16);
9091
}
9192

9293
@Override

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

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
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;
11+
12+
import java.util.regex.Matcher;
13+
import java.util.regex.Pattern;
14+
15+
import static org.hamcrest.Matchers.closeTo;
16+
17+
public class ES920DiskBBQBFloat16VectorsFormatTests extends ES920DiskBBQVectorsFormatTests {
18+
@Override
19+
boolean useBFloat16() {
20+
return true;
21+
}
22+
23+
@Override
24+
public void testEmptyByteVectorData() throws Exception {
25+
// no bytes
26+
}
27+
28+
@Override
29+
public void testMergingWithDifferentByteKnnFields() throws Exception {
30+
// no bytes
31+
}
32+
33+
@Override
34+
public void testByteVectorScorerIteration() throws Exception {
35+
// no bytes
36+
}
37+
38+
@Override
39+
public void testSortedIndexBytes() throws Exception {
40+
// no bytes
41+
}
42+
43+
@Override
44+
public void testMismatchedFields() throws Exception {
45+
// no bytes
46+
}
47+
48+
@Override
49+
public void testRandomBytes() throws Exception {
50+
// no bytes
51+
}
52+
53+
@Override
54+
public void testWriterRamEstimate() throws Exception {
55+
// estimate is different due to bfloat16
56+
}
57+
58+
@Override
59+
public void testRandom() throws Exception {
60+
AssertionError err = expectThrows(AssertionError.class, super::testRandom);
61+
assertFloatsWithinBounds(err);
62+
}
63+
64+
@Override
65+
public void testSparseVectors() throws Exception {
66+
AssertionError err = expectThrows(AssertionError.class, super::testSparseVectors);
67+
assertFloatsWithinBounds(err);
68+
}
69+
70+
@Override
71+
public void testVectorValuesReportCorrectDocs() throws Exception {
72+
AssertionError err = expectThrows(AssertionError.class, super::testVectorValuesReportCorrectDocs);
73+
assertFloatsWithinBounds(err);
74+
}
75+
76+
@Override
77+
public void testRandomWithUpdatesAndGraph() throws Exception {
78+
AssertionError err = expectThrows(AssertionError.class, super::testRandomWithUpdatesAndGraph);
79+
assertFloatsWithinBounds(err);
80+
}
81+
82+
private static final Pattern FLOAT_ASSERTION_FAILURE = Pattern.compile(".*expected:<([0-9.-]+)> but was:<([0-9.-]+)>");
83+
84+
private static void assertFloatsWithinBounds(AssertionError error) {
85+
Matcher m = FLOAT_ASSERTION_FAILURE.matcher(error.getMessage());
86+
if (m.matches() == false) {
87+
throw error; // nothing to do with us, just rethrow
88+
}
89+
90+
// numbers just need to be in the same vicinity
91+
double expected = Double.parseDouble(m.group(1));
92+
double actual = Double.parseDouble(m.group(2));
93+
double allowedError = expected * 0.01; // within 1%
94+
assertThat(error.getMessage(), actual, closeTo(expected, allowedError));
95+
}
96+
}

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

Lines changed: 5 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@
88
*/
99
package org.elasticsearch.index.codec.vectors.diskbbq;
1010

11-
import com.carrotsearch.randomizedtesting.annotations.Name;
12-
import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
1311
import com.carrotsearch.randomizedtesting.generators.RandomPicks;
1412

1513
import org.apache.lucene.codecs.Codec;
@@ -62,17 +60,11 @@ public class ES920DiskBBQVectorsFormatTests extends BaseKnnVectorsFormatTestCase
6260
LogConfigurator.loadLog4jPlugins();
6361
LogConfigurator.configureESLogging(); // native access requires logging to be initialized
6462
}
65-
KnnVectorsFormat format;
6663

67-
private boolean useBFloat16;
64+
private KnnVectorsFormat format;
6865

69-
@ParametersFactory
70-
public static Iterable<Object[]> params() {
71-
return List.of(new Object[] { false }, new Object[] { true });
72-
}
73-
74-
public ES920DiskBBQVectorsFormatTests(@Name("useBFloat16") boolean useBFloat16) {
75-
this.useBFloat16 = useBFloat16;
66+
boolean useBFloat16() {
67+
return false;
7668
}
7769

7870
@Before
@@ -83,15 +75,15 @@ public void setUp() throws Exception {
8375
random().nextInt(2 * MIN_VECTORS_PER_CLUSTER, ES920DiskBBQVectorsFormat.MAX_VECTORS_PER_CLUSTER),
8476
random().nextInt(8, ES920DiskBBQVectorsFormat.MAX_CENTROIDS_PER_PARENT_CLUSTER),
8577
random().nextBoolean(),
86-
useBFloat16
78+
useBFloat16()
8779
);
8880
} else {
8981
// run with low numbers to force many clusters with parents
9082
format = new ES920DiskBBQVectorsFormat(
9183
random().nextInt(MIN_VECTORS_PER_CLUSTER, 2 * MIN_VECTORS_PER_CLUSTER),
9284
random().nextInt(MIN_CENTROIDS_PER_PARENT_CLUSTER, 8),
9385
random().nextBoolean(),
94-
useBFloat16
86+
useBFloat16()
9587
);
9688
}
9789
super.setUp();

‎server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93BinaryQuantizedBFloat16VectorsFormatTests.java‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
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.es93;
11+
12+
import org.apache.lucene.index.VectorEncoding;
13+
14+
import java.util.regex.Matcher;
15+
import java.util.regex.Pattern;
16+
17+
import static org.hamcrest.Matchers.closeTo;
18+
19+
public class ES93BinaryQuantizedBFloat16VectorsFormatTests extends ES93BinaryQuantizedVectorsFormatTests {
20+
@Override
21+
boolean useBFloat16() {
22+
return true;
23+
}
24+
25+
@Override
26+
protected VectorEncoding randomVectorEncoding() {
27+
return VectorEncoding.FLOAT32;
28+
}
29+
30+
@Override
31+
public void testEmptyByteVectorData() throws Exception {
32+
// no bytes
33+
}
34+
35+
@Override
36+
public void testMergingWithDifferentByteKnnFields() throws Exception {
37+
// no bytes
38+
}
39+
40+
@Override
41+
public void testByteVectorScorerIteration() throws Exception {
42+
// no bytes
43+
}
44+
45+
@Override
46+
public void testSortedIndexBytes() throws Exception {
47+
// no bytes
48+
}
49+
50+
@Override
51+
public void testMismatchedFields() throws Exception {
52+
// no bytes
53+
}
54+
55+
@Override
56+
public void testRandomBytes() throws Exception {
57+
// no bytes
58+
}
59+
60+
@Override
61+
public void testWriterRamEstimate() throws Exception {
62+
// estimate is different due to bfloat16
63+
}
64+
65+
@Override
66+
public void testRandom() throws Exception {
67+
AssertionError err = expectThrows(AssertionError.class, super::testRandom);
68+
assertFloatsWithinBounds(err);
69+
}
70+
71+
@Override
72+
public void testSparseVectors() throws Exception {
73+
AssertionError err = expectThrows(AssertionError.class, super::testSparseVectors);
74+
assertFloatsWithinBounds(err);
75+
}
76+
77+
@Override
78+
public void testVectorValuesReportCorrectDocs() throws Exception {
79+
AssertionError err = expectThrows(AssertionError.class, super::testVectorValuesReportCorrectDocs);
80+
assertFloatsWithinBounds(err);
81+
}
82+
83+
private static final Pattern FLOAT_ASSERTION_FAILURE = Pattern.compile(".*expected:<([0-9.-]+)> but was:<([0-9.-]+)>");
84+
85+
private static void assertFloatsWithinBounds(AssertionError error) {
86+
Matcher m = FLOAT_ASSERTION_FAILURE.matcher(error.getMessage());
87+
if (m.matches() == false) {
88+
throw error; // nothing to do with us, just rethrow
89+
}
90+
91+
// numbers just need to be in the same vicinity
92+
double expected = Double.parseDouble(m.group(1));
93+
double actual = Double.parseDouble(m.group(2));
94+
double allowedError = expected * 0.01; // within 1%
95+
assertThat(error.getMessage(), actual, closeTo(expected, allowedError));
96+
}
97+
}

‎server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93BinaryQuantizedVectorsFormatTests.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656
import org.apache.lucene.tests.store.MockDirectoryWrapper;
5757
import org.apache.lucene.tests.util.TestUtil;
5858
import org.elasticsearch.common.logging.LogConfigurator;
59+
import org.elasticsearch.index.codec.vectors.BFloat16;
5960

6061
import java.io.IOException;
6162
import java.util.ArrayList;
@@ -77,9 +78,13 @@ public class ES93BinaryQuantizedVectorsFormatTests extends BaseKnnVectorsFormatT
7778

7879
private KnnVectorsFormat format;
7980

81+
boolean useBFloat16() {
82+
return false;
83+
}
84+
8085
@Override
8186
public void setUp() throws Exception {
82-
format = new ES93BinaryQuantizedVectorsFormat(random().nextBoolean());
87+
format = new ES93BinaryQuantizedVectorsFormat(random().nextBoolean(), useBFloat16());
8388
super.setUp();
8489
}
8590

@@ -239,7 +244,8 @@ public void testSimpleOffHeapSizeImpl(Directory dir, IndexWriterConfig config, b
239244
assertEquals(expectVecOffHeap ? 2 : 1, offHeap.size());
240245
assertTrue(offHeap.get("veb") > 0L);
241246
if (expectVecOffHeap) {
242-
assertEquals(vector.length * Float.BYTES, (long) offHeap.get("vec"));
247+
int bytes = useBFloat16() ? BFloat16.BYTES : Float.BYTES;
248+
assertEquals(vector.length * bytes, (long) offHeap.get("vec"));
243249
}
244250
}
245251
}

‎server/src/test/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswBinaryQuantizedVectorsFormatTests.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ public void testToString() {
7575
FilterCodec customCodec = new FilterCodec("foo", Codec.getDefault()) {
7676
@Override
7777
public KnnVectorsFormat knnVectorsFormat() {
78-
return new ES93HnswBinaryQuantizedVectorsFormat(10, 20, false, 1, null);
78+
return new ES93HnswBinaryQuantizedVectorsFormat(10, 20, false, false, 1, null);
7979
}
8080
};
8181
String expectedPattern = "ES93HnswBinaryQuantizedVectorsFormat(name=ES93HnswBinaryQuantizedVectorsFormat, maxConn=10, beamWidth=20,"
@@ -137,7 +137,7 @@ public void testLimits() {
137137
expectThrows(IllegalArgumentException.class, () -> new ES93HnswBinaryQuantizedVectorsFormat(20, 3201));
138138
expectThrows(
139139
IllegalArgumentException.class,
140-
() -> new ES93HnswBinaryQuantizedVectorsFormat(20, 100, false, 1, new SameThreadExecutorService())
140+
() -> new ES93HnswBinaryQuantizedVectorsFormat(20, 100, false, false, 1, new SameThreadExecutorService())
141141
);
142142
}
143143

‎x-pack/qa/rolling-upgrade/src/test/java/org/elasticsearch/upgrades/SemanticTextUpgradeIT.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,6 @@ private static void checkSupportsBFloat16() {
9999
assumeTrue(
100100
"The old cluster needs to support bfloat16 if it is used",
101101
DENSE_MODEL.getServiceSettings().elementType() != DenseVectorFieldMapper.ElementType.BFLOAT16
102-
|| clusterHasFeature(MapperFeatures.BBQ_BFLOAT16)
103102
);
104103
}
105104

0 commit comments

Comments
 (0)