dataset loader improvements

benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java

@@ -249,7 +249,6 @@ private static void writeCompletedCount(int count) {
     // ---------- Benchmark state ----------
     private RandomAccessVectorValues ravv;
     private List<VectorFloat<?>> queryVectors;
-    private List<VectorFloat<?>> baseVectors;
     private List<? extends List<Integer>> groundTruth;
     private DataSet ds;
     private VectorSimilarityFunction similarityFunction;
@@ -399,7 +398,6 @@ public void setup() throws Exception {
 
                 if (datasetPortion == 1.0) {
                     ravv = ds.getBaseRavv();
-                    baseVectors = ds.getBaseVectors();
                 } else {
                     int totalVectors = ds.getBaseRavv().size();
                     int portionedSize = (int) (totalVectors * datasetPortion);
@@ -408,8 +406,7 @@ public void setup() throws Exception {
                                 "datasetPortion=" + datasetPortion + " yields " + portionedSize
                                         + " vectors, fewer than numPartitions=" + numPartitions);
                     }
-                    baseVectors = ds.getBaseVectors().subList(0, portionedSize);
-                    ravv = new ListRandomAccessVectorValues(baseVectors, ds.getDimension());
+                    ravv = ds.getBaseRavv().range(0, portionedSize);
                 }
 
                 similarityFunction = ds.getSimilarityFunction();
@@ -428,7 +425,6 @@ public void setup() throws Exception {
                 }
             } else {
                 ravv = null;
-                baseVectors = null;
                 dimension = -1;
 
                 if (needsRecallData) {
@@ -469,7 +465,7 @@ public void setup() throws Exception {
             }
 
             if (workloadMode == WorkloadMode.PARTITION || workloadMode == WorkloadMode.PARTITION_AND_COMPACT) {
-                var partitionedData = DataSetPartitioner.partition(baseVectors, numPartitions, splitDistribution);
+                var partitionedData = DataSetPartitioner.partition(ravv, numPartitions, splitDistribution);
                 vectorsPerSourceCount = partitionedData.sizes;
             } else {
                 vectorsPerSourceCount = null;
@@ -479,7 +475,7 @@ public void setup() throws Exception {
                 if (jfrPartitioning) {
                     jfrPartitioningRecorder.start(JFR_DIR, "partitioning-" + jfrParamSuffix() + ".jfr", jfrObjectCount);
                 }
-                buildPartitions(ds, baseVectors);
+                buildPartitions(ravv);
                 if (jfrPartitioningRecorder.isActive()) {
                     jfrPartitioningRecorder.stop();
                 }
@@ -595,7 +591,7 @@ private void verifyPartitionsExist(Path partitionsDir, int numPartitions) {
         }
     }
 
-    private void buildPartitions(DataSet ds, List<VectorFloat<?>> baseVectors) throws Exception {
+    private void buildPartitions(RandomAccessVectorValues baseVectors) throws Exception {
 
         var partitionedData = DataSetPartitioner.partition(baseVectors, numPartitions, splitDistribution);
         vectorsPerSourceCount = partitionedData.sizes;
@@ -604,9 +600,9 @@ private void buildPartitions(DataSet ds, List<VectorFloat<?>> baseVectors) throw
                 numPartitions, partitionsBaseDir.toAbsolutePath(), graphDegree, beamWidth, splitDistribution, vectorsPerSourceCount,
                 indexPrecision, parallelWriteThreads, resolvedVectorizationProvider);
 
-        int dimension = baseVectors.get(0).length();
+        int dimension = baseVectors.dimension();
         for (int i = 0; i < numPartitions; i++) {
-            List<VectorFloat<?>> vectorsPerSource = partitionedData.vectors.get(i);
+            RandomAccessVectorValues ravvPerSource = partitionedData.vectors.get(i);
 
             // Round-robin assignment of partition files to storage paths, but still keep canonical base dir name stable.
             Path baseDirForThisSegment = storagePaths.get(i % storagePaths.size());
@@ -616,9 +612,8 @@ private void buildPartitions(DataSet ds, List<VectorFloat<?>> baseVectors) throw
             }
 
             log.info("Building partition {}/{}: vectors={} -> {}",
-                    i + 1, numPartitions, vectorsPerSource.size(), outputPath.toAbsolutePath());
+                    i + 1, numPartitions, ravvPerSource.size(), outputPath.toAbsolutePath());
 
-            var ravvPerSource = new ListRandomAccessVectorValues(vectorsPerSource, dimension);
             BuildScoreProvider bspPerSource;
             ProductQuantization pq = null;
             PQVectors pqVectors = null;
@@ -655,7 +650,8 @@ private void buildPartitions(DataSet ds, List<VectorFloat<?>> baseVectors) throw
 
             try (var writer = writerBuilder.build()) {
                 var suppliers = new EnumMap<FeatureId, IntFunction<Feature.State>>(FeatureId.class);
-                suppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(ravvPerSource.getVector(ordinal)));
+                var vectors = ravvPerSource.threadLocalSupplier(); // the parallel writer calls suppliers from worker threads
+                suppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(vectors.get().getVector(ordinal)));
 
                 if (indexPrecision == IndexPrecision.FUSEDPQ) {
                     var view = graph.getView();
@@ -713,16 +709,16 @@ private long compactPartitions() throws Exception {
         return compactionTimeMs;
     }
 
-    private long buildFromScratch(List<VectorFloat<?>> baseVectors) throws Exception {
+    private long buildFromScratch(RandomAccessVectorValues baseVectors) throws Exception {
         if (scratchOutputPath.getParent() != null) {
             Files.createDirectories(scratchOutputPath.getParent());
         }
         if (Files.exists(scratchOutputPath)) {
             Files.delete(scratchOutputPath);
         }
 
-        int dimension = baseVectors.get(0).length();
-        var full = new ListRandomAccessVectorValues(baseVectors, dimension);
+        int dimension = baseVectors.dimension();
+        var full = baseVectors;
 
         log.info("Building from scratch: vectors={} dim={} sim={} deg={} bw={} precision={} pwThreads={} vp={} -> {}",
                 full.size(), dimension, similarityFunction,
@@ -761,7 +757,8 @@ private long buildFromScratch(List<VectorFloat<?>> baseVectors) throws Exception
 
         try (var writer = writerBuilder.build()) {
             var suppliers = new EnumMap<FeatureId, IntFunction<Feature.State>>(FeatureId.class);
-            suppliers.put(FeatureId.INLINE_VECTORS, ord -> new InlineVectors.State(full.getVector(ord)));
+            var vectors = full.threadLocalSupplier(); // the parallel writer calls suppliers from worker threads
+            suppliers.put(FeatureId.INLINE_VECTORS, ord -> new InlineVectors.State(vectors.get().getVector(ord)));
 
             if (indexPrecision == IndexPrecision.FUSEDPQ) {
                 var view = graph.getView();
@@ -843,7 +840,7 @@ public void run(Blackhole blackhole, RecallResult recallResult) throws Exception
                     break;
 
                 case BUILD:
-                    durationMs = buildFromScratch(baseVectors);
+                    durationMs = buildFromScratch(ravv);
                     if (measureRecall) {
                         searchStats = runRecall(scratchOutputPath);
                         recall = searchStats.recall;

docs/benchmarking.md

@@ -33,6 +33,21 @@ Datasets are grouped into categories. The categories can be arbitrarily chosen f
 
 Dataset similarity functions are configured in `jvector-examples/yaml-configs/dataset-metadata.yml`.
 
+Each entry may carry a loader *profile* and a list of *wrappers* that change how the base vectors are held after loading. Two forms are accepted:
+
+```yaml
+regression-tests:
+  - cap-1M                        # name only: default profile, base vectors cached in heap memory
+  - cohere-english-v3-1M(mmap)    # sugared: name, then wrappers in parentheses
+  - name: cohere-english-v3-10M   # structured
+    profile: default              # optional; "default" when omitted
+    wrappers:                     # optional; applied left to right
+      - mmap
+      - lru: { grain: 4096, capacityMb: 512 }   # a wrapper with options
+```
+
+The sugared form is `name`, `name:profile`, `name(wrapper,...)` or `name:profile(wrapper,...)`, where a wrapper may carry options in brackets as in `lru[grain=4096,capacityMb=512]`, and is also accepted anywhere a dataset name is (command-line patterns, `dataset:` in an index-parameters file). Built-in wrappers are `memory` (cache the base vectors in heap memory, the default when no wrappers are given), `mmap` (serve them from a memory-mapped fvecs file) and `lru` (a bounded least-recently-used cache of vector grains in front of the mapped file, for datasets larger than memory; options `grain` (vectors per grain) and `capacityMb`, defaulting to `-Djvector.dataset.lru.grain` and `-Djvector.dataset.lru.capacityMb` or 1024 and a quarter of the heap). Loaders that do not understand profiles accept `default` and reject any other profile.
+
 Example `datasets.yml`:
 
 ```yaml

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

@@ -78,6 +78,27 @@ default void getVectorInto(int node, VectorFloat<?> destinationVector, int offse
         destinationVector.copyFrom(getVector(node), 0, offset, dimension());
     }
 
+    /**
+     * Returns a view of the ordinal range {@code [fromOrdinal, toOrdinal)} of this RAVV, re-based so that
+     * ordinal {@code i} of the view reads ordinal {@code fromOrdinal + i} of this RAVV.
+     * <p>
+     * The view does not copy vectors: it shares the underlying storage and inherits the sharing semantics
+     * of {@link #isValueShared()}. Its size is fixed at {@code toOrdinal - fromOrdinal} when this method is
+     * called, even if this RAVV later grows. Passing {@code 0} and {@link #size()} yields a view equivalent
+     * to this RAVV.
+     * <p>
+     * The default implementation returns a {@link RangeRandomAccessVectorValues}; implementations with a cheaper
+     * native ranged read may override this.
+     *
+     * @param fromOrdinal the first ordinal of the range, inclusive; must be &ge; 0
+     * @param toOrdinal the last ordinal of the range, exclusive; must be &ge; {@code fromOrdinal} and &le; {@link #size()}
+     * @return a RAVV of size {@code toOrdinal - fromOrdinal} over the requested range
+     * @throws IndexOutOfBoundsException if the range is not within {@code [0, size()]}
+     */
+    default RandomAccessVectorValues range(int fromOrdinal, int toOrdinal) {
+        return new RangeRandomAccessVectorValues(this, fromOrdinal, toOrdinal);
+    }
+
     /**
      * @return true iff the vector returned by `getVector` is shared.  A shared vector will
      * only be valid until the next call to getVector overwrites it.

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

@@ -0,0 +1,113 @@
+/*
+ * 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.graph;
+
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+import java.util.Objects;
+
+/**
+ * A re-based view over a contiguous ordinal range {@code [fromOrdinal, toOrdinal)} of a backing
+ * {@link RandomAccessVectorValues}. Ordinal {@code i} of the view reads ordinal {@code fromOrdinal + i}
+ * of the backing RAVV.
+ * <p>
+ * The view does not copy vectors. Its size is fixed at construction time, so a backing RAVV that grows
+ * afterwards (for example a {@link ListRandomAccessVectorValues} over a list that is still being appended to)
+ * is not reflected in the view. Value-sharing semantics are inherited from the backing RAVV.
+ * <p>
+ * This is the default implementation returned by {@link RandomAccessVectorValues#range(int, int)}; implementations
+ * with a cheaper native ranged read may override that method instead.
+ */
+public final class RangeRandomAccessVectorValues implements RandomAccessVectorValues {
+    private final RandomAccessVectorValues ravv;
+    private final int fromOrdinal;
+    private final int size;
+
+    /**
+     * Creates a view over ordinals {@code [fromOrdinal, toOrdinal)} of {@code ravv}.
+     *
+     * @param ravv        the backing RAVV
+     * @param fromOrdinal the first backing ordinal of the range, inclusive
+     * @param toOrdinal   the last backing ordinal of the range, exclusive
+     * @throws IndexOutOfBoundsException if the range is not within {@code [0, ravv.size()]}
+     */
+    public RangeRandomAccessVectorValues(RandomAccessVectorValues ravv, int fromOrdinal, int toOrdinal) {
+        Objects.checkFromToIndex(fromOrdinal, toOrdinal, ravv.size());
+        this.ravv = ravv;
+        this.fromOrdinal = fromOrdinal;
+        this.size = toOrdinal - fromOrdinal;
+    }
+
+    /**
+     * @return the backing ordinal that view ordinal {@code 0} maps to
+     */
+    public int fromOrdinal() {
+        return fromOrdinal;
+    }
+
+    /**
+     * @return the exclusive upper bound of the backing ordinal range
+     */
+    public int toOrdinal() {
+        return fromOrdinal + size;
+    }
+
+    @Override
+    public int size() {
+        return size;
+    }
+
+    @Override
+    public int dimension() {
+        return ravv.dimension();
+    }
+
+    @Override
+    public VectorFloat<?> getVector(int nodeId) {
+        return ravv.getVector(fromOrdinal + Objects.checkIndex(nodeId, size));
+    }
+
+    @Override
+    public void getVectorInto(int node, VectorFloat<?> destinationVector, int offset) {
+        ravv.getVectorInto(fromOrdinal + Objects.checkIndex(node, size), destinationVector, offset);
+    }
+
+    @Override
+    public boolean isValueShared() {
+        return ravv.isValueShared();
+    }
+
+    /**
+     * Copies the backing RAVV and wraps the copy in an equivalent view. If the backing RAVV is un-shared
+     * and returns itself from {@link RandomAccessVectorValues#copy()}, this view returns itself as well.
+     */
+    @Override
+    public RandomAccessVectorValues copy() {
+        RandomAccessVectorValues copied = ravv.copy();
+        return copied == ravv ? this : new RangeRandomAccessVectorValues(copied, fromOrdinal, toOrdinal());
+    }
+
+    /**
+     * Narrows this view without adding a level of indirection: the result reads the backing RAVV directly
+     * at {@code [this.fromOrdinal + fromOrdinal, this.fromOrdinal + toOrdinal)}.
+     */
+    @Override
+    public RandomAccessVectorValues range(int fromOrdinal, int toOrdinal) {
+        Objects.checkFromToIndex(fromOrdinal, toOrdinal, size);
+        return new RangeRandomAccessVectorValues(ravv, this.fromOrdinal + fromOrdinal, this.fromOrdinal + toOrdinal);
+    }
+}

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

@@ -126,7 +126,7 @@ public static void main(String[] args) throws IOException {
                     DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow(
                             () -> new RuntimeException("Dataset " + datasetName + " not found")
                     ).getDataSet();
-                    logger.info("Dataset loaded: {} with {} vectors", datasetName, ds.getBaseVectors().size());
+                    logger.info("Dataset loaded: {} with {} vectors", datasetName, ds.getBaseRavv().size());
 
                     String normalizedDatasetName = datasetName;
                     if (normalizedDatasetName.endsWith(".hdf5")) {

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

@@ -23,7 +23,7 @@
 import io.github.jbellis.jvector.example.util.CompactionPartitionSource;
 import io.github.jbellis.jvector.example.yaml.TestDataPartition.Distribution;
 import io.github.jbellis.jvector.graph.GraphSearcher;
-import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
 import io.github.jbellis.jvector.graph.SearchResult;
 import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
 import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexCompactor;
@@ -33,7 +33,6 @@
 import io.github.jbellis.jvector.util.Bits;
 import io.github.jbellis.jvector.util.FixedBitSet;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
-import io.github.jbellis.jvector.vector.types.VectorFloat;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -121,7 +120,7 @@ public static List<BenchResult> run(DataSet ds) throws Exception {
     private static BenchResult runConfig(DataSet ds, PartitionConfig cfg) throws Exception {
         String datasetName = ds.getName();
         logger.info("Compaction bench [{}] config {}: {} vectors",
-                datasetName, cfg.dirName(), ds.getBaseVectors().size());
+                datasetName, cfg.dirName(), ds.getBaseRavv().size());
 
         // 1. Fetch pre-built partitions from S3 (cached locally).
         List<Path> partitionPaths = CompactionPartitionSource.ensurePartitions(
@@ -137,15 +136,14 @@ private static BenchResult runConfig(DataSet ds, PartitionConfig cfg) throws Exc
 
     private static BenchResult compactAndMeasure(DataSet ds, PartitionConfig cfg,
                                                  List<Path> partitionPaths, Path tempDir) throws Exception {
-        List<VectorFloat<?>> baseVectors = ds.getBaseVectors();
-        int dimension = ds.getDimension();
+        RandomAccessVectorValues baseVectors = ds.getBaseRavv();
         VectorSimilarityFunction vsf = ds.getSimilarityFunction();
         String datasetName = ds.getName();
         int numPartitions = cfg.numPartitions;
 
         // Load graphs and set up ordinal mapping: partition i's local ordinals shift by the sum of
         // all prior partition sizes, preserving the original base-vector ordering so global ordinal
-        // k maps back to baseVectors.get(k) by construction.
+        // k maps back to baseVectors ordinal k by construction.
         List<ReaderSupplier> rss = new ArrayList<>(numPartitions);
         List<OnDiskGraphIndex> graphs = new ArrayList<>(numPartitions);
         List<OrdinalMapper> remappers = new ArrayList<>(numPartitions);
@@ -184,8 +182,8 @@ private static BenchResult compactAndMeasure(DataSet ds, PartitionConfig cfg,
             rss.clear();
 
             // Search the compacted graph: measure recall and search latency in one pass.
-            // Global ordinal k maps back to baseVectors.get(k) by construction.
-            SearchStats search = searchCompacted(compactPath, ds, baseVectors, dimension, vsf);
+            // Global ordinal k maps back to baseVectors ordinal k by construction.
+            SearchStats search = searchCompacted(compactPath, ds, baseVectors, vsf);
             logger.info(String.format(
                     "%n" +
                     "  ┌─ Compaction result: %s [%s]%n" +
@@ -244,11 +242,10 @@ static final class SearchStats {
      * mean and p99 per-query latency (ms) and throughput (queries/sec, single-threaded sequential).
      */
     private static SearchStats searchCompacted(Path indexPath, DataSet ds,
-                                               List<VectorFloat<?>> baseVectors,
-                                               int dimension, VectorSimilarityFunction vsf) throws Exception {
+                                               RandomAccessVectorValues ravv,
+                                               VectorSimilarityFunction vsf) throws Exception {
         var queryVectors = ds.getQueryVectors();
         var groundTruth = ds.getGroundTruth();
-        var ravv = new ListRandomAccessVectorValues(baseVectors, dimension);
 
         try (var rs = ReaderSupplierFactory.open(indexPath)) {
             var graph = OnDiskGraphIndex.load(rs);