diff --git a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt7uTests.java b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt7uTests.java index fb18377144014..7435bd7f88470 100644 --- a/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt7uTests.java +++ b/libs/native/src/test/java/org/elasticsearch/nativeaccess/jdk/JDKVectorLibraryInt7uTests.java @@ -140,6 +140,61 @@ public void testInt7uBulkWithOffsets() { assertScoresEquals(expectedScores, bulkScoresSeg); } + // All offsets point to the same vector; every score in the result should be identical. + public void testInt7uBulkWithDuplicateOffsets() { + assumeTrue(notSupportedMsg(), supported()); + final int dims = size; + final int numVecs = randomIntBetween(4, 101); + var vectors = new byte[numVecs][dims]; + var vectorsSegment = arena.allocate((long) dims * numVecs); + for (int i = 0; i < numVecs; i++) { + randomBytesBetween(vectors[i], MIN_INT7_VALUE, MAX_INT7_VALUE); + MemorySegment.copy(vectors[i], 0, vectorsSegment, ValueLayout.JAVA_BYTE, (long) i * dims, dims); + } + int target = randomInt(numVecs - 1); + var offsets = new int[numVecs]; + var offsetsSegment = arena.allocate((long) numVecs * Integer.BYTES); + for (int i = 0; i < numVecs; i++) { + offsets[i] = target; + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, i, target); + } + int queryOrd = randomInt(numVecs - 1); + float[] expectedScores = new float[numVecs]; + scalarSimilarityBulkWithOffsets(vectors[queryOrd], vectors, offsets, expectedScores); + + var nativeQuerySeg = vectorsSegment.asSlice((long) queryOrd * dims, dims); + var bulkScoresSeg = arena.allocate((long) numVecs * Float.BYTES); + similarityBulkWithOffsets(vectorsSegment, nativeQuerySeg, dims, dims, offsetsSegment, numVecs, bulkScoresSeg); + assertScoresEquals(expectedScores, bulkScoresSeg); + + for (int i = 1; i < numVecs; i++) { + assertEquals(bulkScoresSeg.get(JAVA_FLOAT_UNALIGNED, 0L), bulkScoresSeg.get(JAVA_FLOAT_UNALIGNED, (long) i * Float.BYTES), 0f); + } + } + + // Identity offsets (offsets[i] == i) should produce the same results as the non-offset bulk path. + public void testInt7uBulkWithIdentityOffsets() { + assumeTrue(notSupportedMsg(), supported()); + final int dims = size; + final int numVecs = randomIntBetween(2, 101); + var vectors = new byte[numVecs][dims]; + var vectorsSegment = arena.allocate((long) dims * numVecs); + var offsetsSegment = arena.allocate((long) numVecs * Integer.BYTES); + for (int i = 0; i < numVecs; i++) { + randomBytesBetween(vectors[i], MIN_INT7_VALUE, MAX_INT7_VALUE); + MemorySegment.copy(vectors[i], 0, vectorsSegment, ValueLayout.JAVA_BYTE, (long) i * dims, dims); + offsetsSegment.setAtIndex(ValueLayout.JAVA_INT, i, i); + } + int queryOrd = randomInt(numVecs - 1); + float[] expectedScores = new float[numVecs]; + scalarSimilarityBulk(vectors[queryOrd], vectors, expectedScores); + + var nativeQuerySeg = vectorsSegment.asSlice((long) queryOrd * dims, dims); + var bulkScoresSeg = arena.allocate((long) numVecs * Float.BYTES); + similarityBulkWithOffsets(vectorsSegment, nativeQuerySeg, dims, dims, offsetsSegment, numVecs, bulkScoresSeg); + assertScoresEquals(expectedScores, bulkScoresSeg); + } + public void testInt7uBulkWithOffsetsAndPitch() { assumeTrue(notSupportedMsg(), supported()); final int dims = size;