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:
parent
72815c2478
commit
bd220a5339
|
@ -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. */
|
||||
|
|
|
@ -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);
|
||||
}
|
Loading…
Reference in New Issue