From 1f5847331f5448dee5867591981a7698037ee2cd Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Mon, 6 Jul 2026 13:16:16 -0400 Subject: [PATCH] fix for disk usage measurements --- .../github/jbellis/jvector/example/Grid.java | 47 ++++++++++++++----- 1 file changed, 34 insertions(+), 13 deletions(-) diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java index 8f45df2a0..6f13710ff 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java @@ -230,7 +230,6 @@ static void runOneGraph(OnDiskGraphIndexCache cache, // Prepare to collect index construction metrics for reporting.... var constructionMetrics = new ConstructionMetrics(); - // TODO this does not capture disk usage for cached indexes. Need to update // Capture initial memory and disk state try (var diagnostics = new BenchmarkDiagnostics(getDiagnosticLevel())) { diagnostics.startMonitoring("testDirectory", workDirectory); @@ -253,6 +252,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache, (buildCompressorObj == null) ? "None" : String.valueOf(buildCompressorObj); Map, ImmutableGraphIndex> indexes = new HashMap<>(); + Map, Long> indexFileSizes = new HashMap<>(); if (buildCompressorObj == null) { indexes = buildInMemory(featureSets, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph, ds, workDirectory); } else { @@ -270,6 +270,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache, if (cached.isPresent()) { System.out.printf("%s: Using cached graph index for %s%n", key.datasetName, fs); indexes.put(fs, cached.get()); + try { indexFileSizes.put(fs, Files.size(cache.resolve(key).finalPath)); } catch (IOException ignored) {} } else { missing.add(fs); if (cache.isEnabled()) { @@ -287,9 +288,10 @@ static void runOneGraph(OnDiskGraphIndexCache cache, if (!missing.isEmpty()) { // At least one index needs to be built (b/c not in cache or cache is disabled) // We pass the handles map so buildOnDisk knows exactly where to write - var newIndexes = buildOnDisk(missing, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph, + var result = buildOnDisk(missing, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph, ds, outputDir, buildCompressorObj, handles, constructionMetrics); - indexes.putAll(newIndexes); + indexes.putAll(result.indexes); + indexFileSizes.putAll(result.fileSizes); } } @@ -303,6 +305,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache, try { for (var cpSupplier : compressionGrid) { indexes.forEach((features, index) -> { + constructionMetrics.indexFileSizeBytes = indexFileSizes.get(features); final Set featureSetForIndex = index instanceof OnDiskGraphIndex ? ((OnDiskGraphIndex) index).getFeatureSet() : Set.of(); CompressedVectors cv; @@ -364,7 +367,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache, } } - private static Map, ImmutableGraphIndex> buildOnDisk(List> featureSets, + private static BuildOnDiskResult buildOnDisk(List> featureSets, int M, int efConstruction, float neighborOverflow, @@ -464,8 +467,9 @@ private static Map, ImmutableGraphIndex> buildOnDisk(List, ImmutableGraphIndex> indexes = new HashMap<>(); + Map, Long> fileSizes = new HashMap<>(); n = 0; for (var features : featureSets) { Path loadPath = handles.containsKey(features) @@ -473,9 +477,10 @@ private static Map, ImmutableGraphIndex> buildOnDisk(List features, @@ -845,6 +850,7 @@ public static List runAllAndCollectResults( diagnostics.startMonitoring("indexCache", Paths.get(indexCacheDir)); diagnostics.capturePrePhaseSnapshot("Build"); Map, ImmutableGraphIndex> indexes = new HashMap<>(); + Map, Long> indexFileSizes = new HashMap<>(); var compressor = getCompressor(buildCompressor, ds); var searchCompressorObj = getCompressor(searchCompressor, ds); @@ -885,6 +891,7 @@ public static List runAllAndCollectResults( if (cached.isPresent()) { System.out.printf("%s: Using cached graph index for %s%n", key.datasetName, fs); indexes.put(fs, cached.get()); + try { indexFileSizes.put(fs, Files.size(cache.resolve(key).finalPath)); } catch (IOException ignored) {} } else { missing.add(fs); if (cache.isEnabled()) { @@ -903,9 +910,10 @@ public static List runAllAndCollectResults( if (!missing.isEmpty()) { // At least one index needs to be built (b/c not in cache or cache is disabled) // We pass the handles map so buildOnDisk knows exactly where to write - var newIndexes = buildOnDisk(missing, m, ef, neighborOverflow, addHierarchy, refineFinalGraph, + var result = buildOnDisk(missing, m, ef, neighborOverflow, addHierarchy, refineFinalGraph, ds, outputDir, compressor, handles, null); - indexes.putAll(newIndexes); + indexes.putAll(result.indexes); + indexFileSizes.putAll(result.fileSizes); } ImmutableGraphIndex index = indexes.get(features); @@ -914,7 +922,6 @@ public static List runAllAndCollectResults( diagnostics.capturePostPhaseSnapshot("Build"); diagnostics.printDiskStatistics("Graph Index Build"); var buildSnapshot = diagnostics.getLatestSystemSnapshot(); - DiskUsageMonitor.MultiDirectorySnapshot buildDiskSnapshot = diagnostics.getLatestDiskSnapshot(); try (ConfiguredSystem cs = new ConfiguredSystem(ds, index, cvArg, features)) { int queryRuns = 2; @@ -966,10 +973,10 @@ public static List runAllAndCollectResults( allMetrics.put("Total Off-Heap (MB)", buildSnapshot.memoryStats.getTotalOffHeapMemory() / 1024.0 / 1024.0); } - // Add disk metrics if available - if (buildDiskSnapshot != null) { - allMetrics.put("Disk Usage (MB)", buildDiskSnapshot.getTotalBytes() / 1024.0 / 1024.0); - allMetrics.put("File Count", buildDiskSnapshot.getTotalFileCount()); + // Add per-feature-set index file size + Long fileSizeBytes = indexFileSizes.get(features); + if (fileSizeBytes != null) { + allMetrics.put("Disk Usage (MB)", fileSizeBytes / 1024.0 / 1024.0); } results.add(new BenchResult(ds.getName(), params, allMetrics)); @@ -1161,9 +1168,19 @@ private static String quantTypeOf(CompressorParameters cp) { *

Call {@link #appendTo(List)} to emit these as {@link Metric} entries so they can be selected for * console/CSV reporting.

*/ + private static final class BuildOnDiskResult { + final Map, ImmutableGraphIndex> indexes; + final Map, Long> fileSizes; + BuildOnDiskResult(Map, ImmutableGraphIndex> indexes, Map, Long> fileSizes) { + this.indexes = indexes; + this.fileSizes = fileSizes; + } + } + static final class ConstructionMetrics { // null means “not applicable / not measured” public Double indexBuildTimeS; + public Long indexFileSizeBytes; // Index-construction phase (single pass per run) private final Map indexByType = new HashMap<>(); @@ -1193,6 +1210,10 @@ void appendTo(List out) { out.add(Metric.of("construction.index_build_time_s", "Index build time (s)", ".2f", indexBuildTimeS)); } + if (indexFileSizeBytes != null) { + out.add(Metric.of("construction.index_file_size_mb", + "Index file size (MB)", ".2f", indexFileSizeBytes / 1024.0 / 1024.0)); + } // Index-construction quant timings appendQuant(out, "construction.index_quant_time_s", "Idx", indexByType);