standardize ReaderSupplier implementations as inner classes of their …

…respective RandomAccessReaders, and add a Supplier for SimpleReader (#323)

Co-authored-by: Mehmet Dogan <[email protected]>

jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java

@@ -19,7 +19,11 @@
 import java.io.IOException;
 
 public interface ReaderSupplier extends AutoCloseable {
-    RandomAccessReader get();
+    /**
+     * @return a new reader.  It is up to the caller to re-use these readers or close them,
+     * the ReaderSupplier is not responsible for caching them.
+     */
+    RandomAccessReader get() throws IOException;
 
     default void close() throws IOException {
     }

jvector-base/src/main/java/io/github/jbellis/jvector/disk/SimpleMappedReader.java

@@ -85,4 +85,22 @@ public void close() {
     public SimpleMappedReader duplicate() {
         return new SimpleMappedReader((MappedByteBuffer) bb.duplicate());
     }
+
+    public static class Supplier implements ReaderSupplier {
+        private final SimpleMappedReader smr;
+
+        public Supplier(Path path) throws IOException {
+            smr = new SimpleMappedReader(path);
+        }
+
+        @Override
+        public RandomAccessReader get() {
+            return smr.duplicate();
+        }
+
+        @Override
+        public void close() {
+            smr.close();
+        }
+    }
 }

jvector-base/src/main/java/io/github/jbellis/jvector/disk/SimpleReader.java

@@ -94,4 +94,17 @@ public void read(float[] floats, int offset, int count) throws IOException {
     public void close() throws IOException {
         raf.close();
     }
+
+    public static class Supplier implements ReaderSupplier {
+        private final Path path;
+
+        public Supplier(Path path) {
+            this.path = path;
+        }
+
+        @Override
+        public RandomAccessReader get() throws FileNotFoundException {
+            return new SimpleReader(path);
+        }
+    }
 }

jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java

@@ -144,7 +144,11 @@ public String toString() {
     // re-declared to specify type
     @Override
     public View getView() {
-        return new View(readerSupplier.get());
+        try {
+            return new View(readerSupplier.get());
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
     }
 
     public class View implements FeatureSource, ScoringView, RandomAccessVectorValues {

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

@@ -19,8 +19,7 @@
 import io.github.jbellis.jvector.disk.RandomAccessReader;
 import io.github.jbellis.jvector.disk.ReaderSupplier;
 import io.github.jbellis.jvector.disk.SimpleMappedReader;
-import io.github.jbellis.jvector.disk.SimpleMappedReaderSupplier;
-import io.github.jbellis.jvector.example.util.MMapReaderSupplier;
+import io.github.jbellis.jvector.example.util.MMapReader;
 import io.github.jbellis.jvector.example.util.SiftLoader;
 import io.github.jbellis.jvector.graph.GraphIndex;
 import io.github.jbellis.jvector.graph.GraphIndexBuilder;
@@ -157,7 +156,7 @@ public static void siftPersisted(List<VectorFloat<?>> baseVectors, List<VectorFl
 
         // on-disk indexes require a ReaderSupplier (not just a Reader) because we will want it to
         // open additional readers for searching
-        ReaderSupplier rs = new SimpleMappedReaderSupplier(indexPath);
+        ReaderSupplier rs = new SimpleMappedReader.Supplier(indexPath);
         OnDiskGraphIndex index = OnDiskGraphIndex.load(rs);
         // measure our recall against the (exactly computed) ground truth
         Function<VectorFloat<?>, SearchScoreProvider> sspFactory = q -> SearchScoreProvider.exact(q, VectorSimilarityFunction.EUCLIDEAN, ravv);
@@ -190,7 +189,7 @@ public static void siftDiskAnn(List<VectorFloat<?>> baseVectors, List<VectorFloa
             pqv.write(out);
         }
 
-        ReaderSupplier rs = new MMapReaderSupplier(indexPath);
+        ReaderSupplier rs = new MMapReader.Supplier(indexPath);
         OnDiskGraphIndex index = OnDiskGraphIndex.load(rs);
         // load the PQVectors that we just wrote to disk
         try (RandomAccessReader in = new SimpleMappedReader(pqPath)) {
@@ -259,7 +258,7 @@ public static void siftDiskAnnLTM(List<VectorFloat<?>> baseVectors, List<VectorF
         }
 
         // searching the index does not change
-        ReaderSupplier rs = new MMapReaderSupplier(indexPath);
+        ReaderSupplier rs = new MMapReader.Supplier(indexPath);
         OnDiskGraphIndex index = OnDiskGraphIndex.load(rs);
         try (RandomAccessReader in = new SimpleMappedReader(pqPath)) {
             PQVectors pqv = PQVectors.load(in);

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

@@ -17,10 +17,13 @@
 
 import com.indeed.util.mmap.MMapBuffer;
 import io.github.jbellis.jvector.disk.RandomAccessReader;
+import io.github.jbellis.jvector.disk.ReaderSupplier;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.nio.channels.FileChannel;
+import java.nio.file.Path;
 
 public class MMapReader implements RandomAccessReader {
     private final MMapBuffer buffer;
@@ -127,4 +130,22 @@ public void readFully(long[] vector) {
     public void close() {
         // don't close buffer, let the Supplier handle that
     }
+
+    public static class Supplier implements ReaderSupplier {
+        private final MMapBuffer buffer;
+
+        public Supplier(Path path) throws IOException {
+            buffer = new MMapBuffer(path, FileChannel.MapMode.READ_ONLY, ByteOrder.BIG_ENDIAN);
+        }
+
+        @Override
+        public RandomAccessReader get() {
+            return new MMapReader(buffer);
+        }
+
+        @Override
+        public void close() throws IOException {
+            buffer.close();
+        }
+    }
 }

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

@@ -16,7 +16,7 @@
 package io.github.jbellis.jvector.example.util;
 
 import io.github.jbellis.jvector.disk.ReaderSupplier;
-import io.github.jbellis.jvector.disk.SimpleMappedReaderSupplier;
+import io.github.jbellis.jvector.disk.SimpleMappedReader;
 
 import java.io.IOException;
 import java.lang.reflect.Constructor;
@@ -27,7 +27,7 @@
 
 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";
+    private static final String MEMORY_SEGMENT_READER_CLASSNAME = "io.github.jbellis.jvector.disk.MemorySegmentReader.Supplier";
 
     public static ReaderSupplier open(Path path) throws IOException {
         try {
@@ -40,15 +40,15 @@ public static ReaderSupplier open(Path path) throws IOException {
         }
 
         try {
-            return new MMapReaderSupplier(path);
+            return new MMapReader.Supplier(path);
         } catch (UnsatisfiedLinkError|NoClassDefFoundError e) {
             LOG.log(Level.WARNING, "MMapReaderSupplier not available, falling back to SimpleMappedReaderSupplier. More details available at level FINE.");
             LOG.log(Level.FINE, "MMapReaderSupplier instantiation exception:", e);
             if (Files.size(path) > Integer.MAX_VALUE) {
                 throw new RuntimeException("File sizes greater than 2GB are not supported on Windows--contributions welcome");
             }
 
-            return new SimpleMappedReaderSupplier(path);
+            return new SimpleMappedReader.Supplier(path);
         }
     }
 }

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

@@ -142,4 +142,48 @@ public void close() {
     public MemorySegmentReader duplicate() {
         return new MemorySegmentReader(arena, memory);
     }
+
+    public static class Supplier implements ReaderSupplier {
+        private final InternalMemorySegmentReader reader;
+
+        public Supplier(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);
+            }
+        }
+    }
 }

jvector-native/src/test/java/io/github/jbellis/jvector/disk/MemorySegmentReaderTest.java

@@ -102,7 +102,7 @@ public void testReaderClose() throws Exception {
 
     @Test
     public void testSupplierClose() throws Exception {
-        var s = new MemorySegmentReaderSupplier(tempFile);
+        var s = new MemorySegmentReader.Supplier(tempFile);
         var r1 = s.get();
         var r2 = s.get();