From 00a03f50b16741c856a444a74ddb8dc087c566a3 Mon Sep 17 00:00:00 2001 From: dian-lun-lin Date: Wed, 11 Nov 2026 22:42:38 +0000 Subject: [PATCH 1/6] Compaction: cell-walk optimizations for large merges The level-0 cell walk and the passes around it, tuned for merges whose store is larger than memory, and the compactor's settings kept its own. Range sizing. The walk's range was bounded by a fixed 3 GiB of query vectors and a pair budget capped at 8 GiB, with the pairs per node estimated once from the cell sizes; that estimate overstates the count several times, so every range was several times smaller than its budget allowed and the store was swept correspondingly more often. The query budget now scales with the heap (a sixteenth of it, at least 3 GiB), the pair budget's cap is 32 GiB, and after a source's first range the pairs per node come from what that range measured, with the busiest target's pair array kept under half the array limit even if the next range doubles them. Record announce-ahead. A source record is announced to the page cache a fixed number of nodes ahead of its read (ReaderSupplier.announce, POSIX_MADV_WILLNEED on the native reader, a positional read elsewhere): in the level-0 pass for batches whose ordinals are scattered by the region order, and in the pre-encode's wide path, which reads the source at random and had no readahead. Without the announcement each read is a demand fault served one at a time. The hub map by position, sized by its hubs. Level-1 neighbours are stored as positions in the level, resolved once when the hub map is built, so the beams walk positions and never look an id up; a scored hub maps to its cell through one position-indexed table. The two tables that held one entry per node of the merge are gone: ids are resolved by binary search over the level's sorted node list, which the merged upper layers now keep in ascending (source, node) order, and that order also makes the cells and the region ordinals independent of which batch finished first. The beam reuses per-thread arrays and rescans for its worst slot only after replacing it. Release when level 0 is done. The upper layers are written from the merged arrays, so the hub map's level-1 mapping and vectors and the cell position tables go before that pass instead of in the final cleanup; the scratches of all worker threads are closed, not only the calling thread's. Direct I/O for a store larger than memory. Every range of the walk reads the targets' store once from end to end. When the store fits beside the heap the mapping serves the later passes from memory; when it does not, each mapped page read evicts a resident page and the kernel reclaims synchronously inside the reading threads. The walk now decides by size and in the direct case reads each run of cells into a per-thread aligned buffer through a second, direct-I/O descriptor on the output file, decoding rows from the buffer; nothing touches the page cache. The pre-encoded cache learns its path and the aligned run read; the wide decoder decodes from a buffer. The walk's budgets, the announce-ahead distance and the direct-read decision have package-private setters for the tests only: each has a default the compactor derives from the machine or the sources, and a caller is not expected to know better. Claude-Session: https://claude.ai/code/session_016pmoQbpUFKvQKpvRc7Xvjz --- .../jbellis/jvector/disk/ReaderSupplier.java | 11 + .../jvector/graph/disk/CompactionContext.java | 6 +- .../jbellis/jvector/graph/disk/ExactJoin.java | 2 + .../jbellis/jvector/graph/disk/HubMap.java | 50 +- .../jvector/graph/disk/OnDiskGraphIndex.java | 18 +- .../graph/disk/OnDiskGraphIndexCompactor.java | 426 ++++++++++++++---- .../graph/disk/PreEncodedCodeCache.java | 125 ++++- .../disk/QuantizationCompactionStrategy.java | 10 +- .../jvector/graph/disk/WideDecoder.java | 38 ++ .../jvector/disk/MemorySegmentReader.java | 30 +- .../disk/TestOnDiskGraphIndexCompactor.java | 26 +- .../graph/disk/TestPreEncodedCodeCache.java | 38 +- 12 files changed, 657 insertions(+), 123 deletions(-) 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 7ebb3f9b0..e976de0fa 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 @@ -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 { } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java index a678faca7..08f7649be 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java +++ b/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 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; } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ExactJoin.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ExactJoin.java index ce9f21fe4..3d5625615 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ExactJoin.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ExactJoin.java @@ -34,6 +34,8 @@ 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; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/HubMap.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/HubMap.java index e83deb4a0..48f87a5c4 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/HubMap.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/HubMap.java @@ -31,7 +31,6 @@ import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; -import org.agrona.collections.Int2IntHashMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import sun.misc.Unsafe; @@ -66,16 +65,16 @@ private static Unsafe getUnsafe() { final int dimension; final int topLevel; int entryNode; - // per level (index 0 unused): node ids, adjacency with stride degree[level] (-1 padded), and - // the nodes' vectors position-major, all indexed by the node's position in nodes[level] + // per level (index 0 unused): node ids in ascending order, adjacency with stride degree[level] as + // positions in the same level (-1 padded), and the nodes' vectors position-major, all indexed by + // the node's position in nodes[level]. Everything is sized by the level's node count: a node id is + // resolved to its position by binary search over nodes[level] (the beams never need it; they walk + // positions), so no table grows with the total node count of the merge. final int[] degree; final int[][] nodes; final int[][] adjacency; final VectorFloat[] vectors; final float[][] norms; // COSINE only: |v| per position - // node id -> position: a flat array for level 1 (the bulk of the upper layers), maps above it - int[] level1Position; - final Int2IntHashMap[] upperPosition; // levels >= 2: node id -> position (level 1 uses level1Position) // level-1 vectors: read-only file mapping, row = position, in chunks private MappedByteBuffer[] level1Chunks; private FloatBuffer[] level1Floats; @@ -92,7 +91,6 @@ private static Unsafe getUnsafe() { this.adjacency = new int[topLevel + 1][]; this.vectors = new VectorFloat[topLevel + 1]; this.norms = new float[topLevel + 1][]; - this.upperPosition = new Int2IntHashMap[topLevel + 1]; } /** Per-thread query state. */ @@ -139,11 +137,12 @@ Scorer scorer() { return new Scorer(); } + /** Position of node id {@code node} in {@code level}, or -1 if the node is not on that level. */ int position(int level, int node) { - if (level == 1) { - return node >= 0 && node < level1Position.length ? level1Position[node] : -1; - } - return upperPosition[level].get(node); // missing value is -1 + int[] ids = nodes[level]; + if (ids == null) return -1; + int at = Arrays.binarySearch(ids, node); + return at >= 0 ? at : -1; } /** @@ -165,23 +164,19 @@ int descend(Scorer scorer) { while (improved) { improved = false; for (int j = 0; j < stride; j++) { - int neighbor = adj[position * stride + j]; - if (neighbor < 0) { - break; - } - int neighborPosition = position(level, neighbor); + int neighborPosition = adj[position * stride + j]; if (neighborPosition < 0) { - continue; + break; } float score = scorer.score(level, neighborPosition); if (score > currentScore) { currentScore = score; - current = neighbor; position = neighborPosition; improved = true; } } } + current = nodes[level][position]; } return current; } @@ -190,15 +185,13 @@ int descend(Scorer scorer) { * Breadth-first walk over the level-1 graph, restarting at the lowest unvisited position * when a component is exhausted. * - * @return walk position per node id (Integer.MAX_VALUE for nodes not on level 1) + * @return walk position per level-1 position (one entry per level-1 node) */ int[] walkPositions() { - int[] l1 = nodes[1]; - int n = l1.length; + int n = nodes[1].length; int stride = degree[1]; int[] adj = adjacency[1]; - int[] positionOf = new int[level1Position.length]; - Arrays.fill(positionOf, Integer.MAX_VALUE); + int[] walkOf = new int[n]; boolean[] seen = new boolean[n]; int[] queue = new int[n]; int head = 0, tail = 0, emitted = 0, nextUnseen = 0; @@ -214,20 +207,19 @@ int[] walkPositions() { queue[tail++] = nextUnseen; } int x = queue[head++]; - positionOf[l1[x]] = emitted++; + walkOf[x] = emitted++; for (int j = 0; j < stride; j++) { - int neighbor = adj[x * stride + j]; - if (neighbor < 0) { + int y = adj[x * stride + j]; + if (y < 0) { break; } - int y = position(1, neighbor); - if (y >= 0 && !seen[y]) { + if (!seen[y]) { seen[y] = true; queue[tail++] = y; } } } - return positionOf; + return walkOf; } /** Bytes of one row of the level-1 vector mapping. */ 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 f43357c27..4316be74c 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 @@ -454,12 +454,26 @@ 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 blockBytes = l0RecordBytes(); long start = neighborsOffset + blockBytes * minNode; readerSupplier.prefetch(start, blockBytes * (maxNode - minNode + 1)); } + /** + * Announces the L0 record of {@code node} to the page cache ahead of its read; see + * {@link ReaderSupplier#announce(long, long)}. Asynchronous where the reader supports it, so a + * caller that knows which record it will read a few steps from now keeps the device busy instead + * of faulting the records in one at a time. Best-effort no-op when unsupported. + */ + public void announceL0Record(int node) { + long blockBytes = l0RecordBytes(); + readerSupplier.announce(neighborsOffset + blockBytes * node, blockBytes); + } + + private long l0RecordBytes() { + return Integer.BYTES + inlineBlockSize + (long) Integer.BYTES * (layerInfo.get(0).degree + 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 e2f4c6bef..fb9dfe5eb 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 @@ -52,7 +52,6 @@ import io.github.jbellis.jvector.vector.types.FloatArray; import io.github.jbellis.jvector.util.BoundedLongHeap; import io.github.jbellis.jvector.util.NumericUtils; -import org.agrona.collections.Int2IntHashMap; import org.agrona.collections.IntHashSet; import org.slf4j.Logger; @@ -102,13 +101,19 @@ public final class OnDiskGraphIndexCompactor implements Accountable { // The raw-vector store every merge scans: a sidecar strategy whose "code" is the vector. private QuantizationCompactionStrategy vectorStore = QuantizationCompactionStrategy.NONE; // The cell walk scans cells for a range of nodes at once. The range is bounded by this many nodes, by the - // memory its query vectors take (about 3 GB), and by the memory of its (cell, node) pairs: a share of - // the heap, since a node probes many small cells to fill its scan budget. + // memory its query vectors take and by the memory of its (cell, node) pairs, both shares of the heap, + // since a node probes many small cells to fill its scan budget; see exactJoinRange. private static final int EXACT_JOIN_MAX_RANGE = 1 << 22; - private static final long EXACT_JOIN_QUERY_BYTES = 3L << 30; - private static final long EXACT_JOIN_PAIR_BYTES_MIN = 512L << 20, EXACT_JOIN_PAIR_BYTES_MAX = 8L << 30; + private static final long EXACT_JOIN_QUERY_BYTES_MIN = 3L << 30; + private static final long EXACT_JOIN_PAIR_BYTES_MIN = 512L << 20, EXACT_JOIN_PAIR_BYTES_MAX = 32L << 30; private static final int MIN_SEARCH_TOP_K = 2; private static final int SEARCH_TOP_K_MULTIPLIER = 4; + // Fixed cell-walk budgets (bytes); 0 = the heap-derived defaults above (see withCellWalkBudgets). + private long walkPairBytes = 0, walkQueryBytes = 0; + // Records announced to the page cache this many nodes ahead of their read in the level-0 pass and + // the pre-encode (see withRecordAnnounceAhead); 0 = demand faults only. + static final int DEFAULT_RECORD_ANNOUNCE_AHEAD = 64; + private int recordAnnounceAhead = DEFAULT_RECORD_ANNOUNCE_AHEAD; private List sources; // Optional non-fused compressed sidecar, parallel to `sources`. Null when sources carry their @@ -204,7 +209,7 @@ public final class OnDiskGraphIndexCompactor implements Accountable { // Merged before the ordinals exist (they define the cells), written after level 0. private int[][] mergedSrc, mergedNode, mergedNbrSrc, mergedNbrNode; private HubMap cellMap; // the merged upper layers, resident through level 0 - private int[] cellWalkPosition; // hub node id -> cell (walk position); MAX_VALUE if not level 1 + private int[] cellWalkPosition; // level-1 position -> cell (walk position) private int[] cellToPosition; // cell -> level-1 position in the hub map private int[][] cellStart; // [source][cell] -> first new ordinal of that source's nodes in the cell private int cellCount; @@ -236,6 +241,9 @@ private void setupCrossLink() { Arrays.toString(l0ProcessOrder), REVERSE_CANDIDATE_SLOTS, OFFER_BAND_WIDTH); } + // Whether the cell walk reads the store with direct I/O: null = decided by size (see walkReadsDirect). + private Boolean walkDirectSetting; + private Boolean walkDirect; // the decision, made once per compaction /** * Constructs a new OnDiskGraphIndexCompactor for graphs without a non-fused compressed sidecar. * Equivalent to calling the 6-arg constructor with {@code sourceCompressed = null}. @@ -303,6 +311,86 @@ public OnDiskGraphIndexCompactor( this.similarityFunction = similarityFunction; } + // The settings below are the compactor's own: every one has a default it derives from the + // machine or the sources, and a caller is not expected to know better. They are package-private + // so the tests can exercise the paths they select. + + /** + * Fixes the cell walk's range budgets instead of deriving them from the heap: {@code pairBytes} + * bounds the (cell, node) pair table of a range and {@code queryBytes} its query vectors (see + * {@link #exactJoinRange}). A value of 0 keeps the heap-derived default for that budget. Fixed + * budgets make the range, and so the sweep count and the walk's memory, a property of the + * configuration rather than of the machine the merge runs on. + */ + OnDiskGraphIndexCompactor withCellWalkBudgets(long pairBytes, long queryBytes) { + if (pairBytes < 0 || queryBytes < 0) throw new IllegalArgumentException("budgets must be >= 0"); + this.walkPairBytes = pairBytes; + this.walkQueryBytes = queryBytes; + return this; + } + + /** + * Sets how many records ahead of its read a node's source record is announced to the page cache + * in the level-0 pass and the pre-encode ({@code POSIX_MADV_WILLNEED} on the native reader, a + * positional read elsewhere). Records of a region-ordered batch are scattered over the source, so + * without the announcement each read is a demand fault served one at a time. 0 disables it. + */ + OnDiskGraphIndexCompactor withRecordAnnounceAhead(int ahead) { + if (ahead < 0) throw new IllegalArgumentException("ahead must be >= 0"); + this.recordAnnounceAhead = ahead; + return this; + } + + /** + * Whether the cell walk reads the store with direct I/O ({@code true}), through its mapping + * ({@code false}), or by a size test ({@code null}, the default): direct when the store does not + * fit in the memory beside the heap. See {@link #walkReadsDirect}. + */ + OnDiskGraphIndexCompactor withWalkDirectReads(Boolean direct) { + this.walkDirectSetting = direct; + return this; + } + + /** + * Whether the cell walk reads the store with direct I/O instead of through its mapping. Every + * range of the walk reads the targets' store once from end to end. When the store fits in the + * memory left beside the heap, the mapping serves every pass after the first from RAM and is the + * right choice. When it does not, a mapping read must evict a resident page for every page it + * brings in, and the kernel does that reclaim synchronously inside the reading threads (measured + * on a 1.5 TB store against 400 GB: a third of the machine in reclaim, the disks at a fraction of + * their rate); a direct read lands in the thread's own buffer and leaves the page cache alone. + */ + private boolean walkReadsDirect(PreEncodedCodeCache store) { + if (walkDirect != null) { + return walkDirect; + } + long storeBytes = store.sectionBytes(); + long beyondHeap = memoryBeyondHeap(); + boolean direct; + if (!store.supportsDirect()) { + direct = false; + } else if (walkDirectSetting != null) { + direct = walkDirectSetting; + } else { + direct = beyondHeap > 0 && storeBytes > beyondHeap * 9 / 10; + } + log.info("Cell walk: store of {} MB, {} MB of memory beyond the heap: the walk reads the store {}", + storeBytes >> 20, beyondHeap < 0 ? "unknown" : String.valueOf(beyondHeap >> 20), + direct ? "with direct I/O" : "through its mapping"); + walkDirect = direct; + return direct; + } + + /** Physical (or container) memory minus the maximum heap, or -1 when the platform bean is unavailable. */ + @SuppressWarnings("deprecation") + private static long memoryBeyondHeap() { + var bean = java.lang.management.ManagementFactory.getOperatingSystemMXBean(); + if (!(bean instanceof com.sun.management.OperatingSystemMXBean)) { + return -1; + } + return ((com.sun.management.OperatingSystemMXBean) bean).getTotalPhysicalMemorySize() - Runtime.getRuntime().maxMemory(); + } + /** * Validates that all source indexes have compatible configurations and required features * before attempting compaction. Ensures consistent dimensions, max degrees, hierarchical @@ -590,7 +678,7 @@ private QuantizationCompactionStrategy detectSidecarStrategy() { /** Snapshot the compactor's state into a {@link CompactionContext} for strategies to consume. */ private CompactionContext buildContext() { return new CompactionContext(sources, sourceCompressed, liveNodes, remappers, - dimension, maxOrdinal, executor, taskWindowSize); + dimension, maxOrdinal, executor, taskWindowSize, recordAnnounceAhead); } /** @@ -746,9 +834,13 @@ private void compactLevels(CompactWriter writer, int upperMaxCandidateSize = upperMaxPerSourceTopK * sources.size(); int maxCandidateSize = Math.max(baseMaxCandidateSize, upperMaxCandidateSize); int scratchDegree = Math.max(maxDegrees.get(0), Math.max(1, maxUpperDegree)); - final ThreadLocal threadLocalScratch = ThreadLocal.withInitial(() -> - new Scratch(maxCandidateSize, scratchDegree, dimension, sources, pq) - ); + // every worker's scratch is recorded so all of them, not only the calling thread's, are closed at the end + final List openedScratch = Collections.synchronizedList(new ArrayList<>()); + final ThreadLocal threadLocalScratch = ThreadLocal.withInitial(() -> { + Scratch sc = new Scratch(maxCandidateSize, scratchDegree, dimension, sources, pq); + openedScratch.add(sc); + return sc; + }); scratchLocal = threadLocalScratch; setupCrossLink(); @@ -832,10 +924,25 @@ private void compactLevels(CompactWriter writer, // nodes are one contiguous range, so a group of consecutive batches is a range of // new ordinals. The cell walk scans for the whole range before its batches run. OrdinalMapper m = remappers.get(s); - int range = exactJoinRange(s); - log.info("Cell walk: source {} in ranges of {} nodes", s, range); + int range = exactJoinRange(s, 0, 0, 0); + log.info("Cell walk: source {} in ranges of {} nodes (pair budget {} MB{}, query budget {} MB{})", s, range, + (walkPairBytes > 0 ? walkPairBytes : Math.max(EXACT_JOIN_PAIR_BYTES_MIN, Math.min(EXACT_JOIN_PAIR_BYTES_MAX, Runtime.getRuntime().maxMemory() / 8))) >> 20, + walkPairBytes > 0 ? " fixed" : " from the heap", + (walkQueryBytes > 0 ? walkQueryBytes : Math.max(EXACT_JOIN_QUERY_BYTES_MIN, Runtime.getRuntime().maxMemory() / 16)) >> 20, + walkQueryBytes > 0 ? " fixed" : " from the heap"); int gStart = 0; while (gStart < batches.size()) { + ExactJoin previous = exactJoin; + exactJoin = null; // released before the next walk allocates + if (previous != null) { + // size the next range from the pairs the previous one measured + int next = exactJoinRange(s, previous.n, previous.pairsWalked, previous.maxTargetPairs); + if (next != range) { + log.info("Cell walk: source {} ranges of {} nodes ({} pairs per node measured)", + s, next, previous.pairsWalked / Math.max(1, previous.n)); + range = next; + } + } BatchSpec first = batches.get(gStart); int lo = m.oldToNew(first.nodes[first.start]); int gEnd = gStart, hi = lo; @@ -848,9 +955,9 @@ private void compactLevels(CompactWriter writer, } exactJoin = cellWalk(s, lo, hi, params.searchTopK); runBatchesWithBackpressure(batches.subList(gStart, gEnd), ecs, submitOne, writeResults); - exactJoin = null; gStart = gEnd; } + exactJoin = null; } } @@ -870,6 +977,15 @@ private void compactLevels(CompactWriter writer, log.info("Wide code: {} candidate scores from the wide cache", l0WideScores.sum()); candidateStore = null; // the strategy unmaps it with its own cache } + if (cellMap != null) { + // the cells served the walk; the upper layers are written from the merged arrays, so the + // hub map's mapping and vectors and the cell tables go before the upper-layer pass + cellMap.close(); + cellMap = null; + cellWalkPosition = null; + cellToPosition = null; + cellStart = null; + } writer.offsetAfterInline(); } else { @@ -878,8 +994,7 @@ private void compactLevels(CompactWriter writer, } } - Scratch s = threadLocalScratch.get(); - s.close(); + for (Scratch sc : openedScratch) sc.close(); threadLocalScratch.remove(); } @@ -969,20 +1084,24 @@ 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. Under - // similarity ordering the batch's ordinals are scattered, so only prefetch when - // they still form a reasonably dense range. - int lo = Integer.MAX_VALUE; - int hi = -1; - for (int i = bs.start; i < bs.end; i++) { - lo = Math.min(lo, bs.nodes[i]); - hi = Math.max(hi, bs.nodes[i]); - } - if ((long) hi - lo <= 8L * (bs.end - bs.start)) { - sources.get(bs.sourceIdx).prefetchL0Records(lo, hi); + // Each node's own record read (adjacency + vector) is fully predictable, while search reads + // into other sources are data-dependent and stay demand-faulted. A batch whose ordinals form + // a dense range is streamed into the page cache whole; a region-ordered batch is scattered + // over the source, so its records are announced one at a time, a fixed number of nodes ahead + // of their read, which keeps the device busy instead of faulting them in one by one. + int lo = Integer.MAX_VALUE, hi = -1; + for (int i = bs.start; i < bs.end; i++) { + lo = Math.min(lo, bs.nodes[i]); + hi = Math.max(hi, bs.nodes[i]); + } + final boolean scattered = bs.end > bs.start && (long) hi - lo > 8L * (bs.end - bs.start); + if (bs.end > bs.start && !scattered) { + sources.get(bs.sourceIdx).prefetchL0Records(lo, hi); + } else if (scattered && recordAnnounceAhead > 0) { + OnDiskGraphIndex source = sources.get(bs.sourceIdx); + FixedBitSet live = liveNodes.get(bs.sourceIdx); + for (int i = bs.start; i < Math.min(bs.end, bs.start + recordAnnounceAhead); i++) { + if (live.get(bs.nodes[i])) source.announceL0Record(bs.nodes[i]); } } @@ -994,6 +1113,10 @@ private List computeBaseBatch(CompactWriter writer, for (int i = bs.start; i < bs.end; i++) { int node = bs.nodes[i]; if (!liveNodes.get(bs.sourceIdx).get(node)) continue; + if (scattered && recordAnnounceAhead > 0 && i + recordAnnounceAhead < bs.end) { + int ahead = bs.nodes[i + recordAnnounceAhead]; + if (liveNodes.get(bs.sourceIdx).get(ahead)) sources.get(bs.sourceIdx).announceL0Record(ahead); + } out.add(processBaseNode(node, bs.sourceIdx, scratch, writer, params)); } @@ -1003,6 +1126,7 @@ private List computeBaseBatch(CompactWriter writer, + /** * Processes a batch of upper layer nodes from one source index. Similar to base layer * processing but returns only ordinal, neighbors, and optional PQ code (no inline vectors). @@ -1068,6 +1192,10 @@ private void mergeUpperLayers(ProductQuantization pq, boolean fusedPQEnabled) } return out; }), all::addAll); + // ascending (source, node): the hub map resolves ids by binary search over this order, and the + // order fixes the level-1 positions, so the cells and region ordinals no longer depend on which + // batch finished first + all.sort(Comparator.comparingInt((RawUpper r) -> r.src).thenComparingInt(r -> r.node)); int n = all.size(), degree = maxDegrees.get(level); int[] ms = new int[n], mn = new int[n], ns = new int[n * degree], nn = new int[n * degree]; Arrays.fill(ns, -1); Arrays.fill(nn, -1); @@ -1102,19 +1230,34 @@ private HubMap buildMergedHubMap() throws IOException { int n = ms.length, degree = maxDegrees.get(level); int[] nodes = new int[n]; for (int i = 0; i < n; i++) nodes[i] = idOffset[ms[i]] + mn[i]; + for (int i = 1; i < n; i++) if (nodes[i] <= nodes[i - 1]) throw new IllegalStateException("merged level " + level + " is not in ascending id order"); h.nodes[level] = nodes; - if (level == 1) { - h.level1Position = new int[idOffset[numSources]]; - Arrays.fill(h.level1Position, -1); - for (int i = 0; i < n; i++) h.level1Position[nodes[i]] = i; - } else { - h.upperPosition[level] = new Int2IntHashMap(n * 2, 0.65f, -1); - for (int i = 0; i < n; i++) h.upperPosition[level].put(nodes[i], i); - } h.degree[level] = degree; - int[] adj = new int[n * degree]; - int[] ns = mergedNbrSrc[level], nn = mergedNbrNode[level]; - for (int e = 0; e < n * degree; e++) adj[e] = ns[e] < 0 ? -1 : idOffset[ns[e]] + nn[e]; + // adjacency as positions in the level: a neighbour absent from the level (-1) ends the row, as + // the merged rows are written; resolved once here, in parallel, so the beams walk positions + final int[] adj = new int[n * degree]; + final int[] ns = mergedNbrSrc[level], nn = mergedNbrNode[level]; + { + int rows = Math.max(1024, n / (4 * taskWindowSize) + 1); + List> resolve = new ArrayList<>(); + for (int from = 0; from < n; from += rows) { + final int lo = from, hi = Math.min(n, from + rows); + resolve.add(() -> { + for (int i = lo; i < hi; i++) { + int w = i * degree; + for (int j = 0; j < degree; j++) { + int e = i * degree + j; + if (ns[e] < 0) break; + int at = Arrays.binarySearch(nodes, idOffset[ns[e]] + nn[e]); + if (at >= 0) adj[w++] = at; + } + for (; w < (i + 1) * degree; w++) adj[w] = -1; + } + return null; + }); + } + joinAll(resolve); + } final boolean mapped = level == 1; // level 1: vectors to a file instead of the heap if (!mapped && (long) n * dimension > Integer.MAX_VALUE - 8) throw new IllegalStateException("merged hub map: level " + level + " too large"); final VectorFloat levelVectors = mapped ? null : vectorTypeSupport.createFloatVector(n * dimension); @@ -1559,21 +1702,22 @@ private void selectCellsByBeam(Scratch scratch, HubMap.Scorer scorer, int startP if (scratch.l1Visited == null) { scratch.l1Visited = new IntHashSet(2 * beamScoredCapacity(CELL_PROBE_EF, stride)); scratch.scoredKeys = new long[beamScoredCapacity(CELL_PROBE_EF, stride)]; + scratch.beamState = new BeamState(CELL_PROBE_EF, stride); } long[] scored = scratch.scoredKeys; - int scoredCount = beam(cellMap, scorer, 1, startPos, CELL_PROBE_EF, scratch.l1Visited, scored); + int scoredCount = beam(cellMap, scorer, 1, startPos, CELL_PROBE_EF, scratch.l1Visited, scored, scratch.beamState); Arrays.sort(scored, 0, scoredCount); int n = 0; for (int i = scoredCount - 1; i >= 0; i--) { // best first - int cell = cellWalkPosition[cellMap.nodes[1][(int) scored[i]]]; + int cell = cellWalkPosition[(int) scored[i]]; if (cell >= 0 && cell < cellCount) scratch.probeCells[n++] = cell; } scratch.probeCount = n; } /** The best-scoring level-1 position a beam of width {@code ef} from {@code startPos} scores. */ - private static int bestLevel1ByBeam(HubMap map, HubMap.Scorer scorer, int startPos, int ef, IntHashSet visited, long[] scored) { - int count = beam(map, scorer, 1, startPos, ef, visited, scored); + private static int bestLevel1ByBeam(HubMap map, HubMap.Scorer scorer, int startPos, int ef, IntHashSet visited, long[] scored, BeamState state) { + int count = beam(map, scorer, 1, startPos, ef, visited, scored, state); long best = scored[0]; for (int i = 1; i < count; i++) if (scored[i] > best) best = scored[i]; return (int) best; @@ -1586,17 +1730,33 @@ private static int bestLevel1ByBeam(HubMap map, HubMap.Scorer scorer, int startP * expansions are capped at {@code 2 * ef} (the beam has converged long before) and the caller * sizes {@code scored} for that bound ({@link #beamScoredCapacity}). */ - private static int beam(HubMap map, HubMap.Scorer scorer, int level, int startPos, int ef, IntHashSet visited, long[] scored) { + /** A beam's working arrays, allocated once per thread and reused across calls. */ + static final class BeamState { + final int[] pos; + final float[] sc; + final boolean[] expanded; + final int[] nbPos; + + BeamState(int ef, int stride) { + pos = new int[ef]; + sc = new float[ef]; + expanded = new boolean[ef]; + nbPos = new int[stride]; + } + } + + private static int beam(HubMap map, HubMap.Scorer scorer, int level, int startPos, int ef, IntHashSet visited, long[] scored, BeamState state) { final int stride = map.degree[level]; - final int[] adj = map.adjacency[level]; - int[] pos = new int[ef]; - float[] sc = new float[ef]; - boolean[] expanded = new boolean[ef]; - int[] nbPos = new int[stride]; + final int[] adj = map.adjacency[level]; // positions in the level + final int[] pos = state.pos; + final float[] sc = state.sc; + final boolean[] expanded = state.expanded; + final int[] nbPos = state.nbPos; visited.clear(); int size = 0, scoredCount = 0; + int worst = -1; // index of the beam's worst entry once the beam is full; -1 = unknown float s0 = scorer.score(level, startPos); - pos[0] = startPos; sc[0] = s0; size = 1; + pos[0] = startPos; sc[0] = s0; expanded[0] = false; size = 1; visited.add(startPos); scored[scoredCount++] = packScored(s0, startPos); for (int expansions = 0; expansions < 2 * ef; expansions++) { @@ -1610,10 +1770,9 @@ private static int beam(HubMap map, HubMap.Scorer scorer, int level, int startPo // score the unvisited neighbours in one tight loop first (their cache misses overlap) int cnt = 0; for (int j = 0; j < stride; j++) { - int nb = adj[p * stride + j]; - if (nb < 0) break; - int np = map.position(level, nb); - if (np >= 0 && visited.add(np)) nbPos[cnt++] = np; + int np = adj[p * stride + j]; + if (np < 0) break; + if (visited.add(np)) nbPos[cnt++] = np; } for (int j = 0; j < cnt; j++) { scored[scoredCount + j] = packScored(scorer.score(level, nbPos[j]), nbPos[j]); @@ -1624,11 +1783,15 @@ private static int beam(HubMap map, HubMap.Scorer scorer, int level, int startPo float score = unpackScore(key); if (size < ef) { pos[size] = np; sc[size] = score; expanded[size] = false; size++; + worst = -1; } else { - int wi = 0; - for (int i = 1; i < size; i++) if (sc[i] < sc[wi]) wi = i; - if (score > sc[wi]) { - pos[wi] = np; sc[wi] = score; expanded[wi] = false; + if (worst < 0) { // the worst slot is cached and only rescanned after it was replaced + worst = 0; + for (int i = 1; i < size; i++) if (sc[i] < sc[worst]) worst = i; + } + if (score > sc[worst]) { + pos[worst] = np; sc[worst] = score; expanded[worst] = false; + worst = -1; } } } @@ -2194,7 +2357,7 @@ private List buildRegionOrdinalMappers(ProductQuantization pq) { cellStart = new int[numSources][]; cellToPosition = new int[cellCount]; for (int pos = 0; pos < cellCount; pos++) { - int w = walkPosition[hubMap.nodes[1][pos]]; + int w = walkPosition[pos]; if (w >= 0 && w < cellCount) cellToPosition[w] = pos; } } @@ -2223,6 +2386,7 @@ private List buildRegionOrdinalMappers(ProductQuantization pq) { HubMap.Scorer scorer = regionMode ? hubMapRef.scorer() : null; IntHashSet beamVisited = regionMode ? new IntHashSet(2 * beamScoredCapacity(CELL_ASSIGN_EF, hubMapRef.degree[1])) : null; long[] beamScored = regionMode ? new long[beamScoredCapacity(CELL_ASSIGN_EF, hubMapRef.degree[1])] : null; + BeamState beamState = regionMode ? new BeamState(CELL_ASSIGN_EF, hubMapRef.degree[1]) : null; try (var view = (OnDiskGraphIndex.View) source.getView()) { for (int node = lo; node < hi; node++) { if (!alive.get(node)) continue; @@ -2234,9 +2398,9 @@ private List buildRegionOrdinalMappers(ProductQuantization pq) { int landing = hubMapRef.descend(scorer); int lp = hubMapRef.position(1, landing); if (lp >= 0) { - landing = hubMapRef.nodes[1][bestLevel1ByBeam(hubMapRef, scorer, lp, CELL_ASSIGN_EF, beamVisited, beamScored)]; + lp = bestLevel1ByBeam(hubMapRef, scorer, lp, CELL_ASSIGN_EF, beamVisited, beamScored, beamState); } - int pos = landing >= 0 && landing < walkPositionRef.length ? walkPositionRef[landing] : Integer.MAX_VALUE; + int pos = lp >= 0 ? walkPositionRef[lp] : Integer.MAX_VALUE; key = pos & 0xFFFFFFFFL; } else { key = 0; @@ -2380,23 +2544,44 @@ private void prefetchFinalizeRows(BatchSpec bs, Scratch scratch) { /** * Node range size of the cell walk for source {@code src}: at most {@link #EXACT_JOIN_MAX_RANGE} - * nodes, about {@link #EXACT_JOIN_QUERY_BYTES} of query vectors, and (cell, node) pairs within - * an eighth of the heap. A node records one pair per cell it probes, and it probes cells until - * {@link #CELL_BUDGET} members per larger target, so the pairs per node follow from the cell - * sizes: with every source's level-1 nodes as cells that is a few hundred pairs per target. + * nodes, query vectors within a sixteenth of the heap (at least {@link #EXACT_JOIN_QUERY_BYTES_MIN}), + * and (cell, node) pairs within an eighth of the heap, so pairs with their growth slack, the + * per-target sort copy and the queries stay around a third. Every range streams the targets' + * store once, so the range is as large as the heap allows: a store larger than RAM is read + * from the device once per range. + * + *

