diff --git a/libs/simdvec/src/test/java/org/elasticsearch/simdvec/AbstractVectorTestCase.java b/libs/simdvec/src/test/java/org/elasticsearch/simdvec/AbstractVectorTestCase.java index 31c5daa81f92b..3905c462d29ca 100644 --- a/libs/simdvec/src/test/java/org/elasticsearch/simdvec/AbstractVectorTestCase.java +++ b/libs/simdvec/src/test/java/org/elasticsearch/simdvec/AbstractVectorTestCase.java @@ -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 static String platformMsg() { 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. */ diff --git a/libs/simdvec/src/test/java/org/elasticsearch/simdvec/VectorScorerFactoryTests.java b/libs/simdvec/src/test/java/org/elasticsearch/simdvec/Int7SQVectorScorerFactoryTests.java similarity index 96% rename from libs/simdvec/src/test/java/org/elasticsearch/simdvec/VectorScorerFactoryTests.java rename to libs/simdvec/src/test/java/org/elasticsearch/simdvec/Int7SQVectorScorerFactoryTests.java index 070260759d6a0..eca67582b7a16 100644 --- a/libs/simdvec/src/test/java/org/elasticsearch/simdvec/VectorScorerFactoryTests.java +++ b/libs/simdvec/src/test/java/org/elasticsearch/simdvec/Int7SQVectorScorerFactoryTests.java @@ -22,6 +22,7 @@ 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.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 @@ void testSimpleImpl(long maxChunkSize) throws IOException { 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 @@ void testRandomSupplier(long maxChunkSize, Function byteArraySu } 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 @@ QuantizedByteVectorValues vectorValues(int dims, int size, IndexInput in, Vector 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); }