Virtualize and Modularize DataSetLoader logic

- virtualize DataSetLoader
- separate and modularize HDF5, MFD loaders and consolidate helper methods
- organize dataset types into a package together
- update callers, minor cleanups
- use Optional instead of checked IOExceptions
- add javadoc
- organize dataset classes together
- remove bench2d and loader

Author: Jonathan Shook

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

@@ -15,11 +15,9 @@
  */
 package io.github.jbellis.jvector.bench;
 
-import io.github.jbellis.jvector.example.SiftSmall;
 import io.github.jbellis.jvector.example.util.SiftLoader;
 import io.github.jbellis.jvector.graph.*;
 import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider;
-import io.github.jbellis.jvector.util.Bits;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
 import org.openjdk.jmh.annotations.*;

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

@@ -15,7 +15,7 @@
  */
 package io.github.jbellis.jvector.bench;
 
-import io.github.jbellis.jvector.example.SiftSmall;
+import io.github.jbellis.jvector.example.benchmarks.datasets.SiftSmall;
 import io.github.jbellis.jvector.example.util.SiftLoader;
 import io.github.jbellis.jvector.graph.*;
 import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider;

jvector-examples/pom.xml

@@ -163,7 +163,7 @@
                                         <argument>-classpath</argument>
                                         <classpath/>
                                         <argument>-ea</argument>
-                                        <argument>io.github.jbellis.jvector.example.SiftSmall</argument>
+                                        <argument>io.github.jbellis.jvector.example.benchmarks.datasets.SiftSmall</argument>
                                     </arguments>
                                 </configuration>
                             </execution>
@@ -212,7 +212,7 @@
                                         <classpath/>
                                         <argument>--add-modules=jdk.incubator.vector</argument>
                                         <argument>-ea</argument>
-                                        <argument>io.github.jbellis.jvector.example.SiftSmall</argument>
+                                        <argument>io.github.jbellis.jvector.example.benchmarks.datasets.SiftSmall</argument>
                                     </arguments>
                                 </configuration>
                             </execution>
@@ -306,7 +306,7 @@
                                         <argument>--add-modules=jdk.incubator.vector</argument>
                                         <argument>-ea</argument>
                                         <argument>-Djvector.experimental.enable_native_vectorization=true</argument>
-                                        <argument>io.github.jbellis.jvector.example.SiftSmall</argument>
+                                        <argument>io.github.jbellis.jvector.example.benchmarks.datasets.SiftSmall</argument>
                                     </arguments>
                                 </configuration>
                             </execution>

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

@@ -20,8 +20,8 @@
 import io.github.jbellis.jvector.example.util.BenchmarkSummarizer;
 import io.github.jbellis.jvector.example.util.BenchmarkSummarizer.SummaryStats;
 import io.github.jbellis.jvector.example.util.CheckpointManager;
-import io.github.jbellis.jvector.example.util.DataSet;
-import io.github.jbellis.jvector.example.util.DataSetLoader;
+import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
+import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets;
 import io.github.jbellis.jvector.example.yaml.MultiConfig;
 
 import org.slf4j.Logger;