Before the first range of a source the pairs per node are bounded from the cell sizes: a + * node records one pair per cell it probes, and it probes cells until {@link #CELL_BUDGET} + * members per larger target, which overstates the count several times. After the first range + * the measured counts take over ({@code measuredPairs} over {@code measuredNodes}, with + * {@code measuredMaxTarget} the busiest target's pairs), and a target's pair array is kept + * under half the array limit even if the next range doubles them. */ - private int exactJoinRange(int src) { - long byQueries = EXACT_JOIN_QUERY_BYTES / (4L * dimension); - long pairsPerNode = 0; - int maxProbes = beamScoredCapacity(CELL_PROBE_EF, cellMap.degree[1]); - for (int t = 0; t < sources.size(); t++) { - if (sizeRank[t] <= sizeRank[src]) continue; // only larger sources are scanned - long membersPerCell = Math.max(1, numLiveNodesPerSource.get(t) / Math.max(1, cellCount)); - pairsPerNode += Math.min(maxProbes, (CELL_BUDGET + membersPerCell - 1) / membersPerCell); - } - long pairBudget = Math.max(EXACT_JOIN_PAIR_BYTES_MIN, Math.min(EXACT_JOIN_PAIR_BYTES_MAX, Runtime.getRuntime().maxMemory() / 8)); + private int exactJoinRange(int src, long measuredNodes, long measuredPairs, long measuredMaxTarget) { + long heap = Runtime.getRuntime().maxMemory(); + long queryBudget = walkQueryBytes > 0 ? walkQueryBytes : Math.max(EXACT_JOIN_QUERY_BYTES_MIN, heap / 16); + long byQueries = queryBudget / (4L * dimension); + long pairBudget = walkPairBytes > 0 ? walkPairBytes : Math.max(EXACT_JOIN_PAIR_BYTES_MIN, Math.min(EXACT_JOIN_PAIR_BYTES_MAX, heap / 8)); + long pairsPerNode, maxTargetPerNode; + if (measuredNodes > 0) { + pairsPerNode = (measuredPairs + measuredNodes - 1) / measuredNodes; + maxTargetPerNode = (measuredMaxTarget + measuredNodes - 1) / measuredNodes; + } else { + pairsPerNode = 0; + maxTargetPerNode = 0; + int maxProbes = beamScoredCapacity(CELL_PROBE_EF, cellMap.degree[1]); + for (int t = 0; t < sources.size(); t++) { + if (sizeRank[t] <= sizeRank[src]) continue; // only larger sources are scanned + long membersPerCell = Math.max(1, numLiveNodesPerSource.get(t) / Math.max(1, cellCount)); + long perTarget = Math.min(maxProbes, (CELL_BUDGET + membersPerCell - 1) / membersPerCell); + pairsPerNode += perTarget; + maxTargetPerNode = Math.max(maxTargetPerNode, perTarget); + } + } long byPairs = pairsPerNode == 0 ? Long.MAX_VALUE : pairBudget / (Long.BYTES * pairsPerNode); - return (int) Math.max(1 << 16, Math.min(EXACT_JOIN_MAX_RANGE, Math.min(byQueries, byPairs))); + long byTarget = maxTargetPerNode == 0 ? Long.MAX_VALUE : (Integer.MAX_VALUE / 2L) / maxTargetPerNode; + long range = Math.min(Math.min(EXACT_JOIN_MAX_RANGE, byQueries), Math.min(byPairs, byTarget)); + return (int) Math.max(1 << 16, range); } /** @@ -2413,6 +2598,12 @@ private ExactJoin cellWalk(int src, int lo, int hi, int k) { final boolean needNorms = similarityFunction != VectorSimilarityFunction.DOT_PRODUCT; final VectorFloat[] queries = new VectorFloat[n]; final float[] queryNorm2 = new float[n]; + final boolean direct = walkReadsDirect(candidateStore); + final int readRowBytes = candidateStore.codeSize(); + final int readRows = Math.max(1, EXACT_JOIN_READ_BYTES / readRowBytes); // rows per direct read + final int readGapRows = Math.max(1, EXACT_JOIN_READ_GAP_BYTES / readRowBytes); + final int readBufBytes = candidateStore.directBufferBytes(readRows); + final LongAdder directBytes = new LongAdder(); // phase 1: query vectors and probe cells; (cell << 32 | node index) pairs per target, per task final int tasks1 = Math.max(1, Math.min(taskWindowSize * 4, n / 256)); final int chunk = (n + tasks1 - 1) / tasks1; @@ -2432,10 +2623,23 @@ private ExactJoin cellWalk(int src, int lo, int hi, int k) { long[][] pairs = pairsByTask[task]; int[] count = countByTask[task]; for (int t = 0; t < j.targets; t++) pairs[t] = new long[Math.max(64, (to - from) * 6)]; + final ByteBuffer buf = direct ? sc.directBuffer(readBufBytes) : null; + int bufLo = 0, bufHi = 0, bufBase = 0; // rows [bufLo, bufHi) are in buf, row bufLo at bufBase for (int i = from; i < to; i++) { int newOrd = lo + i; VectorFloat q = vectorTypeSupport.createFloatVector(dimension); - sc.wide.decode(newOrd, q); + if (direct) { + if (newOrd >= bufHi) { + int rows = Math.min(readRows, to - i); + bufBase = readRowsDirect(candidateStore, newOrd, rows, buf); + bufLo = newOrd; + bufHi = newOrd + rows; + directBytes.add((long) rows * readRowBytes); + } + sc.wide.decode(buf, bufBase + (newOrd - bufLo) * readRowBytes, q); + } else { + sc.wide.decode(newOrd, q); + } queries[i] = q; queryNorm2[i] = needNorms ? VectorUtil.dotProduct(q, q) : 0f; int ownCell = cellOf(src, newOrd); @@ -2468,7 +2672,7 @@ private ExactJoin cellWalk(int src, int lo, int hi, int k) { joinAll(phase1); final long t1 = System.nanoTime(); // phase 2: per target, sort the pairs by cell and walk the cells in ordinal order - long pairTotal = 0; + long pairTotal = 0, maxTargetPairs = 0; for (int target = 0; target < sources.size(); target++) { final int ti = j.targetIndex[target]; if (ti < 0) continue; @@ -2485,7 +2689,9 @@ private ExactJoin cellWalk(int src, int lo, int hi, int k) { if (total == 0) continue; Arrays.parallelSort(pairs); pairTotal += total; + maxTargetPairs = Math.max(maxTargetPairs, total); final int[] starts = cellStart[target]; + final int rowBytes = candidateStore.codeSize(); // split the sorted pairs into contiguous cell ranges for the workers int tasks2 = Math.max(1, Math.min(taskWindowSize * 8, total / 4096)); List> phase2 = new ArrayList<>(); @@ -2503,6 +2709,15 @@ private ExactJoin cellWalk(int src, int lo, int hi, int k) { final VectorFloat[] chunkQueries = new VectorFloat[EXACT_QUERY_CHUNK]; final float[] chunkScores = new float[EXACT_QUERY_CHUNK]; int[] probers = new int[256]; + // Direct mode reads the task's cells, which are sorted so their store rows form ascending + // runs, a piece at a time into the thread's buffer: a run is planned from the current cell + // while the gaps stay small and the run short (planHi), and nothing touches the page cache. + // Mapped mode decodes straight from the mapping, as a store that fits in RAM wants. + int qAhead = from; + long read = 0; + final ByteBuffer buf = direct ? sc.directBuffer(readBufBytes) : null; + int bufLo = 0, bufHi = 0, bufBase = 0; // rows [bufLo, bufHi) are in buf, row bufLo at bufBase + int planHi = 0; int a = from; while (a < to) { int cell = (int) (pairs[a] >>> 32); @@ -2516,7 +2731,30 @@ private ExactJoin cellWalk(int src, int lo, int hi, int k) { final int cn = Math.min(EXACT_QUERY_CHUNK, count - c0); for (int r = 0; r < cn; r++) chunkQueries[r] = queries[probers[c0 + r]]; for (int newOrd = cLo; newOrd < cHi; newOrd++) { - sc.wide.decode(newOrd, tmp); + if (direct) { + if (newOrd >= bufHi || newOrd < bufLo) { + if (newOrd >= planHi) { + int runHi = cHi; + if (qAhead < b) qAhead = b; + while (qAhead < to && runHi - newOrd < readRows) { + int mc = (int) (pairs[qAhead] >>> 32); + int mLo = starts[mc], mHi = starts[mc + 1]; + if (mLo - runHi > readGapRows) break; + if (mHi > runHi) runHi = mHi; + while (qAhead < to && (int) (pairs[qAhead] >>> 32) == mc) qAhead++; + } + planHi = runHi; + } + int rows = Math.min(readRows, planHi - newOrd); + bufBase = readRowsDirect(candidateStore, newOrd, rows, buf); + bufLo = newOrd; + bufHi = newOrd + rows; + read += (long) rows * readRowBytes; + } + sc.wide.decode(buf, bufBase + (newOrd - bufLo) * readRowBytes, tmp); + } else { + sc.wide.decode(newOrd, tmp); + } VectorUtil.dotProductMulti(tmp, chunkQueries, cn, chunkScores); float vNorm2 = needNorms ? VectorUtil.dotProduct(tmp, tmp) : 0f; for (int r = 0; r < cn; r++) { @@ -2539,17 +2777,31 @@ private ExactJoin cellWalk(int src, int lo, int hi, int k) { } a = b; } + directBytes.add(read); return null; }); } joinAll(phase2); } - log.info("Cell walk: source {} ordinals [{}, {}): probes {} ms, {} (cell, node) pairs over {} targets, walk {} ms", - src, lo, hi, (t1 - t0) / 1_000_000, pairTotal, j.targets, (System.nanoTime() - t1) / 1_000_000); + j.pairsWalked = pairTotal; + j.maxTargetPairs = maxTargetPairs; + log.info("Cell walk: source {} ordinals [{}, {}): probes {} ms, {} (cell, node) pairs over {} targets, walk {} ms{}", + src, lo, hi, (t1 - t0) / 1_000_000, pairTotal, j.targets, (System.nanoTime() - t1) / 1_000_000, + direct ? ", " + (directBytes.sum() >> 20) + " MB of store read directly" : ""); return j; } private static final int EXACT_QUERY_CHUNK = 32; // queries scored together against each vector; 32 x 1.5 KB stays in L1 + private static final int EXACT_JOIN_READ_BYTES = 2 << 20; // store bytes per direct read (direct walk) + private static final int EXACT_JOIN_READ_GAP_BYTES = 64 << 10; // unwanted bytes a direct read may span + + private static int readRowsDirect(PreEncodedCodeCache store, int first, int rows, ByteBuffer buf) { + try { + return store.readRowsDirect(first, rows, buf); + } catch (IOException e) { + throw new UncheckedIOException("reading the merge store", e); + } + } private static final class Scratch implements AutoCloseable { final int[] candSrc, candNode; @@ -2568,6 +2820,7 @@ private static final class Scratch implements AutoCloseable { int probeCount; long[] scoredKeys; // positions scored by the current probe beam, packed with their scores IntHashSet l1Visited; // visited set for the probe beam + BeamState beamState; // the probe beam's working arrays int probeNode = -1, probeSrc = -1; int[] finalizeOrds; // full-precision merge: the batch's candidate rows, for the ordered prefetch @@ -2603,6 +2856,15 @@ public void close() throws IOException { for (var s : gs) s.close(); selectedCache.reset(); } + + ByteBuffer directBuf; // cell walk in direct mode: the thread's aligned read buffer + + ByteBuffer directBuffer(int bytes) { + if (directBuf == null || directBuf.capacity() < bytes) { + directBuf = PreEncodedCodeCache.allocateAligned(bytes); + } + return directBuf; + } } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PreEncodedCodeCache.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PreEncodedCodeCache.java index 64535ccc0..9da556ba5 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PreEncodedCodeCache.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PreEncodedCodeCache.java @@ -26,6 +26,8 @@ import java.nio.ByteBuffer; import java.nio.MappedByteBuffer; import java.nio.channels.FileChannel; +import java.nio.file.StandardOpenOption; +import java.nio.file.Path; /** * A memory-mapped array of pre-encoded fused codes, indexed by new ordinal. @@ -88,13 +90,20 @@ private static Unsafe getUnsafe() { private final ThreadLocal viewsPerThread; private final ThreadLocal warmByte = ThreadLocal.withInitial(() -> ByteBuffer.allocateDirect(1)); + /** Alignment direct I/O needs for offsets, lengths and buffer addresses (the file system block). */ + public static final int DIRECT_ALIGN = 4096; + private final Path path; // null when the cache was mapped without one: no direct reads + private volatile FileChannel direct; // opened on first use with the direct-I/O option + private volatile boolean directUnavailable; + private PreEncodedCodeCache(MappedByteBuffer[] chunks, int codeSize, int codesPerChunk, - FileChannel channel, long baseOffset) { + FileChannel channel, long baseOffset, Path path) { this.chunks = chunks; this.codeSize = codeSize; this.codesPerChunk = codesPerChunk; this.channel = channel; this.baseOffset = baseOffset; + this.path = path; this.viewsPerThread = ThreadLocal.withInitial(() -> { ByteBuffer[] views = new ByteBuffer[chunks.length]; for (int i = 0; i < chunks.length; i++) { @@ -116,12 +125,26 @@ private PreEncodedCodeCache(MappedByteBuffer[] chunks, int codeSize, int codesPe */ public static PreEncodedCodeCache map(FileChannel fc, long offset, int count, int codeSize) throws IOException { - return map(fc, offset, count, codeSize, MAX_CHUNK_BYTES); + return map(fc, null, offset, count, codeSize, MAX_CHUNK_BYTES); + } + + /** + * As {@link #map(FileChannel, long, int, int)}, also recording the file's path so the cache can + * open it a second time for {@link #readRowsDirect direct reads}. + */ + public static PreEncodedCodeCache map(FileChannel fc, Path path, long offset, int count, int codeSize) + throws IOException { + return map(fc, path, offset, count, codeSize, MAX_CHUNK_BYTES); } /** As {@link #map(FileChannel, long, int, int)} with an explicit chunk target, so tests can exercise the multi-chunk path. */ static PreEncodedCodeCache map(FileChannel fc, long offset, int count, int codeSize, int maxChunkBytes) throws IOException { + return map(fc, null, offset, count, codeSize, maxChunkBytes); + } + + static PreEncodedCodeCache map(FileChannel fc, Path path, long offset, int count, int codeSize, int maxChunkBytes) + throws IOException { if (codeSize <= 0) { throw new IllegalArgumentException("codeSize must be positive, got " + codeSize); } @@ -147,7 +170,93 @@ static PreEncodedCodeCache map(FileChannel fc, long offset, int count, int codeS offset + firstCode * codeSize, codesHere * codeSize); } - return new PreEncodedCodeCache(chunks, codeSize, codesPerChunk, fc, offset); + return new PreEncodedCodeCache(chunks, codeSize, codesPerChunk, fc, offset, path); + } + + /** Bytes of the whole section: every row of every chunk. */ + public long sectionBytes() { + long bytes = 0; + for (MappedByteBuffer c : chunks) { + bytes += c.capacity(); + } + return bytes; + } + + /** + * Whether {@link #readRowsDirect} can be used: the cache knows its file and the file system + * accepts the direct-I/O open option (tmpfs, for one, does not). + */ + public boolean supportsDirect() { + return directChannel() != null; + } + + private FileChannel directChannel() { + FileChannel d = direct; + if (d != null || directUnavailable || path == null) { + return d; + } + synchronized (this) { + if (direct == null && !directUnavailable) { + try { + direct = FileChannel.open(path, StandardOpenOption.READ, com.sun.nio.file.ExtendedOpenOption.DIRECT); + } catch (IOException | UnsupportedOperationException | IllegalArgumentException e) { + log.info("Direct reads of {} unavailable ({}); the cache is read through its mapping", path, e.toString()); + directUnavailable = true; + } + } + return direct; + } + } + + /** A direct buffer aligned for {@link #readRowsDirect}, holding at least {@code bytes}. */ + public static ByteBuffer allocateAligned(int bytes) { + return ByteBuffer.allocateDirect(bytes + DIRECT_ALIGN).alignedSlice(DIRECT_ALIGN); + } + + /** Buffer capacity {@link #readRowsDirect} needs for {@code count} rows at once. */ + public int directBufferBytes(int count) { + return (int) Math.min(Integer.MAX_VALUE - DIRECT_ALIGN, (long) count * codeSize + 2L * DIRECT_ALIGN); + } + + /** + * Reads the rows {@code [firstOrdinal, firstOrdinal + count)} straight from the file into + * {@code dst}, bypassing the page cache: the bytes land in the caller's buffer and nothing is + * cached, mapped or evicted on their account. For a pass that reads a store larger than RAM + * once from end to end, this keeps the kernel's memory reclaim out of the walk, where a mapping + * read of the same rows must evict a resident page for every page it brings in. Offsets and + * lengths are widened to {@link #DIRECT_ALIGN} as direct I/O requires, so {@code dst} needs + * {@link #directBufferBytes} capacity and the rows start at the returned offset. + * + * @return the offset in {@code dst} at which row {@code firstOrdinal} begins + */ + public int readRowsDirect(int firstOrdinal, int count, ByteBuffer dst) throws IOException { + FileChannel d = directChannel(); + if (d == null) { + throw new IOException("direct reads unavailable for " + path); + } + long start = baseOffset + (long) firstOrdinal * codeSize; + long end = start + (long) count * codeSize; + long alignedStart = start & -(long) DIRECT_ALIGN; + long alignedEnd = (end + DIRECT_ALIGN - 1) & -(long) DIRECT_ALIGN; + int bytes = (int) (alignedEnd - alignedStart); + if (bytes > dst.capacity()) { + throw new IllegalArgumentException("buffer of " + dst.capacity() + " bytes cannot hold " + count + " rows (" + bytes + " aligned bytes)"); + } + dst.clear(); + dst.limit(bytes); + long pos = alignedStart; + while (dst.hasRemaining()) { + int n = d.read(dst, pos); + if (n < 0) { + break; + } + pos += n; + if (pos % DIRECT_ALIGN != 0) { + // a short read is the end of the file: the aligned tail ran past it, and nothing beyond is a row + break; + } + } + return (int) (start - alignedStart); } /** Bytes of the mapped region for {@code count} codes. */ @@ -155,7 +264,6 @@ public static long sectionBytes(int count, int codeSize) { return (long) count * codeSize; } - /** Bytes per code. */ public int codeSize() { return codeSize; @@ -230,6 +338,15 @@ public void get(int ordinal, byte[] dst) { @Override public void close() { + FileChannel d = direct; + if (d != null) { + try { + d.close(); + } catch (IOException ignored) { + // nothing further to release + } + direct = null; + } if (UNSAFE == null) { return; } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java index d8b6a825a..36361d0cb 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java @@ -276,7 +276,7 @@ protected PreEncodedCodeCache precomputeCodeCache(CompactionContext ctx, Compact final PreEncodedCodeCache cache; try (FileChannel fc = FileChannel.open(writer.getOutputPath(), StandardOpenOption.READ, StandardOpenOption.WRITE)) { fc.write(ByteBuffer.wrap(new byte[]{0}), secondaryOffset + secondaryBytes - 1); - cache = PreEncodedCodeCache.map(fc, tempOffset, codeCount, codeSize); + cache = PreEncodedCodeCache.map(fc, writer.getOutputPath(), tempOffset, codeCount, codeSize); if (secondarySize > 0) { secondaryCache = PreEncodedCodeCache.map(fc, secondaryOffset, codeCount, secondarySize); } @@ -314,8 +314,16 @@ protected PreEncodedCodeCache precomputeCodeCache(CompactionContext ctx, Compact ByteSequence code2 = enc2 == null ? null : vectorTypeSupport.createByteSequence(secondarySize); VectorFloat vec = vectorTypeSupport.createFloatVector(ctx.dimension); long count = 0; + // the source is read at random (new-ordinal order), so each record is announced a fixed + // number of nodes ahead of its read instead of being faulted in when the read reaches it + final int ahead = ctx.recordAnnounceAhead; + int announced = cStart; // first new ordinal not yet announced try (var view = source.getView()) { for (int newOrd = cStart; newOrd < cEnd; newOrd++) { + for (int stop = Math.min(cEnd, newOrd + ahead + 1); ahead > 0 && announced < stop; announced++) { + int oldAhead = mapper.newToOld(announced); + if (oldAhead >= 0 && alive.get(oldAhead)) source.announceL0Record(oldAhead); + } int oldOrd = mapper.newToOld(newOrd); if (oldOrd < 0 || !alive.get(oldOrd)) continue; view.getVectorInto(oldOrd, vec, 0); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/WideDecoder.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/WideDecoder.java index 02309892b..581407adc 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/WideDecoder.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/WideDecoder.java @@ -130,6 +130,44 @@ void decode(int newOrdinal, VectorFloat dst) { return; } cache.get(newOrdinal, codeBytes); + decodeCodeBytes(dst); + } + + private java.nio.ByteBuffer bufSrc; // the buffer the views below were taken from + private java.nio.ByteBuffer bufBytes; + private java.nio.FloatBuffer bufFloats; + + /** + * Decodes the code that starts at byte {@code offset} of {@code src}, a buffer the caller filled + * from the cache's file ({@link PreEncodedCodeCache#readRowsDirect}). Raw rows are little-endian + * floats, read through one float view per buffer so the per-row cost is a bulk copy. + */ + void decode(java.nio.ByteBuffer src, int offset, VectorFloat dst) { + if (src != bufSrc) { + // views over the whole capacity: the buffer's position and limit are the caller's read state + bufSrc = src; + java.nio.ByteBuffer whole = src.duplicate(); + whole.clear(); + bufFloats = whole.order(java.nio.ByteOrder.LITTLE_ENDIAN).asFloatBuffer(); + bufBytes = whole; + } + if (raw) { + if (dst instanceof FloatArray && (offset & 3) == 0) { + bufFloats.position(offset >> 2); + bufFloats.get(((FloatArray) dst).array(), 0, dimension); + } else { + bufBytes.position(offset); + bufBytes.get(codeBytes, 0, codeBytes.length); + RawVectorCode.decodeInto(codeBytes, dst, dimension); + } + return; + } + bufBytes.position(offset); + bufBytes.get(codeBytes, 0, codeBytes.length); + decodeCodeBytes(dst); + } + + private void decodeCodeBytes(VectorFloat dst) { if (codebooks != null && dst instanceof FloatArray) { float[] out = ((FloatArray) dst).array(); if (pairs != null) { 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 0a5fde49e..110e1957a 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 @@ -18,6 +18,7 @@ import java.io.IOException; import java.lang.foreign.Arena; +import java.lang.invoke.MethodHandle; import java.lang.foreign.FunctionDescriptor; import java.lang.foreign.Linker; import java.lang.foreign.MemorySegment; @@ -46,6 +47,8 @@ public class MemorySegmentReader implements RandomAccessReader { private static final Logger logger = LoggerFactory.getLogger(MemorySegmentReader.class); private static final int MADV_RANDOM = 1; // Value for Linux + private static final int MADV_WILLNEED = 3; // POSIX value + private static final long PAGE = 4096; private static final OfInt intLayout = ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN); private static final OfFloat floatLayout = ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN); private static final OfLong longLayout = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN); @@ -149,6 +152,7 @@ public static class Supplier implements ReaderSupplier { private final Arena arena; private final MemorySegment memory; private final Path path; + private final MethodHandle madvise; // posix_madvise, or null when the symbol is missing public Supplier(Path path) throws IOException { this.path = path; @@ -160,13 +164,14 @@ public Supplier(Path path) throws IOException { var linker = Linker.nativeLinker(); var maybeMadvise = linker.defaultLookup().find("posix_madvise"); if (maybeMadvise.isPresent()) { - var madvise = linker.downcallHandle(maybeMadvise.get(), + madvise = linker.downcallHandle(maybeMadvise.get(), FunctionDescriptor.of(ValueLayout.JAVA_INT, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, ValueLayout.JAVA_INT)); int result = (int) madvise.invokeExact(memory, memory.byteSize(), MADV_RANDOM); if (result != 0) { throw new IOException("posix_madvise failed with error code: " + result); } } else { + madvise = null; logger.warn("posix_madvise not found, MADV_RANDOM advice not applied"); } } catch (Throwable e) { @@ -188,6 +193,29 @@ public Supplier(Path path) throws IOException { private static final ThreadLocal PREFETCH_BUF = ThreadLocal.withInitial(() -> java.nio.ByteBuffer.wrap(new byte[64 << 10])); + /** + * {@code POSIX_MADV_WILLNEED} on the pages of {@code [offset, offset + length)}: the kernel + * starts reading them in and returns; a later access through the mapping finds them resident + * or already in flight. The positional-read prefetch without the symbol. + */ + @Override + public void announce(long offset, long length) { + if (length <= 0 || madvise == null) { + prefetch(offset, length); + return; + } + long from = Math.max(0, offset) & -PAGE; + long to = Math.min(memory.byteSize(), offset + length); + if (to <= from) { + return; + } + try { + int ignored = (int) madvise.invokeExact(memory.asSlice(from, to - from), to - from, MADV_WILLNEED); + } catch (Throwable e) { + logger.warn("posix_madvise(WILLNEED) of {} [{}, {}) failed; continuing without the announcement", path, from, to, e); + } + } + @Override public void prefetch(long offset, long length) { if (length <= 0) { diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java index 40be0aecf..f648ad654 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java @@ -1205,7 +1205,22 @@ public void testCellJoinMergeRecallCosine() throws Exception { cellJoinMergeRecall(VectorSimilarityFunction.COSINE, true); } + /** + * The walk's direct reads, the record announce-ahead and fixed budgets: the options change how + * the store and the sources are read, never what is scored, so the merge must come out the same. + */ + @Test + public void testCellJoinMergeDirectReadsAndFixedBudgets() throws Exception { + cellJoinMergeRecall(VectorSimilarityFunction.DOT_PRODUCT, false, + c -> c.withWalkDirectReads(true).withRecordAnnounceAhead(8).withCellWalkBudgets(64L << 20, 64L << 20)); + } + private void cellJoinMergeRecall(VectorSimilarityFunction vsf, boolean fusedSources) throws Exception { + cellJoinMergeRecall(vsf, fusedSources, c -> c); + } + + private void cellJoinMergeRecall(VectorSimilarityFunction vsf, boolean fusedSources, + java.util.function.UnaryOperator configure) throws Exception { int perSource = 3000; int nSrc = 3; final java.util.Random rnd = new java.util.Random(20260916); // fixed: the two merges are compared @@ -1247,7 +1262,7 @@ private void cellJoinMergeRecall(VectorSimilarityFunction vsf, boolean fusedSour } String joinName = "celljoin_" + (fusedSources ? "" : "plain_") + vsf + "_join"; - double joinRecall = mergeAndRecall(graphs, live, perSource, vsf, all, queries, gt, topK, joinName); + double joinRecall = mergeAndRecall(graphs, live, perSource, vsf, all, queries, gt, topK, joinName, configure); System.out.printf("Cell-join merge recall (%s, %s sources): %.4f%n", vsf, fusedSources ? "fused" : "full-precision", joinRecall); try (ReaderSupplier rs = ReaderSupplierFactory.open(testDirectory.resolve(joinName))) { @@ -1290,9 +1305,16 @@ private double mergeAndRecall(List graphs, List l VectorSimilarityFunction vsf, List> all, List> queries, List> gt, int topK, String name) throws Exception { + return mergeAndRecall(graphs, live, perSource, vsf, all, queries, gt, topK, name, c -> c); + } + + private double mergeAndRecall(List graphs, List live, int perSource, + VectorSimilarityFunction vsf, List> all, + List> queries, List> gt, int topK, + String name, java.util.function.UnaryOperator configure) throws Exception { int nSrc = graphs.size(); int total = nSrc * perSource; - var compactor = new OnDiskGraphIndexCompactor(graphs, live, vsf, null); + var compactor = configure.apply(new OnDiskGraphIndexCompactor(graphs, live, vsf, null)); var outputPath = testDirectory.resolve(name); compactor.compact(outputPath); assertTrue("level-0 candidates should come from cell scans", compactor.cellScanCount() > 0); diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPreEncodedCodeCache.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPreEncodedCodeCache.java index 0d2c1e4f4..56fea5d64 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPreEncodedCodeCache.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPreEncodedCodeCache.java @@ -79,7 +79,6 @@ private Path newCacheFile(int count, int codeSize) throws IOException { return p; } - /** Round-trips every ordinal with chunks small enough that most ordinals land in a different one. */ @Test public void testRoundTripAcrossChunkBoundaries() throws IOException { @@ -104,6 +103,43 @@ public void testRoundTripAcrossChunkBoundaries() throws IOException { } } + /** + * Direct reads of row runs return the same bytes as the mapping, across chunk boundaries and from + * an unaligned section offset, with the rows landing at the returned buffer offset. + */ + @Test + public void testReadRowsDirectMatchesMapping() throws IOException { + int codeSize = 40, count = 300, offset = 4096 + 13; + int maxChunkBytes = codeSize * 11; // 11 codes per chunk + Path p = Files.createTempFile("code-cache-", ".bin"); + p.toFile().deleteOnExit(); + try (FileChannel fc = FileChannel.open(p, StandardOpenOption.WRITE)) { + fc.write(ByteBuffer.wrap(new byte[]{0}), offset + PreEncodedCodeCache.sectionBytes(count, codeSize) - 1); + } + try (FileChannel fc = FileChannel.open(p, StandardOpenOption.READ, StandardOpenOption.WRITE); + PreEncodedCodeCache cache = PreEncodedCodeCache.map(fc, p, offset, count, codeSize, maxChunkBytes)) { + for (int i = 0; i < count; i++) { + cache.put(i, codeFor(i, codeSize)); + } + if (!cache.supportsDirect()) { + return; // a file system without direct I/O (tmpfs): the mapped path is the only one + } + byte[] dst = new byte[codeSize]; + int[][] runs = {{0, 1}, {7, 50}, {290, count - 290}, {0, count}}; + for (int[] run : runs) { + int first = run[0], rows = run[1]; + ByteBuffer buf = PreEncodedCodeCache.allocateAligned(cache.directBufferBytes(rows)); + int at = cache.readRowsDirect(first, rows, buf); + assertTrue("rows start inside the aligned head", at >= 0 && at < PreEncodedCodeCache.DIRECT_ALIGN); + for (int r = 0; r < rows; r++) { + buf.position(at + r * codeSize); + buf.get(dst, 0, codeSize); + assertCodeMatches(first + r, codeSize, dst); + } + } + } + } + /** The single-chunk path must stay byte-identical to the chunked one. */ @Test public void testSingleChunkMatchesChunked() throws IOException { From 10c93a11f0eaa266f320a41030204e753a99d6f0 Mon Sep 17 00:00:00 2001 From: dian-lun-lin Date: Wed, 11 Nov 2026 22:44:02 +0000 Subject: [PATCH 2/6] Compaction: sharded compaction sessions sharded() opens a sharded compaction of the compactor's sources: each compactShard(shardLiveNodes, path) call merges one output, with that shard's nodes as its live set, and returns the shard's ordinal mappers, so a writer can merge each shard as it learns its membership. What does not depend on the shard is done once: the inline and sidecar codebooks and the wide diversity code are trained on the whole compaction and adopted by every shard (QuantizationCompactionStrategy.adoptCompressor). A shard keeps 1/M of every source, so a node kept only 1/M of its same-source edges and nothing replaced them: measured on 2 x 1M vectors into 64 shards, a node kept 0.5 of its 32 same-source edges and the output was nearly bipartite. A source whose live nodes are a minority (below half) is therefore joined with itself like another source: one more target of the level-0 cell walk, one more graph search at the upper levels. Prefetching a whole source window is skipped when less than an eighth of it is live. A compaction over a slice of the sources' nodes can also leave the sources' top levels without a live node, and the merge failed with "No live nodes found at maxLevel". The output's hierarchy now ends at the highest merged level that still has a live node, with the entry node taken from it. ShardExp (benchmarks-jmh) is the experiment behind the numbers: it builds N source graphs over slices of a dataset, produces K output shards by per-shard merge, by one session, or by a rebuild of every shard from its vectors, and reports recall@10 of the merged top-10 against the ground truth with wall and CPU time per shard. CompactorBenchmark.md describes how to run it. On tessera 2M (two 1M sources, degree 32) into 64 shards: 627 s and recall@10 0.945 in one session, against 1182 s and 0.875 for 64 per-shard merges and 886 s and 0.947 for 64 rebuilds. Claude-Session: https://claude.ai/code/session_016pmoQbpUFKvQKpvRc7Xvjz --- .../jvector/bench/CompactorBenchmark.md | 18 + .../jbellis/jvector/bench/ShardExp.java | 336 +++++++++++++++++ docs/compaction.md | 35 +- .../jbellis/jvector/graph/disk/ExactJoin.java | 8 +- .../graph/disk/FusedCompactionStrategy.java | 6 + .../graph/disk/OnDiskGraphIndexCompactor.java | 338 +++++++++++++++--- .../disk/QuantizationCompactionStrategy.java | 15 + .../graph/disk/SidecarCompactionStrategy.java | 5 + .../disk/TestOnDiskGraphIndexCompactor.java | 209 +++++++++++ 9 files changed, 923 insertions(+), 47 deletions(-) create mode 100644 benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/ShardExp.java diff --git a/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.md b/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.md index 0e89b0f62..cdd2fb995 100644 --- a/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.md +++ b/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.md @@ -251,3 +251,21 @@ target/benchmark-results/compactor-/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 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`. diff --git a/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/ShardExp.java b/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/ShardExp.java new file mode 100644 index 000000000..a80b9892c --- /dev/null +++ b/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/ShardExp.java @@ -0,0 +1,336 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.jbellis.jvector.bench; + +import io.github.jbellis.jvector.disk.ReaderSupplier; +import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.graph.*; +import io.github.jbellis.jvector.graph.disk.*; +import io.github.jbellis.jvector.graph.disk.feature.*; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; +import io.github.jbellis.jvector.quantization.PQVectors; +import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.util.FixedBitSet; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; + +import java.io.*; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.MappedByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.*; +import java.util.*; +import java.util.function.IntFunction; +import java.util.stream.IntStream; + +/** + * Sharded compaction on a real dataset: the jvector side of Cassandra's {@code ShardedVectorCompactionBench}. + *

+ * A Cassandra compaction with UCS sharding writes its output as K sstables, each covering a token range, + * and every output sstable's vector index is produced by a merge over all source graphs that keeps + * only the nodes of its range. This program reproduces that shape against raw {@code .fvecs} files: + * 

+ *   build-sources               build S source graphs over consecutive slices of the base (the "sstables")
+ *   shard K merge|sharded|rebuild
+ *                               produce K output shards, rows assigned to shards by a hash of the row (the token)
+ *                               merge:   one OnDiskGraphIndexCompactor per shard over ALL sources, liveNodes = the shard
+ *                               sharded: one OnDiskGraphIndexCompactor, one sharded() session, compactShard per shard
+ *                               rebuild: one fresh graph per shard from that shard's vectors (graph merge disabled)
+ *   sharded K                   same as "shard K sharded"
+ *   recall K merge|rebuild|sharded NQ
+ *                               search every shard for NQ queries, merge the top-10, recall@10 against the ground truth
+ *
+ * Every command prints {@code RESULT} lines: per-shard and total wall and CPU time, and recall. Outputs land + * under {@code exp.dir}: the sources in {@code sources/}, the shards of a run in {@code out/-K/} + * with a {@code .rows} file per shard that maps its ordinals back to base rows for the recall pass. + *

+ * System properties: {@code exp.base}, {@code exp.query}, {@code exp.gt} (the dataset's base and query + * {@code .fvecs} and ground-truth {@code .ivecs}; required), {@code exp.dir} (output directory, default + * {@code shardexp}), {@code exp.sources} (2), {@code exp.rows} (all), {@code exp.degree} (32), {@code exp.beam} (100), + * {@code exp.tag} (suffix on the output directory and the RESULT mode, to keep variants apart). + *

+ * Run from the shaded benchmark jar, for example on tessera-2M with two 1M-vector sources: + * 

+ *   J="java -Xmx24g --add-modules jdk.incubator.vector --enable-native-access=ALL-UNNAMED \
+ *      -Dexp.base=base.fvecs -Dexp.query=query.fvecs -Dexp.gt=gt.ivecs \
+ *      -cp benchmarks-jmh/target/compactor-benchmark.jar io.github.jbellis.jvector.bench.ShardExp"
+ *   $J build-sources
+ *   for K in 1 4 16 64; do $J shard $K merge; $J shard $K rebuild; $J sharded $K; done
+ *   for K in 1 4 16 64; do for m in merge rebuild sharded; do $J recall $K $m 10000; done; done
+ *
+ * See the "Compacting into several outputs (sharding)" section of {@code docs/compaction.md} for the + * numbers this produced. + */ +public class ShardExp { + static final VectorTypeSupport VTS = VectorizationProvider.getInstance().getVectorTypeSupport(); + static final VectorSimilarityFunction VSF = VectorSimilarityFunction.COSINE; + static final Path DIR = Path.of(System.getProperty("exp.dir", "shardexp")); + static final int SOURCES = Integer.getInteger("exp.sources", 2); + static final int DEGREE = Integer.getInteger("exp.degree", 32); + static final int BEAM = Integer.getInteger("exp.beam", 100); + + /** A path property a command needs: {@code build-sources} and {@code shard} read the base, {@code recall} the queries and the ground truth. */ + static String required(String property) { + String v = System.getProperty(property); + if (v == null) throw new IllegalArgumentException("-D" + property + "= is required"); + return v; + } + + /** .fvecs (4-byte dim header per row), mapped read-only in chunks of whole rows (a mapping holds at most 2 GB). */ + static final class Fvecs { + final MappedByteBuffer[] chunks; final int dim; final int rows; final long rec; final int rowsPerChunk; + Fvecs(String path) throws IOException { + try (var fc = FileChannel.open(Path.of(path), StandardOpenOption.READ)) { + var head = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN); + if (fc.read(head, 0) != 4) throw new IOException("empty fvecs file " + path); + dim = head.getInt(0); + rec = 4L + 4L * dim; + rows = (int) (fc.size() / rec); + rowsPerChunk = (int) (Integer.MAX_VALUE / rec); + chunks = new MappedByteBuffer[(rows + rowsPerChunk - 1) / rowsPerChunk]; + for (int c = 0; c < chunks.length; c++) { + long from = (long) c * rowsPerChunk * rec; + long to = Math.min(fc.size(), from + (long) rowsPerChunk * rec); + chunks[c] = fc.map(FileChannel.MapMode.READ_ONLY, from, to - from); + chunks[c].order(ByteOrder.LITTLE_ENDIAN); + } + } + } + float[] row(int r) { + float[] a = new float[dim]; + var b = chunks[r / rowsPerChunk]; + int at = (int) ((r % rowsPerChunk) * rec) + 4; + for (int i = 0; i < dim; i++) a[i] = b.getFloat(at + 4 * i); // absolute reads: thread-safe + return a; + } + } + + /** The rows {@code rows[i]} of a mapped file as a vector source with ordinals 0..n-1. */ + static final class RowsRavv implements RandomAccessVectorValues { + final Fvecs f; final int[] rows; + RowsRavv(Fvecs f, int[] rows) { this.f = f; this.rows = rows; } + public int size() { return rows.length; } + public int dimension() { return f.dim; } + public VectorFloat getVector(int o) { return VTS.createFloatVector(f.row(rows[o])); } + public boolean isValueShared() { return false; } + public RandomAccessVectorValues copy() { return this; } + } + + static int rowsTotal(Fvecs f) { return Integer.getInteger("exp.rows", f.rows); } + static int perSource(Fvecs f) { return rowsTotal(f) / SOURCES; } + + /** UCS token-range shard of a row: a 64-bit hash of the key, split into K equal ranges. */ + static int shardOf(int row, int k) { + long z = row * 0x9E3779B97F4A7C15L; + z = (z ^ (z >>> 30)) * 0xBF58476D1CE4E5B9L; + z = (z ^ (z >>> 27)) * 0x94D049BB133111EBL; + z ^= z >>> 31; + return (int) Math.multiplyHigh(z >>> 1, (long) k << 1); // in [0, k) + } + + static long cpuNs() { + return ((com.sun.management.OperatingSystemMXBean) java.lang.management.ManagementFactory.getOperatingSystemMXBean()).getProcessCpuTime(); + } + + /** PQ-scored build + FusedPQ write, as CompactorBenchmark builds partitions (and as Cassandra's CompactionGraph rebuilds). */ + static void buildGraph(RandomAccessVectorValues ravv, Path out) throws IOException { + int dim = ravv.dimension(); + var pq = ProductQuantization.compute(ravv, dim / 8, 256, false); + var pqv = (PQVectors) pq.encodeAll(ravv); + var bsp = BuildScoreProvider.pqBuildScoreProvider(VSF, pqv); + try (var builder = new GraphIndexBuilder(bsp, dim, DEGREE, BEAM, 1.2f, 1.2f, true)) { + var graph = builder.build(ravv); + Files.deleteIfExists(out); + try (var writer = new OnDiskGraphIndexWriter.Builder(graph, out) + .with(new InlineVectors(dim)).with(new FusedPQ(graph.maxDegree(), pq)).build()) { + var suppliers = new EnumMap>(FeatureId.class); + suppliers.put(FeatureId.INLINE_VECTORS, o -> new InlineVectors.State(ravv.getVector(o))); + var view = graph.getView(); + suppliers.put(FeatureId.FUSED_PQ, o -> new FusedPQ.State(view, pqv, o)); + writer.write(suppliers); + } + } + } + + static Path sourcePath(int s) { return DIR.resolve("sources").resolve("source-" + s); } + static final String TAG = System.getProperty("exp.tag", ""); + static Path shardDir(int k, String mode) { return DIR.resolve("out").resolve(mode + TAG + "-K" + k); } + + + static void buildSources() throws IOException { + var f = new Fvecs(required("exp.base")); + int per = perSource(f); + Files.createDirectories(sourcePath(0).getParent()); + for (int s = 0; s < SOURCES; s++) { + int[] rows = IntStream.range(s * per, (s + 1) * per).toArray(); + long t0 = System.nanoTime(), c0 = cpuNs(); + buildGraph(new RowsRavv(f, rows), sourcePath(s)); + System.out.printf("RESULT build-source s=%d n=%d wall_ms=%d cpu_ms=%d bytes=%d%n", s, per, + (System.nanoTime() - t0) / 1_000_000, (cpuNs() - c0) / 1_000_000, Files.size(sourcePath(s))); + } + } + + static void writeRows(Path p, int[] rowOf) throws IOException { + try (var out = new DataOutputStream(new BufferedOutputStream(Files.newOutputStream(p)))) { + out.writeInt(rowOf.length); + for (int r : rowOf) out.writeInt(r); + } + } + + static int[] readRows(Path p) throws IOException { + try (var in = new DataInputStream(new BufferedInputStream(Files.newInputStream(p)))) { + int[] a = new int[in.readInt()]; + for (int i = 0; i < a.length; i++) a[i] = in.readInt(); + return a; + } + } + + /** + * K output shards, a row's shard by {@link #shardOf}. {@code merge}: one compactor per shard over + * all sources with the shard's nodes live; {@code sharded}: one compactor over every node, one + * session, {@code compactShard} per shard; {@code rebuild}: a fresh graph from the shard's vectors. + */ + static void shard(int K, String mode) throws IOException { + var f = new Fvecs(required("exp.base")); + int per = perSource(f); + Path dir = shardDir(K, mode); + Files.createDirectories(dir); + boolean merge = mode.equals("merge"), sharded = mode.equals("sharded"); + if (!merge && !sharded && !mode.equals("rebuild")) throw new IllegalArgumentException("mode must be merge, sharded or rebuild"); + List rss = new ArrayList<>(); + List graphs = new ArrayList<>(); + if (merge || sharded) { + for (int s = 0; s < SOURCES; s++) { + rss.add(ReaderSupplierFactory.open(sourcePath(s))); + graphs.add(OnDiskGraphIndex.load(rss.get(s))); + } + } + OnDiskGraphIndexCompactor.ShardedCompaction session = null; + if (sharded) { + List all = new ArrayList<>(); + for (int s = 0; s < SOURCES; s++) { + var bs = new FixedBitSet(graphs.get(s).getIdUpperBound()); + bs.set(0, per); + all.add(bs); + } + session = new OnDiskGraphIndexCompactor(graphs, all, VSF, null).sharded(); + } + long wallTotal = 0, cpuTotal = 0; + for (int kk = 0; kk < K; kk++) { + final int k = kk; + Path out = dir.resolve("shard-" + k); + Files.deleteIfExists(out); + long t0 = System.nanoTime(), c0 = cpuNs(); + int[] rowOf; + if (merge || sharded) { + List live = new ArrayList<>(); + int n = 0; + for (int s = 0; s < SOURCES; s++) { + var bs = new FixedBitSet(graphs.get(s).getIdUpperBound()); + for (int i = 0; i < per; i++) if (shardOf(s * per + i, K) == k) { bs.set(i); n++; } + live.add(bs); + } + List mappers; + if (sharded) { + mappers = session.compactShard(live, out); + } else { + var compactor = new OnDiskGraphIndexCompactor(graphs, live, VSF, null); + compactor.compact(out); + mappers = compactor.ordinalMappers(); + } + rowOf = new int[n]; + for (int s = 0; s < SOURCES; s++) + for (int i = 0; i < per; i++) + if (live.get(s).get(i)) rowOf[mappers.get(s).oldToNew(i)] = s * per + i; + } else { + rowOf = IntStream.range(0, SOURCES * per).filter(r -> shardOf(r, K) == k).toArray(); + buildGraph(new RowsRavv(f, rowOf), out); + } + long wall = (System.nanoTime() - t0) / 1_000_000, cpu = (cpuNs() - c0) / 1_000_000; + wallTotal += wall; cpuTotal += cpu; + writeRows(dir.resolve("shard-" + k + ".rows"), rowOf); + System.out.printf("RESULT shard mode=%s K=%d k=%d nodes=%d wall_ms=%d cpu_ms=%d%n", mode + TAG, K, k, rowOf.length, wall, cpu); + } + for (var g : graphs) g.close(); + for (var r : rss) r.close(); + System.out.printf("RESULT total mode=%s K=%d wall_ms=%d cpu_ms=%d%n", mode + TAG, K, wallTotal, cpuTotal); + } + + static void recall(int K, String mode, int nq) throws Exception { + var q = new Fvecs(required("exp.query")); + int[][] gt; + try (var in = new DataInputStream(new BufferedInputStream(new FileInputStream(required("exp.gt"))))) { + gt = new int[nq][]; + for (int i = 0; i < nq; i++) { + int k = Integer.reverseBytes(in.readInt()); + gt[i] = new int[k]; + for (int j = 0; j < k; j++) gt[i][j] = Integer.reverseBytes(in.readInt()); + } + } + Path dir = shardDir(K, mode); + List graphs = new ArrayList<>(); + List rss = new ArrayList<>(); + int[][] rowOf = new int[K][]; + for (int k = 0; k < K; k++) { + rss.add(ReaderSupplierFactory.open(dir.resolve("shard-" + k))); + graphs.add(OnDiskGraphIndex.load(rss.get(k))); + rowOf[k] = readRows(dir.resolve("shard-" + k + ".rows")); + } + ThreadLocal searchers = ThreadLocal.withInitial(() -> { + var a = new GraphSearcher[K]; + for (int k = 0; k < K; k++) { a[k] = new GraphSearcher(graphs.get(k)); a[k].usePruning(false); } + return a; + }); + int topK = 10; + long hits = IntStream.range(0, nq).parallel().mapToLong(i -> { + var qv = VTS.createFloatVector(q.row(i)); + var ss = searchers.get(); + var all = new ArrayList(); + int[] owner = new int[K * topK]; int at = 0; + List cand = new ArrayList<>(); + for (int k = 0; k < K; k++) { + var view = (ImmutableGraphIndex.ScoringView) ss[k].getView(); + var ssp = new DefaultSearchScoreProvider(view.approximateScoreFunctionFor(qv, VSF), view.rerankerFor(qv, VSF)); + var r = ss[k].search(ssp, topK, topK, 0f, 0f, Bits.ALL); + for (var ns : r.getNodes()) cand.add(new long[]{Float.floatToIntBits(ns.score), rowOf[k][ns.node]}); + } + cand.sort((a, b) -> Float.compare(Float.intBitsToFloat((int) b[0]), Float.intBitsToFloat((int) a[0]))); + var truth = new HashSet(); + for (int j = 0; j < topK; j++) truth.add(gt[i][j]); + long h = 0; + for (int j = 0; j < Math.min(topK, cand.size()); j++) if (truth.contains((int) cand.get(j)[1])) h++; + return h; + }).sum(); + System.out.printf("RESULT recall mode=%s K=%d nq=%d recall@10=%.4f%n", mode + TAG, K, nq, hits / (double) (nq * topK)); + for (var g : graphs) g.close(); + for (var r : rss) r.close(); + } + + public static void main(String[] args) throws Exception { + switch (args[0]) { + case "build-sources": buildSources(); break; + case "shard": shard(Integer.parseInt(args[1]), args[2]); break; + case "sharded": shard(Integer.parseInt(args[1]), "sharded"); break; + case "recall": recall(Integer.parseInt(args[1]), args[2], Integer.parseInt(args[3])); break; + default: throw new IllegalArgumentException("unknown command " + args[0] + "; see the class comment"); + } + } +} diff --git a/docs/compaction.md b/docs/compaction.md index 776a52c2f..1e4a097d4 100644 --- a/docs/compaction.md +++ b/docs/compaction.md @@ -53,6 +53,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 shard0 = session.compactShard(shard0LiveNodes, output0); // shard0.get(s).oldToNew(n): ordinal of source s's node n in output0, or OMITTED +List 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 +110,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) and any heavily deleted source: a node then keeps `degree × live fraction` of its same-source edges, and the output is close to bipartite between the sources. So 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. + **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 +147,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 diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ExactJoin.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ExactJoin.java index 3d5625615..87dd32c7a 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ExactJoin.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/ExactJoin.java @@ -41,7 +41,11 @@ final class ExactJoin { 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; @@ -52,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)]; 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 ed49becce..1b0907178 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 @@ -83,6 +83,12 @@ public void retrain(VectorSimilarityFunction vsf) { this.retrainedCompressor = (VectorCompressor>) (VectorCompressor) retrainer.retrain(vsf); } + @Override + @SuppressWarnings("unchecked") + public void adoptCompressor(VectorCompressor trained) { + this.retrainedCompressor = (VectorCompressor>) java.util.Objects.requireNonNull(trained, "trained"); + } + @Override public VectorCompressor compressor() { return retrainedCompressor; 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 fb9dfe5eb..cf4643321 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 @@ -46,6 +46,7 @@ import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.quantization.CompressedVectors; import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.quantization.VectorCompressor; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; import io.github.jbellis.jvector.vector.types.ByteSequence; @@ -130,7 +131,7 @@ public final class OnDiskGraphIndexCompactor implements Accountable { * of every target. Callers read the mapping back via {@link #ordinalMappers()}. */ private List remappers; - private final List maxDegrees; + private List maxDegrees; private final int dimension; private int maxOrdinal = -1; @@ -244,6 +245,25 @@ private void setupCrossLink() { // Whether the cell walk reads the store with direct I/O: null = decided by size (see walkReadsDirect). private Boolean walkDirectSetting; private Boolean walkDirect; // the decision, made once per compaction + + // ---- Sparse sources: the same-source join ---- + // A node keeps only its LIVE same-source edges. When most of a source is live that is nearly all + // of them; when a source's live nodes are a minority (a shard of a sharded compaction keeps 1/M + // of every source, a heavily deleted source keeps little) most of a node's edges point at dead + // nodes and nothing replaced them: measured on 2 x 1M vectors split into 64 shards, a node kept + // 0.5 of its 32 same-source edges and the output was nearly bipartite between the sources. Below + // this live fraction a source is therefore searched for its own nodes like any other source: at + // level 0 it is one more target of the cell walk, elsewhere one more graph search. + private static final float DEFAULT_SAME_SOURCE_JOIN_BELOW = 0.5f; + private float sameSourceJoinBelow = DEFAULT_SAME_SOURCE_JOIN_BELOW; + + // Sharded compaction: what is trained once for every shard. Null for an ordinary compaction. + private SharedTraining shared; + + /** Whether so few of source {@code s}'s nodes are live that its own nodes are joined like another source's. */ + private boolean sparseSource(int s) { + return numLiveNodesPerSource.get(s) < (double) sameSourceJoinBelow * sources.get(s).size(0); + } /** * Constructs a new OnDiskGraphIndexCompactor for graphs without a non-fused compressed sidecar. * Equivalent to calling the 6-arg constructor with {@code sourceCompressed = null}. @@ -391,6 +411,18 @@ private static long memoryBeyondHeap() { return ((com.sun.management.OperatingSystemMXBean) bean).getTotalPhysicalMemorySize() - Runtime.getRuntime().maxMemory(); } + /** + * Sets the live fraction below which a source's own nodes are joined like another source's: a + * node of such a source gathers same-source candidates beyond the edges it retained, because + * most of those edges point at dead nodes. The default is one half; 0 never joins a source with + * itself, 1 always does (as long as any of its nodes is dead). + */ + OnDiskGraphIndexCompactor withSameSourceJoinBelow(float liveFraction) { + if (!(liveFraction >= 0f && liveFraction <= 1f)) throw new IllegalArgumentException("liveFraction must be in [0, 1]"); + this.sameSourceJoinBelow = liveFraction; + return this; + } + /** * Validates that all source indexes have compatible configurations and required features * before attempting compaction. Ensures consistent dimensions, max degrees, hierarchical @@ -533,6 +565,149 @@ public List ordinalMappers() { return remappers; } + // ---- Sharded compaction: N sources into M outputs ---- + + /** + * What a sharded compaction does once for all of its shards: the codebooks, trained on the live + * nodes of every shard together. It belongs to the compactor {@link #sharded()} was called on, + * whose live nodes are those of all the shards; each shard's merge adopts the results instead + * of training on its own slice. Training is deferred to the first shard that needs it. + */ + private final class SharedTraining { + final Once> inline = new Once<>("inline compressor", () -> trainedCompressor(detectInlineStrategy())); + final Once> sidecar = new Once<>("sidecar compressor", () -> trainedCompressor(detectSidecarStrategy())); + final Once wide = new Once<>("wide code", OnDiskGraphIndexCompactor.this::trainWide); + + private VectorCompressor trainedCompressor(QuantizationCompactionStrategy strategy) { + strategy.retrain(similarityFunction); + return strategy.compressor(); + } + } + + /** A value computed by the first {@link #get} and shared by the rest. */ + private static final class Once { + private final String what; + private final java.util.function.Supplier compute; + private boolean done; + private T value; + + Once(String what, java.util.function.Supplier compute) { + this.what = what; + this.compute = compute; + } + + synchronized T get() { + if (!done) { + long t0 = System.nanoTime(); + value = compute.get(); + done = true; + if (value != null) { + log.info("Sharded compaction: {} trained once for every shard in {} ms", what, (System.nanoTime() - t0) / 1_000_000); + } + } + return value; + } + } + + /** Trains the strategy's compressor, or hands it the one a sharded compaction already trained. */ + private void train(QuantizationCompactionStrategy strategy, boolean sidecar) { + if (shared == null) { + strategy.retrain(similarityFunction); + } else { + strategy.adoptCompressor((sidecar ? shared.sidecar : shared.inline).get()); + } + } + + /** The wide diversity code (see {@link #wideSubspaces}), trained on the live nodes. */ + private ProductQuantization trainWide() { + return new PQRetrainer(sources, liveNodes, dimension).train(wideSubspaces()); + } + + /** + * A compaction of this compactor's sources into several outputs, one {@link #compactShard} call + * per output. What does not depend on the shard is done once and shared: the codebooks are + * trained on the live nodes this compactor was built with (those of every shard) and each + * shard's merge adopts them. A shard's merge is otherwise an ordinary compaction whose live + * nodes are the shard's, and it finds a node's same-source neighbours again, since most of + * the edges a node had in its source lead to nodes of other shards + * (see {@link #withSameSourceJoinBelow}). + *

+ * Shards may become known one at a time, as when each output is written as its rows arrive: + * nothing is done until the first {@link #compactShard}. Shards are compacted one at a time; + * a session is not thread-safe. + */ + @Experimental + public final class ShardedCompaction { + private final SharedTraining training = new SharedTraining(); + private int shardsCompacted; + + private ShardedCompaction() { + } + + /** + * Merges one shard into {@code graphPath}. + * + * @param shardLiveNodes per source, the nodes of this shard: a subset of the live nodes the + * compactor was built with, in bitsets of the same lengths + * @return the output ordinal of every source node in this shard, one mapper per source: + * {@code oldToNew} of a node of the shard is its ordinal in {@code graphPath}, of + * any other node {@link OrdinalMapper#OMITTED}. Every shard has its own mapping. + */ + public List compactShard(List shardLiveNodes, Path graphPath) throws FileNotFoundException { + return compactShard(shardLiveNodes, graphPath, null); + } + + /** + * Merges one shard of sources that ship a compressed sidecar: the graph into + * {@code graphPath}, the merged compressed vectors into {@code compressedPath}. + */ + public List compactShard(List shardLiveNodes, Path graphPath, Path compressedPath) throws FileNotFoundException { + OnDiskGraphIndexCompactor shard = shardCompactor(shardLiveNodes); + if (compressedPath == null) { + shard.compact(graphPath); + } else { + shard.compact(graphPath, compressedPath); + } + return shard.ordinalMappers(); + } + + private OnDiskGraphIndexCompactor shardCompactor(List shardLiveNodes) { + Objects.requireNonNull(shardLiveNodes, "shardLiveNodes"); + if (shardLiveNodes.size() != sources.size()) { + throw new IllegalArgumentException("shardLiveNodes must have one bitset per source"); + } + for (int s = 0; s < sources.size(); s++) { + FixedBitSet all = liveNodes.get(s), mine = shardLiveNodes.get(s); + if (mine.length() != all.length()) { + throw new IllegalArgumentException("source " + s + ": shard bitset length " + mine.length() + " != " + all.length()); + } + for (int n = mine.nextSetBit(0); n != DocIdSetIterator.NO_MORE_DOCS; n = n + 1 < mine.length() ? mine.nextSetBit(n + 1) : DocIdSetIterator.NO_MORE_DOCS) { + if (!all.get(n)) { + throw new IllegalArgumentException("source " + s + " node " + n + " is in the shard but not live in the compaction"); + } + } + } + OnDiskGraphIndexCompactor shard = new OnDiskGraphIndexCompactor(sources, sourceCompressed, shardLiveNodes, similarityFunction, executor); + shard.shared = training; + shard.walkPairBytes = walkPairBytes; + shard.walkQueryBytes = walkQueryBytes; + shard.recordAnnounceAhead = recordAnnounceAhead; + shard.walkDirectSetting = walkDirectSetting; + shard.sameSourceJoinBelow = sameSourceJoinBelow; + log.info("Sharded compaction: shard {} with {} live nodes", shardsCompacted++, shard.numTotalNodes); + return shard; + } + } + + /** + * Opens a sharded compaction of this compactor's sources; see {@link ShardedCompaction}. The + * live nodes this compactor was built with must be those of every shard together. + */ + @Experimental + public ShardedCompaction sharded() { + return new ShardedCompaction(); + } + /** * Main compaction entry point. Merges all source indexes into a single output index at the * specified path, handling PQ retraining if needed, and writing header, all layers, and footer. @@ -584,7 +759,7 @@ public void compact(Path graphPath, Path compressedPath) throws FileNotFoundExce QuantizationCompactionStrategy sidecarStrategy = detectSidecarStrategy(); QuantizationCompactionStrategy scratch = SidecarCompactionStrategy.scratchVectors(buildContext()); try { - sidecarStrategy.retrain(similarityFunction); + train(sidecarStrategy, true); scratch.retrain(similarityFunction); activeSidecarStrategy = sidecarStrategy; vectorStore = scratch; @@ -693,7 +868,7 @@ private CompactionContext buildContext() { */ private void compactGraphImpl(Path outputPath, QuantizationCompactionStrategy strategy) throws FileNotFoundException { this.spillParent = outputPath.toAbsolutePath().getParent(); - strategy.retrain(similarityFunction); + train(strategy, false); boolean fusedPQEnabled = strategy.writesCodesInline(); ProductQuantization pq = strategy.compressorAsPQ(); @@ -719,7 +894,7 @@ private void compactGraphImpl(Path outputPath, QuantizationCompactionStrategy st if (cellMap != null) { log.info("Level-0 candidates are scanned and scored on the vectors themselves ({} B per node in scratch)", 4 * dimension); - pqWide = new PQRetrainer(sources, liveNodes, dimension).train(wideSubspaces()); + pqWide = shared != null ? shared.wide.get() : trainWide(); if (pqWide != null) { vectorStore.setSecondaryCompressor(pqWide); log.info("Diversity comparisons use a {}-subspace wide code ({} B per node); every candidate score stays exact", @@ -781,14 +956,14 @@ private void compactGraphImpl(Path outputPath, QuantizationCompactionStrategy st * across all sources. */ private int[] resolveEntryNodeSource() { - int maxLevel = sources.stream().mapToInt(OnDiskGraphIndex::getMaxLevel).max().orElse(0); + int maxLevel = maxDegrees.size() - 1; // The on-disk format sets entryNode.level = layerInfo.size() - 1 (i.e. maxLevel). // So the chosen node must actually have neighbors written at maxLevel — meaning it // must exist at maxLevel in its source. Prefer the designated entry node of a // maxLevel source; fall back to any live node that is at maxLevel. for (int s = 0; s < sources.size(); s++) { - if (sources.get(s).getMaxLevel() == maxLevel) { + if (sources.get(s).getMaxLevel() >= maxLevel) { int originalEntry = sources.get(s).getView().entryNode().node; if (liveNodes.get(s).get(originalEntry)) { return new int[]{s, originalEntry}; @@ -823,17 +998,8 @@ private void compactLevels(CompactWriter writer, ProductQuantization pq) throws IOException, ExecutionException, InterruptedException { - int maxUpperDegree = 0; - for (int level = 1; level < maxDegrees.size(); level++) { - maxUpperDegree = Math.max(maxUpperDegree, maxDegrees.get(level)); - } - - int baseSearchTopK = Math.max(MIN_SEARCH_TOP_K, ((maxDegrees.get(0) + sources.size() - 1) / sources.size()) * SEARCH_TOP_K_MULTIPLIER); - int baseMaxCandidateSize = baseSearchTopK * (sources.size() - 1) + maxDegrees.get(0) + REVERSE_CANDIDATE_SLOTS; - int upperMaxPerSourceTopK = maxUpperDegree == 0 ? 0 : Math.max(MIN_SEARCH_TOP_K, ((maxUpperDegree + sources.size() - 1) / sources.size()) * SEARCH_TOP_K_MULTIPLIER); - int upperMaxCandidateSize = upperMaxPerSourceTopK * sources.size(); - int maxCandidateSize = Math.max(baseMaxCandidateSize, upperMaxCandidateSize); - int scratchDegree = Math.max(maxDegrees.get(0), Math.max(1, maxUpperDegree)); + int maxCandidateSize = maxCandidateSize(); + int scratchDegree = scratchDegree(); // every worker's scratch is recorded so all of them, not only the calling thread's, are closed at the end final List openedScratch = Collections.synchronizedList(new ArrayList<>()); final ThreadLocal threadLocalScratch = ThreadLocal.withInitial(() -> { @@ -855,7 +1021,7 @@ private void compactLevels(CompactWriter writer, } for (int level = 0; level < maxDegrees.size(); level++) { - int searchTopK = Math.max(MIN_SEARCH_TOP_K, ((maxDegrees.get(level) + sources.size() - 1) / sources.size()) * SEARCH_TOP_K_MULTIPLIER); + int searchTopK = searchTopK(maxDegrees.get(level)); if (level == 0) log.info("Cross-link search budget: searchTopK={} per target", searchTopK); CompactionParams params = new CompactionParams(fusedPQEnabled, compressedPrecision, searchTopK, pq); @@ -914,8 +1080,9 @@ private void compactLevels(CompactWriter writer, // completed before that source's own nodes read them. for (int s : l0ProcessOrder) { List batches = buildBatchesForSource(s, 0); + // the largest source scans no other source, but a sparse one still scans itself boolean walk = exactMode && cellMap != null && candidateStore != null - && sizeRank[s] < sources.size() - 1 && !batches.isEmpty(); + && (sizeRank[s] < sources.size() - 1 || sparseSource(s)) && !batches.isEmpty(); if (!walk) { runBatchesWithBackpressure(batches, ecs, submitOne, writeResults); continue; @@ -1141,14 +1308,35 @@ private static final class RawUpper { } private ThreadLocal newScratchLocal(ProductQuantization pq) { + int maxCandidateSize = maxCandidateSize(), scratchDegree = scratchDegree(); + return ThreadLocal.withInitial(() -> new Scratch(maxCandidateSize, scratchDegree, dimension, sources, pq)); + } + + private int maxUpperDegree() { int maxUpperDegree = 0; for (int level = 1; level < maxDegrees.size(); level++) maxUpperDegree = Math.max(maxUpperDegree, maxDegrees.get(level)); - int baseSearchTopK = Math.max(MIN_SEARCH_TOP_K, ((maxDegrees.get(0) + sources.size() - 1) / sources.size()) * SEARCH_TOP_K_MULTIPLIER); - int baseMaxCandidateSize = baseSearchTopK * (sources.size() - 1) + maxDegrees.get(0) + REVERSE_CANDIDATE_SLOTS; - int upperMaxPerSourceTopK = maxUpperDegree == 0 ? 0 : Math.max(MIN_SEARCH_TOP_K, ((maxUpperDegree + sources.size() - 1) / sources.size()) * SEARCH_TOP_K_MULTIPLIER); - int maxCandidateSize = Math.max(baseMaxCandidateSize, upperMaxPerSourceTopK * sources.size()); - int scratchDegree = Math.max(maxDegrees.get(0), Math.max(1, maxUpperDegree)); - return ThreadLocal.withInitial(() -> new Scratch(maxCandidateSize, scratchDegree, dimension, sources, pq)); + return maxUpperDegree; + } + + /** The search's top-K per source for a level of the given degree. */ + private int searchTopK(int degree) { + return Math.max(MIN_SEARCH_TOP_K, ((degree + sources.size() - 1) / sources.size()) * SEARCH_TOP_K_MULTIPLIER); + } + + /** + * The most candidates a node gathers at any level: every source can contribute a search's + * worth, the node's own included (the same-source join of a sparse source), on top of the edges + * it retained and, at level 0, the reverse candidates. + */ + private int maxCandidateSize() { + int maxUpperDegree = maxUpperDegree(); + int base = searchTopK(maxDegrees.get(0)) * sources.size() + maxDegrees.get(0) + REVERSE_CANDIDATE_SLOTS; + int upper = maxUpperDegree == 0 ? 0 : searchTopK(maxUpperDegree) * sources.size() + maxUpperDegree; + return Math.max(base, upper); + } + + private int scratchDegree() { + return Math.max(maxDegrees.get(0), Math.max(1, maxUpperDegree())); } /** @@ -1168,7 +1356,7 @@ private void mergeUpperLayers(ProductQuantization pq, boolean fusedPQEnabled) List opened = Collections.synchronizedList(new ArrayList<>()); for (int level = 1; level <= maxLevel; level++) { final int lvl = level; - int searchTopK = Math.max(MIN_SEARCH_TOP_K, ((maxDegrees.get(level) + sources.size() - 1) / sources.size()) * SEARCH_TOP_K_MULTIPLIER); + int searchTopK = searchTopK(maxDegrees.get(level)); CompactionParams params = new CompactionParams(fusedPQEnabled, fusedPQEnabled, searchTopK, pq); List all = Collections.synchronizedList(new ArrayList<>()); ExecutorCompletionService> ecs = new ExecutorCompletionService<>(executor); @@ -1209,6 +1397,19 @@ private void mergeUpperLayers(ProductQuantization pq, boolean fusedPQEnabled) } for (Scratch sc : opened) sc.close(); log.info("Upper layers merged in {} ms", (System.nanoTime() - t0) / 1_000_000); + // A sharded merge keeps a slice of the sources' nodes, and the slice may hold none of the + // nodes of the top levels: drop every merged level that came out empty, so the output's + // hierarchy ends at the highest level that still has a live node. + int top = maxLevel; + while (top > 0 && mergedSrc[top].length == 0) top--; + if (top < maxLevel) { + log.info("Levels {}..{} have no live node; the merged graph tops out at level {}", top + 1, maxLevel, top); + maxDegrees = new ArrayList<>(maxDegrees.subList(0, top + 1)); + mergedSrc = Arrays.copyOf(mergedSrc, top + 1); + mergedNode = Arrays.copyOf(mergedNode, top + 1); + mergedNbrSrc = Arrays.copyOf(mergedNbrSrc, top + 1); + mergedNbrNode = Arrays.copyOf(mergedNbrNode, top + 1); + } } /** @@ -1363,7 +1564,9 @@ private WriteResult processBaseNode( // it also received no reverse candidates its candidate set is exactly its retained // same-source edges — and re-running diversity over an already-diversity-selected edge // set is a fixed point. Skip selection entirely: filter dead neighbors, remap, write. - if (reverseCandidates != null && sizeRank[sourceIdx] == sources.size() - 1) { + // A sparse source is exempt: its retained edges are a remnant, and it gathers same-source + // candidates of its own. + if (reverseCandidates != null && sizeRank[sourceIdx] == sources.size() - 1 && !sparseSource(sourceIdx)) { int newOrdinal = remappers.get(sourceIdx).oldToNew(node); if (reverseCandidates.countAt(sourceIdx, newOrdinal) == 0) { return writeRetainedOnlyRecord(node, sourceIdx, newOrdinal, scratch, writer); @@ -1514,8 +1717,13 @@ private int gatherCandidates( var indexAlive = liveNodes.get(ss); if (ss == sourceIdx) { + int retainedStart = candSize; candSize = gatherFromSameSource(node, level, ss, searchView, indexAlive, baseVec, scratch, candSize); + if (sparseSource(ss)) { + candSize = gatherSameSourceReplacements(node, level, ss, searchView, indexAlive, + baseVec, scratch, retainedStart, candSize, params); + } } else { // Cross-link: at L0 only search LARGER sources; candidates from smaller sources // arrive via reverse propagation (consumed below), offered when those sources' @@ -1578,6 +1786,54 @@ private int gatherFromSameSource(int node, int level, int sourceIdx, } + /** + * Same-source candidates beyond the retained edges, for a source whose live nodes are a + * minority (see {@link #sparseSource}): most of the node's own edges point at dead nodes, so its + * source is searched like any other one - the cell walk's results for the source as its own + * target at level 0, a graph search elsewhere. The node itself and the neighbours it already + * retained are dropped from what the search returns. + */ + private int gatherSameSourceReplacements(int node, int level, int sourceIdx, + OnDiskGraphIndex.View searchView, FixedBitSet indexAlive, + VectorFloat baseVec, Scratch scratch, + int retainedStart, int candSize, CompactionParams params) { + final int from = candSize; + if (level == 0 && cellMap != null) { + ExactJoin j = exactJoin; + if (j == null || j.targetIndex[sourceIdx] < 0) { + return candSize; + } + candSize = gatherFromCellWalk(node, sourceIdx, sourceIdx, scratch, candSize); + } else { + candSize = gatherFromOtherSource(node, sourceIdx, level, sourceIdx, searchView, indexAlive, + baseVec, scratch, candSize, params); + } + int kept = from; + for (int i = from; i < candSize; i++) { + int cand = scratch.candNode[i]; + boolean drop = cand == node; + for (int r = retainedStart; r < from && !drop; r++) { + drop = scratch.candNode[r] == cand; + } + if (drop) continue; + if (kept != i) { + scratch.candSrc[kept] = scratch.candSrc[i]; + scratch.candNode[kept] = cand; + scratch.candScore[kept] = scratch.candScore[i]; + // the decoded vectors are per-slot buffers: swap them, never alias two slots + VectorFloat v = scratch.candVec[kept]; + scratch.candVec[kept] = scratch.candVec[i]; + scratch.candVec[i] = v; + scratch.candHasVec[kept] = scratch.candHasVec[i]; + } + kept++; + } + for (int i = kept; i < candSize; i++) { + scratch.candHasVec[i] = false; // the slots past the kept candidates are free again + } + return kept; + } + /** * Level 0: hands a node the cell walk's results for one target - the best {@code searchTopK} * ordinals with their exact scores - decoding each candidate's diversity representation as it @@ -1920,7 +2176,7 @@ private int appendApproximateResults(NodeQueue queue, * across all source indexes. */ private List computeLayerInfoFromSources() { - int maxLevel = sources.stream().mapToInt(OnDiskGraphIndex::getMaxLevel).max().orElse(0); + int maxLevel = maxDegrees.size() - 1; List layerInfo = new ArrayList<>(maxLevel + 1); for (int level = 0; level <= maxLevel; level++) { int count = 0; @@ -2115,17 +2371,8 @@ private long estimateScratchSpacePerThread() { int OH = RamUsageEstimator.NUM_BYTES_OBJECT_HEADER; int REF = RamUsageEstimator.NUM_BYTES_OBJECT_REF; - // Calculate maxCandidateSize and maxDegree (same logic as in compactLevels) - int maxUpperDegree = 0; - for (int level = 1; level < maxDegrees.size(); level++) { - maxUpperDegree = Math.max(maxUpperDegree, maxDegrees.get(level)); - } - int baseSearchTopK = Math.max(MIN_SEARCH_TOP_K, ((maxDegrees.get(0) + sources.size() - 1) / sources.size()) * SEARCH_TOP_K_MULTIPLIER); - int baseMaxCandidateSize = baseSearchTopK * (sources.size() - 1) + maxDegrees.get(0) + REVERSE_CANDIDATE_SLOTS; - int upperMaxPerSourceTopK = maxUpperDegree == 0 ? 0 : Math.max(MIN_SEARCH_TOP_K, ((maxUpperDegree + sources.size() - 1) / sources.size()) * SEARCH_TOP_K_MULTIPLIER); - int upperMaxCandidateSize = upperMaxPerSourceTopK * sources.size(); - int maxCandidateSize = Math.max(baseMaxCandidateSize, upperMaxCandidateSize); - int scratchDegree = Math.max(maxDegrees.get(0), Math.max(1, maxUpperDegree)); + int maxCandidateSize = maxCandidateSize(); + int scratchDegree = scratchDegree(); long scratchSize = OH + 6L * REF; @@ -2375,12 +2622,16 @@ private List buildRegionOrdinalMappers(ProductQuantization pq) { AtomicInteger fill = new AtomicInteger(); // enough windows to keep every worker busy on small sources, at most 256K nodes each int window = Math.max(4096, Math.min(1 << 18, size / (4 * taskWindowSize) + 1)); + // Streaming a whole window into the page cache pays when most of it is read. A shard of a + // sharded compaction reads 1/M of every source, and prefetching each window for it would + // read the whole source once per shard. + final boolean streamWindows = liveCount >= size / 8; List> tasks = new ArrayList<>(); for (int from = 0; from < size; from += window) { final int lo = from; final int hi = Math.min(size, from + window); tasks.add(() -> { - source.prefetchL0Records(lo, hi - 1); + if (streamWindows) source.prefetchL0Records(lo, hi - 1); VectorFloat vec = vectorTypeSupport.createFloatVector(dimension); ByteSequence code = pq == null ? null : vectorTypeSupport.createByteSequence(pq.getSubspaceCount()); HubMap.Scorer scorer = regionMode ? hubMapRef.scorer() : null; @@ -2571,7 +2822,8 @@ private int exactJoinRange(int src, long measuredNodes, long measuredPairs, long maxTargetPerNode = 0; int maxProbes = beamScoredCapacity(CELL_PROBE_EF, cellMap.degree[1]); for (int t = 0; t < sources.size(); t++) { - if (sizeRank[t] <= sizeRank[src]) continue; // only larger sources are scanned + // only larger sources are scanned, and a sparse source itself + if (t == src ? !sparseSource(src) : sizeRank[t] <= sizeRank[src]) continue; long membersPerCell = Math.max(1, numLiveNodesPerSource.get(t) / Math.max(1, cellCount)); long perTarget = Math.min(maxProbes, (CELL_BUDGET + membersPerCell - 1) / membersPerCell); pairsPerNode += perTarget; @@ -2593,7 +2845,7 @@ private int exactJoinRange(int src, long measuredNodes, long measuredPairs, long */ private ExactJoin cellWalk(int src, int lo, int hi, int k) { final long t0 = System.nanoTime(); - final ExactJoin j = new ExactJoin(src, lo, hi, k, sources.size(), sizeRank); + final ExactJoin j = new ExactJoin(src, lo, hi, k, sources.size(), sizeRank, sparseSource(src)); final int n = j.n; final boolean needNorms = similarityFunction != VectorSimilarityFunction.DOT_PRODUCT; final VectorFloat[] queries = new VectorFloat[n]; @@ -2692,6 +2944,7 @@ private ExactJoin cellWalk(int src, int lo, int hi, int k) { maxTargetPairs = Math.max(maxTargetPairs, total); final int[] starts = cellStart[target]; final int rowBytes = candidateStore.codeSize(); + final boolean self = target == src; // the source as its own target: a node is not its own candidate // split the sorted pairs into contiguous cell ranges for the workers int tasks2 = Math.max(1, Math.min(taskWindowSize * 8, total / 4096)); List> phase2 = new ArrayList<>(); @@ -2759,6 +3012,7 @@ private ExactJoin cellWalk(int src, int lo, int hi, int k) { float vNorm2 = needNorms ? VectorUtil.dotProduct(tmp, tmp) : 0f; for (int r = 0; r < cn; r++) { int i = probers[c0 + r]; + if (self && lo + i == newOrd) continue; float dot = chunkScores[r]; float score; if (!needNorms) { diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java index 36361d0cb..02d413540 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java @@ -82,12 +82,27 @@ public VectorCompressor compressor() { return null; } + @Override + public void adoptCompressor(VectorCompressor trained) { + // no-op: nothing is quantized + } + @Override public String toString() { return "QuantizationCompactionStrategy.NONE"; } }; + /** + * Installs a compressor trained elsewhere in place of {@link #retrain}. A sharded compaction + * trains once, on the live nodes of every shard, and each shard's merge adopts the result + * instead of training again on its own slice. After this call {@link #compressor()} returns + * {@code trained}. + */ + public void adoptCompressor(VectorCompressor trained) { + throw new UnsupportedOperationException(getClass().getSimpleName() + " cannot adopt a compressor"); + } + /** * Trains a fresh compressor on a balanced sample of merged source vectors. May be a no-op * for strategies that don't carry a compressor (e.g. {@link #NONE}). After this call, diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/SidecarCompactionStrategy.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/SidecarCompactionStrategy.java index b3a091c09..290f6320e 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/SidecarCompactionStrategy.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/SidecarCompactionStrategy.java @@ -98,6 +98,11 @@ public void retrain(VectorSimilarityFunction vsf) { this.retrainedCompressor = retrainer.retrain(vsf); } + @Override + public void adoptCompressor(VectorCompressor trained) { + this.retrainedCompressor = java.util.Objects.requireNonNull(trained, "trained"); + } + @Override public VectorCompressor compressor() { return retrainedCompressor; diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java index f648ad654..2b1cc44fb 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java @@ -1215,6 +1215,215 @@ public void testCellJoinMergeDirectReadsAndFixedBudgets() throws Exception { c -> c.withWalkDirectReads(true).withRecordAnnounceAhead(8).withCellWalkBudgets(64L << 20, 64L << 20)); } + /** + * A sharded merge keeps a slice of every source's nodes, and the slice can miss every node of + * the sources' top levels. The output's hierarchy then ends at the highest level that still has + * a live node, with the entry node on it; dropping the top level only, and dropping every upper + * level (a single-layer output), are both merged and searched. + */ + @Test + public void testMergeWithoutLiveNodesAtTopLevels() throws Exception { + VectorSimilarityFunction vsf = VectorSimilarityFunction.COSINE; + int perSource = 3000, nSrc = 2; + java.util.Random rnd = new java.util.Random(20261107); + List> all = new ArrayList<>(); + List graphs = new ArrayList<>(); + List rss = new ArrayList<>(); + for (int sIdx = 0; sIdx < nSrc; sIdx++) { + List> vecs = seededVectors(rnd, perSource, dimension); + Path path = buildPlainSourceGraph(vecs, vsf, true, "notop_src_" + sIdx); + rss.add(ReaderSupplierFactory.open(path)); + graphs.add(OnDiskGraphIndex.load(rss.get(sIdx))); + all.addAll(vecs); + } + int sourceTop = graphs.stream().mapToInt(OnDiskGraphIndex::getMaxLevel).max().orElse(0); + assertTrue("sources must have a hierarchy", sourceTop >= 1); + // keep = the highest level that keeps its live nodes: every node of the levels above is dead + for (int keep : new int[]{sourceTop - 1, 0}) { + List live = new ArrayList<>(); + int liveTotal = 0; + for (var g : graphs) { + var lv = new FixedBitSet(g.getIdUpperBound()); + lv.set(0, g.getIdUpperBound()); + for (int level = keep + 1; level <= g.getMaxLevel(); level++) { + var it = g.getNodes(level); + while (it.hasNext()) lv.clear(it.next()); + } + live.add(lv); + liveTotal += lv.cardinality(); + } + assertTrue("the test must delete something", liveTotal < nSrc * perSource); + var compactor = new OnDiskGraphIndexCompactor(graphs, live, vsf, null); + Path out = testDirectory.resolve("notop_merged_" + keep); + compactor.compact(out); + var mappers = compactor.ordinalMappers(); + List> reordered = new ArrayList<>(Collections.nCopies(liveTotal, (VectorFloat) null)); + for (int sIdx = 0; sIdx < nSrc; sIdx++) { + for (int i = live.get(sIdx).nextSetBit(0); i >= 0 && i < perSource; i = i + 1 < perSource ? live.get(sIdx).nextSetBit(i + 1) : -1) { + reordered.set(mappers.get(sIdx).oldToNew(i), all.get(sIdx * perSource + i)); + } + } + try (ReaderSupplier rs = ReaderSupplierFactory.open(out)) { + var merged = OnDiskGraphIndex.load(rs); + assertEquals("live nodes (keep=" + keep + ")", liveTotal, merged.size(0)); + assertEquals("output top level (keep=" + keep + ")", keep, merged.getMaxLevel()); + assertTrue("the top level has nodes", merged.size(keep) > 0); + try (var view = merged.getView()) { + assertEquals("entry node level", keep, view.entryNode().level); + assertTrue("entry node is a live node", view.entryNode().node >= 0 && view.entryNode().node < liveTotal); + var v = vectorTypeSupport.createFloatVector(dimension); + for (int n = 0; n < liveTotal; n++) { + view.getVectorInto(n, v, 0); + assertVecEquals(reordered.get(n), v, n); + } + } + var ravv = new ListRandomAccessVectorValues(reordered, dimension); + try (GraphSearcher searcher = new GraphSearcher(merged)) { + int selfHits = 0, nq = 100; + for (int qi = 0; qi < nq; qi++) { + int target = rnd.nextInt(liveTotal); + SearchScoreProvider ssp = DefaultSearchScoreProvider.exact(reordered.get(target), vsf, ravv); + SearchResult sr = searcher.search(ssp, 10, Bits.ALL); + for (var ns : sr.getNodes()) if (ns.node == target) selfHits++; + } + assertTrue("self-hits " + selfHits + "/" + nq + " (keep=" + keep + ")", selfHits >= nq * 9 / 10); + } + } + } + for (var r : rss) r.close(); + } + + /** One session, several outputs, full-precision sources. */ + @Test + public void testCompactShardedFullPrecisionSources() throws Exception { + compactSharded(false); + } + + /** One session, several outputs, fused-PQ sources: every shard adopts the codebook trained once. */ + @Test + public void testCompactShardedFusedSources() throws Exception { + compactSharded(true); + } + + /** + * {@code sharded().compactShard} per output: every node of a shard lands in that shard, at the + * ordinal the shard's mappers report, with its vector; a node in no shard and a deleted node + * land nowhere. A shard keeps a quarter of each source, so a node keeps about a quarter of its + * same-source edges: the same-source join must restore them, which the edge mix of the + * outputs shows. + */ + private void compactSharded(boolean fusedSources) throws Exception { + VectorSimilarityFunction vsf = VectorSimilarityFunction.COSINE; + int perSource = 3000, nSrc = 2, shards = 4; + java.util.Random rnd = new java.util.Random(20261109); + List> all = new ArrayList<>(); + List graphs = new ArrayList<>(); + List rss = new ArrayList<>(); + List live = new ArrayList<>(); + for (int sIdx = 0; sIdx < nSrc; sIdx++) { + List> vecs = seededVectors(rnd, perSource, dimension); + String name = "sharded_" + (fusedSources ? "fused" : "plain") + "_src_" + sIdx; + Path path = fusedSources ? buildFusedSourceGraph(vecs, vsf, true, name) : buildPlainSourceGraph(vecs, vsf, true, name); + rss.add(ReaderSupplierFactory.open(path)); + graphs.add(OnDiskGraphIndex.load(rss.get(sIdx))); + var lv = new FixedBitSet(perSource); + for (int i = 0; i < perSource; i++) if (i % 10 != 3) lv.set(i); // every tenth node is deleted + live.add(lv); + all.addAll(vecs); + } + // the live nodes are spread over the shards; a few are in no shard + List> shardNodes = new ArrayList<>(); // [shard][source] + for (int m = 0; m < shards; m++) { + List perShard = new ArrayList<>(); + for (int sIdx = 0; sIdx < nSrc; sIdx++) { + var bits = new FixedBitSet(perSource); + for (int i = 0; i < perSource; i++) { + if (live.get(sIdx).get(i) && i % 97 != 0 && (i * 31 + sIdx) % shards == m) bits.set(i); + } + perShard.add(bits); + } + shardNodes.add(perShard); + } + List outs = new ArrayList<>(); + String outName = "sharded_" + (fusedSources ? "fused" : "plain") + "_out_"; + for (int m = 0; m < shards; m++) outs.add(testDirectory.resolve(outName + m)); + + var compactor = new OnDiskGraphIndexCompactor(graphs, live, vsf, null); + var session = compactor.sharded(); + List> mappers = new ArrayList<>(); + for (int m = 0; m < shards; m++) mappers.add(session.compactShard(shardNodes.get(m), outs.get(m))); + + long sameSourceEdges = 0, edges = 0; + for (int m = 0; m < shards; m++) { + List> expected = new ArrayList<>(); + int count = 0; + for (int sIdx = 0; sIdx < nSrc; sIdx++) count += shardNodes.get(m).get(sIdx).cardinality(); + for (int i = 0; i < count; i++) expected.add(null); + int[] sourceOf = new int[count]; + for (int sIdx = 0; sIdx < nSrc; sIdx++) { + for (int i = 0; i < perSource; i++) { + int ord = mappers.get(m).get(sIdx).oldToNew(i); + if (shardNodes.get(m).get(sIdx).get(i)) { + assertTrue("ordinal " + ord + " of a shard of " + count, ord >= 0 && ord < count); + assertTrue("two nodes at ordinal " + ord, expected.get(ord) == null); + expected.set(ord, all.get(sIdx * perSource + i)); + sourceOf[ord] = sIdx; + } else { + assertEquals("a node outside the shard has no ordinal", OrdinalMapper.OMITTED, ord); + } + } + } + try (ReaderSupplier rs = ReaderSupplierFactory.open(outs.get(m))) { + var merged = OnDiskGraphIndex.load(rs); + assertEquals("nodes of shard " + m, count, merged.size(0)); + assertEquals("quantized feature in shard " + m, fusedSources, merged.getFeatures().containsKey(FeatureId.FUSED_PQ)); + try (var view = merged.getView()) { + var v = vectorTypeSupport.createFloatVector(dimension); + for (int n = 0; n < count; n++) { + view.getVectorInto(n, v, 0); + assertVecEquals(expected.get(n), v, n); + var it = view.getNeighborsIterator(0, n); + while (it.hasNext()) { + int nb = it.nextInt(); + assertTrue("neighbour " + nb + " of " + n, nb >= 0 && nb < count && nb != n); + if (sourceOf[nb] == sourceOf[n]) sameSourceEdges++; + edges++; + } + } + } + var ravv = new ListRandomAccessVectorValues(expected, dimension); + try (GraphSearcher searcher = new GraphSearcher(merged)) { + int selfHits = 0, nq = 100; + for (int qi = 0; qi < nq; qi++) { + int target = rnd.nextInt(count); + SearchScoreProvider ssp = DefaultSearchScoreProvider.exact(expected.get(target), vsf, ravv); + SearchResult sr = searcher.search(ssp, 10, Bits.ALL); + for (var ns : sr.getNodes()) if (ns.node == target) selfHits++; + } + assertTrue("self-hits " + selfHits + "/" + nq + " in shard " + m, selfHits >= nq * 9 / 10); + } + } + } + double sameShare = sameSourceEdges / (double) edges; + System.out.printf("Sharded compaction (%s sources): %.1f%% of the level-0 edges stay within a source%n", + fusedSources ? "fused" : "full-precision", 100 * sameShare); + // The two sources are draws of one distribution, so about half of a node's nearest + // neighbours are of its own source; retained edges alone would leave about a quarter. + assertTrue("same-source share of the edges " + sameShare, sameShare >= 0.40); + + // a shard that is not a subset of the compaction's live nodes is rejected + var notLive = new ArrayList(); + for (int sIdx = 0; sIdx < nSrc; sIdx++) notLive.add(new FixedBitSet(perSource)); + notLive.get(0).set(3); // node 3 of source 0 is deleted + try { + compactor.sharded().compactShard(notLive, testDirectory.resolve(outName + "rejected")); + org.junit.Assert.fail("a shard holding a deleted node must be rejected"); + } catch (IllegalArgumentException expectedException) { + // expected + } + for (var r : rss) r.close(); + } + private void cellJoinMergeRecall(VectorSimilarityFunction vsf, boolean fusedSources) throws Exception { cellJoinMergeRecall(vsf, fusedSources, c -> c); } From 43ee5cb7a86d3f35b933ced72556149b29c0fd4f Mon Sep 17 00:00:00 2001 From: dian-lun-lin Date: Wed, 11 Nov 2026 22:44:02 +0000 Subject: [PATCH 3/6] Compaction: withVersion chooses the output's on-disk format version The compactor wrote OnDiskGraphIndex.CURRENT_VERSION only; CompactWriter had no version parameter. withVersion(int) now names the version of the output, any one GraphIndexFormatFactory knows, and the writer follows it: the header is built through the version's format as before, the footer is written only for footer-based versions (5 and above), and a version that cannot hold the output is rejected up front (fused PQ needs 6, upper layers need 4 or above). The sources' versions are still detected from their files and may differ from the output's. A sharded compaction writes every shard at the session's version. testCompactWithVersion round-trips a hierarchical merge at 4, 5 and 6 and a single-layer merge at 2 to 6 through the versioned loader, and checks the rejections. Claude-Session: https://claude.ai/code/session_016pmoQbpUFKvQKpvRc7Xvjz --- docs/compaction.md | 4 + .../jvector/graph/disk/CompactWriter.java | 35 ++++-- .../graph/disk/OnDiskGraphIndexCompactor.java | 24 +++- .../disk/TestOnDiskGraphIndexCompactor.java | 107 ++++++++++++++++++ 4 files changed, 158 insertions(+), 12 deletions(-) diff --git a/docs/compaction.md b/docs/compaction.md index 1e4a097d4..9cbe38d1b 100644 --- a/docs/compaction.md +++ b/docs/compaction.md @@ -39,6 +39,10 @@ List 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). diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactWriter.java index 73da527d8..bacfe6bcb 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactWriter.java +++ b/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 bufferPerThread; private final ThreadLocal> zeroPQ; private final boolean fusedPQEnabled; + private final boolean usesFooter; private final Path outputPath; private final List configuredLayerInfo; private final List configuredLayerDegrees; @@ -75,7 +76,13 @@ final class CompactWriter implements AutoCloseable { private int pqCodeSize; private ThreadLocal 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; } 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 cf4643321..64a8b2cf1 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 @@ -260,6 +260,9 @@ private void setupCrossLink() { // Sharded compaction: what is trained once for every shard. Null for an ordinary compaction. private SharedTraining shared; + // The on-disk format version of the output. + private int version = OnDiskGraphIndex.CURRENT_VERSION; + /** Whether so few of source {@code s}'s nodes are live that its own nodes are joined like another source's. */ private boolean sparseSource(int s) { return numLiveNodesPerSource.get(s) < (double) sameSourceJoinBelow * sources.get(s).size(0); @@ -331,6 +334,24 @@ public OnDiskGraphIndexCompactor( this.similarityFunction = similarityFunction; } + /** + * Sets the on-disk format version of the output; the default is {@link OnDiskGraphIndex#CURRENT_VERSION}. + * The sources may be of any version the reader supports, and the output is written at this one. + * An output that carries fused PQ needs version 6; one with upper layers needs version 4 or above. + * + * @throws IllegalArgumentException if the version is not one the writer supports + */ + @Experimental + public OnDiskGraphIndexCompactor withVersion(int version) { + try { + GraphIndexFormatFactory.forVersion(version); + } catch (GraphIndexFormatFactory.UnsupportedVersionException e) { + throw new IllegalArgumentException(e.getMessage(), e); + } + this.version = version; + return this; + } + // The settings below are the compactor's own: every one has a default it derives from the // machine or the sources, and a caller is not expected to know better. They are package-private // so the tests can exercise the paths they select. @@ -689,6 +710,7 @@ private OnDiskGraphIndexCompactor shardCompactor(List shardLiveNode } OnDiskGraphIndexCompactor shard = new OnDiskGraphIndexCompactor(sources, sourceCompressed, shardLiveNodes, similarityFunction, executor); shard.shared = training; + shard.version = version; shard.walkPairBytes = walkPairBytes; shard.walkQueryBytes = walkQueryBytes; shard.recordAnnounceAhead = recordAnnounceAhead; @@ -917,7 +939,7 @@ private void compactGraphImpl(Path outputPath, QuantizationCompactionStrategy st log.info("Writing compacted graph : {} total nodes, maxOrdinal={}, dimension={}, degree={}", numTotalNodes, maxOrdinal, dimension, maxDegrees.get(0)); - try (CompactWriter writer = new CompactWriter(outputPath, maxOrdinal, numTotalNodes, 0, layerInfo, entryNode, dimension, maxDegrees, outputFusedFeature)) { + try (CompactWriter writer = new CompactWriter(outputPath, version, maxOrdinal, numTotalNodes, 0, layerInfo, entryNode, dimension, maxDegrees, outputFusedFeature)) { // Header has to be written first so the writer's position is past the header // before any strategy that mmaps past the projected end of the output runs. writer.writeHeader(); diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java index 2b1cc44fb..09dbec5a2 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java @@ -1215,6 +1215,113 @@ public void testCellJoinMergeDirectReadsAndFixedBudgets() throws Exception { c -> c.withWalkDirectReads(true).withRecordAnnounceAhead(8).withCellWalkBudgets(64L << 20, 64L << 20)); } + /** + * {@code withVersion}: the output is written at the requested on-disk format version, whatever + * the sources' version, and reads back through the versioned loader with its vectors, its + * hierarchy and a working search. Versions that cannot hold the output are rejected up front: + * upper layers need 4 or above, fused PQ needs 6. + */ + @Test + public void testCompactWithVersion() throws Exception { + VectorSimilarityFunction vsf = VectorSimilarityFunction.COSINE; + int perSource = 2000, nSrc = 2; + java.util.Random rnd = new java.util.Random(20261111); + for (boolean hierarchy : new boolean[]{true, false}) { + List> all = new ArrayList<>(); + List paths = new ArrayList<>(); + for (int sIdx = 0; sIdx < nSrc; sIdx++) { + List> vecs = seededVectors(rnd, perSource, dimension); + paths.add(buildPlainSourceGraph(vecs, vsf, hierarchy, "version_" + (hierarchy ? "h" : "flat") + "_src_" + sIdx)); + all.addAll(vecs); + } + int[] versions = hierarchy ? new int[]{4, 5, 6} : new int[]{2, 3, 4, 5, 6}; + for (int version : versions) { + compactAtVersion(paths, all, vsf, version, hierarchy, "version_" + (hierarchy ? "h" : "flat") + "_v" + version); + } + if (hierarchy) { + assertRejectedVersion(paths, vsf, 3, "upper layers need version 4 or above"); + } + } + // fused PQ sources: the output carries fused PQ, which only version 6 holds + List fused = new ArrayList<>(); + for (int sIdx = 0; sIdx < nSrc; sIdx++) { + fused.add(buildFusedSourceGraph(seededVectors(rnd, perSource, dimension), vsf, true, "version_fused_src_" + sIdx)); + } + assertRejectedVersion(fused, vsf, 5, "fused PQ needs version 6"); + try { + openCompactor(fused, vsf).withVersion(OnDiskGraphIndex.CURRENT_VERSION + 1); + org.junit.Assert.fail("an unknown version must be rejected"); + } catch (IllegalArgumentException expected) { + // expected + } + } + + private OnDiskGraphIndexCompactor openCompactor(List paths, VectorSimilarityFunction vsf) throws IOException { + List graphs = new ArrayList<>(); + List live = new ArrayList<>(); + for (Path path : paths) { + var g = OnDiskGraphIndex.load(ReaderSupplierFactory.open(path)); + graphs.add(g); + var lv = new FixedBitSet(g.getIdUpperBound()); + lv.set(0, g.getIdUpperBound()); + live.add(lv); + } + return new OnDiskGraphIndexCompactor(graphs, live, vsf, null); + } + + private void compactAtVersion(List paths, List> all, VectorSimilarityFunction vsf, + int version, boolean hierarchy, String name) throws Exception { + var compactor = openCompactor(paths, vsf).withVersion(version); + Path out = testDirectory.resolve(name); + compactor.compact(out); + var mappers = compactor.ordinalMappers(); + int total = all.size(), perSource = total / paths.size(); + try (ReaderSupplier rs = ReaderSupplierFactory.open(out)) { + var merged = OnDiskGraphIndex.load(rs); + assertEquals("version of " + name, version, merged.version); + assertEquals("nodes of " + name, total, merged.size(0)); + assertEquals("hierarchy of " + name, hierarchy, merged.getMaxLevel() >= 1); + assertFalse("fused PQ in " + name, merged.getFeatures().containsKey(FeatureId.FUSED_PQ)); + List> expected = new ArrayList<>(java.util.Collections.nCopies(total, (VectorFloat) null)); + for (int sIdx = 0; sIdx < paths.size(); sIdx++) { + for (int i = 0; i < perSource; i++) { + expected.set(mappers.get(sIdx).oldToNew(i), all.get(sIdx * perSource + i)); + } + } + try (var view = merged.getView()) { + var v = vectorTypeSupport.createFloatVector(dimension); + for (int n = 0; n < total; n++) { + view.getVectorInto(n, v, 0); + assertVecEquals(expected.get(n), v, n); + } + } + var ravv = new ListRandomAccessVectorValues(expected, dimension); + try (GraphSearcher searcher = new GraphSearcher(merged)) { + int selfHits = 0, nq = 100; + java.util.Random rnd = new java.util.Random(version); + for (int qi = 0; qi < nq; qi++) { + int target = rnd.nextInt(total); + SearchScoreProvider ssp = DefaultSearchScoreProvider.exact(expected.get(target), vsf, ravv); + SearchResult sr = searcher.search(ssp, 10, Bits.ALL); + for (var ns : sr.getNodes()) if (ns.node == target) selfHits++; + } + assertTrue("self-hits " + selfHits + "/" + nq + " in " + name, selfHits >= nq * 9 / 10); + } + } + } + + private void assertRejectedVersion(List paths, VectorSimilarityFunction vsf, int version, String why) throws Exception { + var compactor = openCompactor(paths, vsf).withVersion(version); + try { + compactor.compact(testDirectory.resolve("version_rejected_" + version)); + org.junit.Assert.fail(why); + } catch (IllegalArgumentException expected) { + // expected + } catch (RuntimeException e) { + assertTrue(why + ": " + e, e.getCause() instanceof IllegalArgumentException); + } + } + /** * A sharded merge keeps a slice of every source's nodes, and the slice can miss every node of * the sources' top levels. The output's hierarchy then ends at the highest level that still has From 59b83a85e254797767b0fed3a3e8b698950322ff Mon Sep 17 00:00:00 2001 From: dian-lun-lin Date: Thu, 12 Nov 2026 23:14:20 +0000 Subject: [PATCH 4/6] Compaction: direct reads align to the file store's block size The direct-I/O path aligned offsets, lengths and buffers to 4 KB and only checked that the direct open succeeded. The JDK requires alignment to the file store's block size, so on a store with 64 KB blocks (XFS or btrfs on 64 KB-page kernels, HFS+, NFS) the first read failed and the compaction aborted with "Channel position ... is not a multiple of the block size". PreEncodedCodeCache now takes the alignment from the file store (at least 4 KB, a power of two up to 1 MB), probes one aligned read when it opens the direct channel, and falls back to the mapping on any failure; allocateAligned and directBufferBytes use the cache's alignment. Verified on a loop-mounted XFS volume with 64 KB blocks: the old code fails as reported, the new code reads and a full merge runs with direct I/O. --- .../graph/disk/OnDiskGraphIndexCompactor.java | 8 +-- .../graph/disk/PreEncodedCodeCache.java | 68 +++++++++++++++---- .../graph/disk/TestPreEncodedCodeCache.java | 4 +- 3 files changed, 61 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 64a8b2cf1..bacca6270 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 @@ -2897,7 +2897,7 @@ private ExactJoin cellWalk(int src, int lo, int hi, int k) { long[][] pairs = pairsByTask[task]; int[] count = countByTask[task]; for (int t = 0; t < j.targets; t++) pairs[t] = new long[Math.max(64, (to - from) * 6)]; - final ByteBuffer buf = direct ? sc.directBuffer(readBufBytes) : null; + final ByteBuffer buf = direct ? sc.directBuffer(candidateStore, readBufBytes) : null; int bufLo = 0, bufHi = 0, bufBase = 0; // rows [bufLo, bufHi) are in buf, row bufLo at bufBase for (int i = from; i < to; i++) { int newOrd = lo + i; @@ -2990,7 +2990,7 @@ private ExactJoin cellWalk(int src, int lo, int hi, int k) { // Mapped mode decodes straight from the mapping, as a store that fits in RAM wants. int qAhead = from; long read = 0; - final ByteBuffer buf = direct ? sc.directBuffer(readBufBytes) : null; + final ByteBuffer buf = direct ? sc.directBuffer(candidateStore, readBufBytes) : null; int bufLo = 0, bufHi = 0, bufBase = 0; // rows [bufLo, bufHi) are in buf, row bufLo at bufBase int planHi = 0; int a = from; @@ -3135,9 +3135,9 @@ public void close() throws IOException { ByteBuffer directBuf; // cell walk in direct mode: the thread's aligned read buffer - ByteBuffer directBuffer(int bytes) { + ByteBuffer directBuffer(PreEncodedCodeCache store, int bytes) { if (directBuf == null || directBuf.capacity() < bytes) { - directBuf = PreEncodedCodeCache.allocateAligned(bytes); + directBuf = store.allocateAligned(bytes); } return directBuf; } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PreEncodedCodeCache.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PreEncodedCodeCache.java index 9da556ba5..93a99838a 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PreEncodedCodeCache.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PreEncodedCodeCache.java @@ -26,6 +26,7 @@ import java.nio.ByteBuffer; import java.nio.MappedByteBuffer; import java.nio.channels.FileChannel; +import java.nio.file.Files; import java.nio.file.StandardOpenOption; import java.nio.file.Path; @@ -90,10 +91,16 @@ private static Unsafe getUnsafe() { private final ThreadLocal viewsPerThread; private final ThreadLocal warmByte = ThreadLocal.withInitial(() -> ByteBuffer.allocateDirect(1)); - /** Alignment direct I/O needs for offsets, lengths and buffer addresses (the file system block). */ + /** + * The least alignment direct I/O needs for offsets, lengths and buffer addresses; the file's + * store may need more (see {@link #directAlign()}). + */ public static final int DIRECT_ALIGN = 4096; + /** The largest file-store block size direct reads are attempted with. */ + private static final long MAX_DIRECT_ALIGN = 1 << 20; private final Path path; // null when the cache was mapped without one: no direct reads private volatile FileChannel direct; // opened on first use with the direct-I/O option + private volatile int directAlign = DIRECT_ALIGN; // the alignment of direct reads of this file, once opened private volatile boolean directUnavailable; private PreEncodedCodeCache(MappedByteBuffer[] chunks, int codeSize, int codesPerChunk, @@ -183,13 +190,25 @@ public long sectionBytes() { } /** - * Whether {@link #readRowsDirect} can be used: the cache knows its file and the file system - * accepts the direct-I/O open option (tmpfs, for one, does not). + * Whether {@link #readRowsDirect} can be used: the cache knows its file, the file system + * accepts the direct-I/O open option (tmpfs, for one, does not) and an aligned read through + * it succeeds. */ public boolean supportsDirect() { return directChannel() != null; } + /** + * The alignment of this file's direct reads: offsets, lengths and buffer addresses must be + * multiples of it. It is the file store's block size, at least {@link #DIRECT_ALIGN}, once the + * direct channel is open; a file store with 64 KB blocks (some XFS and btrfs volumes on 64 KB + * page kernels, HFS+, NFS) rejects 4 KB alignment. + */ + public int directAlign() { + directChannel(); + return directAlign; + } + private FileChannel directChannel() { FileChannel d = direct; if (d != null || directUnavailable || path == null) { @@ -197,25 +216,46 @@ private FileChannel directChannel() { } synchronized (this) { if (direct == null && !directUnavailable) { + FileChannel opened = null; try { - direct = FileChannel.open(path, StandardOpenOption.READ, com.sun.nio.file.ExtendedOpenOption.DIRECT); + long block = Files.getFileStore(path).getBlockSize(); + if (block <= 0 || block > MAX_DIRECT_ALIGN || Long.bitCount(block) != 1) { + throw new IOException("file store block size " + block + " is not a power of two up to " + MAX_DIRECT_ALIGN); + } + int align = (int) Math.max(DIRECT_ALIGN, block); + opened = FileChannel.open(path, StandardOpenOption.READ, com.sun.nio.file.ExtendedOpenOption.DIRECT); + // an open that succeeds does not guarantee that reads do: probe one aligned read + ByteBuffer probe = ByteBuffer.allocateDirect(2 * align).alignedSlice(align); + probe.limit(align); + opened.read(probe, 0); + directAlign = align; + direct = opened; } catch (IOException | UnsupportedOperationException | IllegalArgumentException e) { log.info("Direct reads of {} unavailable ({}); the cache is read through its mapping", path, e.toString()); directUnavailable = true; + if (opened != null) { + try { + opened.close(); + } catch (IOException ignored) { + // the mapping remains the way to read the cache + } + } } } return direct; } } - /** A direct buffer aligned for {@link #readRowsDirect}, holding at least {@code bytes}. */ - public static ByteBuffer allocateAligned(int bytes) { - return ByteBuffer.allocateDirect(bytes + DIRECT_ALIGN).alignedSlice(DIRECT_ALIGN); + /** A direct buffer aligned for {@link #readRowsDirect} of this cache, holding at least {@code bytes}. */ + public ByteBuffer allocateAligned(int bytes) { + int align = directAlign(); + return ByteBuffer.allocateDirect(bytes + align).alignedSlice(align); } /** Buffer capacity {@link #readRowsDirect} needs for {@code count} rows at once. */ public int directBufferBytes(int count) { - return (int) Math.min(Integer.MAX_VALUE - DIRECT_ALIGN, (long) count * codeSize + 2L * DIRECT_ALIGN); + int align = directAlign(); + return (int) Math.min(Integer.MAX_VALUE - align, (long) count * codeSize + 2L * align); } /** @@ -224,8 +264,9 @@ public int directBufferBytes(int count) { * cached, mapped or evicted on their account. For a pass that reads a store larger than RAM * once from end to end, this keeps the kernel's memory reclaim out of the walk, where a mapping * read of the same rows must evict a resident page for every page it brings in. Offsets and - * lengths are widened to {@link #DIRECT_ALIGN} as direct I/O requires, so {@code dst} needs - * {@link #directBufferBytes} capacity and the rows start at the returned offset. + * lengths are widened to {@link #directAlign()} as direct I/O requires, so {@code dst} needs + * {@link #directBufferBytes} capacity, allocated by {@link #allocateAligned}, and the rows + * start at the returned offset. * * @return the offset in {@code dst} at which row {@code firstOrdinal} begins */ @@ -234,10 +275,11 @@ public int readRowsDirect(int firstOrdinal, int count, ByteBuffer dst) throws IO if (d == null) { throw new IOException("direct reads unavailable for " + path); } + final int align = directAlign; long start = baseOffset + (long) firstOrdinal * codeSize; long end = start + (long) count * codeSize; - long alignedStart = start & -(long) DIRECT_ALIGN; - long alignedEnd = (end + DIRECT_ALIGN - 1) & -(long) DIRECT_ALIGN; + long alignedStart = start & -(long) align; + long alignedEnd = (end + align - 1) & -(long) align; int bytes = (int) (alignedEnd - alignedStart); if (bytes > dst.capacity()) { throw new IllegalArgumentException("buffer of " + dst.capacity() + " bytes cannot hold " + count + " rows (" + bytes + " aligned bytes)"); @@ -251,7 +293,7 @@ public int readRowsDirect(int firstOrdinal, int count, ByteBuffer dst) throws IO break; } pos += n; - if (pos % DIRECT_ALIGN != 0) { + if (pos % align != 0) { // a short read is the end of the file: the aligned tail ran past it, and nothing beyond is a row break; } diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPreEncodedCodeCache.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPreEncodedCodeCache.java index 56fea5d64..c5ec1e6d4 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPreEncodedCodeCache.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPreEncodedCodeCache.java @@ -128,9 +128,9 @@ public void testReadRowsDirectMatchesMapping() throws IOException { int[][] runs = {{0, 1}, {7, 50}, {290, count - 290}, {0, count}}; for (int[] run : runs) { int first = run[0], rows = run[1]; - ByteBuffer buf = PreEncodedCodeCache.allocateAligned(cache.directBufferBytes(rows)); + ByteBuffer buf = cache.allocateAligned(cache.directBufferBytes(rows)); int at = cache.readRowsDirect(first, rows, buf); - assertTrue("rows start inside the aligned head", at >= 0 && at < PreEncodedCodeCache.DIRECT_ALIGN); + assertTrue("rows start inside the aligned head", at >= 0 && at < cache.directAlign()); for (int r = 0; r < rows; r++) { buf.position(at + r * codeSize); buf.get(dst, 0, codeSize); From dee2371a2857896bc8b3019fbcb7bcd948b22124 Mon Sep 17 00:00:00 2001 From: dian-lun-lin Date: Thu, 12 Nov 2026 23:14:20 +0000 Subject: [PATCH 5/6] Compaction: adoptCompressor is abstract A strategy that implemented only retrain and compressor worked in compact() and failed on the first compactShard(); every strategy must now say how it adopts a shared compressor, so the gap shows at compile time. --- .../jvector/graph/disk/QuantizationCompactionStrategy.java | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java index 02d413540..64388219b 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java @@ -97,11 +97,9 @@ public String toString() { * Installs a compressor trained elsewhere in place of {@link #retrain}. A sharded compaction * trains once, on the live nodes of every shard, and each shard's merge adopts the result * instead of training again on its own slice. After this call {@link #compressor()} returns - * {@code trained}. + * {@code trained}; a strategy without a compressor (e.g. {@link #NONE}) ignores it. */ - public void adoptCompressor(VectorCompressor trained) { - throw new UnsupportedOperationException(getClass().getSimpleName() + " cannot adopt a compressor"); - } + public abstract void adoptCompressor(VectorCompressor trained); /** * Trains a fresh compressor on a balanced sample of merged source vectors. May be a no-op From a19ce5cb01e4791108246b6908b1bdcc9195c127 Mon Sep 17 00:00:00 2001 From: dian-lun-lin Date: Thu, 12 Nov 2026 23:23:13 +0000 Subject: [PATCH 6/6] Compaction: the same-source join is a sharded compaction's only The join of a sparse source with itself was on in every compactor, so an ordinary compact() over a source with under half its nodes live also got it: an extra cell-walk target and upper-level search, the largest source losing its retained-only path, and every worker's candidate scratch sized for one more search. It was measured for shards only. An ordinary compaction now never joins a source with itself; a shard's compactor joins sources below half live, as before. The candidate scratch counts the extra search only when some source is joined with itself. The package-private setter, which no test used, is gone. --- docs/compaction.md | 2 +- .../graph/disk/OnDiskGraphIndexCompactor.java | 52 +++++++++---------- 2 files changed, 27 insertions(+), 27 deletions(-) diff --git a/docs/compaction.md b/docs/compaction.md index 9cbe38d1b..f453a0398 100644 --- a/docs/compaction.md +++ b/docs/compaction.md @@ -115,7 +115,7 @@ searchTopK = max(2, ceil(degree / numSources) * 4) **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) and any heavily deleted source: a node then keeps `degree × live fraction` of its same-source edges, and the output is close to bipartite between the sources. So 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. +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. 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 bacca6270..20918255d 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 @@ -248,14 +248,15 @@ private void setupCrossLink() { // ---- Sparse sources: the same-source join ---- // A node keeps only its LIVE same-source edges. When most of a source is live that is nearly all - // of them; when a source's live nodes are a minority (a shard of a sharded compaction keeps 1/M - // of every source, a heavily deleted source keeps little) most of a node's edges point at dead - // nodes and nothing replaced them: measured on 2 x 1M vectors split into 64 shards, a node kept - // 0.5 of its 32 same-source edges and the output was nearly bipartite between the sources. Below - // this live fraction a source is therefore searched for its own nodes like any other source: at - // level 0 it is one more target of the cell walk, elsewhere one more graph search. - private static final float DEFAULT_SAME_SOURCE_JOIN_BELOW = 0.5f; - private float sameSourceJoinBelow = DEFAULT_SAME_SOURCE_JOIN_BELOW; + // of them; when a source's live nodes are a minority, as in a shard of a sharded compaction + // (which keeps 1/M of every source), most of a node's edges point at dead nodes and nothing + // replaced them: measured on 2 x 1M vectors split into 64 shards, a node kept 0.5 of its 32 + // same-source edges and the output was nearly bipartite between the sources. In a shard, a + // source below this live fraction is therefore searched for its own nodes like any other source: + // at level 0 it is one more target of the cell walk, elsewhere one more graph search. An ordinary + // compaction never joins a source with itself (the fraction is 0), whatever its deletions. + private static final float SHARD_SAME_SOURCE_JOIN_BELOW = 0.5f; + private float sameSourceJoinBelow = 0f; // Sharded compaction: what is trained once for every shard. Null for an ordinary compaction. private SharedTraining shared; @@ -267,6 +268,16 @@ private void setupCrossLink() { private boolean sparseSource(int s) { return numLiveNodesPerSource.get(s) < (double) sameSourceJoinBelow * sources.get(s).size(0); } + + /** Whether any source is joined with itself (see {@link #sparseSource}). */ + private boolean anySparseSource() { + for (int s = 0; s < sources.size(); s++) { + if (sparseSource(s)) { + return true; + } + } + return false; + } /** * Constructs a new OnDiskGraphIndexCompactor for graphs without a non-fused compressed sidecar. * Equivalent to calling the 6-arg constructor with {@code sourceCompressed = null}. @@ -432,18 +443,6 @@ private static long memoryBeyondHeap() { return ((com.sun.management.OperatingSystemMXBean) bean).getTotalPhysicalMemorySize() - Runtime.getRuntime().maxMemory(); } - /** - * Sets the live fraction below which a source's own nodes are joined like another source's: a - * node of such a source gathers same-source candidates beyond the edges it retained, because - * most of those edges point at dead nodes. The default is one half; 0 never joins a source with - * itself, 1 always does (as long as any of its nodes is dead). - */ - OnDiskGraphIndexCompactor withSameSourceJoinBelow(float liveFraction) { - if (!(liveFraction >= 0f && liveFraction <= 1f)) throw new IllegalArgumentException("liveFraction must be in [0, 1]"); - this.sameSourceJoinBelow = liveFraction; - return this; - } - /** * Validates that all source indexes have compatible configurations and required features * before attempting compaction. Ensures consistent dimensions, max degrees, hierarchical @@ -651,7 +650,7 @@ private ProductQuantization trainWide() { * shard's merge adopts them. A shard's merge is otherwise an ordinary compaction whose live * nodes are the shard's, and it finds a node's same-source neighbours again, since most of * the edges a node had in its source lead to nodes of other shards - * (see {@link #withSameSourceJoinBelow}). + * (see {@link #sparseSource}). *

* Shards may become known one at a time, as when each output is written as its rows arrive: * nothing is done until the first {@link #compactShard}. Shards are compacted one at a time; @@ -715,7 +714,7 @@ private OnDiskGraphIndexCompactor shardCompactor(List shardLiveNode shard.walkQueryBytes = walkQueryBytes; shard.recordAnnounceAhead = recordAnnounceAhead; shard.walkDirectSetting = walkDirectSetting; - shard.sameSourceJoinBelow = sameSourceJoinBelow; + shard.sameSourceJoinBelow = SHARD_SAME_SOURCE_JOIN_BELOW; log.info("Sharded compaction: shard {} with {} live nodes", shardsCompacted++, shard.numTotalNodes); return shard; } @@ -1346,13 +1345,14 @@ private int searchTopK(int degree) { } /** - * The most candidates a node gathers at any level: every source can contribute a search's - * worth, the node's own included (the same-source join of a sparse source), on top of the edges - * it retained and, at level 0, the reverse candidates. + * The most candidates a node gathers at any level: every other source can contribute a search's + * worth, the node's own too when a source is joined with itself (a sparse source of a shard), on + * top of the edges it retained and, at level 0, the reverse candidates. */ private int maxCandidateSize() { int maxUpperDegree = maxUpperDegree(); - int base = searchTopK(maxDegrees.get(0)) * sources.size() + maxDegrees.get(0) + REVERSE_CANDIDATE_SLOTS; + int level0Targets = anySparseSource() ? sources.size() : sources.size() - 1; + int base = searchTopK(maxDegrees.get(0)) * level0Targets + maxDegrees.get(0) + REVERSE_CANDIDATE_SLOTS; int upper = maxUpperDegree == 0 ? 0 : searchTopK(maxUpperDegree) * sources.size() + maxUpperDegree; return Math.max(base, upper); }