From 7d2e035805bf91a0af30c576c1517a9ae82fe455 Mon Sep 17 00:00:00 2001 From: ChrisHegarty Date: Fri, 20 Mar 2026 17:24:25 +0000 Subject: [PATCH 1/4] Split native and Panama OSQ vector scorer implementations --- .../simdvec/ESNextOSQVectorsScorer.java | 8 +- .../MSBitToInt4ESNextOSQVectorsScorer.java | 239 ++++------------- .../MSD7Q7ESNextOSQVectorsScorer.java | 41 +-- .../MSDibitToInt4ESNextOSQVectorsScorer.java | 239 ++++------------- ...MSInt4SymmetricESNextOSQVectorsScorer.java | 240 ++++-------------- .../MemorySegmentESNextOSQVectorsScorer.java | 108 +++++++- .../vectorization/NativeBitToInt4Scorer.java | 58 +++++ .../NativeDibitToInt4Scorer.java | 58 +++++ .../NativeInt4SymmetricScorer.java | 58 +++++ .../NativeMemorySegmentScorer.java | 175 +++++++++++++ 10 files changed, 615 insertions(+), 609 deletions(-) create mode 100644 libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeBitToInt4Scorer.java create mode 100644 libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeDibitToInt4Scorer.java create mode 100644 libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeInt4SymmetricScorer.java create mode 100644 libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeMemorySegmentScorer.java diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESNextOSQVectorsScorer.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESNextOSQVectorsScorer.java index 34967e5f9572b..efdcfc0051573 100644 --- a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESNextOSQVectorsScorer.java +++ b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESNextOSQVectorsScorer.java @@ -273,9 +273,7 @@ public float scoreBulk( quantizeScoreBulk(q, bulkSize, scores); in.readFloats(lowerIntervals, 0, bulkSize); in.readFloats(upperIntervals, 0, bulkSize); - for (int i = 0; i < bulkSize; i++) { - targetComponentSums[i] = in.readInt(); - } + in.readInts(targetComponentSums, 0, bulkSize); in.readFloats(additionalCorrections, 0, bulkSize); float maxScore = Float.NEGATIVE_INFINITY; for (int i = 0; i < bulkSize; i++) { @@ -326,9 +324,7 @@ public float scoreBulkOffsets( quantizeScoreBulkOffsets(q, offsets, offsetsCount, scores, count); in.readFloats(lowerIntervals, 0, count); in.readFloats(upperIntervals, 0, count); - for (int i = 0; i < count; i++) { - targetComponentSums[i] = in.readInt(); - } + in.readInts(targetComponentSums, 0, count); in.readFloats(additionalCorrections, 0, count); float maxScore = Float.NEGATIVE_INFINITY; int offsetIndex = 0; diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSBitToInt4ESNextOSQVectorsScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSBitToInt4ESNextOSQVectorsScorer.java index 14e2ca8b31cdc..85d5de7b98278 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSBitToInt4ESNextOSQVectorsScorer.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSBitToInt4ESNextOSQVectorsScorer.java @@ -21,16 +21,12 @@ import org.elasticsearch.simdvec.internal.IndexInputUtils; import java.io.IOException; -import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; import java.nio.ByteOrder; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; -import static org.elasticsearch.simdvec.internal.Similarities.dotProductD1Q4; -import static org.elasticsearch.simdvec.internal.Similarities.dotProductD1Q4Bulk; -import static org.elasticsearch.simdvec.internal.Similarities.dotProductD1Q4BulkWithOffsets; /** Panamized scorer for quantized vectors stored as a {@link MemorySegment}. */ final class MSBitToInt4ESNextOSQVectorsScorer extends MemorySegmentESNextOSQVectorsScorer.MemorySegmentScorer { @@ -42,40 +38,16 @@ final class MSBitToInt4ESNextOSQVectorsScorer extends MemorySegmentESNextOSQVect @Override public long quantizeScore(byte[] q) throws IOException { assert q.length == length * 4; - // 128 / 8 == 16 - if (length >= 16) { - if (NATIVE_SUPPORTED) { - return nativeQuantizeScore(q); - } else if (PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { - if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { - return quantizeScore256(q); - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { - return quantizeScore128(q); - } + if (length >= 16 && PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { + if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { + return quantizeScore256(q); + } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { + return quantizeScore128(q); } } return Long.MIN_VALUE; } - private long nativeQuantizeScore(byte[] q) throws IOException { - return IndexInputUtils.withSlice(in, length, this::getScratch, segment -> nativeQuantizeScoreImpl(q, segment, length)); - } - - private static long nativeQuantizeScoreImpl(byte[] q, MemorySegment datasetMemorySegment, int length) { - final long qScore; - if (SUPPORTS_HEAP_SEGMENTS) { - var queryMemorySegment = MemorySegment.ofArray(q); - qScore = dotProductD1Q4(datasetMemorySegment, queryMemorySegment, length); - } else { - try (var arena = Arena.ofConfined()) { - var queryMemorySegment = arena.allocate(q.length, 32); - MemorySegment.copy(q, 0, queryMemorySegment, ValueLayout.JAVA_BYTE, 0, q.length); - qScore = dotProductD1Q4(datasetMemorySegment, queryMemorySegment, length); - } - } - return qScore; - } - private long quantizeScore256(byte[] q) throws IOException { return IndexInputUtils.withSlice(in, length, this::getScratch, segment -> quantizeScore256Impl(q, segment, length)); } @@ -214,44 +186,18 @@ private static long quantizeScore128Impl(byte[] q, MemorySegment memorySegment, @Override public boolean quantizeScoreBulk(byte[] q, int count, float[] scores) throws IOException { assert q.length == length * 4; - // 128 / 8 == 16 - if (length >= 16) { - if (NATIVE_SUPPORTED) { - if (SUPPORTS_HEAP_SEGMENTS) { - var querySegment = MemorySegment.ofArray(q); - var scoresSegment = MemorySegment.ofArray(scores); - nativeQuantizeScoreBulk(querySegment, count, scoresSegment); - } else { - try (var arena = Arena.ofConfined()) { - var querySegment = arena.allocate(q.length, 32); - var scoresSegment = arena.allocate((long) scores.length * Float.BYTES, 32); - MemorySegment.copy(q, 0, querySegment, ValueLayout.JAVA_BYTE, 0, q.length); - nativeQuantizeScoreBulk(querySegment, count, scoresSegment); - MemorySegment.copy(scoresSegment, ValueLayout.JAVA_FLOAT, 0, scores, 0, scores.length); - } - } + if (length >= 16 && PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { + if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { + quantizeScore256Bulk(q, count, scores); + return true; + } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { + quantizeScore128Bulk(q, count, scores); return true; - } else if (PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { - if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { - quantizeScore256Bulk(q, count, scores); - return true; - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { - quantizeScore128Bulk(q, count, scores); - return true; - } } } return false; } - private void nativeQuantizeScoreBulk(MemorySegment querySegment, int count, MemorySegment scoresSegment) throws IOException { - var datasetLengthInBytes = (long) length * count; - IndexInputUtils.withSlice(in, datasetLengthInBytes, this::getScratch, datasetSegment -> { - dotProductD1Q4Bulk(datasetSegment, querySegment, length, count, scoresSegment); - return null; - }); - } - private void quantizeScore128Bulk(byte[] q, int count, float[] scores) throws IOException { var datasetLengthInBytes = (long) length * count; IndexInputUtils.withSlice(in, datasetLengthInBytes, this::getScratch, segment -> { @@ -403,46 +349,11 @@ private static void quantizeScore256BulkImpl(byte[] q, MemorySegment memorySegme @Override public boolean quantizeScoreBulkOffsets(byte[] q, int[] offsets, int offsetsCount, float[] scores, int count) throws IOException { - assert q.length == length * 4; - if (NATIVE_SUPPORTED) { - if (SUPPORTS_HEAP_SEGMENTS) { - var querySegment = MemorySegment.ofArray(q); - var offsetsSegment = MemorySegment.ofArray(offsets); - var scoresSegment = MemorySegment.ofArray(scores); - nativeQuantizeScoreBulkOffsets(querySegment, offsetsSegment, scoresSegment, offsetsCount, count); - } else { - try (var arena = Arena.ofConfined()) { - var querySegment = arena.allocate(q.length, 32); - var offsetsSegment = arena.allocate((long) offsetsCount * Integer.BYTES, 32); - var scoresSegment = arena.allocate((long) scores.length * Float.BYTES, 32); - MemorySegment.copy(q, 0, querySegment, ValueLayout.JAVA_BYTE, 0, q.length); - MemorySegment.copy(offsets, 0, offsetsSegment, ValueLayout.JAVA_INT, 0, offsetsCount); - nativeQuantizeScoreBulkOffsets(querySegment, offsetsSegment, scoresSegment, offsetsCount, count); - MemorySegment.copy(scoresSegment, ValueLayout.JAVA_FLOAT, 0, scores, 0, scores.length); - } - } - repositionScoresMatchingOffsets(offsets, offsetsCount, scores); - return true; - } return false; } - private void nativeQuantizeScoreBulkOffsets( - MemorySegment querySegment, - MemorySegment offsetsSegment, - MemorySegment scoresSegment, - int offsetsCount, - int totalCount - ) throws IOException { - var datasetLengthInBytes = (long) length * totalCount; - IndexInputUtils.withSlice(in, datasetLengthInBytes, this::getScratch, datasetSegment -> { - dotProductD1Q4BulkWithOffsets(datasetSegment, querySegment, length, length, offsetsSegment, offsetsCount, scoresSegment); - return null; - }); - } - @Override - public float scoreBulk( + float scoreBulkOffsets( byte[] q, float queryLowerInterval, float queryUpperInterval, @@ -450,107 +361,55 @@ public float scoreBulk( float queryAdditionalCorrection, VectorSimilarityFunction similarityFunction, float centroidDp, + int[] offsets, + int offsetsCount, float[] scores, - int bulkSize - ) throws IOException { - assert q.length == length * 4; - // 128 / 8 == 16 - if (length >= 16) { - if (PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { - if (NATIVE_SUPPORTED) { - if (SUPPORTS_HEAP_SEGMENTS) { - var querySegment = MemorySegment.ofArray(q); - var scoresSegment = MemorySegment.ofArray(scores); - nativeQuantizeScoreBulk(querySegment, bulkSize, scoresSegment); - return nativeApplyCorrectionsBulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scoresSegment, - bulkSize - ); - } else { - try (var arena = Arena.ofConfined()) { - var querySegment = arena.allocate(q.length, 32); - var scoresSegment = arena.allocate((long) scores.length * Float.BYTES, 32); - MemorySegment.copy(q, 0, querySegment, ValueLayout.JAVA_BYTE, 0, q.length); - nativeQuantizeScoreBulk(querySegment, bulkSize, scoresSegment); - var maxScore = nativeApplyCorrectionsBulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scoresSegment, - bulkSize - ); - MemorySegment.copy(scoresSegment, ValueLayout.JAVA_FLOAT, 0, scores, 0, scores.length); - return maxScore; - } - } - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { - quantizeScore256Bulk(q, bulkSize, scores); - return applyCorrections256Bulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scores, - bulkSize - ); - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { - quantizeScore128Bulk(q, bulkSize, scores); - return applyCorrections128Bulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scores, - bulkSize - ); - } - } - } + int count + ) { return Float.NEGATIVE_INFINITY; } - private float nativeApplyCorrectionsBulk( + @Override + public float scoreBulk( + byte[] q, float queryLowerInterval, float queryUpperInterval, int queryComponentSum, float queryAdditionalCorrection, VectorSimilarityFunction similarityFunction, float centroidDp, - MemorySegment scoresSegment, + float[] scores, int bulkSize ) throws IOException { - return IndexInputUtils.withSlice( - in, - 16L * bulkSize, - this::getScratch, - seg -> ScoreCorrections.nativeApplyCorrectionsBulk( - similarityFunction, - seg, - bulkSize, - dimensions, - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - FOUR_BIT_SCALE, - ONE_BIT_SCALE, - centroidDp, - scoresSegment - ) - ); + assert q.length == length * 4; + if (length >= 16 && PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { + if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { + quantizeScore256Bulk(q, bulkSize, scores); + return applyCorrections256Bulk( + queryLowerInterval, + queryUpperInterval, + queryComponentSum, + queryAdditionalCorrection, + similarityFunction, + centroidDp, + scores, + bulkSize + ); + } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { + quantizeScore128Bulk(q, bulkSize, scores); + return applyCorrections128Bulk( + queryLowerInterval, + queryUpperInterval, + queryComponentSum, + queryAdditionalCorrection, + similarityFunction, + centroidDp, + scores, + bulkSize + ); + } + } + return Float.NEGATIVE_INFINITY; } private float applyCorrections128Bulk( diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSD7Q7ESNextOSQVectorsScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSD7Q7ESNextOSQVectorsScorer.java index ba13c409a0d55..6719b2c61d261 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSD7Q7ESNextOSQVectorsScorer.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSD7Q7ESNextOSQVectorsScorer.java @@ -14,9 +14,7 @@ import org.elasticsearch.simdvec.internal.MemorySegmentES92Int7VectorsScorer; import java.io.IOException; -import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; -import java.lang.foreign.ValueLayout; import static org.elasticsearch.simdvec.internal.Similarities.dotProductI7uBulkWithOffsets; @@ -44,23 +42,11 @@ boolean quantizeScoreBulk(byte[] q, int count, float[] scores) throws IOExceptio @Override public boolean quantizeScoreBulkOffsets(byte[] q, int[] offsets, int offsetsCount, float[] scores, int count) throws IOException { assert q.length == length; - if (NATIVE_SUPPORTED) { - if (SUPPORTS_HEAP_SEGMENTS) { - var querySegment = MemorySegment.ofArray(q); - var offsetsSegment = MemorySegment.ofArray(offsets); - var scoresSegment = MemorySegment.ofArray(scores); - nativeQuantizeScoreBulkOffsets(querySegment, offsetsSegment, scoresSegment, offsetsCount, count); - } else { - try (var arena = Arena.ofConfined()) { - var querySegment = arena.allocate(q.length, 32); - var offsetsSegment = arena.allocate((long) offsetsCount * Integer.BYTES, 32); - var scoresSegment = arena.allocate((long) scores.length * Float.BYTES, 32); - MemorySegment.copy(q, 0, querySegment, ValueLayout.JAVA_BYTE, 0, q.length); - MemorySegment.copy(offsets, 0, offsetsSegment, ValueLayout.JAVA_INT, 0, offsetsCount); - nativeQuantizeScoreBulkOffsets(querySegment, offsetsSegment, scoresSegment, offsetsCount, count); - MemorySegment.copy(scoresSegment, ValueLayout.JAVA_FLOAT, 0, scores, 0, scores.length); - } - } + if (NATIVE_SUPPORTED && SUPPORTS_HEAP_SEGMENTS) { + var querySegment = MemorySegment.ofArray(q); + var offsetsSegment = MemorySegment.ofArray(offsets); + var scoresSegment = MemorySegment.ofArray(scores); + nativeQuantizeScoreBulkOffsets(querySegment, offsetsSegment, scoresSegment, offsetsCount, count); repositionScoresMatchingOffsets(offsets, offsetsCount, scores); return true; } @@ -81,6 +67,23 @@ private void nativeQuantizeScoreBulkOffsets( }); } + @Override + float scoreBulkOffsets( + byte[] q, + float queryLowerInterval, + float queryUpperInterval, + int queryComponentSum, + float queryAdditionalCorrection, + VectorSimilarityFunction similarityFunction, + float centroidDp, + int[] offsets, + int offsetsCount, + float[] scores, + int count + ) { + return Float.NEGATIVE_INFINITY; + } + @Override float scoreBulk( byte[] q, diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSDibitToInt4ESNextOSQVectorsScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSDibitToInt4ESNextOSQVectorsScorer.java index 939dbf3e9a898..d439974bd9291 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSDibitToInt4ESNextOSQVectorsScorer.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSDibitToInt4ESNextOSQVectorsScorer.java @@ -21,16 +21,12 @@ import org.elasticsearch.simdvec.internal.IndexInputUtils; import java.io.IOException; -import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; import java.nio.ByteOrder; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; -import static org.elasticsearch.simdvec.internal.Similarities.dotProductD2Q4; -import static org.elasticsearch.simdvec.internal.Similarities.dotProductD2Q4Bulk; -import static org.elasticsearch.simdvec.internal.Similarities.dotProductD2Q4BulkWithOffsets; /** Panamized scorer for quantized vectors stored as a {@link MemorySegment}. */ final class MSDibitToInt4ESNextOSQVectorsScorer extends MemorySegmentESNextOSQVectorsScorer.MemorySegmentScorer { @@ -42,40 +38,16 @@ final class MSDibitToInt4ESNextOSQVectorsScorer extends MemorySegmentESNextOSQVe @Override public long quantizeScore(byte[] q) throws IOException { assert q.length == length * 2; - // 128 / 8 == 16 - if (length >= 16) { - if (NATIVE_SUPPORTED) { - return nativeQuantizeScore(q); - } else if (PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { - if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { - return quantizeScore256DibitToInt4(q); - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { - return quantizeScore128DibitToInt4(q); - } + if (length >= 16 && PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { + if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { + return quantizeScore256DibitToInt4(q); + } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { + return quantizeScore128DibitToInt4(q); } } return Long.MIN_VALUE; } - private long nativeQuantizeScore(byte[] q) throws IOException { - return IndexInputUtils.withSlice(in, length, this::getScratch, segment -> nativeQuantizeScoreImpl(q, segment, length)); - } - - private static long nativeQuantizeScoreImpl(byte[] q, MemorySegment datasetMemorySegment, int length) { - final long qScore; - if (SUPPORTS_HEAP_SEGMENTS) { - var queryMemorySegment = MemorySegment.ofArray(q); - qScore = dotProductD2Q4(datasetMemorySegment, queryMemorySegment, length); - } else { - try (var arena = Arena.ofConfined()) { - var queryMemorySegment = arena.allocate(q.length, 32); - MemorySegment.copy(q, 0, queryMemorySegment, ValueLayout.JAVA_BYTE, 0, q.length); - qScore = dotProductD2Q4(datasetMemorySegment, queryMemorySegment, length); - } - } - return qScore; - } - private long quantizeScore256DibitToInt4(byte[] q) throws IOException { int lower = (int) quantizeScore256(q); int upper = (int) quantizeScore256(q); @@ -229,44 +201,18 @@ private static long quantizeScore128Impl(byte[] q, MemorySegment memorySegment, @Override public boolean quantizeScoreBulk(byte[] q, int count, float[] scores) throws IOException { assert q.length == length * 2; - // 128 / 8 == 16 - if (length >= 16) { - if (NATIVE_SUPPORTED) { - if (SUPPORTS_HEAP_SEGMENTS) { - var queryMemorySegment = MemorySegment.ofArray(q); - var scoresSegment = MemorySegment.ofArray(scores); - nativeQuantizeScoreBulk(queryMemorySegment, count, scoresSegment); - } else { - try (var arena = Arena.ofConfined()) { - var queryMemorySegment = arena.allocate(q.length, 32); - var scoresSegment = arena.allocate((long) scores.length * Float.BYTES, 32); - MemorySegment.copy(q, 0, queryMemorySegment, ValueLayout.JAVA_BYTE, 0, q.length); - nativeQuantizeScoreBulk(queryMemorySegment, count, scoresSegment); - MemorySegment.copy(scoresSegment, ValueLayout.JAVA_FLOAT, 0, scores, 0, scores.length); - } - } + if (length >= 16 && PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { + if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { + quantizeScore256Bulk(q, count, scores); + return true; + } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { + quantizeScore128Bulk(q, count, scores); return true; - } else if (PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { - if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { - quantizeScore256Bulk(q, count, scores); - return true; - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { - quantizeScore128Bulk(q, count, scores); - return true; - } } } return false; } - private void nativeQuantizeScoreBulk(MemorySegment queryMemorySegment, int count, MemorySegment scoresSegment) throws IOException { - var datasetLengthInBytes = (long) length * count; - IndexInputUtils.withSlice(in, datasetLengthInBytes, this::getScratch, datasetSegment -> { - dotProductD2Q4Bulk(datasetSegment, queryMemorySegment, length, count, scoresSegment); - return null; - }); - } - private void quantizeScore128Bulk(byte[] q, int count, float[] scores) throws IOException { for (int iter = 0; iter < count; iter++) { scores[iter] = quantizeScore128DibitToInt4(q); @@ -281,46 +227,11 @@ private void quantizeScore256Bulk(byte[] q, int count, float[] scores) throws IO @Override public boolean quantizeScoreBulkOffsets(byte[] q, int[] offsets, int offsetsCount, float[] scores, int count) throws IOException { - assert q.length == length * 2; - if (NATIVE_SUPPORTED) { - if (SUPPORTS_HEAP_SEGMENTS) { - var querySegment = MemorySegment.ofArray(q); - var offsetsSegment = MemorySegment.ofArray(offsets); - var scoresSegment = MemorySegment.ofArray(scores); - nativeQuantizeScoreBulkOffsets(querySegment, offsetsSegment, scoresSegment, offsetsCount, count); - } else { - try (var arena = Arena.ofConfined()) { - var querySegment = arena.allocate(q.length, 32); - var offsetsSegment = arena.allocate((long) offsetsCount * Integer.BYTES, 32); - var scoresSegment = arena.allocate((long) scores.length * Float.BYTES, 32); - MemorySegment.copy(q, 0, querySegment, ValueLayout.JAVA_BYTE, 0, q.length); - MemorySegment.copy(offsets, 0, offsetsSegment, ValueLayout.JAVA_INT, 0, offsetsCount); - nativeQuantizeScoreBulkOffsets(querySegment, offsetsSegment, scoresSegment, offsetsCount, count); - MemorySegment.copy(scoresSegment, ValueLayout.JAVA_FLOAT, 0, scores, 0, scores.length); - } - } - repositionScoresMatchingOffsets(offsets, offsetsCount, scores); - return true; - } return false; } - private void nativeQuantizeScoreBulkOffsets( - MemorySegment querySegment, - MemorySegment offsetsSegment, - MemorySegment scoresSegment, - int offsetsCount, - int totalCount - ) throws IOException { - var datasetLengthInBytes = (long) length * totalCount; - IndexInputUtils.withSlice(in, datasetLengthInBytes, this::getScratch, datasetSegment -> { - dotProductD2Q4BulkWithOffsets(datasetSegment, querySegment, length, length, offsetsSegment, offsetsCount, scoresSegment); - return null; - }); - } - @Override - public float scoreBulk( + float scoreBulkOffsets( byte[] q, float queryLowerInterval, float queryUpperInterval, @@ -328,107 +239,55 @@ public float scoreBulk( float queryAdditionalCorrection, VectorSimilarityFunction similarityFunction, float centroidDp, + int[] offsets, + int offsetsCount, float[] scores, - int bulkSize - ) throws IOException { - assert q.length == length * 2; - // 128 / 8 == 16 - if (length >= 16) { - if (PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { - if (NATIVE_SUPPORTED) { - if (SUPPORTS_HEAP_SEGMENTS) { - var querySegment = MemorySegment.ofArray(q); - var scoresSegment = MemorySegment.ofArray(scores); - nativeQuantizeScoreBulk(querySegment, bulkSize, scoresSegment); - return nativeApplyCorrectionsBulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scoresSegment, - bulkSize - ); - } else { - try (var arena = Arena.ofConfined()) { - var querySegment = arena.allocate(q.length, 32); - var scoresSegment = arena.allocate((long) scores.length * Float.BYTES, 32); - MemorySegment.copy(q, 0, querySegment, ValueLayout.JAVA_BYTE, 0, q.length); - nativeQuantizeScoreBulk(querySegment, bulkSize, scoresSegment); - var maxScore = nativeApplyCorrectionsBulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scoresSegment, - bulkSize - ); - MemorySegment.copy(scoresSegment, ValueLayout.JAVA_FLOAT, 0, scores, 0, scores.length); - return maxScore; - } - } - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { - quantizeScore256Bulk(q, bulkSize, scores); - return applyCorrections256Bulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scores, - bulkSize - ); - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { - quantizeScore128Bulk(q, bulkSize, scores); - return applyCorrections128Bulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scores, - bulkSize - ); - } - } - } + int count + ) { return Float.NEGATIVE_INFINITY; } - private float nativeApplyCorrectionsBulk( + @Override + public float scoreBulk( + byte[] q, float queryLowerInterval, float queryUpperInterval, int queryComponentSum, float queryAdditionalCorrection, VectorSimilarityFunction similarityFunction, float centroidDp, - MemorySegment scoresSegment, + float[] scores, int bulkSize ) throws IOException { - return IndexInputUtils.withSlice( - in, - 16L * bulkSize, - this::getScratch, - seg -> ScoreCorrections.nativeApplyCorrectionsBulk( - similarityFunction, - seg, - bulkSize, - dimensions, - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - FOUR_BIT_SCALE, - TWO_BIT_SCALE, - centroidDp, - scoresSegment - ) - ); + assert q.length == length * 2; + if (length >= 16 && PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { + if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { + quantizeScore256Bulk(q, bulkSize, scores); + return applyCorrections256Bulk( + queryLowerInterval, + queryUpperInterval, + queryComponentSum, + queryAdditionalCorrection, + similarityFunction, + centroidDp, + scores, + bulkSize + ); + } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { + quantizeScore128Bulk(q, bulkSize, scores); + return applyCorrections128Bulk( + queryLowerInterval, + queryUpperInterval, + queryComponentSum, + queryAdditionalCorrection, + similarityFunction, + centroidDp, + scores, + bulkSize + ); + } + } + return Float.NEGATIVE_INFINITY; } private float applyCorrections128Bulk( diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSInt4SymmetricESNextOSQVectorsScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSInt4SymmetricESNextOSQVectorsScorer.java index 2793d2016295a..fe235b27914fb 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSInt4SymmetricESNextOSQVectorsScorer.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSInt4SymmetricESNextOSQVectorsScorer.java @@ -21,16 +21,12 @@ import org.elasticsearch.simdvec.internal.IndexInputUtils; import java.io.IOException; -import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; import java.nio.ByteOrder; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; -import static org.elasticsearch.simdvec.internal.Similarities.dotProductD4Q4; -import static org.elasticsearch.simdvec.internal.Similarities.dotProductD4Q4Bulk; -import static org.elasticsearch.simdvec.internal.Similarities.dotProductD4Q4BulkWithOffsets; /** Panamized scorer for quantized vectors stored as a {@link MemorySegment}. */ final class MSInt4SymmetricESNextOSQVectorsScorer extends MemorySegmentESNextOSQVectorsScorer.MemorySegmentScorer { @@ -42,40 +38,16 @@ final class MSInt4SymmetricESNextOSQVectorsScorer extends MemorySegmentESNextOSQ @Override public long quantizeScore(byte[] q) throws IOException { assert q.length == length; - // 128 / 8 == 16 - if (length >= 16) { - if (NATIVE_SUPPORTED) { - return nativeQuantizeScore(q); - } else if (PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { - if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { - return quantizeScoreSymmetric256(q); - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { - return quantizeScoreSymmetric128(q); - } + if (length >= 16 && PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { + if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { + return quantizeScoreSymmetric256(q); + } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { + return quantizeScoreSymmetric128(q); } } return Long.MIN_VALUE; } - private long nativeQuantizeScore(byte[] q) throws IOException { - return IndexInputUtils.withSlice(in, length, this::getScratch, segment -> nativeQuantizeScoreImpl(q, segment, length)); - } - - private static long nativeQuantizeScoreImpl(byte[] q, MemorySegment datasetMemorySegment, int length) { - final long qScore; - if (SUPPORTS_HEAP_SEGMENTS) { - var queryMemorySegment = MemorySegment.ofArray(q); - qScore = dotProductD4Q4(datasetMemorySegment, queryMemorySegment, length); - } else { - try (var arena = Arena.ofConfined()) { - var queryMemorySegment = arena.allocate(q.length, 32); - MemorySegment.copy(q, 0, queryMemorySegment, ValueLayout.JAVA_BYTE, 0, q.length); - qScore = dotProductD4Q4(datasetMemorySegment, queryMemorySegment, length); - } - } - return qScore; - } - private long quantizeScoreSymmetric128(byte[] q) throws IOException { int stripe0 = (int) quantizeScore128(q); int stripe1 = (int) quantizeScore128(q); @@ -233,45 +205,18 @@ private static long quantizeScore128Impl(byte[] q, MemorySegment memorySegment, @Override public boolean quantizeScoreBulk(byte[] q, int count, float[] scores) throws IOException { assert q.length == length; - // 128 / 8 == 16 - if (length >= 16) { - if (NATIVE_SUPPORTED) { - if (SUPPORTS_HEAP_SEGMENTS) { - var queryMemorySegment = MemorySegment.ofArray(q); - var scoresSegment = MemorySegment.ofArray(scores); - nativeQuantizeScoreBulk(queryMemorySegment, count, scoresSegment); - } else { - try (var arena = Arena.ofConfined()) { - var queryMemorySegment = arena.allocate(q.length, 32); - var scoresSegment = arena.allocate((long) scores.length * Float.BYTES, 32); - MemorySegment.copy(q, 0, queryMemorySegment, ValueLayout.JAVA_BYTE, 0, q.length); - nativeQuantizeScoreBulk(queryMemorySegment, count, scoresSegment); - MemorySegment.copy(scoresSegment, ValueLayout.JAVA_FLOAT, 0, scores, 0, scores.length); - } - } + if (length >= 16 && PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { + if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { + quantizeScore256Bulk(q, count, scores); + return true; + } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { + quantizeScore128Bulk(q, count, scores); return true; - } - if (PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { - if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { - quantizeScore256Bulk(q, count, scores); - return true; - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { - quantizeScore128Bulk(q, count, scores); - return true; - } } } return false; } - private void nativeQuantizeScoreBulk(MemorySegment queryMemorySegment, int count, MemorySegment scoresSegment) throws IOException { - var datasetLengthInBytes = (long) length * count; - IndexInputUtils.withSlice(in, datasetLengthInBytes, this::getScratch, datasetSegment -> { - dotProductD4Q4Bulk(datasetSegment, queryMemorySegment, length, count, scoresSegment); - return null; - }); - } - private void quantizeScore128Bulk(byte[] q, int count, float[] scores) throws IOException { for (int iter = 0; iter < count; iter++) { scores[iter] = quantizeScoreSymmetric128(q); @@ -286,46 +231,11 @@ private void quantizeScore256Bulk(byte[] q, int count, float[] scores) throws IO @Override public boolean quantizeScoreBulkOffsets(byte[] q, int[] offsets, int offsetsCount, float[] scores, int count) throws IOException { - assert q.length == length; - if (NATIVE_SUPPORTED) { - if (SUPPORTS_HEAP_SEGMENTS) { - var querySegment = MemorySegment.ofArray(q); - var offsetsSegment = MemorySegment.ofArray(offsets); - var scoresSegment = MemorySegment.ofArray(scores); - nativeQuantizeScoreBulkOffsets(querySegment, offsetsSegment, scoresSegment, offsetsCount, count); - } else { - try (var arena = Arena.ofConfined()) { - var querySegment = arena.allocate(q.length, 32); - var offsetsSegment = arena.allocate((long) offsetsCount * Integer.BYTES, 32); - var scoresSegment = arena.allocate((long) scores.length * Float.BYTES, 32); - MemorySegment.copy(q, 0, querySegment, ValueLayout.JAVA_BYTE, 0, q.length); - MemorySegment.copy(offsets, 0, offsetsSegment, ValueLayout.JAVA_INT, 0, offsetsCount); - nativeQuantizeScoreBulkOffsets(querySegment, offsetsSegment, scoresSegment, offsetsCount, count); - MemorySegment.copy(scoresSegment, ValueLayout.JAVA_FLOAT, 0, scores, 0, scores.length); - } - } - repositionScoresMatchingOffsets(offsets, offsetsCount, scores); - return true; - } return false; } - private void nativeQuantizeScoreBulkOffsets( - MemorySegment querySegment, - MemorySegment offsetsSegment, - MemorySegment scoresSegment, - int offsetsCount, - int totalCount - ) throws IOException { - var datasetLengthInBytes = (long) length * totalCount; - IndexInputUtils.withSlice(in, datasetLengthInBytes, this::getScratch, datasetSegment -> { - dotProductD4Q4BulkWithOffsets(datasetSegment, querySegment, length, length, offsetsSegment, offsetsCount, scoresSegment); - return null; - }); - } - @Override - public float scoreBulk( + float scoreBulkOffsets( byte[] q, float queryLowerInterval, float queryUpperInterval, @@ -333,107 +243,55 @@ public float scoreBulk( float queryAdditionalCorrection, VectorSimilarityFunction similarityFunction, float centroidDp, + int[] offsets, + int offsetsCount, float[] scores, - int bulkSize - ) throws IOException { - assert q.length == length; - // 128 / 8 == 16 - if (length >= 16) { - if (PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { - if (NATIVE_SUPPORTED) { - if (SUPPORTS_HEAP_SEGMENTS) { - var querySegment = MemorySegment.ofArray(q); - var scoresSegment = MemorySegment.ofArray(scores); - nativeQuantizeScoreBulk(querySegment, bulkSize, scoresSegment); - return nativeApplyCorrectionsBulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scoresSegment, - bulkSize - ); - } else { - try (var arena = Arena.ofConfined()) { - var querySegment = arena.allocate(q.length, 32); - var scoresSegment = arena.allocate((long) scores.length * Float.BYTES, 32); - MemorySegment.copy(q, 0, querySegment, ValueLayout.JAVA_BYTE, 0, q.length); - nativeQuantizeScoreBulk(querySegment, bulkSize, scoresSegment); - var maxScore = nativeApplyCorrectionsBulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scoresSegment, - bulkSize - ); - MemorySegment.copy(scoresSegment, ValueLayout.JAVA_FLOAT, 0, scores, 0, scores.length); - return maxScore; - } - } - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { - quantizeScore256Bulk(q, bulkSize, scores); - return applyCorrections256Bulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scores, - bulkSize - ); - } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { - quantizeScore128Bulk(q, bulkSize, scores); - return applyCorrections128Bulk( - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - similarityFunction, - centroidDp, - scores, - bulkSize - ); - } - } - } + int count + ) { return Float.NEGATIVE_INFINITY; } - private float nativeApplyCorrectionsBulk( + @Override + public float scoreBulk( + byte[] q, float queryLowerInterval, float queryUpperInterval, int queryComponentSum, float queryAdditionalCorrection, VectorSimilarityFunction similarityFunction, float centroidDp, - MemorySegment scoresSegment, + float[] scores, int bulkSize ) throws IOException { - return IndexInputUtils.withSlice( - in, - 16L * bulkSize, - this::getScratch, - memorySegment -> ScoreCorrections.nativeApplyCorrectionsBulk( - similarityFunction, - memorySegment, - bulkSize, - dimensions, - queryLowerInterval, - queryUpperInterval, - queryComponentSum, - queryAdditionalCorrection, - FOUR_BIT_SCALE, - FOUR_BIT_SCALE, - centroidDp, - scoresSegment - ) - ); + assert q.length == length; + if (length >= 16 && PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS) { + if (PanamaESVectorUtilSupport.VECTOR_BITSIZE >= 256) { + quantizeScore256Bulk(q, bulkSize, scores); + return applyCorrections256Bulk( + queryLowerInterval, + queryUpperInterval, + queryComponentSum, + queryAdditionalCorrection, + similarityFunction, + centroidDp, + scores, + bulkSize + ); + } else if (PanamaESVectorUtilSupport.VECTOR_BITSIZE == 128) { + quantizeScore128Bulk(q, bulkSize, scores); + return applyCorrections128Bulk( + queryLowerInterval, + queryUpperInterval, + queryComponentSum, + queryAdditionalCorrection, + similarityFunction, + centroidDp, + scores, + bulkSize + ); + } + } + return Float.NEGATIVE_INFINITY; } private float applyCorrections128Bulk( diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java index adf3a7dd64320..1f0bd04601153 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java @@ -27,6 +27,8 @@ import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; import java.nio.ByteOrder; +import java.util.Map; +import java.util.Objects; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; @@ -34,6 +36,30 @@ /** Panamized scorer for quantized vectors stored as a {@link MemorySegment}. */ public final class MemorySegmentESNextOSQVectorsScorer extends ESNextOSQVectorsScorer { + @FunctionalInterface + private interface ScorerFactory { + MemorySegmentScorer create(IndexInput in, int dimensions, int dataLength, int bulkSize); + } + + private static final boolean USE_NATIVE = MemorySegmentScorer.NATIVE_SUPPORTED && MemorySegmentScorer.SUPPORTS_HEAP_SEGMENTS; + + private static final Map NATIVE_SCORER_FACTORIES = Map.of( + key(4, 1), NativeBitToInt4Scorer::new, + key(4, 2), NativeDibitToInt4Scorer::new, + key(4, 4), NativeInt4SymmetricScorer::new + ); + + private static final Map SCORER_FACTORIES = Map.of( + key(4, 1), MSBitToInt4ESNextOSQVectorsScorer::new, + key(4, 2), MSDibitToInt4ESNextOSQVectorsScorer::new, + key(4, 4), MSInt4SymmetricESNextOSQVectorsScorer::new, + key(7, 7), MSD7Q7ESNextOSQVectorsScorer::new + ); + + private static int key(int queryBits, int indexBits) { + return (queryBits << 8) | indexBits; + } + private final MemorySegmentScorer scorer; public MemorySegmentESNextOSQVectorsScorer( @@ -45,17 +71,15 @@ public MemorySegmentESNextOSQVectorsScorer( int bulkSize ) { super(in, queryBits, indexBits, dimensions, dataLength); - if (queryBits == 4 && indexBits == 1) { - this.scorer = new MSBitToInt4ESNextOSQVectorsScorer(in, dimensions, dataLength, bulkSize); - } else if (queryBits == 4 && indexBits == 4) { - this.scorer = new MSInt4SymmetricESNextOSQVectorsScorer(in, dimensions, dataLength, bulkSize); - } else if (queryBits == 4 && indexBits == 2) { - this.scorer = new MSDibitToInt4ESNextOSQVectorsScorer(in, dimensions, dataLength, bulkSize); - } else if (queryBits == 7 && indexBits == 7) { - this.scorer = new MSD7Q7ESNextOSQVectorsScorer(in, dimensions, dataLength, bulkSize); - } else { - throw new IllegalArgumentException("Unsupported query/index bits combination: " + queryBits + "/" + indexBits); + int k = key(queryBits, indexBits); + ScorerFactory factory = USE_NATIVE && dataLength >= 16 ? NATIVE_SCORER_FACTORIES.get(k) : null; + if (factory == null) { + factory = Objects.requireNonNull( + SCORER_FACTORIES.get(k), + () -> "Unsupported query/index bits combination: " + queryBits + "/" + indexBits + ); } + this.scorer = factory.create(in, dimensions, dataLength, bulkSize); } @Override @@ -158,10 +182,54 @@ public float scoreBulk( ); } - abstract static sealed class MemorySegmentScorer permits MSBitToInt4ESNextOSQVectorsScorer, MSDibitToInt4ESNextOSQVectorsScorer, - MSInt4SymmetricESNextOSQVectorsScorer, MSD7Q7ESNextOSQVectorsScorer { + @Override + public float scoreBulkOffsets( + byte[] q, + float queryLowerInterval, + float queryUpperInterval, + int queryComponentSum, + float queryAdditionalCorrection, + VectorSimilarityFunction similarityFunction, + float centroidDp, + int[] offsets, + int offsetsCount, + float[] scores, + int count + ) throws IOException { + float score = scorer.scoreBulkOffsets( + q, + queryLowerInterval, + queryUpperInterval, + queryComponentSum, + queryAdditionalCorrection, + similarityFunction, + centroidDp, + offsets, + offsetsCount, + scores, + count + ); + if (score != Float.NEGATIVE_INFINITY) { + return score; + } + return super.scoreBulkOffsets( + q, + queryLowerInterval, + queryUpperInterval, + queryComponentSum, + queryAdditionalCorrection, + similarityFunction, + centroidDp, + offsets, + offsetsCount, + scores, + count + ); + } + + abstract static sealed class MemorySegmentScorer permits NativeMemorySegmentScorer, MSBitToInt4ESNextOSQVectorsScorer, + MSDibitToInt4ESNextOSQVectorsScorer, MSInt4SymmetricESNextOSQVectorsScorer, MSD7Q7ESNextOSQVectorsScorer { - // TODO: split Panama and Native implementations static final boolean NATIVE_SUPPORTED = NativeAccess.instance().getVectorSimilarityFunctions().isPresent(); static final boolean SUPPORTS_HEAP_SEGMENTS = Runtime.version().feature() >= 22; @@ -283,6 +351,20 @@ abstract float scoreBulk( int bulkSize ) throws IOException; + abstract float scoreBulkOffsets( + byte[] q, + float queryLowerInterval, + float queryUpperInterval, + int queryComponentSum, + float queryAdditionalCorrection, + VectorSimilarityFunction similarityFunction, + float centroidDp, + int[] offsets, + int offsetsCount, + float[] scores, + int count + ) throws IOException; + protected float applyCorrectionsIndividually( MemorySegment memorySegment, float queryAdditionalCorrection, diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeBitToInt4Scorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeBitToInt4Scorer.java new file mode 100644 index 0000000000000..cd721e4cce193 --- /dev/null +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeBitToInt4Scorer.java @@ -0,0 +1,58 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ +package org.elasticsearch.simdvec.internal.vectorization; + +import org.apache.lucene.store.IndexInput; + +import java.lang.foreign.MemorySegment; + +import static org.elasticsearch.simdvec.internal.Similarities.dotProductD1Q4; +import static org.elasticsearch.simdvec.internal.Similarities.dotProductD1Q4Bulk; +import static org.elasticsearch.simdvec.internal.Similarities.dotProductD1Q4BulkWithOffsets; + +/** Native scorer for 1-bit index / 4-bit query quantization. */ +final class NativeBitToInt4Scorer extends NativeMemorySegmentScorer { + + NativeBitToInt4Scorer(IndexInput in, int dimensions, int dataLength, int bulkSize) { + super(in, dimensions, dataLength, bulkSize); + } + + @Override + long dotProduct(MemorySegment dataset, MemorySegment query, int length) { + return dotProductD1Q4(dataset, query, length); + } + + @Override + void dotProductBulk(MemorySegment dataset, MemorySegment query, int length, int count, MemorySegment scores) { + dotProductD1Q4Bulk(dataset, query, length, count, scores); + } + + @Override + void dotProductBulkWithOffsets( + MemorySegment dataset, + MemorySegment query, + int dataLength, + int dataStride, + MemorySegment offsets, + int offsetsCount, + MemorySegment scores + ) { + dotProductD1Q4BulkWithOffsets(dataset, query, dataLength, dataStride, offsets, offsetsCount, scores); + } + + @Override + float queryBitScale() { + return FOUR_BIT_SCALE; + } + + @Override + float indexBitScale() { + return ONE_BIT_SCALE; + } +} diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeDibitToInt4Scorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeDibitToInt4Scorer.java new file mode 100644 index 0000000000000..4c072d0f92482 --- /dev/null +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeDibitToInt4Scorer.java @@ -0,0 +1,58 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ +package org.elasticsearch.simdvec.internal.vectorization; + +import org.apache.lucene.store.IndexInput; + +import java.lang.foreign.MemorySegment; + +import static org.elasticsearch.simdvec.internal.Similarities.dotProductD2Q4; +import static org.elasticsearch.simdvec.internal.Similarities.dotProductD2Q4Bulk; +import static org.elasticsearch.simdvec.internal.Similarities.dotProductD2Q4BulkWithOffsets; + +/** Native scorer for 2-bit index / 4-bit query quantization. */ +final class NativeDibitToInt4Scorer extends NativeMemorySegmentScorer { + + NativeDibitToInt4Scorer(IndexInput in, int dimensions, int dataLength, int bulkSize) { + super(in, dimensions, dataLength, bulkSize); + } + + @Override + long dotProduct(MemorySegment dataset, MemorySegment query, int length) { + return dotProductD2Q4(dataset, query, length); + } + + @Override + void dotProductBulk(MemorySegment dataset, MemorySegment query, int length, int count, MemorySegment scores) { + dotProductD2Q4Bulk(dataset, query, length, count, scores); + } + + @Override + void dotProductBulkWithOffsets( + MemorySegment dataset, + MemorySegment query, + int dataLength, + int dataStride, + MemorySegment offsets, + int offsetsCount, + MemorySegment scores + ) { + dotProductD2Q4BulkWithOffsets(dataset, query, dataLength, dataStride, offsets, offsetsCount, scores); + } + + @Override + float queryBitScale() { + return FOUR_BIT_SCALE; + } + + @Override + float indexBitScale() { + return TWO_BIT_SCALE; + } +} diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeInt4SymmetricScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeInt4SymmetricScorer.java new file mode 100644 index 0000000000000..105c766131cd2 --- /dev/null +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeInt4SymmetricScorer.java @@ -0,0 +1,58 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ +package org.elasticsearch.simdvec.internal.vectorization; + +import org.apache.lucene.store.IndexInput; + +import java.lang.foreign.MemorySegment; + +import static org.elasticsearch.simdvec.internal.Similarities.dotProductD4Q4; +import static org.elasticsearch.simdvec.internal.Similarities.dotProductD4Q4Bulk; +import static org.elasticsearch.simdvec.internal.Similarities.dotProductD4Q4BulkWithOffsets; + +/** Native scorer for 4-bit symmetric quantization. */ +final class NativeInt4SymmetricScorer extends NativeMemorySegmentScorer { + + NativeInt4SymmetricScorer(IndexInput in, int dimensions, int dataLength, int bulkSize) { + super(in, dimensions, dataLength, bulkSize); + } + + @Override + long dotProduct(MemorySegment dataset, MemorySegment query, int length) { + return dotProductD4Q4(dataset, query, length); + } + + @Override + void dotProductBulk(MemorySegment dataset, MemorySegment query, int length, int count, MemorySegment scores) { + dotProductD4Q4Bulk(dataset, query, length, count, scores); + } + + @Override + void dotProductBulkWithOffsets( + MemorySegment dataset, + MemorySegment query, + int dataLength, + int dataStride, + MemorySegment offsets, + int offsetsCount, + MemorySegment scores + ) { + dotProductD4Q4BulkWithOffsets(dataset, query, dataLength, dataStride, offsets, offsetsCount, scores); + } + + @Override + float queryBitScale() { + return FOUR_BIT_SCALE; + } + + @Override + float indexBitScale() { + return FOUR_BIT_SCALE; + } +} diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeMemorySegmentScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeMemorySegmentScorer.java new file mode 100644 index 0000000000000..9bbc34ed64231 --- /dev/null +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeMemorySegmentScorer.java @@ -0,0 +1,175 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ +package org.elasticsearch.simdvec.internal.vectorization; + +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.IndexInput; +import org.elasticsearch.simdvec.internal.IndexInputUtils; + +import java.io.IOException; +import java.lang.foreign.MemorySegment; + +/** + * Native scorer for quantized vectors, optimized for JDK 22+ with heap segment support. + * All methods assume {@code NATIVE_SUPPORTED && SUPPORTS_HEAP_SEGMENTS} and use + * {@link MemorySegment#ofArray} directly (no arena copies). + * + *

Subclasses provide only the native dot-product function and bit-scale constants; + * the scoring logic (quantize, bulk, corrections) is handled by template methods here. + */ +abstract sealed class NativeMemorySegmentScorer extends MemorySegmentESNextOSQVectorsScorer.MemorySegmentScorer + permits NativeBitToInt4Scorer, NativeDibitToInt4Scorer, NativeInt4SymmetricScorer { + + NativeMemorySegmentScorer(IndexInput in, int dimensions, int dataLength, int bulkSize) { + super(in, dimensions, dataLength, bulkSize); + assert dataLength >= 16 : "NativeMemorySegmentScorer requires dataLength >= 16, got " + dataLength; + } + + abstract long dotProduct(MemorySegment dataset, MemorySegment query, int length); + + abstract void dotProductBulk(MemorySegment dataset, MemorySegment query, int length, int count, MemorySegment scores); + + abstract void dotProductBulkWithOffsets( + MemorySegment dataset, + MemorySegment query, + int dataLength, + int dataStride, + MemorySegment offsets, + int offsetsCount, + MemorySegment scores + ); + + abstract float queryBitScale(); + + abstract float indexBitScale(); + + @Override + final long quantizeScore(byte[] q) throws IOException { + return IndexInputUtils.withSlice( + in, + length, + this::getScratch, + segment -> dotProduct(segment, MemorySegment.ofArray(q), length) + ); + } + + @Override + final boolean quantizeScoreBulk(byte[] q, int count, float[] scores) throws IOException { + var qSeg = MemorySegment.ofArray(q); + var sSeg = MemorySegment.ofArray(scores); + IndexInputUtils.withSlice(in, (long) length * count, this::getScratch, dSeg -> { + dotProductBulk(dSeg, qSeg, length, count, sSeg); + return null; + }); + return true; + } + + @Override + final boolean quantizeScoreBulkOffsets(byte[] q, int[] offsets, int offsetsCount, float[] scores, int count) throws IOException { + var qSeg = MemorySegment.ofArray(q); + var offsetsSeg = MemorySegment.ofArray(offsets); + var sSeg = MemorySegment.ofArray(scores); + IndexInputUtils.withSlice(in, (long) length * count, this::getScratch, dSeg -> { + dotProductBulkWithOffsets(dSeg, qSeg, length, length, offsetsSeg, offsetsCount, sSeg); + return null; + }); + repositionScoresMatchingOffsets(offsets, offsetsCount, scores); + return true; + } + + @Override + final float scoreBulk( + byte[] q, + float queryLowerInterval, + float queryUpperInterval, + int queryComponentSum, + float queryAdditionalCorrection, + VectorSimilarityFunction similarityFunction, + float centroidDp, + float[] scores, + int bulkSize + ) throws IOException { + var qSeg = MemorySegment.ofArray(q); + var sSeg = MemorySegment.ofArray(scores); + IndexInputUtils.withSlice(in, (long) length * bulkSize, this::getScratch, dSeg -> { + dotProductBulk(dSeg, qSeg, length, bulkSize, sSeg); + return null; + }); + return IndexInputUtils.withSlice( + in, + 16L * bulkSize, + this::getScratch, + cSeg -> ScoreCorrections.nativeApplyCorrectionsBulk( + similarityFunction, + cSeg, + bulkSize, + dimensions, + queryLowerInterval, + queryUpperInterval, + queryComponentSum, + queryAdditionalCorrection, + queryBitScale(), + indexBitScale(), + centroidDp, + sSeg + ) + ); + } + + @Override + final float scoreBulkOffsets( + byte[] q, + float queryLowerInterval, + float queryUpperInterval, + int queryComponentSum, + float queryAdditionalCorrection, + VectorSimilarityFunction similarityFunction, + float centroidDp, + int[] offsets, + int offsetsCount, + float[] scores, + int count + ) throws IOException { + var qSeg = MemorySegment.ofArray(q); + var offsetsSeg = MemorySegment.ofArray(offsets); + var sSeg = MemorySegment.ofArray(scores); + IndexInputUtils.withSlice(in, (long) length * count, this::getScratch, dSeg -> { + dotProductBulkWithOffsets(dSeg, qSeg, length, length, offsetsSeg, offsetsCount, sSeg); + return null; + }); + repositionScoresMatchingOffsets(offsets, offsetsCount, scores); + IndexInputUtils.withSlice(in, 16L * count, this::getScratch, cSeg -> { + ScoreCorrections.nativeApplyCorrectionsBulk( + similarityFunction, + cSeg, + count, + dimensions, + queryLowerInterval, + queryUpperInterval, + queryComponentSum, + queryAdditionalCorrection, + queryBitScale(), + indexBitScale(), + centroidDp, + sSeg + ); + return null; + }); + float maxScore = Float.NEGATIVE_INFINITY; + for (int i = 0, offsetIdx = 0; i < count; i++) { + if (offsetIdx < offsetsCount && offsets[offsetIdx] == i) { + offsetIdx++; + if (scores[i] > maxScore) maxScore = scores[i]; + } else { + scores[i] = 0.0f; + } + } + return maxScore; + } +} From e1dc29fc6b23777b814c02a244c7a1308418d658 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Fri, 20 Mar 2026 17:36:05 +0000 Subject: [PATCH 2/4] [CI] Auto commit changes from spotless --- .../MemorySegmentESNextOSQVectorsScorer.java | 21 ++++++++++++------- .../NativeMemorySegmentScorer.java | 11 +++------- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java index 1f0bd04601153..9314e154bb4eb 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java @@ -44,16 +44,23 @@ private interface ScorerFactory { private static final boolean USE_NATIVE = MemorySegmentScorer.NATIVE_SUPPORTED && MemorySegmentScorer.SUPPORTS_HEAP_SEGMENTS; private static final Map NATIVE_SCORER_FACTORIES = Map.of( - key(4, 1), NativeBitToInt4Scorer::new, - key(4, 2), NativeDibitToInt4Scorer::new, - key(4, 4), NativeInt4SymmetricScorer::new + key(4, 1), + NativeBitToInt4Scorer::new, + key(4, 2), + NativeDibitToInt4Scorer::new, + key(4, 4), + NativeInt4SymmetricScorer::new ); private static final Map SCORER_FACTORIES = Map.of( - key(4, 1), MSBitToInt4ESNextOSQVectorsScorer::new, - key(4, 2), MSDibitToInt4ESNextOSQVectorsScorer::new, - key(4, 4), MSInt4SymmetricESNextOSQVectorsScorer::new, - key(7, 7), MSD7Q7ESNextOSQVectorsScorer::new + key(4, 1), + MSBitToInt4ESNextOSQVectorsScorer::new, + key(4, 2), + MSDibitToInt4ESNextOSQVectorsScorer::new, + key(4, 4), + MSInt4SymmetricESNextOSQVectorsScorer::new, + key(7, 7), + MSD7Q7ESNextOSQVectorsScorer::new ); private static int key(int queryBits, int indexBits) { diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeMemorySegmentScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeMemorySegmentScorer.java index 9bbc34ed64231..364155e512716 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeMemorySegmentScorer.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/NativeMemorySegmentScorer.java @@ -23,8 +23,8 @@ *

Subclasses provide only the native dot-product function and bit-scale constants; * the scoring logic (quantize, bulk, corrections) is handled by template methods here. */ -abstract sealed class NativeMemorySegmentScorer extends MemorySegmentESNextOSQVectorsScorer.MemorySegmentScorer - permits NativeBitToInt4Scorer, NativeDibitToInt4Scorer, NativeInt4SymmetricScorer { +abstract sealed class NativeMemorySegmentScorer extends MemorySegmentESNextOSQVectorsScorer.MemorySegmentScorer permits + NativeBitToInt4Scorer, NativeDibitToInt4Scorer, NativeInt4SymmetricScorer { NativeMemorySegmentScorer(IndexInput in, int dimensions, int dataLength, int bulkSize) { super(in, dimensions, dataLength, bulkSize); @@ -51,12 +51,7 @@ abstract void dotProductBulkWithOffsets( @Override final long quantizeScore(byte[] q) throws IOException { - return IndexInputUtils.withSlice( - in, - length, - this::getScratch, - segment -> dotProduct(segment, MemorySegment.ofArray(q), length) - ); + return IndexInputUtils.withSlice(in, length, this::getScratch, segment -> dotProduct(segment, MemorySegment.ofArray(q), length)); } @Override From ca42a99cdf6bc7f93707e732911809da95a3b7e9 Mon Sep 17 00:00:00 2001 From: ChrisHegarty Date: Sat, 21 Mar 2026 20:16:31 +0000 Subject: [PATCH 3/4] itr --- .../MemorySegmentESNextOSQVectorsScorer.java | 81 ++++++++++++------- .../PanamaESVectorizationProvider.java | 1 + 2 files changed, 51 insertions(+), 31 deletions(-) diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java index 1f0bd04601153..cff8351e71eb8 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java @@ -27,8 +27,6 @@ import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; import java.nio.ByteOrder; -import java.util.Map; -import java.util.Objects; import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT; @@ -36,28 +34,25 @@ /** Panamized scorer for quantized vectors stored as a {@link MemorySegment}. */ public final class MemorySegmentESNextOSQVectorsScorer extends ESNextOSQVectorsScorer { - @FunctionalInterface - private interface ScorerFactory { - MemorySegmentScorer create(IndexInput in, int dimensions, int dataLength, int bulkSize); - } - private static final boolean USE_NATIVE = MemorySegmentScorer.NATIVE_SUPPORTED && MemorySegmentScorer.SUPPORTS_HEAP_SEGMENTS; - private static final Map NATIVE_SCORER_FACTORIES = Map.of( - key(4, 1), NativeBitToInt4Scorer::new, - key(4, 2), NativeDibitToInt4Scorer::new, - key(4, 4), NativeInt4SymmetricScorer::new - ); - - private static final Map SCORER_FACTORIES = Map.of( - key(4, 1), MSBitToInt4ESNextOSQVectorsScorer::new, - key(4, 2), MSDibitToInt4ESNextOSQVectorsScorer::new, - key(4, 4), MSInt4SymmetricESNextOSQVectorsScorer::new, - key(7, 7), MSD7Q7ESNextOSQVectorsScorer::new - ); - - private static int key(int queryBits, int indexBits) { - return (queryBits << 8) | indexBits; + enum QuantEncoding { + BIT_TO_INT4, + DIBIT_TO_INT4, + INT4_SYMMETRIC, + D7Q7; + + static QuantEncoding of(byte queryBits, byte indexBits) { + return switch ((queryBits << 8) | indexBits) { + case (4 << 8) | 1 -> BIT_TO_INT4; + case (4 << 8) | 2 -> DIBIT_TO_INT4; + case (4 << 8) | 4 -> INT4_SYMMETRIC; + case (7 << 8) | 7 -> D7Q7; + default -> throw new IllegalArgumentException( + "Unsupported query/index bits combination: " + queryBits + "/" + indexBits + ); + }; + } } private final MemorySegmentScorer scorer; @@ -71,15 +66,39 @@ public MemorySegmentESNextOSQVectorsScorer( int bulkSize ) { super(in, queryBits, indexBits, dimensions, dataLength); - int k = key(queryBits, indexBits); - ScorerFactory factory = USE_NATIVE && dataLength >= 16 ? NATIVE_SCORER_FACTORIES.get(k) : null; - if (factory == null) { - factory = Objects.requireNonNull( - SCORER_FACTORIES.get(k), - () -> "Unsupported query/index bits combination: " + queryBits + "/" + indexBits - ); - } - this.scorer = factory.create(in, dimensions, dataLength, bulkSize); + this.scorer = USE_NATIVE + ? createNativeScorer(QuantEncoding.of(queryBits, indexBits), in, dimensions, dataLength, bulkSize) + : createPanamaScorer(QuantEncoding.of(queryBits, indexBits), in, dimensions, dataLength, bulkSize); + } + + private static MemorySegmentScorer createNativeScorer( + QuantEncoding enc, + IndexInput in, + int dimensions, + int dataLength, + int bulkSize + ) { + return switch (enc) { + case BIT_TO_INT4 -> new NativeBitToInt4Scorer(in, dimensions, dataLength, bulkSize); + case DIBIT_TO_INT4 -> new NativeDibitToInt4Scorer(in, dimensions, dataLength, bulkSize); + case INT4_SYMMETRIC -> new NativeInt4SymmetricScorer(in, dimensions, dataLength, bulkSize); + case D7Q7 -> new MSD7Q7ESNextOSQVectorsScorer(in, dimensions, dataLength, bulkSize); + }; + } + + private static MemorySegmentScorer createPanamaScorer( + QuantEncoding enc, + IndexInput in, + int dimensions, + int dataLength, + int bulkSize + ) { + return switch (enc) { + case BIT_TO_INT4 -> new MSBitToInt4ESNextOSQVectorsScorer(in, dimensions, dataLength, bulkSize); + case DIBIT_TO_INT4 -> new MSDibitToInt4ESNextOSQVectorsScorer(in, dimensions, dataLength, bulkSize); + case INT4_SYMMETRIC -> new MSInt4SymmetricESNextOSQVectorsScorer(in, dimensions, dataLength, bulkSize); + case D7Q7 -> new MSD7Q7ESNextOSQVectorsScorer(in, dimensions, dataLength, bulkSize); + }; } @Override diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorizationProvider.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorizationProvider.java index 9be69a5767e2c..ba2adfdcb8448 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorizationProvider.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorizationProvider.java @@ -46,6 +46,7 @@ public ESNextOSQVectorsScorer newESNextOSQVectorsScorer( int bulkSize ) { if (PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS + && dataLength >= 16 && ((queryBits == 4 && (indexBits == 1 || indexBits == 2 || indexBits == 4)) || (queryBits == 7 && indexBits == 7))) { IndexInput unwrappedInput = FilterIndexInput.unwrapOnlyTest(input); unwrappedInput = MemorySegmentAccessInputAccess.unwrap(unwrappedInput); From 74fc0f913246aed9a9e3c09c28f55e48cd61eef2 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Sat, 21 Mar 2026 20:23:30 +0000 Subject: [PATCH 4/4] [CI] Auto commit changes from spotless --- .../MemorySegmentESNextOSQVectorsScorer.java | 20 +++---------------- 1 file changed, 3 insertions(+), 17 deletions(-) diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java index cff8351e71eb8..ff9dfa7e62a2d 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java @@ -48,9 +48,7 @@ static QuantEncoding of(byte queryBits, byte indexBits) { case (4 << 8) | 2 -> DIBIT_TO_INT4; case (4 << 8) | 4 -> INT4_SYMMETRIC; case (7 << 8) | 7 -> D7Q7; - default -> throw new IllegalArgumentException( - "Unsupported query/index bits combination: " + queryBits + "/" + indexBits - ); + default -> throw new IllegalArgumentException("Unsupported query/index bits combination: " + queryBits + "/" + indexBits); }; } } @@ -71,13 +69,7 @@ public MemorySegmentESNextOSQVectorsScorer( : createPanamaScorer(QuantEncoding.of(queryBits, indexBits), in, dimensions, dataLength, bulkSize); } - private static MemorySegmentScorer createNativeScorer( - QuantEncoding enc, - IndexInput in, - int dimensions, - int dataLength, - int bulkSize - ) { + private static MemorySegmentScorer createNativeScorer(QuantEncoding enc, IndexInput in, int dimensions, int dataLength, int bulkSize) { return switch (enc) { case BIT_TO_INT4 -> new NativeBitToInt4Scorer(in, dimensions, dataLength, bulkSize); case DIBIT_TO_INT4 -> new NativeDibitToInt4Scorer(in, dimensions, dataLength, bulkSize); @@ -86,13 +78,7 @@ private static MemorySegmentScorer createNativeScorer( }; } - private static MemorySegmentScorer createPanamaScorer( - QuantEncoding enc, - IndexInput in, - int dimensions, - int dataLength, - int bulkSize - ) { + private static MemorySegmentScorer createPanamaScorer(QuantEncoding enc, IndexInput in, int dimensions, int dataLength, int bulkSize) { return switch (enc) { case BIT_TO_INT4 -> new MSBitToInt4ESNextOSQVectorsScorer(in, dimensions, dataLength, bulkSize); case DIBIT_TO_INT4 -> new MSDibitToInt4ESNextOSQVectorsScorer(in, dimensions, dataLength, bulkSize);