diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/IndexInputUtils.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/IndexInputUtils.java index 41985352a1af0..b61450427fb34 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/IndexInputUtils.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/IndexInputUtils.java @@ -15,7 +15,9 @@ import org.elasticsearch.core.DirectAccessInput; import java.io.IOException; +import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; import java.util.function.IntFunction; /** @@ -74,6 +76,7 @@ public static R withSlice( @SuppressWarnings("unchecked") R[] result = (R[]) new Object[1]; boolean available = dai.withByteBufferSlice(offset, length, bb -> { + assert bb.isDirect(); in.skipBytes(length); result[0] = action.apply(MemorySegment.ofBuffer(bb)); }); @@ -81,7 +84,7 @@ public static R withSlice( return result[0]; } } - return action.apply(copyOnHeap(in, Math.toIntExact(length), scratchSupplier)); + return copyAndApply(in, Math.toIntExact(length), scratchSupplier, action); } /** @@ -100,15 +103,29 @@ public static void checkInputType(IndexInput in) { } } + private static final boolean SUPPORTS_HEAP_SEGMENTS = Runtime.version().feature() >= 22; + /** - * Reads the given number of bytes from the current position of the - * given IndexInput into a heap-backed memory segment. The returned - * segment is sliced to exactly {@code bytesToRead} bytes, even if - * the underlying array is larger. + * Reads bytes from the index input and applies the action to a memory + * segment containing the data. On Java 22+ a heap-backed segment is + * used directly. On Java 21, where heap segments cannot be passed to + * native downcalls, the data is copied into a confined arena. */ - private static MemorySegment copyOnHeap(IndexInput in, int bytesToRead, IntFunction scratchSupplier) throws IOException { + private static R copyAndApply( + IndexInput in, + int bytesToRead, + IntFunction scratchSupplier, + CheckedFunction action + ) throws IOException { byte[] buf = scratchSupplier.apply(bytesToRead); in.readBytes(buf, 0, bytesToRead); - return MemorySegment.ofArray(buf).asSlice(0, bytesToRead); + if (SUPPORTS_HEAP_SEGMENTS) { + return action.apply(MemorySegment.ofArray(buf).asSlice(0, bytesToRead)); + } + try (Arena arena = Arena.ofConfined()) { + MemorySegment nativeSegment = arena.allocate(bytesToRead); + MemorySegment.copy(buf, 0, nativeSegment, ValueLayout.JAVA_BYTE, 0, bytesToRead); + return action.apply(nativeSegment); + } } } diff --git a/libs/simdvec/src/test21/java/org/elasticsearch/simdvec/internal/IndexInputUtilsTests.java b/libs/simdvec/src/test21/java/org/elasticsearch/simdvec/internal/IndexInputUtilsTests.java index 5edfaa16a7af7..683060fd3ad53 100644 --- a/libs/simdvec/src/test21/java/org/elasticsearch/simdvec/internal/IndexInputUtilsTests.java +++ b/libs/simdvec/src/test21/java/org/elasticsearch/simdvec/internal/IndexInputUtilsTests.java @@ -21,7 +21,9 @@ import org.elasticsearch.test.ESTestCase; import java.io.IOException; +import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; import java.nio.ByteBuffer; import java.util.Arrays; @@ -72,6 +74,26 @@ public void testWithSlicePlainIndexInput() throws Exception { } } + public void testPlainInputFallbackProducesNativeSegmentOnJava21() throws Exception { + byte[] data = randomByteArrayOfLength(256); + try (Directory dir = new NIOFSDirectory(createTempDir())) { + writeData(dir, data); + try (IndexInput in = dir.openInput(FILE_NAME, IOContext.DEFAULT)) { + assertFalse(in instanceof MemorySegmentAccessInput); + assertFalse(in instanceof DirectAccessInput); + IndexInputUtils.withSlice(in, data.length, byte[]::new, segment -> { + if (Runtime.version().feature() < 22) { + assertTrue("segment should be native-backed on Java 21", segment.heapBase().isEmpty()); + } + byte[] buf = new byte[(int) segment.byteSize()]; + MemorySegment.ofArray(buf).copyFrom(segment); + assertArrayEquals(data, buf); + return null; + }); + } + } + } + // -- constructor validation tests ----------------------------------------- public void testES92ConstructorAcceptsPlainInput() throws Exception { @@ -167,9 +189,12 @@ static class DirectAccessWrapper extends FilterIndexInput implements DirectAcces @Override public boolean withByteBufferSlice(long offset, long length, CheckedConsumer action) throws IOException { - ByteBuffer bb = ByteBuffer.wrap(data, (int) offset, (int) length).asReadOnlyBuffer(); - action.accept(bb); - return true; + try (Arena arena = Arena.ofConfined()) { + MemorySegment segment = arena.allocate(length); + MemorySegment.copy(data, (int) offset, segment, ValueLayout.JAVA_BYTE, 0, (int) length); + action.accept(segment.asByteBuffer().asReadOnlyBuffer()); + return true; + } } @Override