Add a memory-mapped RandomAccessReader using MemorySegment api (#296)

* Add a memory-mapped RandomAccessReader using MemorySegment api
* Use MemorySegmentReaderSupplier in ReaderSupplierFactory

Author: mdogan

jvector-examples/pom.xml

@@ -214,7 +214,7 @@
             </build>
         </profile>
         <profile>
-            <id>jdk21</id>
+            <id>jdk22</id>
             <activation>
                 <activeByDefault>true</activeByDefault>
             </activation>

jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/ReaderSupplierFactory.java

@@ -19,15 +19,26 @@
 import io.github.jbellis.jvector.disk.SimpleMappedReaderSupplier;
 
 import java.io.IOException;
+import java.lang.reflect.Constructor;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 public class ReaderSupplierFactory {
     private static final Logger LOG = Logger.getLogger(ReaderSupplierFactory.class.getName());
+    private static final String MEMORY_SEGMENT_READER_CLASSNAME = "io.github.jbellis.jvector.disk.MemorySegmentReaderSupplier";
 
     public static ReaderSupplier open(Path path) throws IOException {
+        try {
+            var supplierClass = Class.forName(MEMORY_SEGMENT_READER_CLASSNAME);
+            Constructor<?> ctor = supplierClass.getConstructor(Path.class);
+            return (ReaderSupplier) ctor.newInstance(path);
+        } catch (Exception e) {
+            LOG.log(Level.WARNING, "MemorySegmentReaderSupplier not available, falling back to MMapReaderSupplier. Reason: {0}: {1}",
+                    new Object[]{e.getClass().getName(), e.getMessage()});
+        }
+
         try {
             return new MMapReaderSupplier(path);
         } catch (UnsatisfiedLinkError|NoClassDefFoundError e) {

jvector-native/pom.xml

@@ -35,6 +35,14 @@
                     </nonFilteredFileExtensions>
                 </configuration>
             </plugin>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-surefire-plugin</artifactId>
+                <version>3.1.2</version>
+                <configuration>
+                    <skip>false</skip>
+                </configuration>
+            </plugin>
         </plugins>
     </build>
     <profiles>

jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReader.java

@@ -0,0 +1,145 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.disk;
+
+import java.io.IOException;
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.ValueLayout;
+import java.lang.foreign.ValueLayout.OfFloat;
+import java.lang.foreign.ValueLayout.OfInt;
+import java.lang.foreign.ValueLayout.OfLong;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.channels.FileChannel;
+import java.nio.channels.FileChannel.MapMode;
+import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+
+/**
+ * {@link MemorySegment} based implementation of RandomAccessReader.
+ * MemorySegmentReader doesn't have 2GB file size limitation of {@link SimpleMappedReader}.
+ */
+public class MemorySegmentReader implements RandomAccessReader {
+
+    private static final OfInt intLayout = ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN);
+    private static final OfFloat floatLayout = ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN);
+    private static final OfLong longLayout = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN);
+
+    final Arena arena;
+    final MemorySegment memory;
+    private long position = 0;
+
+    public MemorySegmentReader(Path path) throws IOException {
+        arena = Arena.ofShared();
+        try (var ch = FileChannel.open(path, StandardOpenOption.READ)) {
+            memory = ch.map(MapMode.READ_ONLY, 0L, ch.size(), arena);
+        } catch (Exception e) {
+            arena.close();
+            throw e;
+        }
+    }
+
+    MemorySegmentReader(Arena arena, MemorySegment memory) {
+        this.arena = arena;
+        this.memory = memory;
+    }
+
+    @Override
+    public void seek(long offset) {
+        this.position = offset;
+    }
+
+    @Override
+    public long getPosition() {
+        return position;
+    }
+
+    @Override
+    public void readFully(float[] buffer) {
+        MemorySegment.copy(memory, floatLayout, position, buffer, 0, buffer.length);
+        position += buffer.length * 4L;
+    }
+
+    @Override
+    public void readFully(byte[] b) {
+        MemorySegment.copy(memory, ValueLayout.JAVA_BYTE, position, b, 0, b.length);
+        position += b.length;
+    }
+
+    @Override
+    public void readFully(ByteBuffer buffer) {
+        var remaining = buffer.remaining();
+        var slice = memory.asSlice(position, remaining).asByteBuffer();
+        buffer.put(slice);
+        position += remaining;
+    }
+
+    @Override
+    public void readFully(long[] vector) {
+        MemorySegment.copy(memory, longLayout, position, vector, 0, vector.length);
+        position += vector.length * 8L;
+    }
+
+    @Override
+    public int readInt() {
+        var k = memory.get(intLayout, position);
+        position += 4;
+        return k;
+    }
+
+    @Override
+    public float readFloat() {
+        var f = memory.get(floatLayout, position);
+        position += 4;
+        return f;
+    }
+
+    @Override
+    public void read(int[] ints, int offset, int count) {
+        MemorySegment.copy(memory, intLayout, position, ints, offset, count);
+        position += count * 4L;
+    }
+
+    @Override
+    public void read(float[] floats, int offset, int count) {
+        MemorySegment.copy(memory, floatLayout, position, floats, offset, count);
+        position += count * 4L;
+    }
+
+    /**
+     * Loads the contents of the mapped segment into physical memory.
+     * This is a best-effort mechanism.
+     */
+    public void loadMemory() {
+        memory.load();
+    }
+
+    @Override
+    public void close() {
+        arena.close();
+    }
+
+    /**
+     * Creates a shallow copy of the MemorySegmentReader.
+     * Underlying memory-mapped region will be shared between all copies.
+     * When MemorySegmentReader is closed, all copies will become invalid.
+     */
+    public MemorySegmentReader duplicate() {
+        return new MemorySegmentReader(arena, memory);
+    }
+}

jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReaderSupplier.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.github.jbellis.jvector.disk;
+
+import java.io.IOException;
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemorySegment;
+import java.nio.file.Path;
+
+public class MemorySegmentReaderSupplier implements ReaderSupplier {
+    private final InternalMemorySegmentReader reader;
+
+    public MemorySegmentReaderSupplier(Path path) throws IOException {
+        reader = new InternalMemorySegmentReader(path);
+    }
+
+    @Override
+    public RandomAccessReader get() {
+        return reader.duplicate();
+    }
+
+    @Override
+    public void close() {
+        reader.close();
+    }
+
+    private static class InternalMemorySegmentReader extends MemorySegmentReader {
+
+        private final boolean shouldClose;
+
+        private InternalMemorySegmentReader(Path path) throws IOException {
+            super(path);
+            shouldClose = true;
+        }
+
+        private InternalMemorySegmentReader(Arena arena, MemorySegment memory) {
+            super(arena, memory);
+            shouldClose = false;
+        }
+
+        @Override
+        public void close() {
+            if (shouldClose) {
+                super.close();
+            }
+        }
+
+        @Override
+        public InternalMemorySegmentReader duplicate() {
+            return new InternalMemorySegmentReader(arena, memory);
+        }
+    }
+}