datastax · ashkrisk · Aug 28, 2026
@@ -21,6 +21,7 @@
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSetInfo;
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets;
+import io.github.jbellis.jvector.example.benchmarks.datasets.InMemoryDataSet;
 import io.github.jbellis.jvector.example.reporting.GitInfo;
 import io.github.jbellis.jvector.example.reporting.JfrRecorder;
 import io.github.jbellis.jvector.example.reporting.JsonlWriter;
@@ -251,7 +252,7 @@ private static void writeCompletedCount(int count) {
     private List<VectorFloat<?>> queryVectors;
     private List<VectorFloat<?>> baseVectors;
     private List<? extends List<Integer>> groundTruth;
-    private DataSet ds;
+    private InMemoryDataSet ds;
     private VectorSimilarityFunction similarityFunction;
 
     private final List<OnDiskGraphIndex> graphs = new ArrayList<>();

@@ -22,6 +22,7 @@
 import io.github.jbellis.jvector.example.util.CheckpointManager;
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets;
+import io.github.jbellis.jvector.example.benchmarks.datasets.InMemoryDataSet;
 import io.github.jbellis.jvector.example.yaml.DatasetCollection;
 import io.github.jbellis.jvector.example.yaml.MultiConfig;
 
@@ -123,10 +124,10 @@ public static void main(String[] args) throws IOException {
 
                 logger.info("Loading dataset: {}", datasetName);
                 try {
-                    DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow(
+                    InMemoryDataSet 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")) {

@@ -19,6 +19,7 @@
 import io.github.jbellis.jvector.disk.ReaderSupplier;
 import io.github.jbellis.jvector.disk.ReaderSupplierFactory;
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
+import io.github.jbellis.jvector.example.benchmarks.datasets.InMemoryDataSet;
 import io.github.jbellis.jvector.example.util.AccuracyMetrics;
 import io.github.jbellis.jvector.example.util.CompactionPartitionSource;
 import io.github.jbellis.jvector.example.yaml.TestDataPartition.Distribution;
@@ -97,7 +98,7 @@ private CompactionBench() {}
      * one result per config. A config that fails (e.g. missing partitions) is logged and skipped so
      * the remaining configs still run. Throws if the dataset has no query vectors or ground truth.
      */
-    public static List<BenchResult> run(DataSet ds) throws Exception {
+    public static List<BenchResult> run(InMemoryDataSet ds) throws Exception {
         var queryVectors = ds.getQueryVectors();
         var groundTruth = ds.getGroundTruth();
         if (queryVectors == null || queryVectors.isEmpty()) {
@@ -118,10 +119,10 @@ public static List<BenchResult> run(DataSet ds) throws Exception {
         return results;
     }
 
-    private static BenchResult runConfig(DataSet ds, PartitionConfig cfg) throws Exception {
+    private static BenchResult runConfig(InMemoryDataSet 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(
@@ -135,7 +136,7 @@ private static BenchResult runConfig(DataSet ds, PartitionConfig cfg) throws Exc
         }
     }
 
-    private static BenchResult compactAndMeasure(DataSet ds, PartitionConfig cfg,
+    private static BenchResult compactAndMeasure(InMemoryDataSet ds, PartitionConfig cfg,
                                                  List<Path> partitionPaths, Path tempDir) throws Exception {
         List<VectorFloat<?>> baseVectors = ds.getBaseVectors();
         int dimension = ds.getDimension();
@@ -253,7 +254,6 @@ private static SearchStats searchCompacted(Path indexPath, DataSet ds,
         try (var rs = ReaderSupplierFactory.open(indexPath)) {
             var graph = OnDiskGraphIndex.load(rs);
             try (var searcher = new GraphSearcher(graph)) {
-                searcher.usePruning(false);
                 int n = queryVectors.size();
                 List<SearchResult> results = new ArrayList<>(n);
                 long[] latenciesNanos = new long[n];

@@ -58,6 +58,7 @@
 import io.github.jbellis.jvector.util.ExplicitThreadLocal;
 import io.github.jbellis.jvector.util.PhysicalCoreExecutor;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
+import me.tongfei.progressbar.ProgressBar;
 
 import java.io.FileNotFoundException;
 import java.io.IOException;
@@ -329,7 +330,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache,
                                             "Compressor '%s' was provided but failed to encode vectors for dataset '%s'. " +
                                                     "Aborting to prevent false recall results.", compressor, ds.getName()));
                                 }
-                                System.out.format("%s: %s encoded %d vectors [%.2f MB] in %.2fs%n", ds.getName(), compressor, ds.getBaseVectors().size(), (cv.ramBytesUsed() / 1024f / 1024f), encodingTimeS);
+                                System.out.format("%s: %s encoded %d vectors [%.2f MB] in %.2fs%n", ds.getName(), compressor, ds.getBaseRavv().size(), (cv.ramBytesUsed() / 1024f / 1024f), encodingTimeS);
                             }
                         }
 
@@ -409,27 +410,31 @@ private static Map<Set<FeatureId>, ImmutableGraphIndex> buildOnDisk(List<? exten
             }
         }
         if (scoringWriter == null) {
+            builder.close();
             throw new IllegalStateException("Bench looks for either NVQ_VECTORS or INLINE_VECTORS feature set for scoring compressed builds.");
         }
 
         // build the graph incrementally
         long startTime = System.nanoTime();
         var vv = floatVectors.threadLocalSupplier();
         PhysicalCoreExecutor.pool().submit(() -> {
-            IntStream.range(0, floatVectors.size()).parallel().forEach(node -> {
-                writers.forEach((features, writer) -> {
-                    try {
-                        var stateMap = new EnumMap<FeatureId, Feature.State>(FeatureId.class);
-                        suppliers.get(features).forEach((featureId, supplier) -> {
-                            stateMap.put(featureId, supplier.apply(node));
-                        });
-                        writer.writeInline(node, stateMap);
-                    } catch (IOException e) {
-                        throw new UncheckedIOException(e);
-                    }
+            try (ProgressBar pb = new ProgressBar("build onDiskGraph", floatVectors.size())) {
+                IntStream.range(0, floatVectors.size()).parallel().forEach(node -> {
+                    writers.forEach((features, writer) -> {
+                        try {
+                            var stateMap = new EnumMap<FeatureId, Feature.State>(FeatureId.class);
+                            suppliers.get(features).forEach((featureId, supplier) -> {
+                                stateMap.put(featureId, supplier.apply(node));
+                            });
+                            writer.writeFeaturesInline(node, stateMap);
+                        } catch (IOException e) {
+                            throw new UncheckedIOException(e);
+                        }
+                    });
+                    builder.addGraphNode(node, vv.get().getVector(node));
+                    pb.step();
                 });
-                builder.addGraphNode(node, vv.get().getVector(node));
-            });
+            }
         }).join();
         builder.cleanup();
 
@@ -513,7 +518,8 @@ private static BuilderWithSuppliers builderWithSuppliers(Set<FeatureId> features
                     builder.with(new NVQ(nvq));
                     suppliers.put(FeatureId.NVQ_VECTORS, ordinal -> new NVQ.State(nvq.encode(floatVectors.getVector(ordinal))));
                     break;
-
+                default:
+                    break;
             }
         }
         return new BuilderWithSuppliers(builder, suppliers);
@@ -605,6 +611,7 @@ private static Map<Set<FeatureId>, ImmutableGraphIndex> buildInMemory(List<? ext
             indexes.put(features, index);
         }
         indexBuildTimes.put(ds.getName(), buildTimeS);
+        builder.close();
         return indexes;
     }
 
@@ -867,7 +874,7 @@ public static List<BenchResult> runAllAndCollectResults(
                                                                 searchCompressorObj, ds.getName()));
                                                     }
                                                     System.out.format("%s: %s encoded %d vectors [%.2f MB] for search%n",
-                                                            ds.getName(), searchCompressorObj, ds.getBaseVectors().size(),
+                                                            ds.getName(), searchCompressorObj, ds.getBaseRavv().size(),
                                                             (cvArg.ramBytesUsed() / 1024f / 1024f));
                                                 }
                                             }

@@ -52,12 +52,6 @@ public interface DataSet {
      */
     VectorSimilarityFunction getSimilarityFunction();
 
-    /**
-     * The base vectors as a list.
-     * @return a list of base vectors
-     */
-    List<VectorFloat<?>> getBaseVectors();
-
     /**
      * The query vectors as a list.
      * Each major index corresponds to the self-same index from {@link #getGroundTruth()}.
@@ -69,7 +63,7 @@ public interface DataSet {
     /**
      * The ground truth as a list.
      * Each major index corresponds to the self-same index from {@link #getQueryVectors()}.
-     * Each minor index within represents the corresponding ordinal from {@link #getBaseVectors()} and {@link #getBaseRavv()}.
+     * Each minor index within represents the corresponding ordinal from {@link #getBaseRavv()}.
      * @return a list of query vectors.
      */
     List<? extends List<Integer>> getGroundTruth();

@@ -52,7 +52,7 @@
 public class DataSetInfo implements DataSetProperties {
     private final DataSetFiles dsFiles;
     private final DataSetProperties baseProperties;
-    private volatile DataSet cached;
+    private volatile InMemoryDataSet cached;
 
     /// Creates a new dataset info handle.
     ///
@@ -123,7 +123,7 @@ public boolean isDuplicateVectorFree() {
     /// completes, after which all callers share the same cached instance.
     ///
     /// @return the ready-to-use {@link DataSet}
-    public DataSet getDataSet() {
+    public InMemoryDataSet getDataSet() {
         if (cached == null) {
             synchronized (this) {
                 if (cached == null) {

@@ -43,7 +43,7 @@ public class DataSetUtils {
     /**
      * Processes a dataset using the configured load behavior from the dataset metadata.
      */
-    public static DataSet processDataSet(String pathStr,
+    public static InMemoryDataSet processDataSet(String pathStr,
                                          DataSetProperties props,
                                          List<VectorFloat<?>> baseVectors,
                                          List<VectorFloat<?>> queryVectors,
@@ -76,7 +76,7 @@ public static DataSet getScrubbedDataSet(String pathStr,
         return legacyScrubDataSet(pathStr, vsf, baseVectors, queryVectors, groundTruth);
     }
 
-    private static DataSet legacyScrubDataSet(String pathStr,
+    private static InMemoryDataSet legacyScrubDataSet(String pathStr,
                                               VectorSimilarityFunction vsf,
                                               List<VectorFloat<?>> baseVectors,
                                               List<VectorFloat<?>> queryVectors,

@@ -0,0 +1,32 @@
+/*
+ * 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.util.List;
+
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+/**
+ * A {@link DataSet} which allows for fetching the entire base vector corpus as a list
+ */
+public interface InMemoryDataSet extends DataSet {
+    /**
+     * The base vectors as a list.
+     * @return a list of base vectors
+     */
+    List<VectorFloat<?>> getBaseVectors();
+}
@@ -23,7 +23,7 @@
 
 import java.util.List;
 
-public class SimpleDataSet implements DataSet {
+public class SimpleDataSet implements InMemoryDataSet {
     private final String name;
     private final VectorSimilarityFunction similarityFunction;
     private final List<VectorFloat<?>> baseVectors;

@@ -120,7 +120,7 @@ public static Row fromDataSet(String datasetName,
                 basePath,
                 queryPath,
                 groundTruthPath,
-                ds.getBaseVectors().size(),
+                ds.getBaseRavv().size(),
                 ds.getQueryVectors().size(),
                 ds.getGroundTruth().size(),
                 ds.getDimension(),

@@ -16,7 +16,7 @@
 
 package io.github.jbellis.jvector.example.util;
 
-import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
+import io.github.jbellis.jvector.example.benchmarks.datasets.InMemoryDataSet;
 import io.github.jbellis.jvector.example.yaml.TestDataPartition;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
 
@@ -39,7 +39,7 @@ public PartitionedData(List<List<VectorFloat<?>>> vectors, List<Integer> sizes)
         }
     }
 
-    public static PartitionedData partition(DataSet ds, int numParts, TestDataPartition.Distribution distribution) {
+    public static PartitionedData partition(InMemoryDataSet ds, int numParts, TestDataPartition.Distribution distribution) {
         return partition(ds.getBaseVectors(), numParts, distribution);
     }
 

@@ -305,7 +305,7 @@ public static void main(String[] args) throws IOException {
         DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow(
                 () -> new RuntimeException("Dataset " + datasetName + " not found")
         ).getDataSet();
-        System.out.printf("Loaded %d vectors of dimension %d%n", ds.getBaseVectors().size(), ds.getDimension());
+        System.out.printf("Loaded %d vectors of dimension %d%n", ds.getBaseRavv().size(), ds.getDimension());
 
         var floatVectors = ds.getBaseRavv();
 

@@ -15,9 +15,8 @@
  */
 package io.github.jbellis.jvector.microbench;
 
-
-import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets;
+import io.github.jbellis.jvector.example.benchmarks.datasets.InMemoryDataSet;
 import io.github.jbellis.jvector.graph.GraphIndexBuilder;
 import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
 import org.openjdk.jmh.annotations.Benchmark;
@@ -31,6 +30,7 @@
 import org.openjdk.jmh.annotations.Warmup;
 import org.openjdk.jmh.infra.Blackhole;
 
+import java.io.IOException;
 import java.util.concurrent.TimeUnit;
 
 @Warmup(iterations = 1, time = 5)
@@ -40,7 +40,7 @@ public class GraphBuildBench {
 
     @State(Scope.Benchmark)
     public static class Parameters {
-        final DataSet ds;
+        final InMemoryDataSet ds;
         final ListRandomAccessVectorValues ravv;
 
         public Parameters() {
@@ -54,21 +54,23 @@ public Parameters() {
     @Benchmark
     @BenchmarkMode(Mode.Throughput)
     @OutputTimeUnit(TimeUnit.SECONDS)
-    public void testGraphBuild(Blackhole bh, Parameters p) {
+    public void testGraphBuild(Blackhole bh, Parameters p) throws IOException {
         long start = System.nanoTime();
-        GraphIndexBuilder graphIndexBuilder =  new GraphIndexBuilder(p.ravv, p.ds.getSimilarityFunction(), 8, 60, 1.2f, 1.4f, false);
-        graphIndexBuilder.build(p.ravv);
+        try (GraphIndexBuilder graphIndexBuilder =  new GraphIndexBuilder(p.ravv, p.ds.getSimilarityFunction(), 8, 60, 1.2f, 1.4f, false)) {
+            graphIndexBuilder.build(p.ravv);
+        }
         System.out.format("Build M=%d ef=%d in %.2fs%n",
                 32, 600, (System.nanoTime() - start) / 1_000_000_000.0);
     }
 
     @Benchmark
     @BenchmarkMode(Mode.Throughput)
     @OutputTimeUnit(TimeUnit.SECONDS)
-    public void testGraphBuildWithHierarchy(Blackhole bh, Parameters p) {
+    public void testGraphBuildWithHierarchy(Blackhole bh, Parameters p) throws IOException {
         long start = System.nanoTime();
-        GraphIndexBuilder graphIndexBuilder =  new GraphIndexBuilder(p.ravv, p.ds.getSimilarityFunction(), 8, 60, 1.2f, 1.4f, true);
-        graphIndexBuilder.build(p.ravv);
+        try (GraphIndexBuilder graphIndexBuilder =  new GraphIndexBuilder(p.ravv, p.ds.getSimilarityFunction(), 8, 60, 1.2f, 1.4f, true)) {
+            graphIndexBuilder.build(p.ravv);
+        }
         System.out.format("Build M=%d ef=%d in %.2fs%n",
                 32, 600, (System.nanoTime() - start) / 1_000_000_000.0);
     }