On-disk index cache for the Grid benchmark harness (#612)

* Initial implementation of index cache for Bench / Grid.

* Initial implementation of index cache for Bench / Grid.

* Uncommented datasets that were not working prior to PR613.

* Improved Exception handling.  Added alpha to key signature.  Lazy cached index deletes.  Index cached marked Experimental.

* Integrated index caching into autoBenchYAML and runAllAndCollectResults. useSavedIndexIfExists is No by default.  Added refineFinalGraph to construction parameters.

Author: tlwillke

.gitignore

@@ -6,6 +6,7 @@ local/
 
 ### Bench caches
 pq_cache/
+index_cache/
 
 ### JVM crashes
 hs_err_pid*

jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java

@@ -733,11 +733,12 @@ public long ramBytesUsed() {
     @Override
     public String toString() {
         if (anisotropicThreshold == UNWEIGHTED) {
-            return String.format("ProductQuantization(M=%d, clusters=%d)", M, clusterCount);
+            return String.format("ProductQuantization(M=%d, clusters=%d, centered=%s)", M, clusterCount, globalCentroid != null);
         }
-        return String.format("ProductQuantization(M=%d, clusters=%d, T=%.3f, eta=%.1f)",
+        return String.format("ProductQuantization(M=%d, clusters=%d, centered=%s, anisotropicT=%.3f, eta=%.1f)",
                              M,
                              clusterCount,
+                             globalCentroid != null,
                              anisotropicThreshold,
                              KMeansPlusPlusClusterer.computeParallelCostMultiplier(anisotropicThreshold, originalDimension));
     }

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

@@ -153,11 +153,13 @@ public static void main(String[] args) throws IOException {
                     }
                     logger.info("Using configuration: {}", config);
 
-                    List<BenchResult> datasetResults = Grid.runAllAndCollectResults(ds, 
+                    List<BenchResult> datasetResults = Grid.runAllAndCollectResults(ds,
+                            config.construction.useSavedIndexIfExists,
                             config.construction.outDegree, 
                             config.construction.efConstruction,
                             config.construction.neighborOverflow, 
                             config.construction.addHierarchy,
+                            config.construction.refineFinalGraph,
                             config.construction.getFeatureSets(), 
                             config.construction.getCompressorParameters(),
                             config.search.getCompressorParameters(), 

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

@@ -25,10 +25,7 @@
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
 
 import java.io.IOException;
-import java.util.Arrays;
-import java.util.EnumSet;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.function.Function;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -42,6 +39,10 @@ public class Bench {
     public static void main(String[] args) throws IOException {
         System.out.println("Heap space available is " + Runtime.getRuntime().maxMemory());
 
+        // When enabled, caches built indices for reuse in future runs.
+        // Useful for large indexes and repeated testing.
+        boolean enableIndexCache = false;
+
         var mGrid = List.of(32); // List.of(16, 24, 32, 48, 64, 96, 128);
         var efConstructionGrid = List.of(100); // List.of(60, 80, 100, 120, 160, 200, 400, 600, 800);
         var topKGrid = Map.of(
@@ -81,10 +82,10 @@ public static void main(String[] args) throws IOException {
         // compile regex and do substring matching using find
         var pattern = Pattern.compile(regex);
 
-        execute(pattern, buildCompression, featureSets, searchCompression, mGrid, efConstructionGrid, neighborOverflowGrid, addHierarchyGrid, refineFinalGraphGrid, topKGrid, usePruningGrid);
+        execute(pattern, enableIndexCache, buildCompression, featureSets, searchCompression, mGrid, efConstructionGrid, neighborOverflowGrid, addHierarchyGrid, refineFinalGraphGrid, topKGrid, usePruningGrid);
     }
 
-    private static void execute(Pattern pattern, List<Function<DataSet, CompressorParameters>> buildCompression, List<EnumSet<FeatureId>> featureSets, List<Function<DataSet, CompressorParameters>> compressionGrid, List<Integer> mGrid, List<Integer> efConstructionGrid, List<Float> neighborOverflowGrid, List<Boolean> addHierarchyGrid, List<Boolean> refineFinalGraphGrid, Map<Integer, List<Double>> topKGrid, List<Boolean> usePruningGrid) throws IOException {
+    private static void execute(Pattern pattern, boolean enableIndexCache, List<Function<DataSet, CompressorParameters>> buildCompression, List<EnumSet<FeatureId>> featureSets, List<Function<DataSet, CompressorParameters>> compressionGrid, List<Integer> mGrid, List<Integer> efConstructionGrid, List<Float> neighborOverflowGrid, List<Boolean> addHierarchyGrid, List<Boolean> refineFinalGraphGrid, Map<Integer, List<Double>> topKGrid, List<Boolean> usePruningGrid) throws IOException {
         var datasetCollection = DatasetCollection.load();
         var datasetNames = datasetCollection.getAll().stream().filter(dn -> pattern.matcher(dn).find()).collect(Collectors.toList());
         System.out.println("Executing the following datasets: " + datasetNames);
@@ -93,7 +94,7 @@ private static void execute(Pattern pattern, List<Function<DataSet, CompressorPa
             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);
+            Grid.runAll(ds, enableIndexCache, mGrid, efConstructionGrid, neighborOverflowGrid, addHierarchyGrid, refineFinalGraphGrid, featureSets, buildCompression, compressionGrid, topKGrid, usePruningGrid, null);
         }
     }
 }

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

@@ -94,7 +94,7 @@ public static void main(String[] args) throws IOException {
                     () -> new RuntimeException("Could not load dataset:" + datasetName)
             );
 
-            Grid.runAll(ds, config.construction.outDegree, config.construction.efConstruction,
+            Grid.runAll(ds, config.construction.useSavedIndexIfExists, config.construction.outDegree, config.construction.efConstruction,
                     config.construction.neighborOverflow, config.construction.addHierarchy, config.construction.refineFinalGraph,
                     config.construction.getFeatureSets(), config.construction.getCompressorParameters(),
                     config.search.getCompressorParameters(), config.search.topKOverquery, config.search.useSearchPruning, config.search.benchmarks);