diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java index 30431d6ca..7ebb3f9b0 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java +++ b/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 { } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java index f759bc702..b6a44ece3 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java +++ b/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; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java index 3fb69d967..61d8157fb 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java +++ b/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() { 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..be8ae8c7e 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 @@ -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 computeBaseBatch(CompactWriter writer, CompactionParams params) throws IOException { List 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 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..19bfde85b 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,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> 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 sampleBalanced(int totalSamples) { * cover them efficiently. Each source's view is opened once and reused for all its samples. */ private List> extractVectorsSequential(List 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> extractVectorsSequential(List 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 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. */ 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 diff --git a/jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReader.java b/jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReader.java index c9ac39c21..0a5fde49e 100644 --- a/jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReader.java +++ b/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 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);