Compaction: cell-walk optimizations, sharded compaction sessions, and the output’s format version

benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.md

@@ -251,3 +251,21 @@ target/benchmark-results/compactor-<timestamp>/compactor-results.jsonl
 | `durationMs` | Time in the measured phase only (`compact()` for `COMPACT`); excludes JVM startup and I/O setup |
 | `recall` | Recall@10 (present when `measureRecall=true`) |
 | `avgSearchLatencyMs` / `p99SearchLatencyMs` | Per-query `search()` latency (present when `measureRecall=true`) |
+
+# 9. Sharded compaction: `ShardExp`
+
+`ShardExp`, in the same package, is a plain `main` program (not a JMH benchmark) that measures a
+compaction whose output is split into K shards, as a Cassandra compaction with UCS sharding writes
+K output sstables. It builds S source graphs from a raw `.fvecs` file, then produces the K shards
+three ways and measures recall@10 of each against the dataset's ground truth:
+
+| Command | What it measures |
+|---------|------------------|
+| `shard K merge` | one `OnDiskGraphIndexCompactor` per shard over all sources, live nodes = the shard (Cassandra's merge path today) |
+| `shard K rebuild` | one fresh graph per shard from its vectors (graph merge disabled) |
+| `shard K sharded` (or `sharded K`) | one compactor, one `sharded()` session, `compactShard` per shard |
+| `recall K <mode> NQ` | search every shard, merge the top-10, recall@10 |
+
+The dataset paths are given as system properties (`exp.base`, `exp.query`, `exp.gt`); the class
+comment lists every property and a complete run. Results on tessera-2M are in the "Compacting into
+several outputs (sharding)" section of `docs/compaction.md`.

docs/compaction.md

@@ -39,6 +39,10 @@ List<OrdinalMapper> mapping = compactor.ordinalMappers();
 
 The compactor assigns the output ordinals; the caller does not choose them. Live nodes are numbered `0 .. live - 1`, sources in ascending-size order, and within a source by the region of the merged hierarchy each node descends to, so nodes that are close in vector space get adjacent ordinals: records are written sequentially, every cell's members are contiguous (the cell join below relies on this), and consecutive nodes explore the same neighbourhood of every target. Callers that keep external references to nodes (a row id to ordinal table, for instance) translate them through `ordinalMappers()` after `compact` returns.
 
+### Output format version
+
+The sources may be of any on-disk format version the reader supports (2 to 6); the compactor detects each source's version from its file. The output is written at `OnDiskGraphIndex.CURRENT_VERSION` unless `withVersion(int)` asks for another one, so by default a compaction also upgrades old sources. An output that carries fused PQ needs version 6, and one with upper layers version 4 or above; a version that cannot hold the output is rejected when `compact` starts. A sharded compaction writes every shard at the version of the compactor the session was opened on.
+
 ### Handling Deleted Nodes
 
 Deleted nodes are excluded from the output by marking them as `false` in the corresponding `FixedBitSet`; they get no output ordinal (`ordinalMappers()` reports `OrdinalMapper.OMITTED` for them).
@@ -53,6 +57,32 @@ for (int i = 0; i < source.size(); i++) {
 }
 ```
 
+### Compacting into several outputs (sharding)
+
+A compaction that splits its output, such as one writing each token range to its own file, can run as N separate compactions, one per output, each over all the sources with only that output's nodes marked live. That repeats per output what does not depend on the output: every one of them trains the codebooks again, on its own slice. A sharded compaction session does it once: the compactor is built over everything that may be written, and each shard is merged by its own `compactShard` call, with that shard's nodes as the live set. The codebooks are trained by the first shard that needs them, on the live nodes of every shard together, and reused by the rest.
+
+```java
+// liveNodes: the live nodes of ALL the shards together
+var compactor = new OnDiskGraphIndexCompactor(sources, liveNodes, VectorSimilarityFunction.COSINE, null);
+var session = compactor.sharded();
+
+// shardNLiveNodes: per source, the nodes of shard N, a subset of liveNodes
+List<OrdinalMapper> shard0 = session.compactShard(shard0LiveNodes, output0);   // shard0.get(s).oldToNew(n): ordinal of source s's node n in output0, or OMITTED
+List<OrdinalMapper> shard1 = session.compactShard(shard1LiveNodes, output1);   // every shard has its own mapping
+```
+
+Shards may become known one at a time (each output is closed as its rows arrive); the session does nothing until the first call. Sources that ship a compressed sidecar use the `compactShard` overload that also takes the compressed path. Each shard is otherwise an ordinary merge of its own nodes: its upper layers, cells, ordinals and merge store are its own, and its cost follows its node count. Since a shard holds a small part of every source, each source is sparse in it; see "Sparse sources" below.
+
+Measured on two 1M-node sources (1024 dimensions, cosine, degree 32) split into M shards by a hash of the row, total wall time and recall@10 over all shards (8 cores):
+
+| M | one compaction per shard, retained same-source edges only | one session, same-source join off | one session | rebuild of every shard |
+|---|---|---|---|---|
+| 4 | 454 s, 0.781 | | 566 s, 0.825 | 876 s, 0.836 |
+| 16 | 981 s, 0.823 | 300 s, 0.821 | 567 s, 0.894 | 946 s, 0.903 |
+| 64 | 1203 s, 0.875 | 335 s, 0.875 | 627 s, 0.945 | 886 s, 0.947 |
+
+Training once removes the growth with M (the two codebooks take 43 s once instead of about 14 s per shard at M = 64); the same-source join costs about as much again as the rest of the merge and brings the recall to within a point of a rebuild.
+
 ## Algorithm
 
 ### Ordinal assignment
@@ -84,6 +114,9 @@ searchTopK  = max(2,  ceil(degree / numSources) * 4)
 - *Level L > 0*: the compactor first descends greedily from the source's entry point through each level above L (one `searchOneLayer` call with topK=1 per level, feeding the result into the next via `setEntryPointsFromPreviousLayer()`), then performs the full beam search at level L. This mirrors standard HNSW construction and gives a much better starting point than jumping directly to level L from the global entry node.
 
 
+**Sparse sources: the same-source join**\
+Step 1 keeps only the live neighbours, and nothing replaces the dead ones. That is the right economy while most of a source is live, since the source's own graph already holds its same-source edges. It fails when a source's live nodes are a minority, which is every source of a shard (a shard of M keeps 1/M of each): a node then keeps `degree × live fraction` of its same-source edges, and the output is close to bipartite between the sources. So in a sharded compaction a source with fewer than half of its nodes live is also searched for the node's own neighbours, exactly as another source is: at level 0 it is one more target of the cell walk (the largest source, which otherwise scans nothing, scans itself), at the upper levels and in merges without cells one more graph search. The node itself and the neighbours it retained are dropped from the result, and no reverse offers are made within a source, since every node of it runs the same scan. An ordinary `compact` never joins a source with itself, however many of its nodes are deleted: its cost and output are unchanged by this.
+
 **3. Diversity selection** (Vamana-style)\
 Candidates are sorted by exact score (descending). The compactor selects up to `maxDegree` diverse neighbors using an adaptive alpha. At level 0 the pairwise comparisons between a candidate and the already-selected neighbours run on vectors decoded from the wide code (below), and the candidate's pruning threshold is recomputed against the node from the same decoded vector, so both sides of the inequality carry the same quantization bias; the ordering itself always uses the exact scores.
 
@@ -118,9 +151,9 @@ Processing is batched per source and run in parallel across sources using a `For
 
 ### Entry Node
 
-The entry node of the compacted graph is:
-1. The designated entry node of the first source that reaches the graph's max level, if it is live.
-2. Otherwise, the first live node found at the max level, scanning those sources in order.
+The output's max level is the highest level that still has a live node: a merged level that came out empty (a shard, or a heavily deleted source, can miss every node of the sources' top levels) is dropped. The entry node of the compacted graph is:
+1. The designated entry node of the first source that reaches the output's max level, if it is live.
+2. Otherwise, the first live node found at that level, scanning those sources in order.
 
 ## Benchmarking
 

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

