Refactor VectorScorerFactoryTests to Int7SQVectorScorerFactoryTests.java (#130620)

This commit refactors VectorScorerFactoryTests to Int7SQVectorScorerFactoryTests, in order to make space for other vector scorer benchmarks, namely float32.
This commit is contained in:
Chris Hegarty 2025-07-04 13:41:13 +01:00 committed by GitHub
parent 72815c2478
commit bd220a5339
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 17 deletions

View File

@ -9,7 +9,6 @@
package org.elasticsearch.simdvec;
import org.apache.lucene.util.quantization.ScalarQuantizedVectorSimilarity;
import org.elasticsearch.test.ESTestCase;
import org.junit.BeforeClass;
@ -62,17 +61,9 @@ public abstract class AbstractVectorTestCase extends ESTestCase {
return "JDK=" + jdkVersion + ", os=" + osName + ", arch=" + arch;
}
/** Computes the score using the Lucene implementation. */
public static float luceneScore(
VectorSimilarityType similarityFunc,
byte[] a,
byte[] b,
float correction,
float aOffsetValue,
float bOffsetValue
) {
var scorer = ScalarQuantizedVectorSimilarity.fromVectorSimilarity(VectorSimilarityType.of(similarityFunc), correction, (byte) 7);
return scorer.score(a, aOffsetValue, b, bOffsetValue);
// Support for passing on-heap arrays/segments to native
protected static boolean supportsHeapSegments() {
return Runtime.version().feature() >= 22;
}
/** Converts a float value to a byte array. */

View File

@ -22,6 +22,7 @@ import org.apache.lucene.store.MMapDirectory;
import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer;
import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
import org.apache.lucene.util.quantization.ScalarQuantizedVectorSimilarity;
import org.apache.lucene.util.quantization.ScalarQuantizer;
import java.io.IOException;
@ -47,8 +48,7 @@ import static org.elasticsearch.test.hamcrest.OptionalMatchers.isEmpty;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
// @com.carrotsearch.randomizedtesting.annotations.Repeat(iterations = 100)
public class VectorScorerFactoryTests extends AbstractVectorTestCase {
public class Int7SQVectorScorerFactoryTests extends AbstractVectorTestCase {
// bounds of the range of values that can be seen by int7 scalar quantized vectors
static final byte MIN_INT7_VALUE = 0;
@ -107,7 +107,7 @@ public class VectorScorerFactoryTests extends AbstractVectorTestCase {
scorer.setScoringOrdinal(0);
assertThat(scorer.score(1), equalTo(expected));
if (Runtime.version().feature() >= 22) {
if (supportsHeapSegments()) {
var qScorer = factory.getInt7SQVectorScorer(VectorSimilarityType.of(sim), values, query1).get();
assertThat(qScorer.score(1), equalTo(expected));
}
@ -229,11 +229,11 @@ public class VectorScorerFactoryTests extends AbstractVectorTestCase {
}
public void testRandomScorer() throws IOException {
testRandomScorerImpl(MMapDirectory.DEFAULT_MAX_CHUNK_SIZE, VectorScorerFactoryTests.FLOAT_ARRAY_RANDOM_FUNC);
testRandomScorerImpl(MMapDirectory.DEFAULT_MAX_CHUNK_SIZE, Int7SQVectorScorerFactoryTests.FLOAT_ARRAY_RANDOM_FUNC);
}
public void testRandomScorerMax() throws IOException {
testRandomScorerImpl(MMapDirectory.DEFAULT_MAX_CHUNK_SIZE, VectorScorerFactoryTests.FLOAT_ARRAY_MAX_FUNC);
testRandomScorerImpl(MMapDirectory.DEFAULT_MAX_CHUNK_SIZE, Int7SQVectorScorerFactoryTests.FLOAT_ARRAY_MAX_FUNC);
}
public void testRandomScorerChunkSizeSmall() throws IOException {
@ -461,6 +461,19 @@ public class VectorScorerFactoryTests extends AbstractVectorTestCase {
return new OffHeapQuantizedByteVectorValues.DenseOffHeapVectorValues(dims, size, sq, false, sim, null, slice);
}
/** Computes the score using the Lucene implementation. */
public static float luceneScore(
VectorSimilarityType similarityFunc,
byte[] a,
byte[] b,
float correction,
float aOffsetValue,
float bOffsetValue
) {
var scorer = ScalarQuantizedVectorSimilarity.fromVectorSimilarity(VectorSimilarityType.of(similarityFunc), correction, (byte) 7);
return scorer.score(a, aOffsetValue, b, bOffsetValue);
}
RandomVectorScorerSupplier luceneScoreSupplier(QuantizedByteVectorValues values, VectorSimilarityFunction sim) throws IOException {
return new Lucene99ScalarQuantizedVectorScorer(null).getRandomVectorScorerSupplier(sim, values);
}