Add a RAVV backed by an fvec file

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

@@ -46,9 +46,8 @@ public long getPosition() {
 
     @Override
     public void readFully(float[] buffer) {
-        for (int i = 0; i < buffer.length; i++) {
-            buffer[i] = bb.getFloat();
-        }
+        bb.asFloatBuffer().get(buffer);
+        bb.position(bb.position() + buffer.length * Float.BYTES);
     }
 
     @Override
@@ -68,9 +67,8 @@ public void readFully(ByteBuffer buffer) {
 
     @Override
     public void readFully(long[] vector) {
-        for (int i = 0; i < vector.length; i++) {
-            vector[i] = bb.getLong();
-        }
+        bb.asLongBuffer().get(vector);
+        bb.position(bb.position() + vector.length * Long.BYTES);
     }
 
     @Override
@@ -90,16 +88,14 @@ public float readFloat() {
 
     @Override
     public void read(int[] ints, int offset, int count) {
-        for (int i = 0; i < count; i++) {
-            ints[offset + i] = bb.getInt();
-        }
+        bb.asIntBuffer().get(ints, offset, count);
+        bb.position(bb.position() + count * Integer.BYTES);
     }
 
     @Override
     public void read(float[] floats, int offset, int count) {
-        for (int i = 0; i < count; i++) {
-            floats[offset + i] = bb.getFloat();
-        }
+        bb.asFloatBuffer().get(floats, offset, count);
+        bb.position(bb.position() + count * Float.BYTES);
     }
 
     @Override

jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/FvecRavv.java

@@ -0,0 +1,162 @@
+/*
+ * 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.example.benchmarks.datasets;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.MappedByteBuffer;
+import java.nio.channels.FileChannel;
+import java.nio.channels.FileChannel.MapMode;
+import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import io.github.jbellis.jvector.annotations.VisibleForTesting;
+import io.github.jbellis.jvector.disk.ByteBufferReader;
+import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+
+/**
+ * A {@link RandomAccessVectorValues} over a memory-mapped Fvec file.
+ */
+class FvecRavv implements RandomAccessVectorValues {
+
+    private static final ByteOrder byteOrder = ByteOrder.LITTLE_ENDIAN;
+    private static final VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport();
+
+    private static int readInt(FileChannel ch, long position) throws IOException {
+        ByteBuffer bf = ByteBuffer.allocate(Integer.BYTES).order(byteOrder);
+        ch.read(bf, position);
+        bf.flip();
+        return bf.asIntBuffer().get();
+    }
+
+    private final List<MappedByteBuffer> mbbs;
+    private final int dim;
+    private final int nvecs;
+    private final int vecsPerGroup;
+
+    private final VectorFloat<?> bufVec;
+
+    private FvecRavv(List<MappedByteBuffer> mbbs, int dim, int nvecs, int vecsPerGroup) {
+        this.mbbs = mbbs;
+        this.dim = dim;
+        this.nvecs = nvecs;
+        this.vecsPerGroup = vecsPerGroup;
+        this.bufVec = vts.createFloatVector(dim);
+    }
+
+    static FvecRavv of(Path path) throws IOException {
+        return FvecRavv.of(path, Integer.MAX_VALUE);
+    }
+
+    /** Use {@link #of(Path)} instead */
+    @VisibleForTesting
+    static FvecRavv of(Path path, int maxBytesPerGroup) throws IOException {
+        // This static factory method avoids the ReaderSupplierFactory.open() indirection
+        // used elsewhere in JVector.
+        // ReaderSuppliers and RandomAccessReaders are AutoCloseable, which would require this RAVV
+        // to be AutoCloseable, which complicates upstream integration.
+        // This approach sidesteps the problem simce MappedByteBuffers are valid until garbage-collected,
+        // even after the original FileChannel is closed.
+        try (FileChannel ch = FileChannel.open(path, StandardOpenOption.READ)) {
+            long size = ch.size();
+            int dim = readInt(ch, 0);
+            if (dim <= 0) {
+                throw new RuntimeException("Fvec dimension is negative");
+            }
+
+            int vecBytes = Integer.BYTES + dim * Float.BYTES;
+            int vecsPerGroup = maxBytesPerGroup / vecBytes;
+            int groupSizeBytes = vecsPerGroup * vecBytes;
+
+            int nvecs = Math.toIntExact(size / vecBytes);
+            if (nvecs * (long) vecBytes != size) {
+                throw new RuntimeException("File size is not divisible by row size");
+            }
+
+            var mbbs = new ArrayList<MappedByteBuffer>();
+            for (long i = 0; i < size; i += groupSizeBytes) {
+                int mapSize = Math.toIntExact(Math.min(size - i, groupSizeBytes));
+                var mbb = ch.map(MapMode.READ_ONLY, i, mapSize);
+                mbbs.add(mbb);
+            }
+
+            return new FvecRavv(Collections.unmodifiableList(mbbs), dim, nvecs, vecsPerGroup);
+        }
+    }
+
+	@Override
+	public int size() {
+        return nvecs;
+	}
+
+	@Override
+	public int dimension() {
+        return dim;
+	}
+
+    @Override
+    public void getVectorInto(int node, VectorFloat<?> destinationVector, int offset) {
+        if (node >= nvecs) {
+            throw new IndexOutOfBoundsException(node);
+        }
+        int groupId = node / vecsPerGroup;
+        int inGroupId = node % vecsPerGroup;
+        int inGroupByteOffset = (inGroupId + 1) * Integer.BYTES + inGroupId * dim * Float.BYTES;
+
+        var slice = mbbs.get(groupId)
+            .slice()
+            .position(inGroupByteOffset)
+            .limit(inGroupByteOffset + dim * Float.BYTES)
+            .order(byteOrder);
+
+        // close is expected to be a no-op for ByteBufferReader and therefore cheap
+        try (var reader = new ByteBufferReader(slice)) {
+            vts.readFloatVector(reader, dim, destinationVector, offset);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+	@Override
+	public VectorFloat<?> getVector(int nodeId) {
+        this.getVectorInto(nodeId, bufVec, 0);
+        return bufVec;
+	}
+
+	@Override
+	public boolean isValueShared() {
+        return true;
+	}
+
+	@Override
+	public RandomAccessVectorValues copy() {
+        return new FvecRavv(mbbs, dim, nvecs, vecsPerGroup);
+	}
+
+    @VisibleForTesting
+    public int getNumGroups() {
+        return mbbs.size();
+    }
+}

jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java

@@ -1490,7 +1490,7 @@ private static void writeTestDataFiles(Path dir) throws IOException {
     }
 
     /// Writes vectors in the standard fvecs format.
-    private static void writeTestFvecs(Path path, int dimension, float[][] vectors) throws IOException {
+    static void writeTestFvecs(Path path, int dimension, float[][] vectors) throws IOException {
         int bytesPerVector = Integer.BYTES + dimension * Float.BYTES;
         var buf = ByteBuffer.allocate(vectors.length * bytesPerVector).order(ByteOrder.LITTLE_ENDIAN);
         for (float[] vec : vectors) {

jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/FvecRavvTest.java

@@ -0,0 +1,107 @@
+/*
+* 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.example.benchmarks.datasets;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+
+import java.io.IOException;
+import java.nio.file.Path;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+
+public class FvecRavvTest {
+
+    private static final VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport();
+
+    private static float[] toArray(VectorFloat<?> vec) {
+        var arr = new float[vec.length()];
+        for (int i = 0; i < arr.length; i++) {
+            arr[i] = vec.get(i);
+        }
+        return arr;
+    }
+
+    @TempDir Path testDir;
+    Path fvecPath;
+    float[][] referenceVecs;
+    int dim;
+
+    @BeforeEach
+    void createTestData() throws IOException {
+        dim = 4;
+        referenceVecs = new float[][]{
+            new float[]{0.0f, 0.1f, 0.2f, 0.3f},
+            new float[]{1.0f, 1.1f, 1.2f, 1.3f},
+            new float[]{2.0f, 2.1f, 2.2f, 2.3f},
+            new float[]{3.0f, 3.1f, 3.2f, 3.3f},
+            new float[]{4.0f, 4.1f, 4.2f, 4.3f}
+        };
+        fvecPath = testDir.resolve("fvecs");
+        DataSetLoaderSimpleMFDTest.writeTestFvecs(fvecPath, dim, referenceVecs);
+    }
+
+
+    @Test
+    void testGetVector() throws IOException {
+        var ravv = FvecRavv.of(fvecPath);
+        assertEquals(referenceVecs.length, ravv.size());
+        assertEquals(dim, ravv.dimension());
+        for (int i = 0; i < referenceVecs.length; i++) {
+            assertArrayEquals(referenceVecs[i], toArray(ravv.getVector(i)));
+        }
+    }
+
+    @Test
+    void testGetVectorInto() throws IOException {
+        var vec = vts.createFloatVector(dim);
+        var ravv = FvecRavv.of(fvecPath);
+        ravv.getVectorInto(3, vec, 0);
+        assertArrayEquals(referenceVecs[3], toArray(vec));
+    }
+
+    @Test
+    void testCopyReturnsDistinctVec() throws IOException {
+        var ravv = FvecRavv.of(fvecPath);
+        var a = ravv.getVector(2);
+        var ravvCopy = ravv.copy();
+        var b = ravvCopy.getVector(4);
+
+        assertNotSame(a, b);
+        assertArrayEquals(referenceVecs[2], toArray(a));
+        assertArrayEquals(referenceVecs[4], toArray(b));
+    }
+
+    @Test
+    void testMultipleChunks() throws IOException {
+        // little more than the space needed for 2 complete vecs per chunk, but less that 3
+        int maxBytesPerChunk = (Float.BYTES * dim + Integer.BYTES) * 2 + 2;
+        var ravv = FvecRavv.of(fvecPath, maxBytesPerChunk);
+
+        assertEquals(3, ravv.getNumGroups());
+        for (int i = 0; i < referenceVecs.length; i++) {
+            assertArrayEquals(referenceVecs[i], toArray(ravv.getVector(i)));
+        }
+    }
+}