@@ -41,6 +41,17 @@ public interface ReaderSupplier extends AutoCloseable {
     default void prefetch(long offset, long length) {
     }
 
+    /**
+     * Announces that the byte range {@code [offset, offset + length)} is about to be read, so the
+     * implementation can start bringing it into the page cache without blocking the caller
+     * ({@code POSIX_MADV_WILLNEED} on a mapping). Meant for one record at a time, a fixed number of
+     * records ahead of the read loop that will touch it. Best-effort; by default the synchronous
+     * {@link #prefetch(long, long)}.
+     */
+    default void announce(long offset, long length) {
+        prefetch(offset, length);
+    }
+
     default void close() throws IOException {
     }
 }

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

@@ -60,6 +60,7 @@ final class CompactWriter implements AutoCloseable {
     private final ThreadLocal<ByteBuffer> bufferPerThread;
     private final ThreadLocal<ByteSequence<?>> zeroPQ;
     private final boolean fusedPQEnabled;
+    private final boolean usesFooter;
     private final Path outputPath;
     private final List<CommonHeader.LayerInfo> configuredLayerInfo;
     private final List<Integer> configuredLayerDegrees;
@@ -75,7 +76,13 @@ final class CompactWriter implements AutoCloseable {
     private int pqCodeSize;
     private ThreadLocal<byte[]> pqCodeBufPerThread;
 
+    /**
+     * @param version the on-disk format version to write, one {@link GraphIndexFormatFactory} knows.
+     *                Fused PQ needs a version that supports it (6), an output with upper layers one
+     *                that supports multi-layer graphs (4 and above).
+     */
     CompactWriter(Path outputPath,
+                  int version,
                   int maxOrdinal,
                   int numBaseLayerNodes,
                   long startOffset,
@@ -87,18 +94,19 @@ final class CompactWriter implements AutoCloseable {
             throws IOException {
         this.fusedFeature = fusedFeature;
         this.fusedPQEnabled = fusedFeature != null;
-        // This writer only ever targets the current on-disk format -- unlike AbstractGraphIndexWriter,
-        // there is no builder/version parameter that lets a caller pick an older version. Because of
-        // that, `fusedPQEnabled` alone is a safe stand-in for "the format supports fused PQ" everywhere
-        // below: there's no other version this class could be writing. We still check the format here,
-        // once, instead of just assuming it, so that if this class is ever changed to support writing
-        // older versions, or a future version drops fused PQ support, this fails loudly at construction
-        // instead of silently omitting the fused-PQ footer/packed-neighbor data later on.
-        int version = OnDiskGraphIndex.CURRENT_VERSION;
-        if (fusedPQEnabled && !GraphIndexFormatFactory.forVersion(version).supportsFeature(FeatureId.FUSED_PQ)) {
-            throw new IllegalStateException(
+        // The format is checked once, here, against what this output needs. Everywhere below,
+        // `fusedPQEnabled` alone then stands for "the format supports fused PQ", and `usesFooter`
+        // for "the format is footer-based" (5 and above; the header is written at the start either way).
+        GraphIndexFormat format = GraphIndexFormatFactory.forVersion(version);
+        if (fusedPQEnabled && !format.supportsFeature(FeatureId.FUSED_PQ)) {
+            throw new IllegalArgumentException(
                     "Fused PQ requires a format version that supports it; version " + version + " does not");
         }
+        if (layerInfo.size() > 1 && !format.supportsMultiLayer()) {
+            throw new IllegalArgumentException(
+                    "A graph with upper layers requires a format version that supports them; version " + version + " does not");
+        }
+        this.usesFooter = format.usesFooter();
         this.outputPath = outputPath;
         this.writer = new BufferedRandomAccessWriter(outputPath);
         this.startOffset = startOffset;
@@ -199,6 +207,9 @@ void writeFooter() throws IOException {
                 vectorTypeSupport.writeByteSequence(writer, entryNodePqCode);
             }
         }
+        if (!usesFooter) {
+            return;   // a header-based format: the header written at the start is the metadata
+        }
         long headerOffset = writer.position();
         header.write(writer);
         writer.writeLong(headerOffset);
@@ -270,7 +281,9 @@ public long projectedOutputSize() {
             }
         }
         // Footer trailer: header copy + headerOffset (long) + magic (int).
-        total += headerSize + FOOTER_SIZE;
+        if (usesFooter) {
+            total += headerSize + FOOTER_SIZE;
+        }
         return total;
     }
 

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

@@ -44,6 +44,8 @@ public final class CompactionContext {
     public final int maxOrdinal;
     public final ForkJoinPool executor;
     public final int taskWindowSize;
+    /** Records announced to the page cache this many nodes ahead of their read (0 = none); see OnDiskGraphIndexCompactor#withRecordAnnounceAhead. */
+    public final int recordAnnounceAhead;
 
     public CompactionContext(
             List<OnDiskGraphIndex> sources,
@@ -53,7 +55,8 @@ public CompactionContext(
             int dimension,
             int maxOrdinal,
             ForkJoinPool executor,
-            int taskWindowSize) {
+            int taskWindowSize,
+            int recordAnnounceAhead) {
         this.sources = Collections.unmodifiableList(sources);
         this.sourceCompressed = sourceCompressed == null ? null : Collections.unmodifiableList(sourceCompressed);
         this.liveNodes = Collections.unmodifiableList(liveNodes);
@@ -62,5 +65,6 @@ public CompactionContext(
         this.maxOrdinal = maxOrdinal;
         this.executor = executor;
         this.taskWindowSize = taskWindowSize;
+        this.recordAnnounceAhead = recordAnnounceAhead;
     }
 }

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

@@ -34,12 +34,18 @@ final class ExactJoin {
     final long[] heaps;                // [(targetIndex * n + i) * k + j]
     final int[] sizes;                 // [targetIndex * n + i]
     final float[] thresholds;          // score of the k-th best so far, -inf until the heap is full
+    long pairsWalked;                  // (cell, node) pairs the walk scored, all targets; sizes the next range
+    long maxTargetPairs;               // the busiest target's share of them
     final Object[] locks = new Object[1024];
 
     private final int numSources;
     private final int[] sizeRank;
 
-    ExactJoin(int src, int lo, int hi, int k, int numSources, int[] sizeRank) {
+    /**
+     * @param selfTarget whether the searching source is also scanned for its own nodes (a sparse
+     *                   source, whose retained same-source edges mostly point at dead nodes)
+     */
+    ExactJoin(int src, int lo, int hi, int k, int numSources, int[] sizeRank, boolean selfTarget) {
         this.numSources = numSources;
         this.sizeRank = sizeRank;
         this.src = src;
@@ -50,7 +56,7 @@ final class ExactJoin {
         Arrays.fill(targetIndex, -1);
         int t = 0;
         for (int i = 0; i < numSources; i++) {
-            if (i != src && sizeRank[i] > sizeRank[src]) targetIndex[i] = t++;
+            if (i == src ? selfTarget : sizeRank[i] > sizeRank[src]) targetIndex[i] = t++;
         }
         targets = t;
         heaps = new long[Math.multiplyExact(Math.multiplyExact(targets, n), k)];

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

@@ -83,6 +83,12 @@ public void retrain(VectorSimilarityFunction vsf) {
         this.retrainedCompressor = (VectorCompressor<ByteSequence<?>>) (VectorCompressor<?>) retrainer.retrain(vsf);
     }
 
+    @Override
+    @SuppressWarnings("unchecked")
+    public void adoptCompressor(VectorCompressor<?> trained) {
+        this.retrainedCompressor = (VectorCompressor<ByteSequence<?>>) java.util.Objects.requireNonNull(trained, "trained");
+    }
+
     @Override
     public VectorCompressor<?> compressor() {
         return retrainedCompressor;