Compaction: eliminating disk scan and updating PQ retrain

jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java

@@ -30,6 +30,17 @@ public interface ReaderSupplier extends AutoCloseable {
      */
     RandomAccessReader get() throws IOException;
 
+    /**
+     * Streams the byte range {@code [offset, offset + length)} of the underlying storage into
+     * the page cache, if the implementation supports it. Synchronous and best-effort; a no-op
+     * by default. Sized for repeated windowed calls — a caller that knows which records a bulk
+     * phase is about to read (e.g. graph compaction, whose readers otherwise advise
+     * {@code MADV_RANDOM} and forgo kernel readahead) can warm exactly those, keeping transient
+     * cache demand proportional to the window rather than the file.
+     */
+    default void prefetch(long offset, long length) {
+    }
+
     default void close() throws IOException {
     }
 }

jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java

@@ -230,6 +230,9 @@ private void precomputeCodes(CompactWriter writer) throws IOException {
                 final int cStart = chunkStart;
                 final int cEnd = Math.min(chunkStart + chunkSize, upper);
                 tasks.add(() -> {
+                    // Stream this chunk's records into the page cache before the encode loop;
+                    // the mapping's MADV_RANDOM otherwise faults them one page at a time.
+                    source.prefetchL0Records(cStart, cEnd - 1);
                     ByteSequence<?> code = vectorTypeSupport.createByteSequence(cs);
                     VectorFloat<?> vec = vectorTypeSupport.createFloatVector(ctx.dimension);
                     long count = 0;

jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java

@@ -443,6 +443,22 @@ public int maxDegree() {
         return layerInfo.stream().mapToInt(li -> li.degree).max().orElseThrow();
     }
 
+    /**
+     * Streams the L0 records of ordinals {@code [minNode, maxNode]} (inclusive) into the page
+     * cache; see {@link ReaderSupplier#prefetch(long, long)}. An L0 record holds the node's id,
+     * inline features (vector, fused codes), and adjacency, so warming it covers every read a
+     * bulk scan makes for that node. Best-effort no-op when unsupported.
+     */
+    public void prefetchL0Records(int minNode, int maxNode) {
+        if (maxNode < minNode) {
+            return;
+        }
+        long blockBytes = Integer.BYTES + inlineBlockSize
+                + (long) Integer.BYTES * (layerInfo.get(0).degree + 1);
+        long start = neighborsOffset + blockBytes * minNode;
+        readerSupplier.prefetch(start, blockBytes * (maxNode - minNode + 1));
+    }
+
     // re-declared to specify type
     @Override
     public View getView() {

jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java

@@ -84,7 +84,19 @@ public final class OnDiskGraphIndexCompactor implements Accountable {
     private final ForkJoinPool executor;
     private final int taskWindowSize;
     private final VectorSimilarityFunction similarityFunction;
+    private boolean refineAfterCompaction = false;
 
+    /**
+     * Whether to run the second-pass neighbor refinement after the merged graph is written
+     * (default false). Refinement is a navigability pass: it has no measurable effect on
+     * recall, but it improves query latency on the merged index at the cost of a significant
+     * fraction of total compaction time. Enable it when search latency matters more than
+     * compaction throughput.
+     */
+    @Experimental
+    public void setRefineAfterCompaction(boolean refineAfterCompaction) {
+        this.refineAfterCompaction = refineAfterCompaction;
+    }
     /**
      * Constructs a new OnDiskGraphIndexCompactor for graphs without a non-fused compressed sidecar.
      * Equivalent to calling the 6-arg constructor with {@code sourceCompressed = null}.
@@ -298,7 +310,9 @@ public void compact(Path outputPath) throws FileNotFoundException {
         try {
             compactGraphImpl(outputPath, strategy);
             releaseSourcesBeforeRefine(strategy);
-            refineCompactedGraph(outputPath, strategy);
+            if (refineAfterCompaction) {
+                refineCompactedGraph(outputPath, strategy);
+            }
         } finally {
             // Delayed until after refinement so refineCompactedGraph can read from the pre-encoded
             // code cache appended past the projected EOF; onAfterClose unmaps it and truncates.
@@ -330,7 +344,9 @@ public void compact(Path graphPath, Path compressedPath) throws FileNotFoundExce
         try {
             sidecarStrategy.retrain(similarityFunction);
             compactGraphImpl(graphPath, inlineStrategy);
-            refineCompactedGraph(graphPath, inlineStrategy);
+            if (refineAfterCompaction) {
+                refineCompactedGraph(graphPath, inlineStrategy);
+            }
             sidecarStrategy.writeSidecar(compressedPath);
         } catch (IOException e) {
             throw new RuntimeException("Sidecar compaction failed", e);
@@ -997,6 +1013,12 @@ private List<WriteResult> computeBaseBatch(CompactWriter writer,
                                               CompactionParams params) throws IOException {
 
         List<WriteResult> out = new ArrayList<>(bs.end - bs.start);
+        if (bs.end > bs.start) {
+            // Stream this batch's own records into the page cache before processing. Search
+            // reads into other sources are data-dependent and stay demand-faulted, but each
+            // node's own record read (adjacency + vector) is fully predictable.
+            sources.get(bs.sourceIdx).prefetchL0Records(bs.nodes[bs.start], bs.nodes[bs.end - 1]);
+        }
 
         for (int i = bs.start; i < bs.end; i++) {
             int node = bs.nodes[i];
@@ -1217,8 +1239,11 @@ private int gatherFromOtherSource(int node, int level, int sourceIdx,
         );
 
         if (level == 0) {
+            // rerankK = searchTopK, not beamWidth: the wider beam's extra candidates are largely
+            // pruned by diversity selection, so the doubled approximate-phase cost buys almost
+            // no recall.
             SearchResult results = scratch.gs[sourceIdx].search(
-                    ssp, params.searchTopK, params.beamWidth, 0f, 0f, indexAlive
+                    ssp, params.searchTopK, params.searchTopK, 0f, 0f, indexAlive
             );
 
             for (var r : results.getNodes()) {
@@ -1354,11 +1379,18 @@ private List<CommonHeader.LayerInfo> 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)));

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,26 @@ 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.
+        long t0 = System.nanoTime();
         List<VectorFloat<?>> trainingVectors = extractVectorsSequential(samples);
-        var ravv = new ListRandomAccessVectorValues(trainingVectors, dimension);
+        log.info("Extracted {} vectors in {}ms; starting PQ refinement",
+                 trainingVectors.size(), (System.nanoTime() - t0) / 1_000_000L);
 
-        boolean center = similarityFunction == VectorSimilarityFunction.EUCLIDEAN;
+        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;
     }
 
     /**
@@ -211,6 +217,8 @@ private List<SampleRef> sampleBalanced(int totalSamples) {
      * cover them efficiently. Each source's view is opened once and reused for all its samples.
      */
     private List<VectorFloat<?>> extractVectorsSequential(List<SampleRef> samples) {
+        prefetchSampleRanges(samples);
+
         OnDiskGraphIndex.View[] views = new OnDiskGraphIndex.View[sources.size()];
         for (int s = 0; s < sources.size(); s++) {
             views[s] = (OnDiskGraphIndex.View) sources.get(s).getView();
@@ -225,6 +233,33 @@ private List<VectorFloat<?>> extractVectorsSequential(List<SampleRef> samples) {
         return vectors;
     }
 
+    // Merging samples closer than this many ordinals into one range keeps the prefetch mostly
+    // sequential without dragging in long runs of unsampled records.
+    private static final int PREFETCH_MERGE_GAP = 256;
+
+    /**
+     * Streams the records of the (source, node)-sorted sample list into the page cache before
+     * extraction. The mappings advise {@code MADV_RANDOM}, so without this the extraction loop
+     * faults one page at a time on a cold cache; total demand is bounded by the training-set
+     * size, not the file size.
+     */
+    private void prefetchSampleRanges(List<SampleRef> samples) {
+        int i = 0;
+        while (i < samples.size()) {
+            SampleRef first = samples.get(i);
+            int last = first.node;
+            int j = i + 1;
+            while (j < samples.size()
+                    && samples.get(j).source == first.source
+                    && samples.get(j).node - last <= PREFETCH_MERGE_GAP) {
+                last = samples.get(j).node;
+                j++;
+            }
+            sources.get(first.source).prefetchL0Records(first.node, last);
+            i = j;
+        }
+    }
+
     /**
      * Reference to a sampled vector from a specific source index.
      */

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

@@ -60,7 +60,7 @@ public class ProductQuantization implements VectorCompressor<ByteSequence<?>>, 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

jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReader.java

@@ -148,12 +148,14 @@ public void close() {
     public static class Supplier implements ReaderSupplier {
         private final Arena arena;
         private final MemorySegment memory;
+        private final Path path;
 
         public Supplier(Path path) throws IOException {
+            this.path = path;
             this.arena = Arena.ofShared();
             try (var ch = FileChannel.open(path, StandardOpenOption.READ)) {
                 this.memory = ch.map(MapMode.READ_ONLY, 0L, ch.size(), arena);
-                
+
                 // Apply MADV_RANDOM advice
                 var linker = Linker.nativeLinker();
                 var maybeMadvise = linker.defaultLookup().find("posix_madvise");
@@ -176,6 +178,39 @@ public Supplier(Path path) throws IOException {
             }
         }
 
+        // Windowed prefetch streams ranges through a separate read-ahead-enabled file descriptor:
+        // the mapping itself advises MADV_RANDOM (disables kernel readahead), and MADV_WILLNEED
+        // proved to be a no-op hint for large ranges. The page cache is shared per file, so
+        // subsequent random access through the mapping hits the populated pages. The bytes read
+        // here are discarded, so the buffer is sized to stay L2-resident rather than for syscall
+        // amortization (cold throughput is device-bound at any size >= 32KB, and the warm case
+        // is within noise of larger buffers at 64KB).
+        private static final ThreadLocal<java.nio.ByteBuffer> PREFETCH_BUF =
+                ThreadLocal.withInitial(() -> java.nio.ByteBuffer.wrap(new byte[64 << 10]));
+
+        @Override
+        public void prefetch(long offset, long length) {
+            if (length <= 0) {
+                return;
+            }
+            var buf = PREFETCH_BUF.get();
+            long end = Math.min(offset + length, memory.byteSize());
+            try (var ch = FileChannel.open(path, StandardOpenOption.READ)) {
+                long pos = Math.max(0, offset);
+                while (pos < end) {
+                    buf.clear().limit((int) Math.min(buf.capacity(), end - pos));
+                    int n = ch.read(buf, pos);
+                    if (n < 0) {
+                        break;
+                    }
+                    pos += n;
+                }
+            } catch (IOException e) {
+                logger.warn("ranged prefetch of {} [{}, {}) failed; continuing without warm cache",
+                            path, offset, end, e);
+            }
+        }
+
         @Override
         public MemorySegmentReader get() {
             return new MemorySegmentReader(memory);