@@ -130,7 +130,9 @@ public static void main(String[] args) throws IOException {
 
                 logger.info("Loading dataset: {}", datasetName);
                 try {
-                    DataSet ds = DataSetLoader.loadDataSet(datasetName);
+                    DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow(
+                            () -> new IllegalStateException("Dataset " + datasetName + " not found")
+                    );
                     logger.info("Dataset loaded: {} with {} vectors", datasetName, ds.getBaseVectors().size());
 
                     String normalizedDatasetName = datasetName;

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

@@ -18,8 +18,8 @@
 
 import io.github.jbellis.jvector.example.util.CompressorParameters;
 import io.github.jbellis.jvector.example.util.CompressorParameters.PQParameters;
-import io.github.jbellis.jvector.example.util.DataSet;
-import io.github.jbellis.jvector.example.util.DataSetLoader;
+import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
+import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets;
 import io.github.jbellis.jvector.example.yaml.DatasetCollection;
 import io.github.jbellis.jvector.graph.disk.feature.FeatureId;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
@@ -90,7 +90,9 @@ private static void execute(Pattern pattern, List<Function<DataSet, CompressorPa
         System.out.println("Executing the following datasets: " + datasetNames);
 
         for (var datasetName : datasetNames) {
-            DataSet ds = DataSetLoader.loadDataSet(datasetName);
+            DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow(
+                    () -> new RuntimeException("Dataset " + datasetName + " not found")
+            );
             Grid.runAll(ds, mGrid, efConstructionGrid, neighborOverflowGrid, addHierarchyGrid, refineFinalGraphGrid, featureSets, buildCompression, compressionGrid, topKGrid, usePruningGrid);
         }
     }

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

@@ -16,8 +16,8 @@
 
 package io.github.jbellis.jvector.example;
 
-import io.github.jbellis.jvector.example.util.DataSet;
-import io.github.jbellis.jvector.example.util.DataSetLoader;
+import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
+import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets;
 import io.github.jbellis.jvector.example.yaml.DatasetCollection;
 import io.github.jbellis.jvector.example.yaml.MultiConfig;
 
@@ -52,13 +52,16 @@ public static void main(String[] args) throws IOException {
         if (!datasetNames.isEmpty()) {
             System.out.println("Executing the following datasets: " + datasetNames);
 
-            for (var datasetName : datasetNames) {
-                DataSet ds = DataSetLoader.loadDataSet(datasetName);
+            String hdf5 = ".hdf5";
+            for (var rawname : datasetNames) {
+                String datasetName =
+                        rawname.endsWith(hdf5) ? rawname.substring(0, rawname.length() - hdf5.length() -1) : rawname;
+                // pre-loading and early error phase
+                DataSets.loadDataSet(datasetName).orElseThrow(
+                        () -> new RuntimeException("Could not load dataset:" + datasetName)
+                );
 
-                if (datasetName.endsWith(".hdf5")) {
-                    datasetName = datasetName.substring(0, datasetName.length() - ".hdf5".length());
-                }
-                MultiConfig config = MultiConfig.getDefaultConfig(datasetName);
+                MultiConfig config = MultiConfig.getDefaultConfig(rawname);
                 allConfigs.add(config);
             }
         }
@@ -76,7 +79,9 @@ public static void main(String[] args) throws IOException {
         for (var config : allConfigs) {
             String datasetName = config.dataset;
 
-            DataSet ds = DataSetLoader.loadDataSet(datasetName);
+            DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow(
+                    () -> new RuntimeException("Could not load dataset:" + datasetName)
+            );
 
             Grid.runAll(ds, config.construction.outDegree, config.construction.efConstruction,
                     config.construction.neighborOverflow, config.construction.addHierarchy, config.construction.refineFinalGraph,

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

@@ -27,7 +27,7 @@
 import io.github.jbellis.jvector.example.benchmarks.ThroughputBenchmark;
 import io.github.jbellis.jvector.example.benchmarks.diagnostics.DiagnosticLevel;
 import io.github.jbellis.jvector.example.util.CompressorParameters;
-import io.github.jbellis.jvector.example.util.DataSet;
+import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
 import io.github.jbellis.jvector.example.util.FilteredForkJoinPool;
 import io.github.jbellis.jvector.graph.ImmutableGraphIndex;
 import io.github.jbellis.jvector.graph.GraphIndexBuilder;

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

@@ -16,8 +16,7 @@
 
 package io.github.jbellis.jvector.example;
 
-import io.github.jbellis.jvector.example.util.DataSet;
-import io.github.jbellis.jvector.example.util.DownloadHelper;
+import io.github.jbellis.jvector.example.benchmarks.datasets.DataSetLoaderMFD;
 import io.github.jbellis.jvector.example.yaml.MultiConfig;
 
 import java.io.IOException;
@@ -28,14 +27,10 @@
 public class HelloVectorWorld {
     public static void main(String[] args) throws IOException {
         System.out.println("Heap space available is " + Runtime.getRuntime().maxMemory());
-
         String datasetName = "ada002-100k";
-
-        var mfd = DownloadHelper.maybeDownloadFvecs(datasetName);
-        DataSet ds = mfd.load();
-
+        var ds = new DataSetLoaderMFD().loadDataSet(datasetName)
+                .orElseThrow(() -> new RuntimeException("dataset " + datasetName + " not found"));
         MultiConfig config = MultiConfig.getConfig(datasetName);
-
         Grid.runAll(ds, config.construction.outDegree, config.construction.efConstruction,
                 config.construction.neighborOverflow, config.construction.addHierarchy, config.construction.refineFinalGraph,
                 config.construction.getFeatureSets(), config.construction.getCompressorParameters(),

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

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.github.jbellis.jvector.example.util;
+package io.github.jbellis.jvector.example.benchmarks.datasets;
 
 import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;