From 3c814fac81cc1c09328c3bef8eed7848fd6d2e35 Mon Sep 17 00:00:00 2001 From: "Dian-Lun (Aaron) Lin" Date: Thu, 2 Jul 2026 22:14:58 +0000 Subject: [PATCH] Compaction: fix cold compaction time by eliminating disk scan in computeLayerInfoFromSources MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit computeLayerInfoFromSources called getNodes(0) on each source graph to count live nodes at level 0. getNodes(0) sequentially seeks through every node record on disk to filter out deleted entries. On a cold page cache this touches large amounts of source data before compaction even begins, significantly delaying the start of actual graph merging. Since every live node is present at level 0 by the HNSW invariant, the count is simply liveNodes.get(s).cardinality() — an in-memory popcount requiring no I/O. Also switch PQ retraining from ProductQuantization.compute() (full k-means++ init) to basePQ.refine() (Lloyd's iterations only, warm-started from the existing codebook). The source codebooks are already trained on the same distribution, so warm-starting converges in far fewer passes with no recall loss. --- .../graph/disk/OnDiskGraphIndexCompactor.java | 17 +++++++--- .../jvector/graph/disk/PQRetrainer.java | 33 +++++++++++-------- .../quantization/ProductQuantization.java | 2 +- 3 files changed, 33 insertions(+), 19 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index deffb719f..ddf5d3cc5 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -1354,11 +1354,18 @@ private List computeLayerInfoFromSources() { int count = 0; for (int s = 0; s < sources.size(); s++) { if (level > sources.get(s).getMaxLevel()) continue; - NodesIterator it = sources.get(s).getNodes(level); - FixedBitSet alive = liveNodes.get(s); - while (it.hasNext()) { - int node = it.next(); - if (alive.get(node)) count++; + if (level == 0) { + // Every live node is present at level 0 (HNSW base layer invariant), + // so count directly from the in-memory bitset instead of scanning node + // records on disk (which touches gigabytes of source data on a cold cache). + count += liveNodes.get(s).cardinality(); + } else { + NodesIterator it = sources.get(s).getNodes(level); + FixedBitSet alive = liveNodes.get(s); + while (it.hasNext()) { + int node = it.next(); + if (alive.get(node)) count++; + } } } layerInfo.add(new CommonHeader.LayerInfo(count, maxDegrees.get(level))); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java index 280d75c9c..c45120f70 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java @@ -22,6 +22,7 @@ import io.github.jbellis.jvector.quantization.ProductQuantization; import io.github.jbellis.jvector.util.DocIdSetIterator; import io.github.jbellis.jvector.util.FixedBitSet; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.types.VectorFloat; @@ -32,6 +33,7 @@ import java.util.ArrayList; import java.util.Comparator; import java.util.List; +import java.util.concurrent.ForkJoinPool; import java.util.concurrent.ThreadLocalRandom; /** @@ -96,22 +98,27 @@ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, log.info("Collected {} training samples", samples.size()); - // Extract vectors sequentially in sorted (source, node) order so disk reads are - // purely sequential and the OS read-ahead can cover them efficiently. We do this - // here rather than letting ProductQuantization.compute() drive the reads via its - // parallel stream, which would scatter page faults across a potentially very large - // file and cause I/O that scales with dataset size rather than sample count. List> trainingVectors = extractVectorsSequential(samples); - var ravv = new ListRandomAccessVectorValues(trainingVectors, dimension); - boolean center = similarityFunction == VectorSimilarityFunction.EUCLIDEAN; + long t0 = System.nanoTime(); + log.info("Extracted {} vectors in {}ms; starting PQ refinement", + trainingVectors.size(), (System.nanoTime() - t0) / 1_000_000L); + + var ravv = new ListRandomAccessVectorValues(trainingVectors, dimension); - return ProductQuantization.compute( - ravv, - basePQ.getSubspaceCount(), - basePQ.getClusterCount(), - center - ); + // Warm-start from the existing codebook via Lloyd's-only refinement rather than + // re-running k-means++ from scratch. k-means++ initialization visits every point + // once per centroid (256 passes for k=256), which dominates training time. + // Since the source codebooks are already trained on data from the same underlying + // distribution, this warm-start converges in far fewer passes with no recall loss. + long t1 = System.nanoTime(); + ProductQuantization result = basePQ.refine(ravv, + ProductQuantization.K_MEANS_ITERATIONS, + -1.0f, // UNWEIGHTED / isotropic + PhysicalCoreExecutor.pool(), + ForkJoinPool.commonPool()); + log.info("PQ refinement complete in {}ms", (System.nanoTime() - t1) / 1_000_000L); + return result; } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java index c84b7b955..6d9d23879 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java @@ -60,7 +60,7 @@ public class ProductQuantization implements VectorCompressor>, A private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); static final int DEFAULT_CLUSTERS = 256; // number of clusters per subspace = one byte's worth - static final int K_MEANS_ITERATIONS = 6; + public static final int K_MEANS_ITERATIONS = 6; public static final int MAX_PQ_TRAINING_SET_SIZE = 128000; final VectorFloat[] codebooks; // array of codebooks, where each codebook is a VectorFloat consisting of k contiguous subvectors each of length M