From e479701349a6648dc8e5c53a3f265cf2de511b8b Mon Sep 17 00:00:00 2001 From: Ashwin Krishna Kumar Date: Fri, 28 Aug 2026 17:37:16 +0530 Subject: [PATCH] Add a RAVV backed by an fvec file This can be used to create `DataSet` objects that are not held entirely in memory. --- .../jvector/disk/ByteBufferReader.java | 20 +-- .../example/benchmarks/datasets/FvecRavv.java | 162 ++++++++++++++++++ .../datasets/DataSetLoaderSimpleMFDTest.java | 2 +- .../benchmarks/datasets/FvecRavvTest.java | 107 ++++++++++++ 4 files changed, 278 insertions(+), 13 deletions(-) create mode 100644 jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/FvecRavv.java create mode 100644 jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/FvecRavvTest.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ByteBufferReader.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ByteBufferReader.java index 9d02269d7..dcec610b4 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ByteBufferReader.java +++ b/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 diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/FvecRavv.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/FvecRavv.java new file mode 100644 index 000000000..118ba0d5c --- /dev/null +++ b/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 mbbs; + private final int dim; + private final int nvecs; + private final int vecsPerGroup; + + private final VectorFloat bufVec; + + private FvecRavv(List 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(); + 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(); + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java index 33379dd57..c0a76d2a6 100644 --- a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderSimpleMFDTest.java +++ b/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) { diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/FvecRavvTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/benchmarks/datasets/FvecRavvTest.java new file mode 100644 index 000000000..77ed7c47d --- /dev/null +++ b/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))); + } + } +}