refKeys = sources.get(0).getFeatures().keySet();
+ boolean sameFeatures = sources.stream()
+ .skip(1)
+ .map(s -> s.getFeatures().keySet())
+ .allMatch(refKeys::equals);
+
+ if (!sameFeatures) {
+ throw new IllegalArgumentException("Each source must have the same features");
+ }
+ if (!refKeys.contains(FeatureId.INLINE_VECTORS)) {
+ throw new IllegalArgumentException("Each source must have the INLINE_VECTORS feature");
+ }
+ }
+
+ /**
+ * 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.
+ */
+ public void compact(Path outputPath) throws FileNotFoundException {
+ QuantizationCompactionStrategy strategy = detectInlineStrategy();
+ try {
+ compactGraphImpl(outputPath, strategy);
+ releaseSourcesBeforeRefine(strategy);
+ refineCompactedGraph(outputPath, strategy);
+ } finally {
+ // Delayed until after refinement so refineCompactedGraph can read from the pre-encoded
+ // code cache appended past the projected EOF; onAfterClose unmaps it and truncates.
+ strategy.onAfterClose(outputPath);
+ if (ownsExecutor) executor.shutdown();
+ }
+ }
+
+ /**
+ * Compaction entry point for graphs that ship a non-fused compressed sidecar (e.g.
+ * {@link io.github.jbellis.jvector.quantization.PQVectors}). Writes the merged graph to
+ * {@code graphPath} and the merged compressed vectors to {@code compressedPath}.
+ *
+ * The compressor is retrained on a balanced sample of merged source vectors, then every live
+ * node is re-encoded against the new codebook. Requires that {@code sourceCompressed} was
+ * supplied to the constructor.
+ */
+ public void compact(Path graphPath, Path compressedPath) throws FileNotFoundException {
+ if (sourceCompressed == null) {
+ throw new IllegalStateException(
+ "compact(graphPath, compressedPath) requires sourceCompressed to be supplied to the constructor");
+ }
+ Objects.requireNonNull(compressedPath, "compressedPath");
+
+ // Graph compaction proceeds without fused-PQ retrain (validateCompressed forbids
+ // FUSED_PQ when sourceCompressed is set), then the sidecar is written below.
+ QuantizationCompactionStrategy inlineStrategy = detectInlineStrategy();
+ QuantizationCompactionStrategy sidecarStrategy = detectSidecarStrategy();
+ try {
+ sidecarStrategy.retrain(similarityFunction);
+ compactGraphImpl(graphPath, inlineStrategy);
+ refineCompactedGraph(graphPath, inlineStrategy);
+ sidecarStrategy.writeSidecar(compressedPath);
+ } catch (IOException e) {
+ throw new RuntimeException("Sidecar compaction failed", e);
+ } finally {
+ inlineStrategy.onAfterClose(graphPath);
+ if (ownsExecutor) executor.shutdown();
+ }
+ }
+
+ /**
+ * For compaction use. Drops the compactor's strong references to the source graphs and their
+ * per-source live-node / remapper sidecars, and tells the strategy to release its
+ * {@link CompactionContext} hold on the same. Called between {@code compactGraphImpl} and
+ * {@code refineCompactedGraph} so the source graphs' in-heap upper-layer adjacency and feature
+ * buffers become GC-eligible before refinement loads a second full graph — the peak that was
+ * OOM-ing on memory-tight hosts. The underlying {@code ReaderSupplier}s are still owned and
+ * closed by the caller (per {@link OnDiskGraphIndex#close()}'s contract), so we only drop
+ * references, never close. Not used by the sidecar {@code compact(graphPath, compressedPath)}
+ * path: {@code SidecarCompactionStrategy.writeSidecar} re-reads source vectors after refinement.
+ */
+ private void releaseSourcesBeforeRefine(QuantizationCompactionStrategy strategy) {
+ strategy.releaseSources();
+ sources = null;
+ liveNodes = null;
+ remappers = null;
+ }
+
+ /**
+ * Pick the inline-codes strategy by asking the source's fused feature (if any) for its
+ * compaction strategy. Returns {@link QuantizationCompactionStrategy#NONE} when no fused feature is
+ * present. New fused quantization types extend the compactor purely by implementing
+ * {@link FusedFeature#createCompactionStrategy}.
+ */
+ private QuantizationCompactionStrategy detectInlineStrategy() {
+ for (var feature : sources.get(0).getFeatures().values()) {
+ if (feature instanceof FusedFeature) {
+ return ((FusedFeature) feature).createCompactionStrategy(buildContext());
+ }
+ }
+ return QuantizationCompactionStrategy.NONE;
+ }
+
+ /**
+ * Pick the sidecar-codes strategy by delegating to the first {@link CompressedVectors}'
+ * own factory. Returns {@link QuantizationCompactionStrategy#NONE} when no sidecar input was supplied
+ * to the constructor. New sidecar quantization types extend the compactor purely by
+ * implementing {@link CompressedVectors#createCompactionStrategy}.
+ */
+ private QuantizationCompactionStrategy detectSidecarStrategy() {
+ if (sourceCompressed == null) {
+ return QuantizationCompactionStrategy.NONE;
+ }
+ return sourceCompressed.get(0).createCompactionStrategy(buildContext());
+ }
+
+ /** 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);
+ }
+
+ /**
+ * Internal graph-compaction body. Performs the full graph write but does not shut
+ * down {@link #executor}; the public {@code compact(...)} entry points own that lifecycle so
+ * follow-on passes (e.g. a sidecar write via {@link SidecarCompactionStrategy}) can keep using
+ * the executor.
+ *
+ * Quantization-aware steps (codebook retrain, pre-encode caches, entry-node tail records,
+ * mmap cleanup) are delegated to {@code strategy}. For sources with no inline quantization,
+ * pass {@link QuantizationCompactionStrategy#NONE} for a fully no-op strategy hook set.
+ */
+ private void compactGraphImpl(Path outputPath, QuantizationCompactionStrategy strategy) throws FileNotFoundException {
+ strategy.retrain(similarityFunction);
+
+ boolean fusedPQEnabled = strategy.writesCodesInline();
+ ProductQuantization pq = strategy.compressorAsPQ();
+ boolean compressedPrecision = fusedPQEnabled;
+ int maxBaseDegree = java.util.Collections.max(maxDegrees);
+ io.github.jbellis.jvector.graph.disk.feature.FusedFeature outputFusedFeature =
+ strategy.outputFusedFeature(maxBaseDegree);
+
+ List layerInfo = computeLayerInfoFromSources();
+ int[] entryNodeSource = resolveEntryNodeSource(); // {sourceIdx, originalOrdinal}
+ int entryNode = remappers.get(entryNodeSource[0]).oldToNew(entryNodeSource[1]);
+
+ 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)) {
+ // 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();
+ strategy.onAfterHeader(writer);
+
+ compactLevels(writer, similarityFunction, fusedPQEnabled, compressedPrecision, pq);
+
+ strategy.onAfterLevels(writer, entryNodeSource, maxDegrees);
+
+ writer.writeFooter();
+ log.info("Compaction complete: {}", outputPath);
+ } catch (IOException | ExecutionException | InterruptedException e) {
+ throw new RuntimeException(e);
+ }
+ // strategy.onAfterClose is deferred to the public compact() entry points so refinement
+ // can read from the still-mapped pre-encode cache section past the projected EOF.
+ }
+
+ /**
+ * Second pass over the just-written compacted graph. Mirrors
+ * {@link io.github.jbellis.jvector.graph.GraphIndexBuilder}'s {@code cleanup()} refinement
+ * step: when the merged graph has a hierarchy, iterates only level-1 nodes (which are also
+ * in L0); for each node, descends greedily through upper layers and beam-searches level 0
+ * carrying entry points layer-to-layer, then rewrites the L0 neighbor list (and the inline
+ * per-neighbor PQ codes for fused-PQ outputs) in place. When the merged graph has no
+ * hierarchy, falls back to iterating all live L0 nodes.
+ *
+ * The refinement search uses approximate PQ scoring with an exact reranker when fused-PQ is
+ * available (matching the cross-source path in {@code compactLevels}); otherwise it falls
+ * back to exact-only scoring backed by inline vectors.
+ *
+ * For fused-PQ outputs the per-neighbor code write is a memcpy from the
+ * {@link QuantizationCompactionStrategy#getCodeCache() pre-encode cache} keyed by new
+ * ordinal — no per-neighbor {@code encodeTo} call. The cache lives in the same file past
+ * the projected EOF and is truncated away by {@code onAfterClose} once refinement returns.
+ *
+ * Only L0 records are written. Upper-layer neighbor lists live in an in-memory map after
+ * load and have no addressable file offset, so they're left as written by compactLevels.
+ */
+ private void refineCompactedGraph(Path outputPath, QuantizationCompactionStrategy strategy) {
+ log.info("Refining compacted graph: {}", outputPath);
+ long t0 = System.nanoTime();
+
+ final int baseDegree = maxDegrees.get(0);
+ final boolean hasFusedPQ = strategy.writesCodesInline();
+ @SuppressWarnings("unchecked")
+ final VectorCompressor> compressor =
+ hasFusedPQ ? (VectorCompressor>) (VectorCompressor) strategy.compressor() : null;
+ final int pqCodeSize = hasFusedPQ ? compressor.compressedVectorSize() : 0;
+
+ final int searchTopK = Math.max(MIN_SEARCH_TOP_K,
+ baseDegree * SEARCH_TOP_K_MULTIPLIER);
+ final int beamWidth = Math.max(baseDegree, searchTopK) * BEAM_WIDTH_MULTIPLIER;
+
+ // Code cache may or may not be present; capture once so refineOneNode can take the fast path.
+ // The cache is shared across threads; refineOneNode duplicates per call (cheap; no per-thread
+ // state to track and the duplicates are tiny GC-friendly ByteBuffer wrappers).
+ final java.nio.MappedByteBuffer codeCache = hasFusedPQ ? strategy.getCodeCache() : null;
+ final int cacheCodeSize = hasFusedPQ ? strategy.getCacheCodeSize() : 0;
+
+ try (var supplier = new SimpleReader.Supplier(outputPath);
+ FileChannel fc = FileChannel.open(outputPath, StandardOpenOption.WRITE, StandardOpenOption.READ)) {
+
+ // useFooter=false because the file's logical EOF (where the v6 footer trailer sits) is
+ // before the still-attached pre-encode cache section. loadFromFooter() would seek to
+ // the actual file length and read garbage as the magic.
+ OnDiskGraphIndex mergedGraph = OnDiskGraphIndex.load(supplier, 0, false);
+
+ // Pick the iteration set: when there's a hierarchy, refine only L1 nodes (each also
+ // lives in L0, so their L0 record is what we rewrite). Mirrors GraphIndexBuilder's
+ // cleanup() which gates improveConnections() on `graph.getMaxLevel() > 0` and iterates
+ // `nodeStream(1)`. When there's no hierarchy, fall back to all L0 nodes.
+ int[] liveOrdinals;
+ int iterationLevel = mergedGraph.getMaxLevel() > 0 ? 1 : 0;
+ try (var collectView = mergedGraph.getView()) {
+ NodesIterator it = mergedGraph.getNodes(iterationLevel);
+ liveOrdinals = new int[it.size()];
+ int n = 0;
+ while (it.hasNext()) liveOrdinals[n++] = it.next();
+ }
+
+ final ThreadLocal tls = ThreadLocal.withInitial(() ->
+ new RefineScratch(mergedGraph, baseDegree, dimension, searchTopK, pqCodeSize));
+
+ ExecutorCompletionService ecs = new ExecutorCompletionService<>(executor);
+
+ int total = liveOrdinals.length;
+ int targetBatches = Math.max(taskWindowSize * 4, 16);
+ int batchSize = Math.max(1, (total + targetBatches - 1) / targetBatches);
+
+ final int[] ords = liveOrdinals;
+ final boolean fpq = hasFusedPQ;
+ final int codeSize = pqCodeSize;
+ final VectorCompressor> cmp = compressor;
+ final int bw = beamWidth;
+ final java.nio.MappedByteBuffer cache = codeCache;
+ final int cacheSz = cacheCodeSize;
+ final OnDiskGraphIndex graphRef = mergedGraph;
+
+ log.info("Refining {} live nodes at level {} (hierarchy maxLevel={}, fusedPQ={}, codeCache={})",
+ total, iterationLevel, mergedGraph.getMaxLevel(), fpq, cache != null);
+
+ int submitted = 0;
+ for (int start = 0; start < total; start += batchSize) {
+ final int s = start;
+ final int e = Math.min(start + batchSize, total);
+ ecs.submit(() -> {
+ RefineScratch scratch = tls.get();
+ for (int i = s; i < e; i++) {
+ int node = ords[i];
+ refineOneNode(node, scratch, fc, baseDegree, fpq, codeSize, cmp, bw,
+ graphRef, cache, cacheSz);
+ }
+ return e - s;
+ });
+ submitted++;
+ }
+
+ int completed = 0;
+ int nodesDone = 0;
+ int progressStep = Math.max(1, total / 10);
+ int nextProgress = progressStep;
+ while (completed < submitted) {
+ nodesDone += ecs.take().get();
+ completed++;
+ if (nodesDone >= nextProgress) {
+ log.info("Refinement progress: {}/{} nodes", nodesDone, total);
+ nextProgress += progressStep;
+ }
+ }
+
+ // Per-thread scratches live in worker-thread ThreadLocals; closing the supplier in
+ // try-with-resources tears down the underlying mapping, so any later access would
+ // fail anyway. The references will be GC'd when the worker threads die.
+ } catch (IOException | InterruptedException | ExecutionException e) {
+ throw new RuntimeException("Refinement failed", e);
+ }
+
+ log.info("Refinement complete in {} ms", (System.nanoTime() - t0) / 1_000_000);
+ }
+
+ /**
+ * Refines a single node by mirroring {@code GraphIndexBuilder.improveConnections}:
+ * descend greedily through upper layers carrying entry points layer-to-layer, then beam
+ * search at L0. Diversity selection + in-place L0 record rewrite happen at the end.
+ *
+ * The {@code SearchScoreProvider} uses approximate PQ scoring with an exact reranker when
+ * fused-PQ is available; otherwise exact-only via the inline-vector reranker. Diversity
+ * always runs over exact scores (so we rescore approximate results after the L0 beam).
+ */
+ private void refineOneNode(int node,
+ RefineScratch scratch,
+ FileChannel fc,
+ int baseDegree,
+ boolean hasFusedPQ,
+ int pqCodeSize,
+ VectorCompressor> compressor,
+ int beamWidth,
+ OnDiskGraphIndex mergedGraph,
+ java.nio.MappedByteBuffer codeCache,
+ int cacheCodeSize) {
+ OnDiskGraphIndex.View view = scratch.view;
+ view.getVectorInto(node, scratch.queryVec, 0);
+
+ // Build score provider for this query. Reranker reads the candidate's inline FP vector
+ // (via view.getVectorInto into a worker-private tmp) and computes exact similarity.
+ ScoreFunction.ExactScoreFunction reranker = node2 -> {
+ view.getVectorInto(node2, scratch.tmpVec, 0);
+ return similarityFunction.compare(scratch.queryVec, scratch.tmpVec);
+ };
+ SearchScoreProvider ssp;
+ if (hasFusedPQ) {
+ FusedPQ fpq = (FusedPQ) mergedGraph.getFeatures().get(FeatureId.FUSED_PQ);
+ var asf = fpq.approximateScoreFunctionFor(scratch.queryVec, similarityFunction, view, reranker);
+ ssp = new DefaultSearchScoreProvider(asf, reranker);
+ } else {
+ ssp = new DefaultSearchScoreProvider(reranker);
+ }
+
+ Bits excludeSelf = idx -> idx != node;
+
+ // Per-layer descent. Mirrors GraphSearcher.internalSearch: greedy single-best through
+ // each upper layer, then a beam search at layer 0. Entry points carry forward via
+ // setEntryPointsFromPreviousLayer so the L0 beam starts from the best-known region
+ // rather than the global entry node — much cheaper than the previous full search().
+ GraphSearcher gs = scratch.searcher;
+ var entry = view.entryNode();
+ gs.initializeInternal(ssp, entry, excludeSelf);
+ for (int lvl = entry.level; lvl > 0; lvl--) {
+ gs.searchOneLayer(ssp, 1, 0f, lvl, excludeSelf);
+ gs.setEntryPointsFromPreviousLayer();
+ }
+ gs.searchOneLayer(ssp, beamWidth, 0f, 0, excludeSelf);
+
+ // Collect candidates. Start with the node's existing L0 edges (rescored exact) so
+ // refinement never drops an edge that the search happened to miss — matches the
+ // existing+search union pattern from GraphIndexBuilder.insertDiverse.
+ scratch.candSize = 0;
+ var existing = view.getNeighborsIterator(0, node);
+ while (existing.hasNext()) {
+ int nb = existing.nextInt();
+ if (nb == node) continue;
+ view.getVectorInto(nb, scratch.tmpVec, 0);
+ scratch.candNode[scratch.candSize] = nb;
+ scratch.candScore[scratch.candSize] = similarityFunction.compare(scratch.queryVec, scratch.tmpVec);
+ scratch.candSize++;
+ }
+ // Pull search results from approximateResults. When fused-PQ is on the scores there are
+ // approximate; rescore exact for correct diversity comparison against existing edges.
+ final boolean rescore = hasFusedPQ;
+ gs.approximateResults().foreach((nb, approxScore) -> {
+ if (nb == node) return;
+ for (int k = 0; k < scratch.candSize; k++) {
+ if (scratch.candNode[k] == nb) return; // de-dupe against existing edges
+ }
+ if (scratch.candSize >= scratch.candNode.length) return;
+ float s;
+ if (rescore) {
+ view.getVectorInto(nb, scratch.tmpVec, 0);
+ s = similarityFunction.compare(scratch.queryVec, scratch.tmpVec);
+ } else {
+ s = approxScore;
+ }
+ scratch.candNode[scratch.candSize] = nb;
+ scratch.candScore[scratch.candSize] = s;
+ scratch.candSize++;
+ });
+
+ if (scratch.candSize == 0) {
+ // No live neighbors found — leave the existing record alone.
+ return;
+ }
+
+ // Sort candidates by descending score.
+ int[] order = scratch.order;
+ for (int k = 0; k < scratch.candSize; k++) order[k] = k;
+ sortOrderByScoreDesc(order, scratch.candScore, scratch.candSize);
+
+ // Vamana diversity selection with progressively-relaxed alpha.
+ int selectedSize = retainDiverseSingleSource(
+ view, order, scratch.candNode, scratch.candScore, scratch.candSize,
+ baseDegree, scratch.selectedNodes, scratch.selectedVecs, scratch.tmpVec);
+
+ // Build the trailing-section bytes (PQ codes block — if any — followed by count + neighbors).
+ ByteBuffer rec = scratch.recordBuffer;
+ rec.clear();
+
+ long writeOffset;
+ if (hasFusedPQ) {
+ // PQ codes block sits between the inline vector and the neighbor count.
+ writeOffset = view.offsetFor(node, FeatureId.FUSED_PQ);
+ if (codeCache != null) {
+ // Memcpy from the pre-encoded cache (indexed by new ordinal). Avoids one FP
+ // vector read AND one PQ encode per selected neighbor. duplicate() gives this
+ // call its own position cursor without racing other workers.
+ ByteBuffer cacheView = codeCache.duplicate();
+ byte[] codeBuf = scratch.pqCodeBytes;
+ for (int k = 0; k < selectedSize; k++) {
+ int newOrd = scratch.selectedNodes[k];
+ cacheView.position(newOrd * cacheCodeSize);
+ cacheView.get(codeBuf, 0, cacheCodeSize);
+ rec.put(codeBuf, 0, cacheCodeSize);
+ }
+ } else {
+ // Fallback: re-encode from the selected neighbor's inline vector. Same as before
+ // the cache-reuse optimization. Used when the cache wasn't built (graph too large
+ // for a single mapping, or pre-encode failure).
+ ByteSequence codeOut = scratch.pqCode;
+ for (int k = 0; k < selectedSize; k++) {
+ view.getVectorInto(scratch.selectedNodes[k], scratch.tmpVec, 0);
+ codeOut.zero();
+ compressor.encodeTo(scratch.tmpVec, codeOut);
+ for (int b = 0; b < pqCodeSize; b++) {
+ rec.put(codeOut.get(b));
+ }
+ }
+ }
+ // Pad remaining slots with zero codes (matches CompactWriter's zeroPQ behavior).
+ int padSlots = baseDegree - selectedSize;
+ for (int s = 0; s < padSlots; s++) {
+ for (int b = 0; b < pqCodeSize; b++) rec.put((byte) 0);
+ }
+ } else {
+ writeOffset = view.neighborsOffsetFor(0, node);
+ }
+
+ // Neighbor count + ordinals (-1 padding for unused slots).
+ rec.putInt(selectedSize);
+ for (int k = 0; k < selectedSize; k++) rec.putInt(scratch.selectedNodes[k]);
+ for (int k = selectedSize; k < baseDegree; k++) rec.putInt(-1);
+
+ rec.flip();
+ try {
+ while (rec.hasRemaining()) {
+ int n = fc.write(rec, writeOffset);
+ writeOffset += n;
+ }
+ } catch (IOException e) {
+ throw new UncheckedIOException(e);
+ }
+ }
+
+ /**
+ * Single-source Vamana diversity selection. Mirrors {@link CompactVamanaDiversityProvider}
+ * but operates on one merged graph rather than per-source views, so candidates are bare
+ * (node, score) pairs.
+ *
+ * @return the number of selected neighbors written into {@code selectedNodes}.
+ */
+ private int retainDiverseSingleSource(OnDiskGraphIndex.View view,
+ int[] order, int[] candNode, float[] candScore, int candSize,
+ int maxDegree, int[] selectedNodes,
+ VectorFloat[] selectedVecs, VectorFloat tmp) {
+ if (candSize == 0) return 0;
+ int nSelected = 0;
+ float currentAlpha = 1.0f;
+ final float alpha = 1.2f;
+ while (currentAlpha <= alpha + 1E-6 && nSelected < maxDegree) {
+ for (int i = 0; i < candSize && nSelected < maxDegree; i++) {
+ int ci = order[i];
+ int cNode = candNode[ci];
+ float cScore = candScore[ci];
+
+ view.getVectorInto(cNode, tmp, 0);
+
+ boolean diverse = true;
+ for (int j = 0; j < nSelected; j++) {
+ if (selectedNodes[j] == cNode) { diverse = false; break; }
+ if (similarityFunction.compare(tmp, selectedVecs[j]) > cScore * currentAlpha) {
+ diverse = false;
+ break;
+ }
+ }
+ if (diverse) {
+ selectedNodes[nSelected] = cNode;
+ selectedVecs[nSelected].copyFrom(tmp, 0, 0, tmp.length());
+ nSelected++;
+ }
+ }
+ currentAlpha += DIVERSITY_ALPHA_STEP;
+ }
+ return nSelected;
+ }
+
+ /** Per-thread scratch space for refinement. One per worker thread, populated lazily via ThreadLocal. */
+ private static final class RefineScratch {
+ final OnDiskGraphIndex.View view;
+ final GraphSearcher searcher;
+ final VectorFloat queryVec;
+ final VectorFloat tmpVec;
+ final int[] candNode;
+ final float[] candScore;
+ final int[] order;
+ int candSize;
+ final int[] selectedNodes;
+ final VectorFloat[] selectedVecs;
+ final ByteSequence pqCode;
+ // Heap byte buffer for memcpy from the precomputed code cache into the record buffer.
+ final byte[] pqCodeBytes;
+ final ByteBuffer recordBuffer;
+
+ RefineScratch(OnDiskGraphIndex mergedGraph, int baseDegree, int dimension, int searchTopK, int pqCodeSize) {
+ this.view = mergedGraph.getView();
+ this.searcher = new GraphSearcher(mergedGraph);
+ this.searcher.usePruning(false);
+ this.queryVec = vectorTypeSupport.createFloatVector(dimension);
+ this.tmpVec = vectorTypeSupport.createFloatVector(dimension);
+ // Candidates = existing neighbors (up to baseDegree) ∪ search results (up to searchTopK).
+ int cap = searchTopK + baseDegree + 16;
+ this.candNode = new int[cap];
+ this.candScore = new float[cap];
+ this.order = new int[cap];
+ this.selectedNodes = new int[baseDegree];
+ this.selectedVecs = new VectorFloat[baseDegree];
+ for (int i = 0; i < baseDegree; i++) {
+ this.selectedVecs[i] = vectorTypeSupport.createFloatVector(dimension);
+ }
+ this.pqCode = pqCodeSize > 0 ? vectorTypeSupport.createByteSequence(pqCodeSize) : null;
+ this.pqCodeBytes = pqCodeSize > 0 ? new byte[pqCodeSize] : null;
+ // Trailing section to rewrite: optional PQ codes block + count + neighbor ids.
+ int recordBytes = (pqCodeSize > 0 ? baseDegree * pqCodeSize : 0) + Integer.BYTES + baseDegree * Integer.BYTES;
+ this.recordBuffer = ByteBuffer.allocate(recordBytes).order(java.nio.ByteOrder.BIG_ENDIAN);
+ }
+ }
+
+ /**
+ * Returns {sourceIdx, originalOrdinal} for the entry node of the compacted graph.
+ * The chosen node must exist at maxLevel (since the on-disk format sets entryNode.level =
+ * maxLevel). Prefers the designated entry node of any source whose maxLevel equals the global
+ * maxLevel; if all such entry nodes are deleted, falls back to the first live node at maxLevel
+ * across all sources.
+ */
+ private int[] resolveEntryNodeSource() {
+ int maxLevel = sources.stream().mapToInt(OnDiskGraphIndex::getMaxLevel).max().orElse(0);
+
+ // 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) {
+ int originalEntry = sources.get(s).getView().entryNode().node;
+ if (liveNodes.get(s).get(originalEntry)) {
+ return new int[]{s, originalEntry};
+ }
+ }
+ }
+
+ // Entry nodes were all deleted: scan for any live node that exists at maxLevel.
+ for (int s = 0; s < sources.size(); s++) {
+ if (sources.get(s).getMaxLevel() < maxLevel) continue;
+ NodesIterator it = sources.get(s).getNodes(maxLevel);
+ while (it.hasNext()) {
+ int node = it.next();
+ if (liveNodes.get(s).get(node)) {
+ return new int[]{s, node};
+ }
+ }
+ }
+
+ throw new IllegalStateException("No live nodes found at maxLevel=" + maxLevel);
+ }
+
+ /**
+ * Compacts all hierarchical levels of the graph, processing each level in batches.
+ * For level 0 (base layer), writes inline vectors and neighbors. For upper layers,
+ * writes only graph structure and optional PQ codes.
+ */
+ private void compactLevels(CompactWriter writer,
+ VectorSimilarityFunction similarityFunction,
+ boolean fusedPQEnabled,
+ boolean compressedPrecision,
+ 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);
+ 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));
+ final ThreadLocal threadLocalScratch = ThreadLocal.withInitial(() ->
+ new Scratch(maxCandidateSize, scratchDegree, dimension, sources, pq)
+ );
+
+ for (int level = 0; level < maxDegrees.size(); level++) {
+ List batches = buildBatches(level);
+ int searchTopK = Math.max(MIN_SEARCH_TOP_K, ((maxDegrees.get(level) + sources.size() - 1) / sources.size()) * SEARCH_TOP_K_MULTIPLIER);
+ int beamWidth = Math.max(maxDegrees.get(level), searchTopK) * BEAM_WIDTH_MULTIPLIER;
+
+ CompactionParams params = new CompactionParams(fusedPQEnabled, compressedPrecision, searchTopK, beamWidth, pq);
+
+ if (level == 0) {
+ log.info("Compacting level 0 (base layer)");
+
+ ExecutorCompletionService> ecs =
+ new ExecutorCompletionService<>(executor);
+
+ java.util.function.Consumer submitOne = (bs) -> {
+ ecs.submit(() -> {
+ Scratch scratch = threadLocalScratch.get();
+ return computeBaseBatch(writer, bs, scratch, params);
+ });
+ };
+
+ var wropts = EnumSet.of(StandardOpenOption.WRITE, StandardOpenOption.READ);
+ try (FileChannel fc = FileChannel.open(writer.getOutputPath(), wropts)) {
+
+ runBatchesWithBackpressure(
+ batches,
+ ecs,
+ submitOne,
+ (results) -> {
+ try {
+ for (WriteResult r : results) {
+ ByteBuffer b = r.data;
+ long pos = r.fileOffset;
+ while (b.hasRemaining()) {
+ int n = fc.write(b, pos);
+ pos += n;
+ }
+ }
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ );
+ }
+
+ writer.offsetAfterInline();
+
+ } else {
+ final int lvl = level;
+ log.info("Compacting upper layer {}", level);
+
+ ExecutorCompletionService> ecs =
+ new ExecutorCompletionService<>(executor);
+
+ java.util.function.Consumer submitOne = (bs) -> {
+ ecs.submit(() -> {
+ Scratch scratch = threadLocalScratch.get();
+ return computeUpperBatchForLevel(bs, lvl, scratch, params);
+ });
+ };
+
+ runBatchesWithBackpressure(
+ batches,
+ ecs,
+ submitOne,
+ (results) -> {
+ try {
+ for (UpperLayerWriteResult r : results) {
+ writer.writeUpperLayerNode(
+ lvl,
+ r.ordinal,
+ r.neighbors,
+ r.pqCode
+ );
+ }
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ );
+ }
+ }
+
+ Scratch s = threadLocalScratch.get();
+ s.close();
+ threadLocalScratch.remove();
+ }
+
+ /**
+ * Divides nodes at a given level across all source indexes into processing batches
+ * for parallel execution. Each batch contains a subset of nodes from one source.
+ */
+ private List buildBatches(int level) {
+ List batches = new ArrayList<>();
+
+ for (int s = 0; s < sources.size(); ++s) {
+ var source = sources.get(s);
+ if (level > source.getMaxLevel()) continue;
+ NodesIterator sourceNodes = source.getNodes(level);
+ int numNodes = sourceNodes.size();
+ int[] nodes = new int[numNodes];
+ int i = 0;
+ while (sourceNodes.hasNext()) {
+ nodes[i++] = sourceNodes.next();
+ }
+
+ int numBatches = max(TARGET_BATCHES_PER_SOURCE, (numNodes + TARGET_NODES_PER_BATCH - 1) / TARGET_NODES_PER_BATCH);
+ if (numBatches > numNodes) numBatches = numNodes;
+ int batchSize = (numNodes + numBatches - 1) / numBatches;
+ for (int b = 0; b < numBatches; ++b) {
+ int start = min(numNodes, batchSize * b);
+ int end = min(numNodes, batchSize * (b + 1));
+ batches.add(new BatchSpec(s, nodes, start, end));
+ }
+ }
+
+ return batches;
+ }
+
+ /**
+ * Processes a batch of base layer (level 0) nodes from one source index. For each live node,
+ * gathers candidates from all sources, applies diversity selection, and creates write results
+ * containing the full node record data.
+ */
+ private List computeBaseBatch(CompactWriter writer,
+ BatchSpec bs,
+ Scratch scratch,
+ CompactionParams params) throws IOException {
+
+ List out = new ArrayList<>(bs.end - bs.start);
+
+ for (int i = bs.start; i < bs.end; i++) {
+ int node = bs.nodes[i];
+ if (!liveNodes.get(bs.sourceIdx).get(node)) continue;
+
+ out.add(processBaseNode(node, bs.sourceIdx, scratch, writer, params));
+ }
+
+ return out;
+ }
+
+ /**
+ * 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).
+ */
+ private List computeUpperBatchForLevel(
+ BatchSpec bs,
+ int level,
+ Scratch scratch,
+ CompactionParams params
+ ) {
+ List results =
+ new ArrayList<>(bs.end - bs.start);
+
+ for (int i = bs.start; i < bs.end; i++) {
+ int node = bs.nodes[i];
+
+ if (!liveNodes.get(bs.sourceIdx).get(node)) continue;
+
+ results.add(processUpperNode(node, bs.sourceIdx, level, scratch, params));
+ }
+
+ return results;
+ }
+
+ /**
+ * Processes a single base layer node: retrieves its vector, gathers diverse candidates from
+ * all sources, selects best neighbors using diversity criteria, remaps ordinals, and returns
+ * the complete write result for this node.
+ */
+ private WriteResult processBaseNode(
+ int node,
+ int sourceIdx,
+ Scratch scratch,
+ CompactWriter writer,
+ CompactionParams params
+ ) throws IOException {
+
+ var sourceView = (OnDiskGraphIndex.View) scratch.gs[sourceIdx].getView();
+ sourceView.getVectorInto(node, scratch.baseVec, 0);
+
+ int candSize = gatherCandidates(node, 0, sourceIdx, scratch, scratch.baseVec, params);
+
+ int[] order = IntStream.range(0, candSize).toArray();
+ sortOrderByScoreDesc(order, scratch.candScore, candSize);
+
+ var selected = scratch.selectedCache;
+
+ new CompactVamanaDiversityProvider(similarityFunction, 1.2f)
+ .retainDiverse(
+ scratch.candSrc,
+ scratch.candNode,
+ scratch.candScore,
+ order,
+ candSize,
+ maxDegrees.get(0),
+ selected,
+ scratch.tmpVec,
+ scratch.gs
+ );
+
+ // remap
+ for (int k = 0; k < selected.size; k++) {
+ selected.nodes[k] =
+ remappers.get(selected.sourceIdx[k])
+ .oldToNew(selected.nodes[k]);
+ }
+
+ int newOrdinal = remappers.get(sourceIdx).oldToNew(node);
+
+ return writer.writeInlineNodeRecord(
+ newOrdinal,
+ scratch.baseVec,
+ selected,
+ scratch.pqCode
+ );
+ }
+
+ /**
+ * Processes a single upper layer node: similar to base layer processing but only returns
+ * graph structure (ordinal and neighbors) and optional PQ encoding for level 1.
+ */
+ private UpperLayerWriteResult processUpperNode(
+ int node,
+ int sourceIdx,
+ int level,
+ Scratch scratch,
+ CompactionParams params
+ ) {
+ var sourceView = (OnDiskGraphIndex.View) scratch.gs[sourceIdx].getView();
+ sourceView.getVectorInto(node, scratch.baseVec, 0);
+
+ int candSize = gatherCandidates(node, level, sourceIdx, scratch, scratch.baseVec, params);
+
+ int[] order = IntStream.range(0, candSize).toArray();
+ sortOrderByScoreDesc(order, scratch.candScore, candSize);
+
+ var selected = scratch.selectedCache;
+
+ new CompactVamanaDiversityProvider(similarityFunction, 1.2f)
+ .retainDiverse(
+ scratch.candSrc,
+ scratch.candNode,
+ scratch.candScore,
+ order,
+ candSize,
+ maxDegrees.get(level),
+ selected,
+ scratch.tmpVec,
+ scratch.gs
+ );
+
+ // remap
+ for (int k = 0; k < selected.size; k++) {
+ selected.nodes[k] =
+ remappers.get(selected.sourceIdx[k])
+ .oldToNew(selected.nodes[k]);
+ }
+
+ int newOrdinal = remappers.get(sourceIdx).oldToNew(node);
+
+ ByteSequence pqCode = maybeEncodePQ(level, scratch, params);
+
+ return new UpperLayerWriteResult(newOrdinal, selected, pqCode);
+ }
+
+ /**
+ * Encodes a vector using Product Quantization if enabled and the level is 1.
+ * Returns null otherwise.
+ */
+ private ByteSequence maybeEncodePQ(int level, Scratch scratch, CompactionParams params) {
+ if (!params.fusedPQEnabled || level != 1) {
+ return null;
+ }
+
+ scratch.pqCode.zero();
+ params.pq.encodeTo(scratch.baseVec, scratch.pqCode);
+ return scratch.pqCode.copy();
+ }
+
+ /**
+ * Collects neighbor candidates for a node from all source indexes. For the source containing
+ * the node, uses existing neighbors; for other sources, performs graph search. Returns the
+ * total number of candidates gathered.
+ */
+ private int gatherCandidates(
+ int node,
+ int level,
+ int sourceIdx,
+ Scratch scratch,
+ VectorFloat baseVec,
+ CompactionParams params
+ ) {
+ int candSize = 0;
+
+ for (int ss = 0; ss < sources.size(); ss++) {
+ var searchView = (OnDiskGraphIndex.View) scratch.gs[ss].getView();
+ var indexAlive = liveNodes.get(ss);
+
+ if (ss == sourceIdx) {
+ candSize = gatherFromSameSource(node, level, ss, searchView, indexAlive,
+ baseVec, scratch, candSize);
+ } else {
+ candSize = gatherFromOtherSource(node, level, ss, searchView, indexAlive,
+ baseVec, scratch, candSize, params);
+ }
+ }
+
+ return candSize;
+ }
+
+ /**
+ * Gathers candidates from the same source index that contains the node.
+ * Simply iterates through existing neighbors.
+ */
+ private int gatherFromSameSource(int node, int level, int sourceIdx,
+ OnDiskGraphIndex.View searchView, FixedBitSet indexAlive,
+ VectorFloat baseVec, Scratch scratch, int candSize) {
+ var it = searchView.getNeighborsIterator(level, node);
+ while (it.hasNext()) {
+ int nb = it.nextInt();
+ if (!indexAlive.get(nb)) continue;
+
+ searchView.getVectorInto(nb, scratch.tmpVec, 0);
+
+ scratch.candSrc[candSize] = sourceIdx;
+ scratch.candNode[candSize] = nb;
+ scratch.candScore[candSize] = similarityFunction.compare(baseVec, scratch.tmpVec);
+ candSize++;
+ }
+ return candSize;
+ }
+
+ /**
+ * Gathers candidates from a different source index via graph search.
+ */
+ private int gatherFromOtherSource(int node, int level, int sourceIdx,
+ OnDiskGraphIndex.View searchView, FixedBitSet indexAlive,
+ VectorFloat baseVec, Scratch scratch, int candSize,
+ CompactionParams params) {
+ SearchScoreProvider ssp = buildCrossSourceScoreProvider(
+ params.compressedPrecision,
+ sources.get(sourceIdx),
+ searchView,
+ baseVec,
+ scratch.tmpVec,
+ similarityFunction
+ );
+
+ if (level == 0) {
+ SearchResult results = scratch.gs[sourceIdx].search(
+ ssp, params.searchTopK, params.beamWidth, 0f, 0f, indexAlive
+ );
+
+ for (var r : results.getNodes()) {
+ scratch.candSrc[candSize] = sourceIdx;
+ scratch.candNode[candSize] = r.node;
+ scratch.candScore[candSize] =
+ params.fusedPQEnabled
+ ? rescore(searchView, r.node, baseVec, scratch.tmpVec)
+ : r.score;
+ candSize++;
+ }
+ } else {
+ var entry = searchView.entryNode();
+ if (level > entry.level) return candSize;
+ scratch.gs[sourceIdx].initializeInternal(ssp, entry, Bits.ALL);
+
+ // Descend greedily through levels above the target level, so the search at
+ // `level` starts from the best-known region rather than the global entry node.
+ // This mirrors how GraphSearcher.searchInternal navigates the hierarchy.
+ for (int l = entry.level; l > level; l--) {
+ scratch.gs[sourceIdx].searchOneLayer(ssp, 1, 0f, l, Bits.ALL);
+ scratch.gs[sourceIdx].setEntryPointsFromPreviousLayer();
+ }
+
+ scratch.gs[sourceIdx].searchOneLayer(
+ ssp, params.searchTopK, 0f, level, indexAlive
+ );
+
+ int prev_candSize = candSize;
+ candSize = appendApproximateResults(
+ scratch.gs[sourceIdx].approximateResults(),
+ sourceIdx,
+ scratch,
+ candSize
+ );
+
+ if (params.fusedPQEnabled) {
+ for (int i = prev_candSize; i < candSize; i++) {
+ scratch.candScore[i] = rescore(
+ searchView,
+ scratch.candNode[i],
+ baseVec,
+ scratch.tmpVec
+ );
+ }
+ }
+ }
+
+ return candSize;
+ }
+
+ /**
+ * Recomputes exact similarity score between the base vector and a node's vector,
+ * used to refine approximate PQ-based search results.
+ */
+ private float rescore(OnDiskGraphIndex.View view,
+ int node,
+ VectorFloat base,
+ VectorFloat tmp) {
+ view.getVectorInto(node, tmp, 0);
+ return similarityFunction.compare(base, tmp);
+ }
+
+ /**
+ * Executes batches with controlled concurrency using a sliding window approach. Prevents
+ * overwhelming memory by limiting the number of in-flight tasks while maintaining high
+ * throughput via the completion service.
+ */
+ private void runBatchesWithBackpressure(
+ List batches,
+ ExecutorCompletionService> ecs,
+ java.util.function.Consumer submitOne,
+ java.util.function.Consumer> onComplete
+ ) throws InterruptedException, ExecutionException {
+
+ final int total = batches.size();
+ int nextToSubmit = 0;
+ int inFlight = 0;
+
+ // initial window
+ while (inFlight < taskWindowSize && nextToSubmit < total) {
+ submitOne.accept(batches.get(nextToSubmit++));
+ inFlight++;
+ }
+
+ int completed = 0;
+ while (completed < total) {
+ List results = ecs.take().get();
+ onComplete.accept(results);
+
+ completed++;
+ inFlight--;
+
+ if (nextToSubmit < total) {
+ submitOne.accept(batches.get(nextToSubmit++));
+ inFlight++;
+ }
+ if (completed % 10 == 0) {
+ log.debug("Compaction I/O progress: {}/{} batches written to disk", completed, total);
+ }
+ }
+ }
+
+ /**
+ * Appends search results from a NodeQueue to the candidate arrays, returning the updated
+ * candidate count.
+ */
+ private int appendApproximateResults(NodeQueue queue,
+ int sourceIdx,
+ Scratch scratch,
+ int candSize) {
+ final int ss = sourceIdx;
+ final int[] idx = new int[] { candSize };
+
+ queue.foreach((nb, score) -> {
+ scratch.candSrc[idx[0]] = ss;
+ scratch.candNode[idx[0]] = nb;
+ scratch.candScore[idx[0]] = score;
+ idx[0]++;
+ });
+
+ return idx[0];
+ }
+
+ /**
+ * Computes layer metadata for the compacted graph by counting live nodes at each level
+ * across all source indexes.
+ */
+ private List computeLayerInfoFromSources() {
+ int maxLevel = sources.stream().mapToInt(OnDiskGraphIndex::getMaxLevel).max().orElse(0);
+ List layerInfo = new ArrayList<>(maxLevel + 1);
+ for (int level = 0; level <= maxLevel; level++) {
+ int count = 0;
+ for (int s = 0; s < sources.size(); s++) {
+ if (level > sources.get(s).getMaxLevel()) continue;
+ NodesIterator it = sources.get(s).getNodes(level);
+ FixedBitSet alive = liveNodes.get(s);
+ while (it.hasNext()) {
+ int node = it.next();
+ if (alive.get(node)) count++;
+ }
+ }
+ layerInfo.add(new CommonHeader.LayerInfo(count, maxDegrees.get(level)));
+ }
+ return layerInfo;
+ }
+
+ /**
+ * Creates a score provider for searching across different source indexes. Uses approximate
+ * PQ-based scoring if compressedPrecision is enabled, otherwise uses exact scoring.
+ */
+ private SearchScoreProvider buildCrossSourceScoreProvider(boolean compressedPrecision,
+ OnDiskGraphIndex searchSource,
+ OnDiskGraphIndex.View searchView,
+ VectorFloat baseVec,
+ VectorFloat tmpVec,
+ VectorSimilarityFunction similarityFunction) {
+ if (compressedPrecision) {
+ ScoreFunction.ExactScoreFunction reranker =
+ node2 -> {
+ searchView.getVectorInto(node2, tmpVec, 0);
+ return similarityFunction.compare(baseVec, tmpVec);
+ };
+ var asf = ((FusedPQ) searchSource.getFeatures().get(FeatureId.FUSED_PQ)).approximateScoreFunctionFor(baseVec, similarityFunction, searchView, reranker);
+
+ return new DefaultSearchScoreProvider(asf);
+ }
+
+ var sf = new ScoreFunction.ExactScoreFunction() {
+ @Override
+ public float similarityTo(int node2) {
+ searchView.getVectorInto(node2, tmpVec, 0);
+ return similarityFunction.compare(baseVec, tmpVec);
+ }
+ };
+ return new DefaultSearchScoreProvider(sf);
+ }
+
+ /**
+ * Estimates the RAM usage of this compactor instance.
+ * Accounts for data structures used during compaction including bitsets, remappers,
+ * executor overhead, and per-thread scratch space.
+ */
+ @Override
+ public long ramBytesUsed() {
+ int OH = RamUsageEstimator.NUM_BYTES_OBJECT_HEADER;
+ int REF = RamUsageEstimator.NUM_BYTES_OBJECT_REF;
+
+ // Shallow size of this object (header + fields)
+ // Current fields: sources, liveNodes, numLiveNodesPerSource, remappers, maxDegrees,
+ // dimension(int), maxOrdinal(int), numTotalNodes(int),
+ // ownsExecutor(boolean), executor, taskWindowSize(int), similarityFunction
+ long size = OH + 8L * REF + Integer.BYTES * 4 + 1;
+
+ // liveNodes: FixedBitSet per source. May be null after releaseSourcesBeforeRefine().
+ if (liveNodes != null) {
+ for (var entry : liveNodes) {
+ size += entry.ramBytesUsed();
+ }
+ }
+
+ // numLiveNodesPerSource: ArrayList of Integers
+ size += OH + REF + (long) numLiveNodesPerSource.size() * (OH + Integer.BYTES);
+
+ // remappers: each MapMapper holds an oldToNew HashMap and newToOld Int2IntHashMap.
+ // May be null after releaseSourcesBeforeRefine().
+ if (remappers != null) {
+ for (var mapper : remappers) {
+ // Object overhead + two maps with int key/value pairs
+ // HashMap entry: ~32 bytes each; Int2IntHashMap: ~16 bytes per entry
+ if (mapper instanceof OrdinalMapper.MapMapper) {
+ // rough estimate: the mapper stores two maps over all mapped ordinals
+ size += OH + (long) (maxOrdinal + 1) * 48;
+ }
+ }
+ }
+
+ // maxDegrees: small list of integers
+ size += OH + REF + (long) maxDegrees.size() * (OH + Integer.BYTES);
+
+ // executor: ForkJoinPool overhead (if owned)
+ // Estimate based on number of threads
+ int numThreads = ownsExecutor ? Runtime.getRuntime().availableProcessors() : taskWindowSize;
+ if (ownsExecutor) {
+ size += OH + REF;
+ }
+
+ // Scratch space: ThreadLocal instances (one per active thread)
+ // Each Scratch contains:
+ // - candSrc, candNode, candScore arrays
+ // - SelectedVecCache (with its own arrays and vector copies)
+ // - tmpVec, baseVec (VectorFloat instances)
+ // - GraphSearcher array (one per source)
+ // - pqCode ByteSequence
+ size += estimateScratchSpacePerThread() * numThreads;
+
+ return size;
+ }
+
+ /**
+ * Estimates the RAM usage of a single Scratch instance.
+ */
+ 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);
+ 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));
+
+ long scratchSize = OH + 6L * REF;
+
+ // candSrc, candNode, candScore arrays
+ scratchSize += (long) maxCandidateSize * Integer.BYTES; // candSrc
+ scratchSize += (long) maxCandidateSize * Integer.BYTES; // candNode
+ scratchSize += (long) maxCandidateSize * Float.BYTES; // candScore
+
+ // SelectedVecCache
+ scratchSize += OH + 5L * REF + Integer.BYTES; // SelectedVecCache object
+ scratchSize += (long) scratchDegree * Integer.BYTES; // sourceIdx array
+ scratchSize += (long) scratchDegree * REF; // views array
+ scratchSize += (long) scratchDegree * Integer.BYTES; // nodes array
+ scratchSize += (long) scratchDegree * Float.BYTES; // scores array
+ scratchSize += (long) scratchDegree * REF; // vecs array
+ scratchSize += (long) scratchDegree * (OH + dimension * Float.BYTES); // VectorFloat instances
+
+ // tmpVec and baseVec
+ scratchSize += 2L * (OH + dimension * Float.BYTES);
+
+ // GraphSearcher array (one per source)
+ scratchSize += (long) sources.size() * REF;
+ // Each GraphSearcher has internal state - rough estimate
+ scratchSize += (long) sources.size() * (OH + 10L * REF);
+
+ // Per-thread scratch ByteSequence holding one code's worth of bytes, for each fused
+ // feature carried by the graph. Generalized over fused types so new quantizations
+ // (e.g. FUSED_ASH) don't need an edit here.
+ for (var feature : sources.get(0).getFeatures().values()) {
+ if (feature instanceof FusedFeature) {
+ scratchSize += OH + ((FusedFeature) feature).codeSize();
+ }
+ }
+
+ return scratchSize;
+ }
+
+ /**
+ * Encapsulates common parameters used throughout the compaction process.
+ */
+ private static final class CompactionParams {
+ final boolean fusedPQEnabled;
+ final boolean compressedPrecision;
+ final int searchTopK;
+ final int beamWidth;
+ final ProductQuantization pq;
+
+ CompactionParams(boolean fusedPQEnabled, boolean compressedPrecision,
+ int searchTopK, int beamWidth, ProductQuantization pq) {
+ this.fusedPQEnabled = fusedPQEnabled;
+ this.compressedPrecision = compressedPrecision;
+ this.searchTopK = searchTopK;
+ this.beamWidth = beamWidth;
+ this.pq = pq;
+ }
+ }
+
+ /**
+ * Sorts an index array by descending score values using quicksort.
+ */
+ private static void sortOrderByScoreDesc(int[] order, float[] score, int size) {
+ quicksort(order, score, 0, size - 1);
+ }
+
+ /**
+ * Tail-recursive quicksort implementation for sorting by score in descending order.
+ */
+ private static void quicksort(int[] order, float[] score, int lo, int hi) {
+ while (lo < hi) {
+ int p = partition(order, score, lo, hi);
+ // recurse smaller side first (limits stack)
+ if (p - lo < hi - p) {
+ quicksort(order, score, lo, p - 1);
+ lo = p + 1;
+ } else {
+ quicksort(order, score, p + 1, hi);
+ hi = p - 1;
+ }
+ }
+ }
+
+ /**
+ * Partitions the order array for quicksort using descending score comparison.
+ */
+ private static int partition(int[] order, float[] score, int lo, int hi) {
+ float pivot = score[order[hi]];
+ int i = lo;
+ for (int j = lo; j < hi; j++) {
+ if (score[order[j]] > pivot) { // DESC
+ int t = order[i];
+ order[i] = order[j];
+ order[j] = t;
+ i++;
+ }
+ }
+ int t = order[i];
+ order[i] = order[hi];
+ order[hi] = t;
+ return i;
+ }
+
+ static final class WriteResult {
+ final int newOrdinal;
+ final long fileOffset;
+ final ByteBuffer data;
+
+ WriteResult(int newOrdinal, long fileOffset, ByteBuffer data) {
+ this.newOrdinal = newOrdinal;
+ this.fileOffset = fileOffset;
+ this.data = data;
+ }
+ };
+
+ private static final class UpperLayerWriteResult {
+ final int ordinal;
+ final int[] neighbors;
+ final ByteSequence pqCode;
+
+ UpperLayerWriteResult(int ordinal, SelectedVecCache cache, ByteSequence pqCode) {
+ this.ordinal = ordinal;
+ this.neighbors = Arrays.copyOf(cache.nodes, cache.size);
+ this.pqCode = pqCode == null ? null : pqCode.copy();
+ }
+ };
+
+
+ /**
+ * Thread-local scratch space containing reusable buffers and search state for processing nodes.
+ */
+ private static final class Scratch implements AutoCloseable {
+
+ final int[] candSrc, candNode;
+ final float[] candScore;
+ final SelectedVecCache selectedCache;
+ final VectorFloat tmpVec, baseVec;
+ final GraphSearcher[] gs;
+ final ByteSequence pqCode;
+
+ /**
+ * Constructs scratch space with buffers sized for the maximum expected candidates and degree.
+ */
+ Scratch(int maxCandidateSize, int maxDegree, int dimension, List sources, ProductQuantization pq) {
+ this.candSrc = new int[maxCandidateSize];
+ this.candNode = new int[maxCandidateSize];
+ this.candScore = new float[maxCandidateSize];
+ this.selectedCache = new SelectedVecCache(maxDegree, dimension);
+ this.tmpVec = vectorTypeSupport.createFloatVector(dimension);
+ this.baseVec = vectorTypeSupport.createFloatVector(dimension);
+ this.pqCode = (pq == null) ? null : vectorTypeSupport.createByteSequence(pq.getSubspaceCount());
+
+ this.gs = new GraphSearcher[sources.size()];
+ for (int i = 0; i < sources.size(); i++) {
+ gs[i] = new GraphSearcher(sources.get(i));
+ gs[i].usePruning(false);
+ }
+ }
+
+ /**
+ * Closes all graph searchers and resets the cache.
+ */
+ @Override
+ public void close() throws IOException {
+ for (var s : gs) s.close();
+ selectedCache.reset();
+ }
+ }
+
+ /**
+ * Specification for a batch of nodes to be processed from one source index.
+ */
+ private static final class BatchSpec {
+ final int sourceIdx;
+ final int[] nodes; // materialized node ids for this source
+ final int start;
+ final int end;
+
+ BatchSpec(int sourceIdx, int[] nodes, int start, int end) {
+ this.sourceIdx = sourceIdx;
+ this.nodes = nodes;
+ this.start = start;
+ this.end = end;
+ }
+ }
+
+ /**
+ * Provides Vamana-style diversity filtering for neighbor selection during compaction.
+ */
+ private static final class CompactVamanaDiversityProvider {
+ /**
+ * the diversity threshold; 1.0 is equivalent to HNSW; Vamana uses 1.2 or more
+ */
+ public final float alpha;
+
+ /**
+ * used to compute diversity
+ */
+ public final VectorSimilarityFunction vsf;
+
+ /**
+ * Create a new diversity provider
+ */
+ public CompactVamanaDiversityProvider(VectorSimilarityFunction vsf, float alpha) {
+ this.vsf = vsf;
+ this.alpha = alpha;
+ }
+
+ /**
+ * Selects diverse neighbors from candidates using gradually increasing alpha threshold.
+ * Update `selected` with the diverse members of `neighbors`. `neighbors` is not modified
+ * It assumes that the i-th neighbor with 0 {@literal <=} i {@literal <} diverseBefore is already diverse.
+ */
+ public void retainDiverse(int[] candSrc, int[] candNode, float[] candScore, int[] order, int orderSize, int maxDegree, SelectedVecCache selectedCache, VectorFloat tmp, GraphSearcher[] gs) {
+ selectedCache.reset();
+ if (orderSize == 0) return;
+ int nSelected = 0;
+
+ // add diverse candidates, gradually increasing alpha to the threshold
+ // (so that the nearest candidates are prioritized)
+ float currentAlpha = 1.0f;
+ while (currentAlpha <= alpha + 1E-6 && nSelected < maxDegree) {
+ for (int i = 0; i < orderSize && nSelected < maxDegree; i++) {
+ int ci = order[i];
+ int cSrc = candSrc[ci];
+ int cNode = candNode[ci];
+ float cScore = candScore[ci];
+
+ OnDiskGraphIndex.View cView = (OnDiskGraphIndex.View) gs[cSrc].getView();
+ cView.getVectorInto(cNode, tmp, 0);
+ if (isDiverse(cView, cNode, tmp, cScore, currentAlpha, selectedCache)) {
+ selectedCache.add(cSrc, cView, cNode, cScore, tmp);
+ nSelected++;
+ }
+ }
+
+ currentAlpha += DIVERSITY_ALPHA_STEP;
+ }
+ }
+
+ /**
+ * Checks if a candidate is diverse enough by ensuring it's closer to the base node
+ * than to any already-selected neighbor (scaled by alpha threshold).
+ */
+ private boolean isDiverse(OnDiskGraphIndex.View cView, int cNode, VectorFloat cVec, float cScore, float alpha, SelectedVecCache selectedCache) {
+ for (int j = 0; j < selectedCache.size; j++) {
+ if (selectedCache.views[j] == cView && selectedCache.nodes[j] == cNode) {
+ return false; // already selected; don't add a duplicate
+ }
+ if (vsf.compare(cVec, selectedCache.vecs[j]) > cScore * alpha) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ }
+
+ /**
+ * Cache for storing selected diverse neighbors along with their metadata and vector copies.
+ */
+ static final class SelectedVecCache {
+ int[] sourceIdx;
+ OnDiskGraphIndex.View[] views;
+ int[] nodes;
+ float[] scores;
+ VectorFloat[] vecs;
+ int size;
+
+ /**
+ * Constructs a cache with the specified capacity and vector dimension.
+ */
+ SelectedVecCache(int capacity, int dimension) {
+ sourceIdx = new int[capacity];
+ views = new OnDiskGraphIndex.View[capacity];
+ nodes = new int[capacity];
+ scores = new float[capacity];
+ vecs = new VectorFloat[capacity];
+ for(int c = 0; c < capacity; ++c) {
+ vecs[c] = vectorTypeSupport.createFloatVector(dimension);
+ }
+ size = 0;
+ }
+
+ /**
+ * Resets the cache for reuse.
+ */
+ void reset() {
+ size = 0;
+ }
+
+ /**
+ * Adds a selected neighbor to the cache, copying its vector.
+ */
+ void add(int source, OnDiskGraphIndex.View view, int node, float score, VectorFloat vec) {
+ sourceIdx[size] = source;
+ views[size] = view;
+ nodes[size] = node;
+ scores[size] = score;
+ vecs[size].copyFrom(vec, 0, 0, vec.length());
+ size++;
+ }
+ }
+
+}
+
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OrdinalMapper.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OrdinalMapper.java
index 526241eff..445255980 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OrdinalMapper.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OrdinalMapper.java
@@ -106,4 +106,37 @@ public int newToOld(int newOrdinal) {
return newToOld.get(newOrdinal);
}
}
+
+ /**
+ * A mapper that applies a fixed offset to ordinals.
+ * Used for sequential mapping where local ordinal i maps to globalOffset + i.
+ */
+ class OffsetMapper implements OrdinalMapper {
+ private final int offset;
+ private final int size;
+
+ public OffsetMapper(int offset, int size) {
+ this.offset = offset;
+ this.size = size;
+ }
+
+ @Override
+ public int maxOrdinal() {
+ return offset + size - 1;
+ }
+
+ @Override
+ public int oldToNew(int oldOrdinal) {
+ return oldOrdinal + offset;
+ }
+
+ @Override
+ public int newToOld(int newOrdinal) {
+ int oldOrdinal = newOrdinal - offset;
+ if (oldOrdinal < 0 || oldOrdinal >= size) {
+ return OMITTED;
+ }
+ return oldOrdinal;
+ }
+ }
}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java
new file mode 100644
index 000000000..280d75c9c
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java
@@ -0,0 +1,241 @@
+/*
+ * 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.graph.disk;
+
+import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.disk.feature.FeatureId;
+import io.github.jbellis.jvector.graph.disk.feature.FusedPQ;
+import io.github.jbellis.jvector.quantization.ProductQuantization;
+import io.github.jbellis.jvector.util.DocIdSetIterator;
+import io.github.jbellis.jvector.util.FixedBitSet;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
+
+/**
+ * Handles Product Quantization retraining for graph index compaction.
+ * Performs balanced sampling across multiple source indexes and trains
+ * a new PQ codebook optimized for the combined dataset.
+ */
+public class PQRetrainer {
+ private static final Logger log = LoggerFactory.getLogger(PQRetrainer.class);
+ private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport();
+ private static final int MIN_SAMPLES_PER_SOURCE = 1000;
+ // Number of consecutive nodes to read per chunk before jumping to another location.
+ // Keeping reads sequential within each chunk lets the OS read-ahead cover them,
+ // avoiding the random I/O that would happen with per-node random sampling.
+ private static final int SAMPLE_CHUNK_SIZE = 32;
+
+ private final List sources;
+ private final List liveNodes;
+ private final List numLiveNodesPerSource;
+ private final int dimension;
+ private final int numTotalNodes;
+
+ public PQRetrainer(List sources, List liveNodes, int dimension) {
+ this.sources = sources;
+ this.liveNodes = liveNodes;
+ this.dimension = dimension;
+
+ this.numLiveNodesPerSource = new ArrayList<>(sources.size());
+ int total = 0;
+ for (int s = 0; s < sources.size(); s++) {
+ int numLiveNodes = liveNodes.get(s).cardinality();
+ total += numLiveNodes;
+ this.numLiveNodesPerSource.add(numLiveNodes);
+ }
+ this.numTotalNodes = total;
+ }
+
+ /**
+ * Trains a new Product Quantization codebook using balanced sampling across all source indexes.
+ * The base PQ parameters are taken from the FUSED_PQ feature on the first source.
+ * All sampled vectors are read into memory up front, so ProductQuantization.compute() itself
+ * performs no I/O.
+ */
+ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction) {
+ FusedPQ fpq = (FusedPQ) sources.get(0).getFeatures().get(FeatureId.FUSED_PQ);
+ return retrain(similarityFunction, fpq.getPQ());
+ }
+
+ /**
+ * Trains a new Product Quantization codebook using balanced sampling across all source indexes
+ * and the supplied base PQ for subspace/cluster parameters. Used when the base PQ comes from a
+ * non-fused source (e.g. a sidecar {@code CompressedVectors}) rather than the FUSED_PQ feature.
+ */
+ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, ProductQuantization basePQ) {
+ log.info("Training PQ using balanced sampling across sources");
+
+ List samples = sampleBalanced(ProductQuantization.MAX_PQ_TRAINING_SET_SIZE);
+
+ // Sort by (source, node) so extractVectorsSequential reads each source's file
+ // in ascending order, enabling OS read-ahead instead of random page faults.
+ samples.sort(Comparator.comparingInt((SampleRef r) -> r.source).thenComparingInt(r -> r.node));
+
+ log.info("Collected {} training samples", samples.size());
+
+ // Extract vectors sequentially in sorted (source, node) order so disk reads are
+ // purely sequential and the OS read-ahead can cover them efficiently. We do this
+ // here rather than letting ProductQuantization.compute() drive the reads via its
+ // parallel stream, which would scatter page faults across a potentially very large
+ // file and cause I/O that scales with dataset size rather than sample count.
+ List> trainingVectors = extractVectorsSequential(samples);
+ var ravv = new ListRandomAccessVectorValues(trainingVectors, dimension);
+
+ boolean center = similarityFunction == VectorSimilarityFunction.EUCLIDEAN;
+
+ return ProductQuantization.compute(
+ ravv,
+ basePQ.getSubspaceCount(),
+ basePQ.getClusterCount(),
+ center
+ );
+ }
+
+ /**
+ * Performs balanced sampling across all source indexes to ensure proportional representation.
+ * Guarantees minimum samples per source while respecting total sample budget.
+ */
+ private List sampleBalanced(int totalSamples) {
+ // If total live nodes <= totalSamples, return ALL
+ if (numTotalNodes <= totalSamples) {
+ List all = new ArrayList<>(numTotalNodes);
+
+ for (int s = 0; s < sources.size(); s++) {
+ FixedBitSet live = liveNodes.get(s);
+
+ for (int node = live.nextSetBit(0);
+ node != DocIdSetIterator.NO_MORE_DOCS;
+ node = live.nextSetBit(node + 1)) {
+ all.add(new SampleRef(s, node));
+ }
+ }
+
+ return all;
+ }
+
+ final int MIN_PER_SOURCE = Math.min(MIN_SAMPLES_PER_SOURCE, totalSamples / sources.size());
+
+ int[] quota = new int[sources.size()];
+ int assigned = 0;
+
+ // Proportional allocation
+ for (int s = 0; s < sources.size(); s++) {
+ quota[s] = Math.max(
+ MIN_PER_SOURCE,
+ (int) ((long) totalSamples * numLiveNodesPerSource.get(s) / numTotalNodes)
+ );
+ assigned += quota[s];
+ }
+
+ // Normalize down
+ while (assigned > totalSamples) {
+ for (int s = 0; s < sources.size() && assigned > totalSamples; s++) {
+ if (quota[s] > MIN_PER_SOURCE) {
+ quota[s]--;
+ assigned--;
+ }
+ }
+ }
+
+ // Normalize up
+ while (assigned < totalSamples) {
+ for (int s = 0; s < sources.size() && assigned < totalSamples; s++) {
+ quota[s]++;
+ assigned++;
+ }
+ }
+
+ List samples = new ArrayList<>(totalSamples);
+ ThreadLocalRandom rand = ThreadLocalRandom.current();
+
+ for (int s = 0; s < sources.size(); s++) {
+ FixedBitSet live = liveNodes.get(s);
+ int max = live.length();
+ int numChunks = (max + SAMPLE_CHUNK_SIZE - 1) / SAMPLE_CHUNK_SIZE;
+
+ // Build a shuffled chunk order so samples are representative but
+ // each chunk is read sequentially to minimize page faults.
+ // Fisher-Yates shuffle
+ int[] chunkOrder = new int[numChunks];
+ for (int i = 0; i < numChunks; i++) chunkOrder[i] = i;
+ for (int i = numChunks - 1; i > 0; i--) {
+ int j = rand.nextInt(i + 1);
+ int tmp = chunkOrder[i];
+ chunkOrder[i] = chunkOrder[j];
+ chunkOrder[j] = tmp;
+ }
+
+ int count = 0;
+ outer:
+ for (int ci = 0; ci < numChunks; ci++) {
+ int start = chunkOrder[ci] * SAMPLE_CHUNK_SIZE;
+ int end = Math.min(max, start + SAMPLE_CHUNK_SIZE);
+ for (int node = start; node < end; node++) {
+ if (live.get(node)) {
+ samples.add(new SampleRef(s, node));
+ if (++count >= quota[s]) break outer;
+ }
+ }
+ }
+ }
+
+ return samples;
+ }
+
+ /**
+ * Reads sampled vectors in the order provided. The caller must pre-sort {@code samples}
+ * by (source, node) so reads within each source are ascending, letting the OS read-ahead
+ * cover them efficiently. Each source's view is opened once and reused for all its samples.
+ */
+ private List> extractVectorsSequential(List samples) {
+ OnDiskGraphIndex.View[] views = new OnDiskGraphIndex.View[sources.size()];
+ for (int s = 0; s < sources.size(); s++) {
+ views[s] = (OnDiskGraphIndex.View) sources.get(s).getView();
+ }
+
+ List> vectors = new ArrayList<>(samples.size());
+ VectorFloat tmp = vectorTypeSupport.createFloatVector(dimension);
+ for (SampleRef ref : samples) {
+ views[ref.source].getVectorInto(ref.node, tmp, 0);
+ vectors.add(tmp.copy());
+ }
+ return vectors;
+ }
+
+ /**
+ * Reference to a sampled vector from a specific source index.
+ */
+ private static final class SampleRef {
+ final int source;
+ final int node;
+
+ SampleRef(int source, int node) {
+ this.source = source;
+ this.node = node;
+ }
+ }
+
+}
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
new file mode 100644
index 000000000..5d9c7311c
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java
@@ -0,0 +1,186 @@
+/*
+ * 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.graph.disk;
+
+import io.github.jbellis.jvector.graph.disk.feature.FusedFeature;
+import io.github.jbellis.jvector.quantization.ProductQuantization;
+import io.github.jbellis.jvector.quantization.VectorCompressor;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+
+import java.io.IOException;
+import java.nio.MappedByteBuffer;
+import java.nio.file.Path;
+import java.util.List;
+
+/**
+ * Encapsulates the quantization-aware steps the compactor needs to run during a single
+ * {@code compact()} invocation. Pulling these behind a strategy lets the compactor body stay
+ * scheme-agnostic: it asks the strategy whether to write inline codes, hands it pre/post hooks,
+ * and (for sidecar strategies) defers the merged sidecar write to the strategy.
+ *
+ * One strategy instance per compaction run. Strategies are stateful — they hold the retrained
+ * compressor produced by {@link #retrain} and any transient resources (e.g. memory-mapped
+ * pre-encode caches) until {@link #onAfterClose} releases them.
+ *
+ * Two concrete implementations cover all quantization schemes:
+ *
+ * - {@link FusedCompactionStrategy} — sources carry a {@link FusedFeature} with inline codes;
+ * the strategy is parameterized by a {@link VectorCompressorRetrainer} and the source's
+ * feature (used as a factory for the merged output's feature via
+ * {@link FusedFeature#withCompressor}). No PQ- or ASH-specific code lives in the strategy.
+ * - {@link SidecarCompactionStrategy} — sources ship codes as a non-fused
+ * {@code CompressedVectors} sidecar; the strategy is parameterized by a retrainer plus the
+ * source's {@code CompressedVectors} (used as a format handle).
+ *
+ * Adding a new quantization type (e.g. ASH) requires no strategy classes; the new {@code FusedASH}
+ * and {@code ASHVectors} just return appropriately-parameterized instances of these two strategies.
+ */
+public abstract class QuantizationCompactionStrategy {
+
+ /**
+ * Singleton strategy for sources that ship no quantization at all (no FUSED_PQ, no sidecar).
+ * All hooks are no-ops and {@link #compressor()} returns {@code null}.
+ */
+ public static final QuantizationCompactionStrategy NONE = new QuantizationCompactionStrategy() {
+ @Override
+ public void retrain(VectorSimilarityFunction vsf) {
+ // no-op
+ }
+
+ @Override
+ public VectorCompressor compressor() {
+ return null;
+ }
+
+ @Override
+ public String toString() {
+ return "QuantizationCompactionStrategy.NONE";
+ }
+ };
+
+ /**
+ * 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,
+ * {@link #compressor()} returns the retrained compressor.
+ */
+ public abstract void retrain(VectorSimilarityFunction vsf);
+
+ /** The retrained compressor produced by {@link #retrain}. {@code null} before retrain or for NONE. */
+ public abstract VectorCompressor compressor();
+
+ /**
+ * Whether this strategy writes codes inline in the graph file (FusedPQ-style). When true, the
+ * compactor passes the compressor to {@link CompactWriter} and the strategy expects to drive
+ * per-node code emission via the writer's inline-code path.
+ */
+ public boolean writesCodesInline() {
+ return false;
+ }
+
+ /**
+ * Whether this strategy writes codes to a separate sidecar file (PQVectors-style). When true,
+ * the compactor calls {@link #writeSidecar} after the graph file is closed.
+ */
+ public boolean writesCodesSidecar() {
+ return false;
+ }
+
+ /**
+ * Hook invoked once after {@link CompactWriter#writeHeader()} but before {@code compactLevels}.
+ * Inline strategies can use this to pre-encode every live node's code into a transient cache
+ * that the writer will copy from during inline writes. No-op by default.
+ */
+ public void onAfterHeader(CompactWriter writer) throws IOException {
+ // no-op
+ }
+
+ /**
+ * Hook invoked once after {@code compactLevels} returns but before
+ * {@link CompactWriter#writeFooter()}. Inline strategies that need to emit a per-graph tail
+ * record (e.g. the entry-node PQ code for FusedPQ when there is no hierarchy) do so here.
+ * No-op by default.
+ */
+ public void onAfterLevels(CompactWriter writer, int[] entryNodeSource, List maxDegrees) throws IOException {
+ // no-op
+ }
+
+ /**
+ * Hook invoked once after the graph file is closed (in {@code finally}). Strategies can
+ * release transient resources (e.g. unmap a pre-encode cache and truncate the output file
+ * back to its expected size). No-op by default.
+ */
+ public void onAfterClose(Path graphPath) {
+ // no-op
+ }
+
+ /**
+ * Writes the merged compressed-vectors sidecar file. Called by the compactor's
+ * {@code compact(graphPath, compressedPath)} entry point after the graph is fully written.
+ * Throws {@link UnsupportedOperationException} by default; sidecar strategies override.
+ */
+ public void writeSidecar(Path compressedPath) throws IOException {
+ throw new UnsupportedOperationException(this + " does not write a sidecar");
+ }
+
+ /**
+ * Returns the {@link FusedFeature} the compactor should put in the merged output graph for
+ * an inline strategy. {@code null} for non-inline strategies (NONE and any sidecar strategy).
+ * Called after {@link #retrain} so the strategy can build the output feature from the
+ * retrained compressor.
+ */
+ public FusedFeature outputFusedFeature(int maxDegree) {
+ return null;
+ }
+
+ /**
+ * For compaction use. Returns the precomputed code cache built by {@link #onAfterHeader},
+ * indexed by new ordinal so refinement can memcpy neighbor codes instead of re-encoding them.
+ * Returns {@code null} when no cache is held (non-fused strategy, NONE, or graph too large for
+ * a single mapping). The returned buffer is shared; callers must {@code .duplicate()} per
+ * thread before using.
+ */
+ public MappedByteBuffer getCodeCache() {
+ return null;
+ }
+
+ /** For compaction use. Bytes per code in {@link #getCodeCache()}, or {@code 0} when no cache. */
+ public int getCacheCodeSize() {
+ return 0;
+ }
+
+ /**
+ * For compaction use. Drops the strategy's hold on the {@link CompactionContext} (and thus the
+ * source graphs) once {@code compactGraphImpl} no longer needs it, so the source graphs become
+ * GC-eligible before the refinement pass loads a second full graph. Called only from the
+ * non-sidecar {@code compact(Path)} path. No-op by default; strategies that retain a context
+ * override. Must not be called before {@code onAfterHeader}/{@code onAfterLevels} have run, and
+ * implementations must keep {@code onAfterClose} working without the context.
+ */
+ public void releaseSources() {
+ // no-op
+ }
+
+ /**
+ * Convenience: returns {@link #compressor()} cast to {@link ProductQuantization}, or
+ * {@code null} if no compressor is held. Kept for backward compat with code paths that still
+ * thread a typed {@code ProductQuantization} through {@link CompactWriter}.
+ */
+ protected ProductQuantization compressorAsPQ() {
+ VectorCompressor c = compressor();
+ return (c instanceof ProductQuantization) ? (ProductQuantization) c : null;
+ }
+}
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
new file mode 100644
index 000000000..1317ca88a
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/SidecarCompactionStrategy.java
@@ -0,0 +1,170 @@
+/*
+ * 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.graph.disk;
+
+import io.github.jbellis.jvector.disk.BufferedRandomAccessWriter;
+import io.github.jbellis.jvector.quantization.CompressedVectors;
+import io.github.jbellis.jvector.quantization.VectorCompressor;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.ByteSequence;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+
+/**
+ * Generic compaction strategy for any non-fused {@link CompressedVectors} sidecar. Parameterized
+ * by:
+ *
+ * - a {@link VectorCompressorRetrainer} that produces the merged compressor on retrain (the
+ * only scheme-specific knowledge this strategy carries),
+ * - a {@code formatHandle} {@link CompressedVectors} from the sources, used only to invoke
+ * {@link CompressedVectors#writeSidecarHeader} and {@link CompressedVectors#sidecarVectorsPerChunk}
+ * — the format hooks that decide the on-disk layout for the merged sidecar.
+ *
+ * No PQ-specific (or ASH-specific) code lives here. Adding a new quantization type that ships
+ * a sidecar means implementing those two hooks on its {@code CompressedVectors} class plus a
+ * retrainer; this strategy and the compactor stay untouched.
+ */
+public final class SidecarCompactionStrategy extends QuantizationCompactionStrategy {
+ private static final Logger log = LoggerFactory.getLogger(SidecarCompactionStrategy.class);
+ private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport();
+
+ private final CompactionContext ctx;
+ private final CompressedVectors formatHandle;
+ private final VectorCompressorRetrainer retrainer;
+ private VectorCompressor retrainedCompressor;
+
+ public SidecarCompactionStrategy(CompactionContext ctx,
+ CompressedVectors formatHandle,
+ VectorCompressorRetrainer retrainer) {
+ this.ctx = ctx;
+ this.formatHandle = formatHandle;
+ this.retrainer = retrainer;
+ }
+
+ @Override
+ public void retrain(VectorSimilarityFunction vsf) {
+ log.info("Retraining sidecar compressor ({}) on merged sources",
+ formatHandle.getClass().getSimpleName());
+ this.retrainedCompressor = retrainer.retrain(vsf);
+ }
+
+ @Override
+ public VectorCompressor compressor() {
+ return retrainedCompressor;
+ }
+
+ @Override
+ public boolean writesCodesSidecar() {
+ return true;
+ }
+
+ @Override
+ public void writeSidecar(Path compressedPath) throws IOException {
+ if (retrainedCompressor == null) {
+ throw new IllegalStateException("retrain() must be called before writeSidecar()");
+ }
+ final int vectorsPerChunk = formatHandle.sidecarVectorsPerChunk();
+ final int codeSize = retrainedCompressor.compressedVectorSize();
+ final int count = ctx.maxOrdinal + 1;
+ final int chunkCount = (count + vectorsPerChunk - 1) / vectorsPerChunk;
+
+ log.info("Streaming {} merged ordinals to {} ({} chunks of up to {} entries each)",
+ count, compressedPath, chunkCount, vectorsPerChunk);
+
+ try (var out = new BufferedRandomAccessWriter(compressedPath)) {
+ formatHandle.writeSidecarHeader(out, retrainedCompressor, count);
+
+ int parallelism = Math.max(ctx.taskWindowSize, 1);
+ for (int batchStart = 0; batchStart < chunkCount; batchStart += parallelism) {
+ int batchEnd = Math.min(batchStart + parallelism, chunkCount);
+ List>> tasks = new ArrayList<>(batchEnd - batchStart);
+ for (int c = batchStart; c < batchEnd; c++) {
+ final int chunkStart = c * vectorsPerChunk;
+ final int chunkEnd = Math.min(chunkStart + vectorsPerChunk, count);
+ tasks.add(() -> encodeChunk(chunkStart, chunkEnd, codeSize, retrainedCompressor));
+ }
+ for (var f : ctx.executor.invokeAll(tasks)) {
+ vectorTypeSupport.writeByteSequence(out, f.get());
+ }
+ }
+ } catch (InterruptedException | ExecutionException e) {
+ throw new IOException("Failed to write compressed sidecar to " + compressedPath, e);
+ }
+ log.info("Wrote compacted compressed sidecar to {}", compressedPath);
+ }
+
+ @SuppressWarnings("unchecked")
+ private ByteSequence encodeChunk(int chunkStart, int chunkEnd, int codeSize, VectorCompressor compressor) throws IOException {
+ int chunkBytes = (chunkEnd - chunkStart) * codeSize;
+ ByteSequence chunk = vectorTypeSupport.createByteSequence(chunkBytes);
+ chunk.zero();
+
+ // Cast once; valid for all VectorCompressor implementations that produce ByteSequence codes
+ // (PQ, future ASH, etc.). VectorCompressor's encode/encodeTo contract guarantees T is the
+ // encoded type and for our supported quantization schemes T = ByteSequence.
+ VectorCompressor> byteCompressor = (VectorCompressor>) compressor;
+
+ OnDiskGraphIndex.View[] views = new OnDiskGraphIndex.View[ctx.sources.size()];
+ try {
+ VectorFloat vec = vectorTypeSupport.createFloatVector(ctx.dimension);
+ ByteSequence code = vectorTypeSupport.createByteSequence(codeSize);
+ for (int newOrd = chunkStart; newOrd < chunkEnd; newOrd++) {
+ int[] resolved = resolveSourceForNewOrd(newOrd);
+ if (resolved == null) continue; // hole; slot stays zero
+ int srcIdx = resolved[0];
+ int oldOrd = resolved[1];
+ if (views[srcIdx] == null) {
+ views[srcIdx] = (OnDiskGraphIndex.View) ctx.sources.get(srcIdx).getView();
+ }
+ views[srcIdx].getVectorInto(oldOrd, vec, 0);
+ code.zero();
+ byteCompressor.encodeTo(vec, code);
+ int slotOffset = (newOrd - chunkStart) * codeSize;
+ for (int b = 0; b < codeSize; b++) {
+ chunk.set(slotOffset + b, code.get(b));
+ }
+ }
+ } finally {
+ for (var v : views) {
+ if (v != null) {
+ try { v.close(); } catch (Exception ignore) {}
+ }
+ }
+ }
+ return chunk;
+ }
+
+ private int[] resolveSourceForNewOrd(int newOrd) {
+ for (int s = 0; s < ctx.remappers.size(); s++) {
+ int oldOrd = ctx.remappers.get(s).newToOld(newOrd);
+ if (oldOrd != OrdinalMapper.OMITTED) {
+ return new int[]{s, oldOrd};
+ }
+ }
+ return null;
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/VectorCompressorRetrainer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/VectorCompressorRetrainer.java
new file mode 100644
index 000000000..5110cbe52
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/VectorCompressorRetrainer.java
@@ -0,0 +1,33 @@
+/*
+ * 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.graph.disk;
+
+import io.github.jbellis.jvector.quantization.VectorCompressor;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+
+/**
+ * Trains a fresh {@link VectorCompressor} on the merged source vectors during a compaction run.
+ *
+ * One implementation per quantization scheme (today: {@link PQRetrainer} wrapped behind a lambda;
+ * future: {@code ASHRetrainer}, etc.). Both fused and sidecar generic strategies receive a
+ * retrainer via their constructor and invoke it at {@code retrain(vsf)} time — the strategies
+ * stay quantization-agnostic and the retrainer encapsulates the scheme-specific training math.
+ */
+@FunctionalInterface
+public interface VectorCompressorRetrainer {
+ VectorCompressor retrain(VectorSimilarityFunction vsf);
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedFeature.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedFeature.java
index d54630999..64334310a 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedFeature.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedFeature.java
@@ -18,7 +18,11 @@
import io.github.jbellis.jvector.disk.IndexWriter;
import io.github.jbellis.jvector.disk.RandomAccessReader;
+import io.github.jbellis.jvector.graph.disk.CompactionContext;
+import io.github.jbellis.jvector.graph.disk.QuantizationCompactionStrategy;
+import io.github.jbellis.jvector.quantization.VectorCompressor;
import io.github.jbellis.jvector.util.Accountable;
+import io.github.jbellis.jvector.vector.types.ByteSequence;
import java.io.IOException;
@@ -38,4 +42,42 @@ default boolean isFused() {
interface InlineSource extends Accountable {}
InlineSource loadSourceFeature(RandomAccessReader in) throws IOException;
+
+ /**
+ * For compaction use: bytes occupied on disk by a single stored code (one neighbor's payload).
+ * For fused features {@code featureSize() == codeSize() * maxDegree}. Called by the compactor
+ * (and {@link io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexCompactor#ramBytesUsed}) to
+ * size per-thread scratch buffers and by {@code FusedCompactionStrategy} to size the streaming
+ * pre-encode cache.
+ */
+ int codeSize();
+
+ /**
+ * For compaction use: returns the underlying compressor that produced the inline codes carried
+ * by this feature. Returned typed as {@code VectorCompressor>} so generic
+ * compaction code (the pre-encode pass, the per-write encoding fallback in {@code CompactWriter})
+ * can call {@code encodeTo(VectorFloat, ByteSequence)} without knowing the concrete
+ * quantization scheme.
+ */
+ VectorCompressor> getCompressor();
+
+ /**
+ * For compaction use: returns a fresh {@link FusedFeature} of this same scheme but
+ * parameterized by a new compressor and max degree. Called by
+ * {@code FusedCompactionStrategy.outputFusedFeature} to construct the merged output's fused
+ * feature from a retrained compressor — every {@link FusedFeature} implementation acts as a
+ * factory for itself in this way so the compactor never references concrete subtypes.
+ */
+ FusedFeature withCompressor(VectorCompressor> newCompressor, int maxDegree);
+
+ /**
+ * For compaction use: returns the {@link QuantizationCompactionStrategy} the compactor should
+ * run when merging graphs that carry this fused feature. One strategy instance per
+ * compaction; it owns any transient state (retrained codebook, pre-encode caches) until the
+ * compactor releases it via {@link QuantizationCompactionStrategy#onAfterClose}.
+ *
+ * Implementations must return a fresh strategy on every call — feature instances themselves
+ * are read-mostly objects that may be shared by concurrent readers of the source graph.
+ */
+ QuantizationCompactionStrategy createCompactionStrategy(CompactionContext ctx);
}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java
index 840650ba5..3dc404481 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java
@@ -20,11 +20,17 @@
import io.github.jbellis.jvector.disk.RandomAccessReader;
import io.github.jbellis.jvector.graph.ImmutableGraphIndex;
import io.github.jbellis.jvector.graph.disk.CommonHeader;
+import io.github.jbellis.jvector.graph.disk.CompactionContext;
+import io.github.jbellis.jvector.graph.disk.FusedCompactionStrategy;
import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
+import io.github.jbellis.jvector.graph.disk.PQRetrainer;
+import io.github.jbellis.jvector.graph.disk.QuantizationCompactionStrategy;
+import io.github.jbellis.jvector.graph.disk.VectorCompressorRetrainer;
import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
import io.github.jbellis.jvector.quantization.FusedPQDecoder;
import io.github.jbellis.jvector.quantization.PQVectors;
import io.github.jbellis.jvector.quantization.ProductQuantization;
+import io.github.jbellis.jvector.quantization.VectorCompressor;
import io.github.jbellis.jvector.util.ExplicitThreadLocal;
import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
import io.github.jbellis.jvector.vector.VectorizationProvider;
@@ -67,6 +73,23 @@ public ProductQuantization getPQ() {
return pq;
}
+ /** For compaction use. See {@link FusedFeature#getCompressor}. */
+ @Override
+ @SuppressWarnings("unchecked")
+ public VectorCompressor> getCompressor() {
+ return (VectorCompressor>) (VectorCompressor) pq;
+ }
+
+ /** For compaction use. See {@link FusedFeature#withCompressor}. */
+ @Override
+ public FusedFeature withCompressor(VectorCompressor> newCompressor, int maxDegree) {
+ if (!(newCompressor instanceof ProductQuantization)) {
+ throw new IllegalArgumentException(
+ "FusedPQ requires ProductQuantization; got " + newCompressor.getClass().getSimpleName());
+ }
+ return new FusedPQ(maxDegree, (ProductQuantization) newCompressor);
+ }
+
@Override
public int headerSize() {
return pq.compressorSize();
@@ -77,6 +100,12 @@ public int featureSize() {
return pq.compressedVectorSize() * maxDegree;
}
+ /** For compaction use. See {@link FusedFeature#codeSize}. */
+ @Override
+ public int codeSize() {
+ return pq.compressedVectorSize();
+ }
+
static FusedPQ load(CommonHeader header, RandomAccessReader reader) {
try {
return new FusedPQ(header.layerInfo.get(0).degree, ProductQuantization.load(reader));
@@ -96,6 +125,16 @@ public ScoreFunction.ApproximateScoreFunction approximateScoreFunctionFor(Vector
return FusedPQDecoder.newDecoder(neighbors, pq, hierarchyCachedFeatures, queryVector, reusableNeighborCodes.get(), reusableResults.get(), vsf, esf);
}
+ /** For compaction use. See {@link FusedFeature#createCompactionStrategy}. */
+ @Override
+ public QuantizationCompactionStrategy createCompactionStrategy(CompactionContext ctx) {
+ ProductQuantization basePQ = this.pq;
+ VectorCompressorRetrainer retrainer =
+ vsf -> new PQRetrainer(ctx.sources, ctx.liveNodes, ctx.dimension)
+ .retrain(vsf, basePQ);
+ return new FusedCompactionStrategy(ctx, this, retrainer);
+ }
+
@Override
public void writeHeader(IndexWriter out) throws IOException {
pq.write(out, OnDiskGraphIndex.CURRENT_VERSION);
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/CompressedVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/CompressedVectors.java
index 767659148..020397eba 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/CompressedVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/CompressedVectors.java
@@ -17,7 +17,9 @@
package io.github.jbellis.jvector.quantization;
import io.github.jbellis.jvector.disk.IndexWriter;
+import io.github.jbellis.jvector.graph.disk.CompactionContext;
import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
+import io.github.jbellis.jvector.graph.disk.QuantizationCompactionStrategy;
import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
import io.github.jbellis.jvector.util.Accountable;
import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
@@ -66,4 +68,46 @@ default ScoreFunction.ApproximateScoreFunction approximateScoreFunctionFor(Vecto
/** the number of vectors */
int count();
+
+ /**
+ * For compaction use: returns the {@link QuantizationCompactionStrategy} the compactor should
+ * run when merging graphs whose non-fused compressed sidecars are this kind of
+ * {@code CompressedVectors}. One strategy instance per compaction; it retrains the compressor
+ * on the merged source vectors and streams the merged sidecar to disk.
+ *
+ * Called by {@code OnDiskGraphIndexCompactor.detectSidecarStrategy()}. Named to mirror
+ * {@code FusedFeature.createCompactionStrategy} — same verb, receiver type disambiguates
+ * whether the returned strategy drives the inline-fused or sidecar workflow. Default throws —
+ * implementations supporting compaction must override.
+ */
+ default QuantizationCompactionStrategy createCompactionStrategy(CompactionContext ctx) {
+ throw new UnsupportedOperationException(
+ getClass().getSimpleName() + " does not support sidecar compaction");
+ }
+
+ // ---- For compaction use: sidecar-streaming-write hooks. Called by the generic
+ // SidecarCompactionStrategy to produce a merged-format-compatible sidecar without that
+ // strategy knowing the format. ----
+
+ /**
+ * For compaction use: writes the format-specific sidecar header (compressor params + vector
+ * count + any extras the reader expects between count and the chunk stream). Called once at
+ * the start of a streaming sidecar write by {@code SidecarCompactionStrategy.writeSidecar},
+ * after which the strategy emits chunks of {@code sidecarVectorsPerChunk()} codes each.
+ * Default throws — implementations supporting sidecar compaction must override.
+ */
+ default void writeSidecarHeader(IndexWriter out, VectorCompressor mergedCompressor, int count) throws IOException {
+ throw new UnsupportedOperationException(
+ getClass().getSimpleName() + " does not support sidecar compaction");
+ }
+
+ /**
+ * For compaction use: vectors per chunk for streaming sidecar writes. The chunk size must
+ * match the format the reader expects (e.g. {@code PQVectors} uses 1024 to align with
+ * {@code MutablePQVectors}'s on-disk layout). Read by
+ * {@code SidecarCompactionStrategy.writeSidecar} to size each emitted chunk.
+ */
+ default int sidecarVectorsPerChunk() {
+ return 1024;
+ }
}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java
index 538632da0..760eb38ee 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java
@@ -19,6 +19,8 @@
import io.github.jbellis.jvector.disk.IndexWriter;
import io.github.jbellis.jvector.disk.RandomAccessReader;
import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.disk.CompactionContext;
+import io.github.jbellis.jvector.graph.disk.QuantizationCompactionStrategy;
import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
import io.github.jbellis.jvector.util.RamUsageEstimator;
import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
@@ -416,6 +418,37 @@ public ProductQuantization getCompressor() {
return pq;
}
+ /** For compaction use. See {@link CompressedVectors#createCompactionStrategy}. */
+ @Override
+ public QuantizationCompactionStrategy createCompactionStrategy(CompactionContext ctx) {
+ ProductQuantization basePQ = this.pq;
+ io.github.jbellis.jvector.graph.disk.VectorCompressorRetrainer retrainer =
+ vsf -> new io.github.jbellis.jvector.graph.disk.PQRetrainer(ctx.sources, ctx.liveNodes, ctx.dimension)
+ .retrain(vsf, basePQ);
+ return new io.github.jbellis.jvector.graph.disk.SidecarCompactionStrategy(ctx, this, retrainer);
+ }
+
+ /** For compaction use. See {@link CompressedVectors#writeSidecarHeader}. */
+ @Override
+ public void writeSidecarHeader(IndexWriter out, VectorCompressor mergedCompressor, int count) throws IOException {
+ if (!(mergedCompressor instanceof ProductQuantization)) {
+ throw new IllegalArgumentException(
+ "PQVectors sidecar header requires ProductQuantization; got "
+ + mergedCompressor.getClass().getSimpleName());
+ }
+ ProductQuantization mergedPQ = (ProductQuantization) mergedCompressor;
+ mergedPQ.write(out, io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex.CURRENT_VERSION);
+ out.writeInt(count);
+ out.writeInt(mergedPQ.getSubspaceCount());
+ }
+
+ /** For compaction use. See {@link CompressedVectors#sidecarVectorsPerChunk}. */
+ @Override
+ public int sidecarVectorsPerChunk() {
+ // Match MutablePQVectors so the on-disk layout is identical to what PQVectors.load reconstructs.
+ return 1024;
+ }
+
@Override
public long ramBytesUsed() {
int REF_BYTES = RamUsageEstimator.NUM_BYTES_OBJECT_REF;
@@ -444,7 +477,7 @@ public String toString() {
* This is emulative of modern Java records, but keeps to J11 standards.
* This class consolidates the layout calculations for PQ data into one place
*/
- static class PQLayout {
+ public static class PQLayout {
/**
* total number of vectors
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorizationProvider.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorizationProvider.java
index 1ec46443d..5ab5664d1 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorizationProvider.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorizationProvider.java
@@ -77,6 +77,28 @@ protected VectorizationProvider() {
// visible for tests
static VectorizationProvider lookup(boolean testMode) {
+ String forcedProvider = System.getProperty("jvector.vectorization_provider");
+ if (forcedProvider != null) {
+ switch (forcedProvider.toLowerCase(Locale.ROOT)) {
+ case "default":
+ return new DefaultVectorizationProvider();
+ case "panama":
+ try {
+ return (VectorizationProvider) Class.forName("io.github.jbellis.jvector.vector.PanamaVectorizationProvider").getConstructor().newInstance();
+ } catch (Throwable e) {
+ throw new RuntimeException("Failed to load forced PanamaVectorizationProvider", e);
+ }
+ case "native":
+ try {
+ return (VectorizationProvider) Class.forName("io.github.jbellis.jvector.vector.NativeVectorizationProvider").getConstructor().newInstance();
+ } catch (Throwable e) {
+ throw new RuntimeException("Failed to load forced NativeVectorizationProvider", e);
+ }
+ default:
+ throw new IllegalArgumentException("Unknown vectorization provider: " + forcedProvider);
+ }
+ }
+
final int runtimeVersion = Runtime.version().feature();
if (runtimeVersion >= 20) {
// is locale sane (only buggy in Java 20)
diff --git a/jvector-examples/pom.xml b/jvector-examples/pom.xml
index 9daf7b8cf..731c1d769 100644
--- a/jvector-examples/pom.xml
+++ b/jvector-examples/pom.xml
@@ -85,16 +85,16 @@
gson
2.10.1
-
+
org.slf4j
slf4j-api
2.0.9
- ch.qos.logback
- logback-classic
- 1.4.11
+ org.apache.logging.log4j
+ log4j-slf4j2-impl
+ 2.24.3
software.amazon.awssdk
@@ -112,6 +112,11 @@
aws-crt-client
${awssdk.version}
+
+ software.amazon.awssdk
+ ec2
+ ${awssdk.version}
+
software.amazon.awssdk
s3
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/GitInfo.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/GitInfo.java
new file mode 100644
index 000000000..92979ac5d
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/GitInfo.java
@@ -0,0 +1,54 @@
+/*
+ * 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.example.reporting;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Detects the current git commit hash for tagging benchmark results.
+ */
+public final class GitInfo {
+ private static final Logger log = LoggerFactory.getLogger(GitInfo.class);
+
+ private GitInfo() {}
+
+ // Lazy holder pattern — computed once on first access
+ private static class Holder {
+ static final String SHORT_HASH;
+ static {
+ String hash;
+ try {
+ var process = new ProcessBuilder("git", "rev-parse", "HEAD").redirectErrorStream(true).start();
+ hash = new String(process.getInputStream().readAllBytes()).trim();
+ process.waitFor();
+ if (hash.length() >= 8) {
+ hash = hash.substring(hash.length() - 8);
+ }
+ } catch (Exception e) {
+ log.warn("Could not determine git hash", e);
+ hash = "unknown";
+ }
+ SHORT_HASH = hash;
+ }
+ }
+
+ /** Returns the last 8 characters of {@code git rev-parse HEAD}, or {@code "unknown"} on failure. */
+ public static String getShortHash() {
+ return Holder.SHORT_HASH;
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/JfrRecorder.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/JfrRecorder.java
new file mode 100644
index 000000000..ab5c4b581
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/JfrRecorder.java
@@ -0,0 +1,105 @@
+/*
+ * 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.example.reporting;
+
+import jdk.jfr.Configuration;
+import jdk.jfr.Recording;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.text.ParseException;
+import java.time.Duration;
+
+/**
+ * Manages the lifecycle of a JFR (Java Flight Recorder) recording for benchmarks.
+ */
+public final class JfrRecorder {
+ private static final Logger log = LoggerFactory.getLogger(JfrRecorder.class);
+
+ private Recording recording;
+ private String fileName;
+
+ /**
+ * Creates the output directory, configures a "profile" recording, starts it, and returns the absolute path.
+ *
+ * @param outputDir directory to write the JFR file into
+ * @param fileName name of the JFR file (e.g. {@code "compactor-foo.jfr"})
+ * @return the absolute path of the recording file
+ * @throws IOException if the directory cannot be created
+ * @throws ParseException if the JFR "profile" configuration cannot be loaded
+ */
+ public Path start(Path outputDir, String fileName) throws IOException, ParseException {
+ return start(outputDir, fileName, false);
+ }
+
+ /**
+ * Creates the output directory, configures a "profile" recording, starts it, and returns the absolute path.
+ *
+ * @param outputDir directory to write the JFR file into
+ * @param fileName name of the JFR file
+ * @param objectCount whether to enable periodic 'jdk.ObjectCount' events
+ * @return the absolute path of the recording file
+ */
+ public Path start(Path outputDir, String fileName, boolean objectCount) throws IOException, ParseException {
+ Files.createDirectories(outputDir);
+ Path jfrPath = outputDir.resolve(fileName).toAbsolutePath();
+ recording = new Recording(Configuration.getConfiguration("profile"));
+ recording.setToDisk(true);
+ recording.setDestination(jfrPath);
+
+ // Enable heap occupancy snapshots and old object sampling
+ var settings = recording.getSettings();
+ if (objectCount) {
+ settings.put("jdk.ObjectCount#enabled", "true");
+ settings.put("jdk.ObjectCount#period", "10s"); // Every 10 seconds
+ }
+ settings.put("jdk.OldObjectSample#enabled", "true");
+ // Flush to disk every minute so data is available for inspection during long benchmarks
+ settings.put("flush-interval", Duration.ofMinutes(1).toMillis() + "ms");
+ recording.setSettings(settings);
+ recording.start();
+ this.fileName = fileName;
+ System.out.println("JFR recording started, saving to: " + jfrPath);
+ log.info("JFR recording started, saving to: {}", jfrPath);
+ return jfrPath;
+ }
+
+ /** Stops and closes the recording, logging the saved path. */
+ public void stop() {
+ if (recording != null) {
+ Path jfrPath = recording.getDestination();
+ recording.stop();
+ recording.close();
+ recording = null;
+ System.out.println("JFR recording saved to: " + jfrPath);
+ log.info("JFR recording saved to: {}", jfrPath);
+ }
+ }
+
+ /** Returns {@code true} if a recording is currently in progress. */
+ public boolean isActive() {
+ return recording != null;
+ }
+
+ /** Returns the current file name, or {@code null} if no recording has been started. */
+ public String getFileName() {
+ return fileName;
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/JsonlWriter.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/JsonlWriter.java
new file mode 100644
index 000000000..a025dd207
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/JsonlWriter.java
@@ -0,0 +1,56 @@
+/*
+ * 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.example.reporting;
+
+import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.util.Map;
+
+/**
+ * Append-only JSONL file writer that serializes one map per line using GSON.
+ */
+public final class JsonlWriter {
+ private static final Logger log = LoggerFactory.getLogger(JsonlWriter.class);
+ private static final Gson GSON = new GsonBuilder()
+ .disableHtmlEscaping()
+ .serializeNulls()
+ .create(); // No pretty printing for JSONL
+
+ private final Path outputFile;
+
+ public JsonlWriter(Path outputFile) {
+ this.outputFile = outputFile;
+ }
+
+ /** Serializes the map as a single JSON line and appends it to the output file. */
+ public void writeLine(Map result) {
+ String json = GSON.toJson(result) + "\n";
+ try {
+ Files.writeString(outputFile, json,
+ StandardOpenOption.CREATE, StandardOpenOption.APPEND);
+ } catch (IOException e) {
+ log.error("Failed to persist result to {}", outputFile, e);
+ }
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/SystemStatsCollector.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/SystemStatsCollector.java
new file mode 100644
index 000000000..571782510
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/SystemStatsCollector.java
@@ -0,0 +1,196 @@
+/*
+ * 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.example.reporting;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.time.Instant;
+import java.time.ZoneOffset;
+import java.time.format.DateTimeFormatter;
+import java.util.HashSet;
+import java.util.List;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.regex.Pattern;
+
+/**
+ * Background collector of {@code /proc} system metrics (CPU topology, load, memory, disk I/O).
+ * Reads /proc files directly in Java and appends JSONL lines to a file every 30 seconds.
+ */
+public final class SystemStatsCollector {
+ private static final Logger log = LoggerFactory.getLogger(SystemStatsCollector.class);
+ private static final Path PROC_CPUINFO = Path.of("/proc/cpuinfo");
+ private static final Path PROC_LOADAVG = Path.of("/proc/loadavg");
+ private static final Path PROC_MEMINFO = Path.of("/proc/meminfo");
+ private static final Path PROC_DISKSTATS = Path.of("/proc/diskstats");
+ private static final Pattern DISK_DEVICE_PATTERN = Pattern.compile("sd[a-z]+|nvme[0-9]+n[0-9]+|vd[a-z]+|xvd[a-z]+");
+ private static final DateTimeFormatter TS_FORMAT = DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss'Z'").withZone(ZoneOffset.UTC);
+
+ private ScheduledExecutorService scheduler;
+ private BufferedWriter writer;
+ private String fileName;
+ private int cpuSockets;
+ private int cpuCores;
+ private int cpuThreads;
+
+ public Path start(Path outputDir, String fileName) throws IOException {
+ if (!Files.exists(PROC_CPUINFO)) {
+ log.warn("/proc filesystem not available (not Linux?), system stats collection disabled");
+ return null;
+ }
+
+ Files.createDirectories(outputDir);
+ Path sysStatsPath = outputDir.resolve(fileName).toAbsolutePath();
+ this.fileName = fileName;
+
+ parseCpuTopology();
+
+ this.writer = Files.newBufferedWriter(sysStatsPath,
+ StandardOpenOption.CREATE, StandardOpenOption.APPEND);
+
+ scheduler = Executors.newSingleThreadScheduledExecutor(r -> {
+ Thread t = new Thread(r, "sys-stats-collector");
+ t.setDaemon(true);
+ return t;
+ });
+ scheduler.scheduleAtFixedRate(() -> {
+ try {
+ String line = collectSnapshot();
+ writer.write(line);
+ writer.newLine();
+ writer.flush();
+ } catch (Exception e) {
+ log.warn("Failed to collect system stats", e);
+ }
+ }, 0, 30, TimeUnit.SECONDS);
+
+ log.info("System stats collection started, saving to: {}", sysStatsPath);
+ return sysStatsPath;
+ }
+
+ public void stop(Path outputDir) throws InterruptedException {
+ if (scheduler != null) {
+ scheduler.shutdown();
+ scheduler.awaitTermination(5, TimeUnit.SECONDS);
+ scheduler = null;
+ try {
+ if (writer != null) {
+ writer.close();
+ writer = null;
+ }
+ } catch (IOException e) {
+ log.warn("Failed to close stats writer", e);
+ }
+ log.info("System stats collection stopped, saved to: {}", outputDir.resolve(fileName).toAbsolutePath());
+ }
+ }
+
+ public boolean isActive() {
+ return scheduler != null && !scheduler.isShutdown();
+ }
+
+ public String getFileName() {
+ return fileName;
+ }
+
+ private void parseCpuTopology() throws IOException {
+ List lines = Files.readAllLines(PROC_CPUINFO);
+ int threads = 0;
+ var physicalIds = new HashSet();
+ var coreKeys = new HashSet();
+ String currentPhysicalId = "0";
+
+ for (String line : lines) {
+ if (line.startsWith("processor")) {
+ threads++;
+ } else if (line.startsWith("physical id")) {
+ currentPhysicalId = line.substring(line.indexOf(':') + 1).trim();
+ physicalIds.add(currentPhysicalId);
+ } else if (line.startsWith("core id")) {
+ String coreId = line.substring(line.indexOf(':') + 1).trim();
+ coreKeys.add(currentPhysicalId + "-" + coreId);
+ }
+ }
+
+ this.cpuThreads = threads;
+ this.cpuSockets = physicalIds.isEmpty() ? 1 : physicalIds.size();
+ this.cpuCores = coreKeys.isEmpty() ? cpuThreads : coreKeys.size();
+ }
+
+ private String collectSnapshot() throws IOException {
+ String ts = TS_FORMAT.format(Instant.now());
+
+ // /proc/loadavg: "0.50 0.35 0.25 2/150 12345"
+ String loadLine = Files.readString(PROC_LOADAVG).trim();
+ String[] loadParts = loadLine.split("\\s+");
+ String load1 = loadParts[0];
+ String load5 = loadParts[1];
+ String load15 = loadParts[2];
+ String[] runProcs = loadParts[3].split("/");
+ String running = runProcs[0];
+ String total = runProcs[1];
+
+ // /proc/meminfo
+ long memTotal = 0, memFree = 0, memAvail = 0, buffers = 0, cached = 0, swapTotal = 0, swapFree = 0;
+ for (String line : Files.readAllLines(PROC_MEMINFO)) {
+ if (line.startsWith("MemTotal:")) memTotal = parseMemValue(line);
+ else if (line.startsWith("MemFree:")) memFree = parseMemValue(line);
+ else if (line.startsWith("MemAvailable:")) memAvail = parseMemValue(line);
+ else if (line.startsWith("Buffers:")) buffers = parseMemValue(line);
+ else if (line.startsWith("Cached:")) cached = parseMemValue(line);
+ else if (line.startsWith("SwapTotal:")) swapTotal = parseMemValue(line);
+ else if (line.startsWith("SwapFree:")) swapFree = parseMemValue(line);
+ }
+
+ // /proc/diskstats
+ StringBuilder disks = new StringBuilder();
+ for (String line : Files.readAllLines(PROC_DISKSTATS)) {
+ String[] f = line.trim().split("\\s+");
+ if (f.length < 14) continue;
+ String dev = f[2];
+ if (!DISK_DEVICE_PATTERN.matcher(dev).matches()) continue;
+ if (disks.length() > 0) disks.append(',');
+ disks.append(String.format(
+ "{\"device\":\"%s\",\"readsCompleted\":%s,\"readsMerged\":%s,\"sectorsRead\":%s,\"readTimeMs\":%s,"
+ + "\"writesCompleted\":%s,\"writesMerged\":%s,\"sectorsWritten\":%s,\"writeTimeMs\":%s,"
+ + "\"ioInProgress\":%s,\"ioTimeMs\":%s,\"weightedIoTimeMs\":%s}",
+ dev, f[3], f[4], f[5], f[6], f[7], f[8], f[9], f[10], f[11], f[12], f[13]));
+ }
+
+ return String.format(
+ "{\"timestamp\":\"%s\",\"cpuSockets\":%d,\"cpuCores\":%d,\"cpuThreads\":%d,"
+ + "\"loadAvg1\":%s,\"loadAvg5\":%s,\"loadAvg15\":%s,\"runningProcs\":%s,\"totalProcs\":%s,"
+ + "\"memTotalKB\":%d,\"memFreeKB\":%d,\"memAvailableKB\":%d,\"buffersKB\":%d,\"cachedKB\":%d,"
+ + "\"swapTotalKB\":%d,\"swapFreeKB\":%d,\"diskStats\":[%s]}",
+ ts, cpuSockets, cpuCores, cpuThreads,
+ load1, load5, load15, running, total,
+ memTotal, memFree, memAvail, buffers, cached, swapTotal, swapFree,
+ disks);
+ }
+
+ private static long parseMemValue(String line) {
+ String[] parts = line.split("\\s+");
+ return Long.parseLong(parts[1]);
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/ThreadAllocTracker.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/ThreadAllocTracker.java
new file mode 100644
index 000000000..207d67045
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/ThreadAllocTracker.java
@@ -0,0 +1,203 @@
+/*
+ * 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.example.reporting;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.lang.management.ManagementFactory;
+import java.lang.management.ThreadInfo;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.time.Instant;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Periodically samples per-thread heap allocation via
+ * {@link com.sun.management.ThreadMXBean#getThreadAllocatedBytes(long[])}
+ * and writes JSONL output with per-thread deltas and cumulative totals.
+ *
+ * Lifecycle mirrors {@link SystemStatsCollector}: {@link #start(Path, String)},
+ * {@link #stop()}, {@link #isActive()}, {@link #getFileName()}.
+ */
+public final class ThreadAllocTracker {
+ private static final Logger log = LoggerFactory.getLogger(ThreadAllocTracker.class);
+
+ private static final long DEFAULT_INTERVAL_SECONDS = 10;
+
+ private final com.sun.management.ThreadMXBean threadMXBean;
+ private final long intervalSeconds;
+
+ private volatile Thread samplerThread;
+ private volatile boolean running;
+ private String fileName;
+
+ /// Creates a tracker with the default 10-second sampling interval.
+ public ThreadAllocTracker() {
+ this(DEFAULT_INTERVAL_SECONDS);
+ }
+
+ /// Creates a tracker with a custom sampling interval.
+ ///
+ /// @param intervalSeconds seconds between each sample
+ public ThreadAllocTracker(long intervalSeconds) {
+ this.threadMXBean = (com.sun.management.ThreadMXBean) ManagementFactory.getThreadMXBean();
+ this.intervalSeconds = intervalSeconds;
+ }
+
+ /// Creates the output directory, enables thread allocated memory tracking,
+ /// and spawns a daemon thread that periodically writes JSONL samples.
+ ///
+ /// @param outputDir directory to write the JSONL file into
+ /// @param fileName name of the output file
+ /// @return the absolute path of the output file
+ /// @throws IOException if the directory cannot be created
+ public Path start(Path outputDir, String fileName) throws IOException {
+ Files.createDirectories(outputDir);
+ Path outputPath = outputDir.resolve(fileName).toAbsolutePath();
+ this.fileName = fileName;
+
+ threadMXBean.setThreadAllocatedMemoryEnabled(true);
+
+ running = true;
+ samplerThread = new Thread(() -> sampleLoop(outputPath), "thread-alloc-tracker");
+ samplerThread.setDaemon(true);
+ samplerThread.start();
+
+ log.info("Thread allocation tracking started, saving to: {}", outputPath);
+ return outputPath;
+ }
+
+ /// Stops the sampler thread and writes a final cumulative summary line.
+ public void stop() {
+ running = false;
+ if (samplerThread != null) {
+ samplerThread.interrupt();
+ try {
+ samplerThread.join(5000);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ samplerThread = null;
+ log.info("Thread allocation tracking stopped, saved to: {}", fileName);
+ }
+ }
+
+ /// Returns {@code true} if the sampler thread is currently running.
+ public boolean isActive() {
+ return samplerThread != null && running;
+ }
+
+ /// Returns the current file name, or {@code null} if tracking has not been started.
+ public String getFileName() {
+ return fileName;
+ }
+
+ private void sampleLoop(Path outputPath) {
+ // Track cumulative allocations per thread (by id) for delta computation
+ var previousAllocations = new HashMap();
+
+ try (var writer = Files.newBufferedWriter(outputPath)) {
+ while (running) {
+ try {
+ Thread.sleep(intervalSeconds * 1000);
+ } catch (InterruptedException e) {
+ // On interrupt (from stop()), write final summary and exit
+ break;
+ }
+ writeSample(writer, previousAllocations, false);
+ }
+ // Write final summary with cumulative totals
+ writeSample(writer, previousAllocations, true);
+ } catch (IOException e) {
+ log.error("Failed to write thread allocation sample", e);
+ }
+ }
+
+ private void writeSample(BufferedWriter writer, Map previousAllocations, boolean isSummary)
+ throws IOException {
+ long[] threadIds = threadMXBean.getAllThreadIds();
+ long[] allocatedBytes = threadMXBean.getThreadAllocatedBytes(threadIds);
+ ThreadInfo[] threadInfos = threadMXBean.getThreadInfo(threadIds);
+
+ var sb = new StringBuilder();
+ sb.append("{\"timestamp\":\"").append(Instant.now().toString()).append('"');
+ if (isSummary) {
+ sb.append(",\"event\":\"summary\"");
+ }
+ sb.append(",\"threads\":[");
+
+ long totalAllocated = 0;
+ long totalDelta = 0;
+ boolean first = true;
+
+ for (int i = 0; i < threadIds.length; i++) {
+ if (threadInfos[i] == null || allocatedBytes[i] < 0) {
+ continue;
+ }
+ long id = threadIds[i];
+ long allocated = allocatedBytes[i];
+ long previous = previousAllocations.getOrDefault(id, 0L);
+ long delta = allocated - previous;
+ previousAllocations.put(id, allocated);
+
+ totalAllocated += allocated;
+ totalDelta += delta;
+
+ if (!first) {
+ sb.append(',');
+ }
+ first = false;
+
+ sb.append("{\"id\":").append(id)
+ .append(",\"name\":\"").append(escapeJson(threadInfos[i].getThreadName())).append('"')
+ .append(",\"allocatedBytes\":").append(allocated)
+ .append(",\"deltaBytes\":").append(delta)
+ .append('}');
+ }
+
+ sb.append("],\"totalAllocatedBytes\":").append(totalAllocated)
+ .append(",\"totalDeltaBytes\":").append(totalDelta)
+ .append('}');
+
+ writer.write(sb.toString());
+ writer.newLine();
+ writer.flush();
+ }
+
+ private static String escapeJson(String value) {
+ if (value == null) {
+ return "";
+ }
+ var sb = new StringBuilder(value.length());
+ for (int i = 0; i < value.length(); i++) {
+ char c = value.charAt(i);
+ switch (c) {
+ case '"': sb.append("\\\""); break;
+ case '\\': sb.append("\\\\"); break;
+ case '\n': sb.append("\\n"); break;
+ case '\r': sb.append("\\r"); break;
+ case '\t': sb.append("\\t"); break;
+ default: sb.append(c);
+ }
+ }
+ return sb.toString();
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/DataSetPartitioner.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/DataSetPartitioner.java
new file mode 100644
index 000000000..1e6a83f40
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/DataSetPartitioner.java
@@ -0,0 +1,60 @@
+/*
+ * 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.example.util;
+
+import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
+import io.github.jbellis.jvector.example.yaml.TestDataPartition;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Utility for partitioning a DataSet into multiple segments based on a distribution.
+ */
+public final class DataSetPartitioner {
+ private DataSetPartitioner() {}
+
+ public static final class PartitionedData {
+ public final List>> vectors;
+ public final List sizes;
+
+ public PartitionedData(List>> vectors, List sizes) {
+ this.vectors = vectors;
+ this.sizes = sizes;
+ }
+ }
+
+ public static PartitionedData partition(DataSet ds, int numParts, TestDataPartition.Distribution distribution) {
+ return partition(ds.getBaseVectors(), numParts, distribution);
+ }
+
+ public static PartitionedData partition(List> baseVectors, int numParts, TestDataPartition.Distribution distribution) {
+ List sizes = distribution.computeSplitSizes(baseVectors.size(), numParts);
+ List>> parts = new ArrayList<>(numParts);
+
+ int runningStart = 0;
+ for (int size : sizes) {
+ int start = runningStart;
+ int end = start + size;
+ runningStart = end;
+ parts.add(new ArrayList<>(baseVectors.subList(start, end)));
+ }
+
+ return new PartitionedData(parts, sizes);
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/CloudStorageLayoutUtil.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/CloudStorageLayoutUtil.java
new file mode 100644
index 000000000..9530b8815
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/CloudStorageLayoutUtil.java
@@ -0,0 +1,331 @@
+/*
+ * 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.example.util.storage;
+
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * Cloud wrapper that chooses AWS or GCP storage inspection and maps the provider-specific classes
+ * into cloud-agnostic storage tiers.
+ */
+public final class CloudStorageLayoutUtil {
+ private CloudStorageLayoutUtil() {
+ }
+
+ public enum CloudProvider {
+ AWS_EC2,
+ GCP_GCE,
+ LOCAL_OR_UNKNOWN
+ }
+
+ public enum StorageClass {
+ BLOCK_HDD_COLD,
+ BLOCK_HDD_THROUGHPUT,
+ BLOCK_HDD_STANDARD,
+ BLOCK_SSD_BALANCED,
+ BLOCK_SSD_GENERAL,
+ BLOCK_SSD_HIGH_IOPS,
+ LOCAL_SSD,
+ LOCAL_NVME,
+ NETWORK_FILESYSTEM,
+ MEMORY_TMPFS,
+ PSEUDO_FILESYSTEM,
+ UNKNOWN
+ }
+
+ public static final class StorageSnapshot {
+ private final T cloudSpecificSnapshot;
+ private final CloudProvider provider;
+ private final boolean runningInCloud;
+ private final String instanceId;
+ private final String instanceTypeOrMachineType;
+ private final String regionOrZone;
+ private final Map mountsByMountPoint;
+
+ public StorageSnapshot(T cloudSpecificSnapshot,
+ CloudProvider provider,
+ boolean runningInCloud,
+ String instanceId,
+ String instanceTypeOrMachineType,
+ String regionOrZone,
+ Map mountsByMountPoint) {
+ this.cloudSpecificSnapshot = cloudSpecificSnapshot;
+ this.provider = Objects.requireNonNull(provider, "provider");
+ this.runningInCloud = runningInCloud;
+ this.instanceId = instanceId;
+ this.instanceTypeOrMachineType = instanceTypeOrMachineType;
+ this.regionOrZone = regionOrZone;
+ this.mountsByMountPoint = Objects.requireNonNull(mountsByMountPoint, "mountsByMountPoint");
+ }
+
+ public T cloudSpecificSnapshot() {
+ return cloudSpecificSnapshot;
+ }
+
+ public CloudProvider provider() {
+ return provider;
+ }
+
+ public boolean runningInCloud() {
+ return runningInCloud;
+ }
+
+ public String instanceId() {
+ return instanceId;
+ }
+
+ public String instanceTypeOrMachineType() {
+ return instanceTypeOrMachineType;
+ }
+
+ public String regionOrZone() {
+ return regionOrZone;
+ }
+
+ public Map mountsByMountPoint() {
+ return mountsByMountPoint;
+ }
+ }
+
+ public static final class MountStorageInfo {
+ private final String mountPoint;
+ private final String source;
+ private final String filesystemType;
+ private final StorageClass storageClass;
+ private final String providerSpecificClass;
+
+ public MountStorageInfo(String mountPoint,
+ String source,
+ String filesystemType,
+ StorageClass storageClass,
+ String providerSpecificClass) {
+ this.mountPoint = mountPoint;
+ this.source = source;
+ this.filesystemType = filesystemType;
+ this.storageClass = Objects.requireNonNull(storageClass, "storageClass");
+ this.providerSpecificClass = providerSpecificClass;
+ }
+
+ public String mountPoint() {
+ return mountPoint;
+ }
+
+ public String source() {
+ return source;
+ }
+
+ public String filesystemType() {
+ return filesystemType;
+ }
+
+ public StorageClass storageClass() {
+ return storageClass;
+ }
+
+ public String providerSpecificClass() {
+ return providerSpecificClass;
+ }
+ }
+
+ public static StorageSnapshot inspectStorage() {
+ var awsSnapshot = StorageLayoutUtil.inspectStorage();
+ if (awsSnapshot.runningOnEc2()) {
+ return fromAws(awsSnapshot, CloudProvider.AWS_EC2, true);
+ }
+
+ var gcpSnapshot = GcpStorageLayoutUtil.inspectStorage();
+ if (gcpSnapshot.runningOnGcp()) {
+ return fromGcp(gcpSnapshot);
+ }
+
+ // Not in a detected cloud environment. Use OS-specific local storage inspection.
+ var localSnapshot = LocalStorageLayoutUtil.inspectStorage();
+ return fromLocal(localSnapshot);
+ }
+
+ public static Map storageClassByMountPoint() {
+ var snapshot = inspectStorage();
+ var byMountPoint = new LinkedHashMap(snapshot.mountsByMountPoint().size());
+ for (var entry : snapshot.mountsByMountPoint().entrySet()) {
+ byMountPoint.put(entry.getKey(), entry.getValue().storageClass());
+ }
+ return Collections.unmodifiableMap(byMountPoint);
+ }
+
+ private static StorageSnapshot fromAws(StorageLayoutUtil.StorageSnapshot snapshot,
+ CloudProvider provider,
+ boolean runningInCloud) {
+ var byMountPoint = new LinkedHashMap(snapshot.mountsByMountPoint().size());
+ for (var entry : snapshot.mountsByMountPoint().entrySet()) {
+ var mount = entry.getValue();
+ byMountPoint.put(
+ entry.getKey(),
+ new MountStorageInfo(
+ mount.mountPoint(),
+ mount.source(),
+ mount.filesystemType(),
+ mapAwsClass(mount.storageClass()),
+ mount.storageClass().name()
+ )
+ );
+ }
+
+ return new StorageSnapshot<>(
+ snapshot,
+ provider,
+ runningInCloud,
+ snapshot.instanceId(),
+ snapshot.instanceType(),
+ snapshot.region(),
+ Collections.unmodifiableMap(byMountPoint)
+ );
+ }
+
+ private static StorageSnapshot fromGcp(GcpStorageLayoutUtil.StorageSnapshot snapshot) {
+ var byMountPoint = new LinkedHashMap(snapshot.mountsByMountPoint().size());
+ for (var entry : snapshot.mountsByMountPoint().entrySet()) {
+ var mount = entry.getValue();
+ byMountPoint.put(
+ entry.getKey(),
+ new MountStorageInfo(
+ mount.mountPoint(),
+ mount.source(),
+ mount.filesystemType(),
+ mapGcpClass(mount.storageClass()),
+ mount.storageClass().name()
+ )
+ );
+ }
+
+ return new StorageSnapshot<>(
+ snapshot,
+ CloudProvider.GCP_GCE,
+ true,
+ snapshot.instanceId(),
+ snapshot.machineType(),
+ snapshot.zone(),
+ Collections.unmodifiableMap(byMountPoint)
+ );
+ }
+
+ private static StorageSnapshot fromLocal(LocalStorageLayoutUtil.StorageSnapshot snapshot) {
+ var byMountPoint = new LinkedHashMap(snapshot.mountsByMountPoint().size());
+ for (var entry : snapshot.mountsByMountPoint().entrySet()) {
+ var mount = entry.getValue();
+ byMountPoint.put(
+ entry.getKey(),
+ new MountStorageInfo(
+ mount.mountPoint(),
+ mount.source(),
+ mount.filesystemType(),
+ mapLocalClass(mount.storageClass()),
+ mount.storageClass().name()
+ )
+ );
+ }
+
+ return new StorageSnapshot<>(
+ snapshot,
+ CloudProvider.LOCAL_OR_UNKNOWN,
+ false,
+ null,
+ snapshot.osName(),
+ snapshot.osName(),
+ Collections.unmodifiableMap(byMountPoint)
+ );
+ }
+
+ private static StorageClass mapAwsClass(StorageLayoutUtil.StorageClass storageClass) {
+ switch (storageClass) {
+ case EBS_COLD_HDD:
+ return StorageClass.BLOCK_HDD_COLD;
+ case EBS_THROUGHPUT_HDD:
+ return StorageClass.BLOCK_HDD_THROUGHPUT;
+ case EBS_MAGNETIC:
+ return StorageClass.BLOCK_HDD_STANDARD;
+ case EBS_GP2:
+ return StorageClass.BLOCK_SSD_BALANCED;
+ case EBS_GP3:
+ return StorageClass.BLOCK_SSD_GENERAL;
+ case EBS_PROVISIONED_IOPS_SSD:
+ return StorageClass.BLOCK_SSD_HIGH_IOPS;
+ case INSTANCE_STORE_SSD:
+ return StorageClass.LOCAL_SSD;
+ case INSTANCE_STORE_NVME:
+ return StorageClass.LOCAL_NVME;
+ case NETWORK_FILESYSTEM:
+ return StorageClass.NETWORK_FILESYSTEM;
+ case MEMORY_TMPFS:
+ return StorageClass.MEMORY_TMPFS;
+ case PSEUDO_FILESYSTEM:
+ return StorageClass.PSEUDO_FILESYSTEM;
+ case UNKNOWN:
+ default:
+ return StorageClass.UNKNOWN;
+ }
+ }
+
+ private static StorageClass mapGcpClass(GcpStorageLayoutUtil.StorageClass storageClass) {
+ switch (storageClass) {
+ case PD_STANDARD_HDD:
+ return StorageClass.BLOCK_HDD_STANDARD;
+ case PD_THROUGHPUT_OPTIMIZED:
+ return StorageClass.BLOCK_HDD_THROUGHPUT;
+ case PD_BALANCED_SSD:
+ return StorageClass.BLOCK_SSD_BALANCED;
+ case PD_SSD:
+ return StorageClass.BLOCK_SSD_GENERAL;
+ case PD_EXTREME_SSD:
+ return StorageClass.BLOCK_SSD_HIGH_IOPS;
+ case LOCAL_SSD:
+ return StorageClass.LOCAL_SSD;
+ case LOCAL_NVME:
+ return StorageClass.LOCAL_NVME;
+ case NETWORK_FILESYSTEM:
+ return StorageClass.NETWORK_FILESYSTEM;
+ case MEMORY_TMPFS:
+ return StorageClass.MEMORY_TMPFS;
+ case PSEUDO_FILESYSTEM:
+ return StorageClass.PSEUDO_FILESYSTEM;
+ case UNKNOWN:
+ default:
+ return StorageClass.UNKNOWN;
+ }
+ }
+
+ private static StorageClass mapLocalClass(LocalStorageLayoutUtil.StorageClass storageClass) {
+ switch (storageClass) {
+ case LOCAL_HDD:
+ return StorageClass.BLOCK_HDD_STANDARD;
+ case LOCAL_SSD:
+ return StorageClass.LOCAL_SSD;
+ case LOCAL_NVME:
+ return StorageClass.LOCAL_NVME;
+ case NETWORK_FILESYSTEM:
+ return StorageClass.NETWORK_FILESYSTEM;
+ case MEMORY_TMPFS:
+ return StorageClass.MEMORY_TMPFS;
+ case PSEUDO_FILESYSTEM:
+ return StorageClass.PSEUDO_FILESYSTEM;
+ case UNKNOWN:
+ default:
+ return StorageClass.UNKNOWN;
+ }
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/GcpStorageLayoutUtil.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/GcpStorageLayoutUtil.java
new file mode 100644
index 000000000..a4ecb6d09
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/GcpStorageLayoutUtil.java
@@ -0,0 +1,695 @@
+/*
+ * 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.example.util.storage;
+
+import java.io.IOException;
+import java.net.URI;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.net.http.HttpResponse;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.regex.Pattern;
+import java.util.stream.Stream;
+
+/**
+ * Detects GCE runtime context via metadata service and classifies storage for each mounted filesystem.
+ */
+public final class GcpStorageLayoutUtil {
+ private static final String GCE_METADATA_HOST_ENV = "GCE_METADATA_HOST";
+ private static final String METADATA_HOST_DEFAULT = "metadata.google.internal";
+ private static final String METADATA_PREFIX = "/computeMetadata/v1/";
+ private static final String METADATA_FLAVOR_HEADER = "Metadata-Flavor";
+ private static final String METADATA_FLAVOR_VALUE = "Google";
+ private static final Duration METADATA_TIMEOUT = Duration.ofMillis(300);
+
+ private static final Pattern NVME_PARTITION_SUFFIX = Pattern.compile("p\\d+$");
+ private static final Pattern GENERIC_PARTITION_SUFFIX = Pattern.compile("\\d+$");
+ private static final Set NETWORK_FILESYSTEM_TYPES = Set.of("nfs", "nfs4", "efs", "cifs", "smbfs", "fuse.sshfs");
+
+ private GcpStorageLayoutUtil() {
+ }
+
+ public enum StorageClass {
+ PD_STANDARD_HDD,
+ PD_THROUGHPUT_OPTIMIZED,
+ PD_BALANCED_SSD,
+ PD_SSD,
+ PD_EXTREME_SSD,
+ LOCAL_SSD,
+ LOCAL_NVME,
+ NETWORK_FILESYSTEM,
+ MEMORY_TMPFS,
+ PSEUDO_FILESYSTEM,
+ UNKNOWN
+ }
+
+ public static final class StorageSnapshot {
+ private final boolean runningOnGcp;
+ private final String instanceId;
+ private final String machineType;
+ private final String zone;
+ private final Map mountsByMountPoint;
+
+ public StorageSnapshot(boolean runningOnGcp,
+ String instanceId,
+ String machineType,
+ String zone,
+ Map mountsByMountPoint) {
+ this.runningOnGcp = runningOnGcp;
+ this.instanceId = instanceId;
+ this.machineType = machineType;
+ this.zone = zone;
+ this.mountsByMountPoint = Objects.requireNonNull(mountsByMountPoint, "mountsByMountPoint");
+ }
+
+ public boolean runningOnGcp() {
+ return runningOnGcp;
+ }
+
+ public String instanceId() {
+ return instanceId;
+ }
+
+ public String machineType() {
+ return machineType;
+ }
+
+ public String zone() {
+ return zone;
+ }
+
+ public Map mountsByMountPoint() {
+ return mountsByMountPoint;
+ }
+ }
+
+ public static final class MountStorageInfo {
+ private final String mountPoint;
+ private final String source;
+ private final String filesystemType;
+ private final StorageClass storageClass;
+ private final String deviceName;
+ private final String diskKind;
+ private final String interfaceType;
+
+ public MountStorageInfo(String mountPoint,
+ String source,
+ String filesystemType,
+ StorageClass storageClass,
+ String deviceName,
+ String diskKind,
+ String interfaceType) {
+ this.mountPoint = mountPoint;
+ this.source = source;
+ this.filesystemType = filesystemType;
+ this.storageClass = Objects.requireNonNull(storageClass, "storageClass");
+ this.deviceName = deviceName;
+ this.diskKind = diskKind;
+ this.interfaceType = interfaceType;
+ }
+
+ public String mountPoint() {
+ return mountPoint;
+ }
+
+ public String source() {
+ return source;
+ }
+
+ public String filesystemType() {
+ return filesystemType;
+ }
+
+ public StorageClass storageClass() {
+ return storageClass;
+ }
+
+ public String deviceName() {
+ return deviceName;
+ }
+
+ public String diskKind() {
+ return diskKind;
+ }
+
+ public String interfaceType() {
+ return interfaceType;
+ }
+ }
+
+ public static StorageSnapshot inspectStorage() {
+ var identity = fetchGcpIdentity();
+ var mounts = readMountEntries();
+ var diskData = identity.map(GcpStorageLayoutUtil::fetchGcpDiskData).orElse(GcpDiskData.empty());
+
+ mounts.sort(Comparator.comparing(MountEntry::mountPoint));
+ var byMountPoint = new LinkedHashMap(mounts.size());
+ for (var mount : mounts) {
+ var diskResolution = resolveDisk(mount.source(), diskData);
+ var storageClass = classify(mount, diskResolution);
+ byMountPoint.put(
+ mount.mountPoint(),
+ new MountStorageInfo(
+ mount.mountPoint(),
+ mount.source(),
+ mount.filesystemType(),
+ storageClass,
+ diskResolution.deviceName(),
+ diskResolution.diskKind(),
+ diskResolution.interfaceType()
+ )
+ );
+ }
+
+ return new StorageSnapshot(
+ identity.isPresent(),
+ identity.map(GcpIdentity::instanceId).orElse(null),
+ identity.map(GcpIdentity::machineType).orElse(null),
+ identity.map(GcpIdentity::zone).orElse(null),
+ Collections.unmodifiableMap(byMountPoint)
+ );
+ }
+
+ public static Map storageClassByMountPoint() {
+ var snapshot = inspectStorage();
+ var byMountPoint = new LinkedHashMap(snapshot.mountsByMountPoint().size());
+ for (var entry : snapshot.mountsByMountPoint().entrySet()) {
+ byMountPoint.put(entry.getKey(), entry.getValue().storageClass());
+ }
+ return Collections.unmodifiableMap(byMountPoint);
+ }
+
+ private static Optional fetchGcpIdentity() {
+ var client = HttpClient.newBuilder()
+ .connectTimeout(METADATA_TIMEOUT)
+ .build();
+
+ var instanceId = readMetadata(client, "instance/id");
+ if (instanceId == null || instanceId.isBlank()) {
+ return Optional.empty();
+ }
+
+ var machineType = parseLeafResource(readMetadata(client, "instance/machine-type"));
+ var zone = parseLeafResource(readMetadata(client, "instance/zone"));
+ return Optional.of(new GcpIdentity(instanceId.trim(), machineType, zone));
+ }
+
+ private static GcpDiskData fetchGcpDiskData(GcpIdentity ignoredIdentity) {
+ var byDeviceName = fetchDisksByDeviceNameFromMetadata();
+ var aliasesByNormalizedDevice = mapGoogleAliasesByNormalizedDevice();
+ return new GcpDiskData(byDeviceName, aliasesByNormalizedDevice);
+ }
+
+ private static Map fetchDisksByDeviceNameFromMetadata() {
+ var client = HttpClient.newBuilder()
+ .connectTimeout(METADATA_TIMEOUT)
+ .build();
+
+ var indexListing = readMetadata(client, "instance/disks/");
+ if (indexListing == null || indexListing.isBlank()) {
+ return Map.of();
+ }
+
+ var byDeviceName = new LinkedHashMap();
+ for (var rawLine : indexListing.split("\n")) {
+ var line = rawLine.trim();
+ if (line.isEmpty()) {
+ continue;
+ }
+ var index = line.endsWith("/") ? line.substring(0, line.length() - 1) : line;
+ var deviceName = readMetadata(client, "instance/disks/" + index + "/device-name");
+ if (deviceName == null || deviceName.isBlank()) {
+ continue;
+ }
+
+ var diskKind = safeLower(readMetadata(client, "instance/disks/" + index + "/type"));
+ var interfaceType = safeUpper(readMetadata(client, "instance/disks/" + index + "/interface"));
+ var diskTypeHint = readMetadata(client, "instance/disks/" + index + "/disk-type");
+ byDeviceName.put(deviceName.trim(), new GcpDiskInfo(deviceName.trim(), diskKind, interfaceType, safeLower(diskTypeHint)));
+ }
+ return byDeviceName;
+ }
+
+ private static Map> mapGoogleAliasesByNormalizedDevice() {
+ var byIdDir = Path.of("/dev/disk/by-id");
+ if (!Files.isDirectory(byIdDir)) {
+ return Map.of();
+ }
+
+ var aliasesByDevice = new LinkedHashMap>();
+ try (Stream entries = Files.list(byIdDir)) {
+ entries.filter(Files::isSymbolicLink).forEach(link -> {
+ var alias = link.getFileName().toString();
+ if (!alias.startsWith("google-")) {
+ return;
+ }
+ try {
+ var target = normalizeDevice(link.toRealPath().toString());
+ aliasesByDevice.computeIfAbsent(target, unused -> new ArrayList<>()).add(alias);
+ } catch (IOException ignored) {
+ // continue
+ }
+ });
+ } catch (IOException ignored) {
+ return Map.of();
+ }
+
+ for (var aliases : aliasesByDevice.values()) {
+ aliases.sort(String::compareTo);
+ }
+ return aliasesByDevice;
+ }
+
+ private static DiskResolution resolveDisk(String mountSource, GcpDiskData diskData) {
+ if (mountSource == null || !mountSource.startsWith("/dev/")) {
+ return DiskResolution.empty();
+ }
+
+ var normalized = normalizeDevice(mountSource);
+ var aliases = diskData.aliasesByNormalizedDevice().getOrDefault(normalized, List.of());
+ var primaryAlias = aliases.isEmpty() ? null : aliases.get(0);
+ var inferredDeviceName = primaryAlias == null ? null : stripGooglePrefix(primaryAlias);
+ GcpDiskInfo info = inferredDeviceName == null ? null : diskData.byDeviceName().get(inferredDeviceName);
+
+ // Try all aliases in case the first one doesn't match a metadata device-name.
+ if (info == null) {
+ for (var alias : aliases) {
+ var candidate = stripGooglePrefix(alias);
+ if (candidate == null) {
+ continue;
+ }
+ info = diskData.byDeviceName().get(candidate);
+ if (info != null) {
+ inferredDeviceName = candidate;
+ break;
+ }
+ }
+ }
+
+ var rotational = readRotationalFlag(normalized);
+ if (info == null) {
+ return new DiskResolution(normalized, inferredDeviceName, null, null, null, rotational);
+ }
+ return new DiskResolution(
+ normalized,
+ inferredDeviceName,
+ info.diskKind(),
+ info.interfaceType(),
+ info.diskTypeHint(),
+ rotational
+ );
+ }
+
+ private static StorageClass classify(MountEntry mount, DiskResolution diskResolution) {
+ var fsType = safeLower(mount.filesystemType());
+ var source = mount.source();
+ var sourceLower = safeLower(source);
+
+ if ("tmpfs".equals(fsType)) {
+ return StorageClass.MEMORY_TMPFS;
+ }
+ if (NETWORK_FILESYSTEM_TYPES.contains(fsType)) {
+ return StorageClass.NETWORK_FILESYSTEM;
+ }
+ if (isPseudoFileSystem(fsType, sourceLower)) {
+ return StorageClass.PSEUDO_FILESYSTEM;
+ }
+
+ if ("scratch".equals(diskResolution.diskKind())) {
+ if ("NVME".equals(diskResolution.interfaceType()) || sourceLower.contains("nvme")) {
+ return StorageClass.LOCAL_NVME;
+ }
+ return StorageClass.LOCAL_SSD;
+ }
+ if ("persistent".equals(diskResolution.diskKind())) {
+ return classifyPersistentDisk(diskResolution);
+ }
+
+ // Best-effort fallback based on device name hints and local block characteristics.
+ var hints = safeLower(diskResolution.deviceName()) + " "
+ + safeLower(diskResolution.diskTypeHint()) + " "
+ + sourceLower;
+ if (hints.contains("local-ssd")) {
+ return sourceLower.contains("nvme") ? StorageClass.LOCAL_NVME : StorageClass.LOCAL_SSD;
+ }
+ if (source != null && source.startsWith("/dev/")) {
+ if (sourceLower.contains("nvme")) {
+ return StorageClass.LOCAL_NVME;
+ }
+ if (Boolean.TRUE.equals(diskResolution.rotational())) {
+ return StorageClass.PD_STANDARD_HDD;
+ }
+ return StorageClass.LOCAL_SSD;
+ }
+ return StorageClass.UNKNOWN;
+ }
+
+ private static StorageClass classifyPersistentDisk(DiskResolution diskResolution) {
+ var hints = safeLower(diskResolution.deviceName()) + " " + safeLower(diskResolution.diskTypeHint());
+ if (hints.contains("extreme")) {
+ return StorageClass.PD_EXTREME_SSD;
+ }
+ if (hints.contains("throughput")) {
+ return StorageClass.PD_THROUGHPUT_OPTIMIZED;
+ }
+ if (hints.contains("balanced")) {
+ return StorageClass.PD_BALANCED_SSD;
+ }
+ if (hints.contains("pd-ssd") || hints.contains("ssd")) {
+ return StorageClass.PD_SSD;
+ }
+ if (hints.contains("standard")) {
+ return StorageClass.PD_STANDARD_HDD;
+ }
+
+ if (Boolean.TRUE.equals(diskResolution.rotational())) {
+ return StorageClass.PD_STANDARD_HDD;
+ }
+ return StorageClass.PD_BALANCED_SSD;
+ }
+
+ private static String readMetadata(HttpClient client, String relativePath) {
+ var host = Optional.ofNullable(System.getenv(GCE_METADATA_HOST_ENV)).orElse(METADATA_HOST_DEFAULT);
+ var uri = URI.create("http://" + host + METADATA_PREFIX + relativePath);
+ try {
+ var request = HttpRequest.newBuilder(uri)
+ .timeout(METADATA_TIMEOUT)
+ .header(METADATA_FLAVOR_HEADER, METADATA_FLAVOR_VALUE)
+ .GET()
+ .build();
+ var response = client.send(request, HttpResponse.BodyHandlers.ofString());
+ if (response.statusCode() != 200) {
+ return null;
+ }
+ var flavorHeader = response.headers().firstValue(METADATA_FLAVOR_HEADER).orElse("");
+ if (!METADATA_FLAVOR_VALUE.equalsIgnoreCase(flavorHeader)) {
+ return null;
+ }
+ return response.body();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return null;
+ } catch (IOException e) {
+ return null;
+ }
+ }
+
+ private static String parseLeafResource(String value) {
+ if (value == null) {
+ return null;
+ }
+ var trimmed = value.trim();
+ if (trimmed.isEmpty()) {
+ return null;
+ }
+ var idx = trimmed.lastIndexOf('/');
+ if (idx < 0 || idx == trimmed.length() - 1) {
+ return trimmed;
+ }
+ return trimmed.substring(idx + 1);
+ }
+
+ private static List readMountEntries() {
+ var mountsPath = Files.isReadable(Path.of("/proc/self/mounts"))
+ ? Path.of("/proc/self/mounts")
+ : Path.of("/proc/mounts");
+
+ if (!Files.isReadable(mountsPath)) {
+ return new ArrayList<>();
+ }
+
+ var entries = new ArrayList();
+ try (Stream lines = Files.lines(mountsPath)) {
+ lines.forEach(line -> {
+ var parts = line.split(" ");
+ if (parts.length < 3) {
+ return;
+ }
+ var source = decodeMountToken(parts[0]);
+ var mountPoint = decodeMountToken(parts[1]);
+ var filesystemType = decodeMountToken(parts[2]);
+ entries.add(new MountEntry(source, mountPoint, filesystemType));
+ });
+ } catch (IOException ignored) {
+ return new ArrayList<>();
+ }
+ return entries;
+ }
+
+ private static Boolean readRotationalFlag(String normalizedDevice) {
+ if (normalizedDevice == null || !normalizedDevice.startsWith("/dev/")) {
+ return null;
+ }
+ var blockName = normalizedDevice.substring("/dev/".length());
+ var rotaPath = Path.of("/sys/class/block", blockName, "queue", "rotational");
+ if (!Files.isReadable(rotaPath)) {
+ return null;
+ }
+ try {
+ var value = Files.readString(rotaPath).trim();
+ if ("1".equals(value)) {
+ return Boolean.TRUE;
+ }
+ if ("0".equals(value)) {
+ return Boolean.FALSE;
+ }
+ } catch (IOException ignored) {
+ return null;
+ }
+ return null;
+ }
+
+ private static boolean isPseudoFileSystem(String fsType, String sourceLower) {
+ return fsType.equals("proc")
+ || fsType.equals("sysfs")
+ || fsType.equals("devpts")
+ || fsType.equals("devtmpfs")
+ || fsType.equals("cgroup")
+ || fsType.equals("cgroup2")
+ || fsType.equals("autofs")
+ || fsType.equals("mqueue")
+ || fsType.equals("tracefs")
+ || fsType.equals("pstore")
+ || fsType.equals("securityfs")
+ || fsType.equals("debugfs")
+ || fsType.equals("configfs")
+ || fsType.equals("fusectl")
+ || fsType.equals("binfmt_misc")
+ || fsType.equals("rpc_pipefs")
+ || sourceLower.equals("proc")
+ || sourceLower.equals("sysfs")
+ || sourceLower.equals("tmpfs");
+ }
+
+ private static String normalizeDevice(String device) {
+ if (device == null) {
+ return null;
+ }
+ if (!device.startsWith("/dev/")) {
+ return device;
+ }
+ if (device.startsWith("/dev/nvme")) {
+ return NVME_PARTITION_SUFFIX.matcher(device).replaceAll("");
+ }
+ return GENERIC_PARTITION_SUFFIX.matcher(device).replaceAll("");
+ }
+
+ private static String decodeMountToken(String token) {
+ return token
+ .replace("\\040", " ")
+ .replace("\\011", "\t")
+ .replace("\\012", "\n")
+ .replace("\\134", "\\");
+ }
+
+ private static String stripGooglePrefix(String alias) {
+ if (alias == null || !alias.startsWith("google-") || alias.length() <= "google-".length()) {
+ return null;
+ }
+ return alias.substring("google-".length());
+ }
+
+ private static String safeLower(String value) {
+ return value == null ? "" : value.toLowerCase(Locale.ROOT);
+ }
+
+ private static String safeUpper(String value) {
+ return value == null ? null : value.trim().toUpperCase(Locale.ROOT);
+ }
+
+ private static final class MountEntry {
+ private final String source;
+ private final String mountPoint;
+ private final String filesystemType;
+
+ private MountEntry(String source, String mountPoint, String filesystemType) {
+ this.source = source;
+ this.mountPoint = mountPoint;
+ this.filesystemType = filesystemType;
+ }
+
+ private String source() {
+ return source;
+ }
+
+ private String mountPoint() {
+ return mountPoint;
+ }
+
+ private String filesystemType() {
+ return filesystemType;
+ }
+ }
+
+ private static final class GcpIdentity {
+ private final String instanceId;
+ private final String machineType;
+ private final String zone;
+
+ private GcpIdentity(String instanceId, String machineType, String zone) {
+ this.instanceId = instanceId;
+ this.machineType = machineType;
+ this.zone = zone;
+ }
+
+ private String instanceId() {
+ return instanceId;
+ }
+
+ private String machineType() {
+ return machineType;
+ }
+
+ private String zone() {
+ return zone;
+ }
+ }
+
+ private static final class GcpDiskInfo {
+ private final String deviceName;
+ private final String diskKind;
+ private final String interfaceType;
+ private final String diskTypeHint;
+
+ private GcpDiskInfo(String deviceName, String diskKind, String interfaceType, String diskTypeHint) {
+ this.deviceName = deviceName;
+ this.diskKind = diskKind;
+ this.interfaceType = interfaceType;
+ this.diskTypeHint = diskTypeHint;
+ }
+
+ private String deviceName() {
+ return deviceName;
+ }
+
+ private String diskKind() {
+ return diskKind;
+ }
+
+ private String interfaceType() {
+ return interfaceType;
+ }
+
+ private String diskTypeHint() {
+ return diskTypeHint;
+ }
+ }
+
+ private static final class GcpDiskData {
+ private final Map byDeviceName;
+ private final Map> aliasesByNormalizedDevice;
+
+ private GcpDiskData(Map byDeviceName, Map> aliasesByNormalizedDevice) {
+ this.byDeviceName = Objects.requireNonNull(byDeviceName, "byDeviceName");
+ this.aliasesByNormalizedDevice = Objects.requireNonNull(aliasesByNormalizedDevice, "aliasesByNormalizedDevice");
+ }
+
+ private Map byDeviceName() {
+ return byDeviceName;
+ }
+
+ private Map> aliasesByNormalizedDevice() {
+ return aliasesByNormalizedDevice;
+ }
+
+ private static GcpDiskData empty() {
+ return new GcpDiskData(Map.of(), Map.of());
+ }
+ }
+
+ private static final class DiskResolution {
+ private final String normalizedDevice;
+ private final String deviceName;
+ private final String diskKind;
+ private final String interfaceType;
+ private final String diskTypeHint;
+ private final Boolean rotational;
+
+ private DiskResolution(String normalizedDevice,
+ String deviceName,
+ String diskKind,
+ String interfaceType,
+ String diskTypeHint,
+ Boolean rotational) {
+ this.normalizedDevice = normalizedDevice;
+ this.deviceName = deviceName;
+ this.diskKind = diskKind;
+ this.interfaceType = interfaceType;
+ this.diskTypeHint = diskTypeHint;
+ this.rotational = rotational;
+ }
+
+ private static DiskResolution empty() {
+ return new DiskResolution(null, null, null, null, null, null);
+ }
+
+ private String normalizedDevice() {
+ return normalizedDevice;
+ }
+
+ private String deviceName() {
+ return deviceName;
+ }
+
+ private String diskKind() {
+ return diskKind;
+ }
+
+ private String interfaceType() {
+ return interfaceType;
+ }
+
+ private String diskTypeHint() {
+ return diskTypeHint;
+ }
+
+ private Boolean rotational() {
+ return rotational;
+ }
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/LocalStorageLayoutUtil.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/LocalStorageLayoutUtil.java
new file mode 100644
index 000000000..23dad5a18
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/LocalStorageLayoutUtil.java
@@ -0,0 +1,524 @@
+/*
+ * 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.example.util.storage;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.regex.Pattern;
+import java.util.stream.Stream;
+
+/**
+ * Best-effort storage inspection utility for non-cloud environments.
+ * Supports Linux, macOS, and Windows using local OS signals and common mount metadata.
+ */
+public final class LocalStorageLayoutUtil {
+ private static final Pattern LINUX_NVME_PARTITION_SUFFIX = Pattern.compile("p\\d+$");
+ private static final Pattern GENERIC_PARTITION_SUFFIX = Pattern.compile("\\d+$");
+ private static final Pattern MAC_MOUNT_PATTERN = Pattern.compile("^(.+) on (.+) \\((.+)\\)$");
+ private static final Pattern MAC_DISK_SLICE_SUFFIX = Pattern.compile("s\\d+$");
+ private static final Set NETWORK_FILESYSTEM_TYPES =
+ Set.of("nfs", "nfs4", "efs", "cifs", "smbfs", "fuse.sshfs", "afpfs", "webdav", "davfs");
+
+ private LocalStorageLayoutUtil() {
+ }
+
+ public enum StorageClass {
+ LOCAL_HDD,
+ LOCAL_SSD,
+ LOCAL_NVME,
+ NETWORK_FILESYSTEM,
+ MEMORY_TMPFS,
+ PSEUDO_FILESYSTEM,
+ UNKNOWN
+ }
+
+ public static final class StorageSnapshot {
+ private final String osName;
+ private final Map mountsByMountPoint;
+
+ public StorageSnapshot(String osName, Map mountsByMountPoint) {
+ this.osName = osName;
+ this.mountsByMountPoint = Objects.requireNonNull(mountsByMountPoint, "mountsByMountPoint");
+ }
+
+ public String osName() {
+ return osName;
+ }
+
+ public Map mountsByMountPoint() {
+ return mountsByMountPoint;
+ }
+ }
+
+ public static final class MountStorageInfo {
+ private final String mountPoint;
+ private final String source;
+ private final String filesystemType;
+ private final StorageClass storageClass;
+ private final String osHint;
+
+ public MountStorageInfo(String mountPoint,
+ String source,
+ String filesystemType,
+ StorageClass storageClass,
+ String osHint) {
+ this.mountPoint = mountPoint;
+ this.source = source;
+ this.filesystemType = filesystemType;
+ this.storageClass = Objects.requireNonNull(storageClass, "storageClass");
+ this.osHint = osHint;
+ }
+
+ public String mountPoint() {
+ return mountPoint;
+ }
+
+ public String source() {
+ return source;
+ }
+
+ public String filesystemType() {
+ return filesystemType;
+ }
+
+ public StorageClass storageClass() {
+ return storageClass;
+ }
+
+ public String osHint() {
+ return osHint;
+ }
+ }
+
+ public static StorageSnapshot inspectStorage() {
+ var os = safeLower(System.getProperty("os.name"));
+ List mounts;
+ if (isLinux(os)) {
+ mounts = readLinuxMountEntries();
+ } else if (isMac(os)) {
+ mounts = readMacMountEntries();
+ } else if (isWindows(os)) {
+ mounts = readWindowsMountEntries();
+ } else {
+ mounts = readGenericMountEntries();
+ }
+
+ mounts.sort(Comparator.comparing(MountEntry::mountPoint));
+ var byMountPoint = new LinkedHashMap(mounts.size());
+ for (var mount : mounts) {
+ StorageClass storageClass;
+ String osHint;
+ if (isLinux(os)) {
+ storageClass = classifyLinux(mount);
+ osHint = "linux";
+ } else if (isMac(os)) {
+ storageClass = classifyMac(mount);
+ osHint = "macos";
+ } else if (isWindows(os)) {
+ storageClass = classifyWindows(mount);
+ osHint = "windows";
+ } else {
+ storageClass = classifyGeneric(mount);
+ osHint = "generic";
+ }
+
+ byMountPoint.put(
+ mount.mountPoint(),
+ new MountStorageInfo(
+ mount.mountPoint(),
+ mount.source(),
+ mount.filesystemType(),
+ storageClass,
+ osHint
+ )
+ );
+ }
+
+ return new StorageSnapshot(
+ System.getProperty("os.name"),
+ Collections.unmodifiableMap(byMountPoint)
+ );
+ }
+
+ public static Map storageClassByMountPoint() {
+ var snapshot = inspectStorage();
+ var byMountPoint = new LinkedHashMap(snapshot.mountsByMountPoint().size());
+ for (var entry : snapshot.mountsByMountPoint().entrySet()) {
+ byMountPoint.put(entry.getKey(), entry.getValue().storageClass());
+ }
+ return Collections.unmodifiableMap(byMountPoint);
+ }
+
+ private static List readLinuxMountEntries() {
+ var mountsPath = Files.isReadable(Path.of("/proc/self/mounts"))
+ ? Path.of("/proc/self/mounts")
+ : Path.of("/proc/mounts");
+ if (!Files.isReadable(mountsPath)) {
+ return new ArrayList<>();
+ }
+
+ var entries = new ArrayList();
+ try (Stream lines = Files.lines(mountsPath)) {
+ lines.forEach(line -> {
+ var parts = line.split(" ");
+ if (parts.length < 3) {
+ return;
+ }
+ entries.add(new MountEntry(
+ decodeMountToken(parts[0]),
+ decodeMountToken(parts[1]),
+ decodeMountToken(parts[2])
+ ));
+ });
+ } catch (IOException ignored) {
+ return new ArrayList<>();
+ }
+ return entries;
+ }
+
+ private static List readMacMountEntries() {
+ var entries = new ArrayList();
+ for (String line : runCommandLines("mount")) {
+ var matcher = MAC_MOUNT_PATTERN.matcher(line);
+ if (!matcher.matches()) {
+ continue;
+ }
+ var source = matcher.group(1).trim();
+ var mountPoint = matcher.group(2).trim();
+ var options = matcher.group(3).trim();
+ var fsType = options.split(",")[0].trim();
+ entries.add(new MountEntry(source, mountPoint, fsType));
+ }
+ if (entries.isEmpty()) {
+ return readGenericMountEntries();
+ }
+ return entries;
+ }
+
+ private static List readWindowsMountEntries() {
+ var entries = new ArrayList();
+ var roots = File.listRoots();
+ if (roots == null) {
+ return entries;
+ }
+ for (var root : roots) {
+ if (root == null) {
+ continue;
+ }
+ var path = root.toPath();
+ String fsType = "unknown";
+ try {
+ fsType = Files.getFileStore(path).type();
+ } catch (IOException ignored) {
+ // keep default
+ }
+ entries.add(new MountEntry(root.getPath(), root.getPath(), fsType));
+ }
+ return entries;
+ }
+
+ private static List readGenericMountEntries() {
+ var entries = new ArrayList();
+ var roots = File.listRoots();
+ if (roots == null) {
+ return entries;
+ }
+ for (var root : roots) {
+ if (root == null) {
+ continue;
+ }
+ String fsType = "unknown";
+ try {
+ fsType = Files.getFileStore(root.toPath()).type();
+ } catch (IOException ignored) {
+ // keep default
+ }
+ entries.add(new MountEntry(root.getPath(), root.getPath(), fsType));
+ }
+ return entries;
+ }
+
+ private static StorageClass classifyLinux(MountEntry mount) {
+ var fsType = safeLower(mount.filesystemType());
+ var source = mount.source();
+ var sourceLower = safeLower(source);
+
+ if ("tmpfs".equals(fsType) || "ramfs".equals(fsType)) {
+ return StorageClass.MEMORY_TMPFS;
+ }
+ if (NETWORK_FILESYSTEM_TYPES.contains(fsType) || sourceLower.startsWith("//")) {
+ return StorageClass.NETWORK_FILESYSTEM;
+ }
+ if (isPseudoFileSystem(fsType, sourceLower)) {
+ return StorageClass.PSEUDO_FILESYSTEM;
+ }
+
+ if (source != null && source.startsWith("/dev/")) {
+ var normalized = normalizeLinuxDevice(sourceLower);
+ if (normalized.contains("nvme")) {
+ return StorageClass.LOCAL_NVME;
+ }
+
+ Boolean rotational = readLinuxRotationalFlag(normalized);
+ if (Boolean.TRUE.equals(rotational)) {
+ return StorageClass.LOCAL_HDD;
+ }
+ if (Boolean.FALSE.equals(rotational)) {
+ return StorageClass.LOCAL_SSD;
+ }
+ return StorageClass.UNKNOWN;
+ }
+ return StorageClass.UNKNOWN;
+ }
+
+ private static StorageClass classifyMac(MountEntry mount) {
+ var fsType = safeLower(mount.filesystemType());
+ var source = mount.source();
+ var sourceLower = safeLower(source);
+
+ if ("devfs".equals(fsType) || "autofs".equals(fsType) || "procfs".equals(fsType)) {
+ return StorageClass.PSEUDO_FILESYSTEM;
+ }
+ if ("tmpfs".equals(fsType) || "ramfs".equals(fsType)) {
+ return StorageClass.MEMORY_TMPFS;
+ }
+ if (NETWORK_FILESYSTEM_TYPES.contains(fsType) || sourceLower.startsWith("//")) {
+ return StorageClass.NETWORK_FILESYSTEM;
+ }
+
+ if (source != null && source.startsWith("/dev/")) {
+ var diskInfo = readMacDiskInfo(source);
+ if (diskInfo.protocolNvme) {
+ return StorageClass.LOCAL_NVME;
+ }
+ if (diskInfo.solidState != null) {
+ return diskInfo.solidState ? StorageClass.LOCAL_SSD : StorageClass.LOCAL_HDD;
+ }
+ if (sourceLower.contains("nvme")) {
+ return StorageClass.LOCAL_NVME;
+ }
+ return StorageClass.UNKNOWN;
+ }
+ return StorageClass.UNKNOWN;
+ }
+
+ private static StorageClass classifyWindows(MountEntry mount) {
+ var fsType = safeLower(mount.filesystemType());
+ var source = mount.source();
+ var sourceLower = safeLower(source);
+
+ if (NETWORK_FILESYSTEM_TYPES.contains(fsType)
+ || fsType.contains("smb")
+ || fsType.contains("cifs")
+ || sourceLower.startsWith("\\\\")) {
+ return StorageClass.NETWORK_FILESYSTEM;
+ }
+ if (fsType.contains("tmp") || fsType.contains("ram")) {
+ return StorageClass.MEMORY_TMPFS;
+ }
+
+ // Generic stub: fixed drives are treated as local SSD class when media specifics are unavailable.
+ if (source != null && source.matches("^[A-Za-z]:\\\\.*")) {
+ return StorageClass.LOCAL_SSD;
+ }
+ return StorageClass.UNKNOWN;
+ }
+
+ private static StorageClass classifyGeneric(MountEntry mount) {
+ var fsType = safeLower(mount.filesystemType());
+ if ("tmpfs".equals(fsType) || "ramfs".equals(fsType)) {
+ return StorageClass.MEMORY_TMPFS;
+ }
+ if (NETWORK_FILESYSTEM_TYPES.contains(fsType)) {
+ return StorageClass.NETWORK_FILESYSTEM;
+ }
+ return StorageClass.UNKNOWN;
+ }
+
+ private static boolean isPseudoFileSystem(String fsType, String sourceLower) {
+ return fsType.equals("proc")
+ || fsType.equals("sysfs")
+ || fsType.equals("devpts")
+ || fsType.equals("devtmpfs")
+ || fsType.equals("cgroup")
+ || fsType.equals("cgroup2")
+ || fsType.equals("autofs")
+ || fsType.equals("mqueue")
+ || fsType.equals("tracefs")
+ || fsType.equals("pstore")
+ || fsType.equals("securityfs")
+ || fsType.equals("debugfs")
+ || fsType.equals("configfs")
+ || fsType.equals("fusectl")
+ || fsType.equals("binfmt_misc")
+ || fsType.equals("rpc_pipefs")
+ || sourceLower.equals("proc")
+ || sourceLower.equals("sysfs")
+ || sourceLower.equals("tmpfs");
+ }
+
+ private static String normalizeLinuxDevice(String device) {
+ if (!device.startsWith("/dev/")) {
+ return device;
+ }
+ if (device.startsWith("/dev/nvme")) {
+ return LINUX_NVME_PARTITION_SUFFIX.matcher(device).replaceAll("");
+ }
+ return GENERIC_PARTITION_SUFFIX.matcher(device).replaceAll("");
+ }
+
+ private static Boolean readLinuxRotationalFlag(String normalizedDevice) {
+ if (normalizedDevice == null || !normalizedDevice.startsWith("/dev/")) {
+ return null;
+ }
+ var blockName = normalizedDevice.substring("/dev/".length());
+ var rotaPath = Path.of("/sys/class/block", blockName, "queue", "rotational");
+ if (!Files.isReadable(rotaPath)) {
+ return null;
+ }
+ try {
+ var value = Files.readString(rotaPath).trim();
+ if ("1".equals(value)) {
+ return Boolean.TRUE;
+ }
+ if ("0".equals(value)) {
+ return Boolean.FALSE;
+ }
+ } catch (IOException ignored) {
+ return null;
+ }
+ return null;
+ }
+
+ private static MacDiskInfo readMacDiskInfo(String sourceDevice) {
+ var base = sourceDevice;
+ var slash = sourceDevice.lastIndexOf('/');
+ if (slash >= 0 && slash + 1 < sourceDevice.length()) {
+ var leaf = sourceDevice.substring(slash + 1);
+ if (leaf.startsWith("disk")) {
+ leaf = MAC_DISK_SLICE_SUFFIX.matcher(leaf).replaceAll("");
+ base = "/dev/" + leaf;
+ }
+ }
+
+ Boolean solidState = null;
+ boolean protocolNvme = false;
+ for (String line : runCommandLines("diskutil", "info", base)) {
+ var trimmed = line.trim();
+ var lower = safeLower(trimmed);
+ if (lower.startsWith("solid state:")) {
+ solidState = lower.endsWith("yes");
+ } else if (lower.startsWith("protocol:")) {
+ protocolNvme = lower.contains("nvme");
+ } else if (lower.startsWith("device / media name:") && lower.contains("nvme")) {
+ protocolNvme = true;
+ }
+ }
+ return new MacDiskInfo(solidState, protocolNvme);
+ }
+
+ private static List runCommandLines(String... command) {
+ var lines = new ArrayList();
+ var pb = new ProcessBuilder(command);
+ pb.redirectErrorStream(true);
+ try {
+ var process = pb.start();
+ try (var reader = new BufferedReader(new InputStreamReader(process.getInputStream()))) {
+ String line;
+ while ((line = reader.readLine()) != null) {
+ lines.add(line);
+ }
+ }
+ process.waitFor();
+ } catch (IOException | InterruptedException ignored) {
+ if (ignored instanceof InterruptedException) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ return lines;
+ }
+
+ private static String decodeMountToken(String token) {
+ return token
+ .replace("\\040", " ")
+ .replace("\\011", "\t")
+ .replace("\\012", "\n")
+ .replace("\\134", "\\");
+ }
+
+ private static boolean isLinux(String osNameLower) {
+ return osNameLower.contains("linux");
+ }
+
+ private static boolean isMac(String osNameLower) {
+ return osNameLower.contains("mac") || osNameLower.contains("darwin");
+ }
+
+ private static boolean isWindows(String osNameLower) {
+ return osNameLower.contains("win");
+ }
+
+ private static String safeLower(String value) {
+ return value == null ? "" : value.toLowerCase(Locale.ROOT);
+ }
+
+ private static final class MountEntry {
+ private final String source;
+ private final String mountPoint;
+ private final String filesystemType;
+
+ private MountEntry(String source, String mountPoint, String filesystemType) {
+ this.source = source;
+ this.mountPoint = mountPoint;
+ this.filesystemType = filesystemType;
+ }
+
+ private String source() {
+ return source;
+ }
+
+ private String mountPoint() {
+ return mountPoint;
+ }
+
+ private String filesystemType() {
+ return filesystemType;
+ }
+ }
+
+ private static final class MacDiskInfo {
+ private final Boolean solidState;
+ private final boolean protocolNvme;
+
+ private MacDiskInfo(Boolean solidState, boolean protocolNvme) {
+ this.solidState = solidState;
+ this.protocolNvme = protocolNvme;
+ }
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/StorageLayoutUtil.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/StorageLayoutUtil.java
new file mode 100644
index 000000000..bfbaed234
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/StorageLayoutUtil.java
@@ -0,0 +1,584 @@
+/*
+ * 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.example.util.storage;
+
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.ec2.Ec2Client;
+import software.amazon.awssdk.services.ec2.model.DescribeInstancesRequest;
+import software.amazon.awssdk.services.ec2.model.DescribeVolumesRequest;
+import software.amazon.awssdk.services.ec2.model.InstanceBlockDeviceMapping;
+import software.amazon.awssdk.services.ec2.model.Volume;
+
+import java.io.IOException;
+import java.net.URI;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.net.http.HttpResponse;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.regex.Pattern;
+import java.util.stream.Stream;
+
+/**
+ * Detects EC2 runtime context via IMDSv2 and classifies storage for each mounted filesystem.
+ */
+public final class StorageLayoutUtil {
+ private static final String AWS_EC2_METADATA_DISABLED = "AWS_EC2_METADATA_DISABLED";
+ private static final URI IMDS_TOKEN_URI = URI.create("http://169.254.169.254/latest/api/token");
+ private static final URI IMDS_IDENTITY_URI = URI.create("http://169.254.169.254/latest/dynamic/instance-identity/document");
+ private static final Duration IMDS_TIMEOUT = Duration.ofMillis(300);
+ private static final String IMDS_TOKEN_HEADER = "X-aws-ec2-metadata-token";
+ private static final String IMDS_TOKEN_TTL_HEADER = "X-aws-ec2-metadata-token-ttl-seconds";
+
+ private static final Pattern JSON_FIELD_PATTERN = Pattern.compile("\"([^\"]+)\"\\s*:\\s*\"([^\"]+)\"");
+ private static final Pattern VOL_ID_PATTERN = Pattern.compile("vol-?[0-9a-fA-F]+");
+ private static final Pattern NVME_PARTITION_SUFFIX = Pattern.compile("p\\d+$");
+ private static final Pattern GENERIC_PARTITION_SUFFIX = Pattern.compile("\\d+$");
+ private static final Set NETWORK_FILESYSTEM_TYPES = Set.of("nfs", "nfs4", "efs", "cifs", "smbfs", "fuse.sshfs");
+
+ private StorageLayoutUtil() {
+ }
+
+ public enum StorageClass {
+ // Slowest EBS tiers
+ EBS_COLD_HDD,
+ EBS_THROUGHPUT_HDD,
+ EBS_MAGNETIC,
+
+ // Faster EBS SSD tiers
+ EBS_GP2,
+ EBS_GP3,
+ EBS_PROVISIONED_IOPS_SSD,
+
+ // Local instance storage
+ INSTANCE_STORE_SSD,
+ INSTANCE_STORE_NVME,
+
+ // Non-block storage
+ NETWORK_FILESYSTEM,
+ MEMORY_TMPFS,
+ PSEUDO_FILESYSTEM,
+ UNKNOWN
+ }
+
+ public static final class StorageSnapshot {
+ private final boolean runningOnEc2;
+ private final String instanceId;
+ private final String instanceType;
+ private final String region;
+ private final Map mountsByMountPoint;
+
+ public StorageSnapshot(boolean runningOnEc2,
+ String instanceId,
+ String instanceType,
+ String region,
+ Map mountsByMountPoint) {
+ this.runningOnEc2 = runningOnEc2;
+ this.instanceId = instanceId;
+ this.instanceType = instanceType;
+ this.region = region;
+ this.mountsByMountPoint = Objects.requireNonNull(mountsByMountPoint, "mountsByMountPoint");
+ }
+
+ public boolean runningOnEc2() {
+ return runningOnEc2;
+ }
+
+ public String instanceId() {
+ return instanceId;
+ }
+
+ public String instanceType() {
+ return instanceType;
+ }
+
+ public String region() {
+ return region;
+ }
+
+ public Map mountsByMountPoint() {
+ return mountsByMountPoint;
+ }
+ }
+
+ public static final class MountStorageInfo {
+ private final String mountPoint;
+ private final String source;
+ private final String filesystemType;
+ private final StorageClass storageClass;
+ private final String volumeId;
+ private final String volumeType;
+
+ public MountStorageInfo(String mountPoint,
+ String source,
+ String filesystemType,
+ StorageClass storageClass,
+ String volumeId,
+ String volumeType) {
+ this.mountPoint = mountPoint;
+ this.source = source;
+ this.filesystemType = filesystemType;
+ this.storageClass = Objects.requireNonNull(storageClass, "storageClass");
+ this.volumeId = volumeId;
+ this.volumeType = volumeType;
+ }
+
+ public String mountPoint() {
+ return mountPoint;
+ }
+
+ public String source() {
+ return source;
+ }
+
+ public String filesystemType() {
+ return filesystemType;
+ }
+
+ public StorageClass storageClass() {
+ return storageClass;
+ }
+
+ public String volumeId() {
+ return volumeId;
+ }
+
+ public String volumeType() {
+ return volumeType;
+ }
+ }
+
+ public static StorageSnapshot inspectStorage() {
+ var identity = fetchEc2Identity();
+ var mounts = readMountEntries();
+ var ec2Data = identity.map(StorageLayoutUtil::fetchEc2VolumeData).orElse(Ec2VolumeData.empty());
+
+ mounts.sort(Comparator.comparing(MountEntry::mountPoint));
+ var byMountPoint = new LinkedHashMap(mounts.size());
+ for (var mount : mounts) {
+ var resolvedVolumeId = resolveVolumeId(mount.source(), ec2Data);
+ var volumeType = resolvedVolumeId == null ? null : ec2Data.volumeTypeById().get(resolvedVolumeId);
+ var storageClass = classify(mount, resolvedVolumeId, volumeType);
+ byMountPoint.put(
+ mount.mountPoint(),
+ new MountStorageInfo(
+ mount.mountPoint(),
+ mount.source(),
+ mount.filesystemType(),
+ storageClass,
+ resolvedVolumeId,
+ volumeType
+ )
+ );
+ }
+
+ return new StorageSnapshot(
+ identity.isPresent(),
+ identity.map(Ec2Identity::instanceId).orElse(null),
+ identity.map(Ec2Identity::instanceType).orElse(null),
+ identity.map(Ec2Identity::region).orElse(null),
+ Collections.unmodifiableMap(byMountPoint)
+ );
+ }
+
+ public static Map storageClassByMountPoint() {
+ var snapshot = inspectStorage();
+ var byMountPoint = new LinkedHashMap(snapshot.mountsByMountPoint().size());
+ for (var entry : snapshot.mountsByMountPoint().entrySet()) {
+ byMountPoint.put(entry.getKey(), entry.getValue().storageClass());
+ }
+ return Collections.unmodifiableMap(byMountPoint);
+ }
+
+ private static Optional fetchEc2Identity() {
+ var imdsDisabled = System.getenv(AWS_EC2_METADATA_DISABLED);
+ if (imdsDisabled != null && "true".equalsIgnoreCase(imdsDisabled)) {
+ return Optional.empty();
+ }
+
+ var client = HttpClient.newBuilder()
+ .connectTimeout(IMDS_TIMEOUT)
+ .build();
+ try {
+ var tokenRequest = HttpRequest.newBuilder(IMDS_TOKEN_URI)
+ .timeout(IMDS_TIMEOUT)
+ .header(IMDS_TOKEN_TTL_HEADER, "60")
+ .method("PUT", HttpRequest.BodyPublishers.noBody())
+ .build();
+ var tokenResponse = client.send(tokenRequest, HttpResponse.BodyHandlers.ofString());
+ if (tokenResponse.statusCode() != 200) {
+ return Optional.empty();
+ }
+
+ var token = tokenResponse.body();
+ if (token == null || token.isBlank()) {
+ return Optional.empty();
+ }
+
+ var identityRequest = HttpRequest.newBuilder(IMDS_IDENTITY_URI)
+ .timeout(IMDS_TIMEOUT)
+ .header(IMDS_TOKEN_HEADER, token)
+ .GET()
+ .build();
+ var identityResponse = client.send(identityRequest, HttpResponse.BodyHandlers.ofString());
+ if (identityResponse.statusCode() != 200) {
+ return Optional.empty();
+ }
+
+ return parseIdentity(identityResponse.body());
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return Optional.empty();
+ } catch (IOException e) {
+ return Optional.empty();
+ }
+ }
+
+ private static Optional parseIdentity(String json) {
+ if (json == null || json.isBlank()) {
+ return Optional.empty();
+ }
+ var values = new LinkedHashMap();
+ var matcher = JSON_FIELD_PATTERN.matcher(json);
+ while (matcher.find()) {
+ values.put(matcher.group(1), matcher.group(2));
+ }
+
+ var instanceId = values.get("instanceId");
+ var instanceType = values.get("instanceType");
+ var region = values.get("region");
+ if (instanceId == null || instanceType == null || region == null) {
+ return Optional.empty();
+ }
+ return Optional.of(new Ec2Identity(instanceId, instanceType, region));
+ }
+
+ private static Ec2VolumeData fetchEc2VolumeData(Ec2Identity identity) {
+ var deviceNameToVolumeId = new LinkedHashMap();
+ var volumeTypeById = new LinkedHashMap();
+ var nvmeDeviceToVolumeId = mapNvmeDevicesToVolumeIds();
+
+ try (var ec2 = Ec2Client.builder().region(Region.of(identity.region())).build()) {
+ var instanceRequest = DescribeInstancesRequest.builder()
+ .instanceIds(identity.instanceId())
+ .build();
+ var instanceResponse = ec2.describeInstances(instanceRequest);
+ var reservations = instanceResponse.reservations();
+ if (reservations != null) {
+ for (var reservation : reservations) {
+ for (var instance : reservation.instances()) {
+ for (InstanceBlockDeviceMapping mapping : instance.blockDeviceMappings()) {
+ if (mapping.ebs() == null || mapping.ebs().volumeId() == null || mapping.deviceName() == null) {
+ continue;
+ }
+ deviceNameToVolumeId.put(normalizeDevice(mapping.deviceName()), mapping.ebs().volumeId());
+ }
+ }
+ }
+ }
+
+ if (!deviceNameToVolumeId.isEmpty()) {
+ var volumeResponse = ec2.describeVolumes(DescribeVolumesRequest.builder()
+ .volumeIds(deviceNameToVolumeId.values())
+ .build());
+ for (Volume volume : volumeResponse.volumes()) {
+ if (volume.volumeId() != null && volume.volumeType() != null) {
+ volumeTypeById.put(volume.volumeId(), volume.volumeTypeAsString());
+ }
+ }
+ }
+ } catch (RuntimeException ignored) {
+ // If IAM permissions or service calls fail, we still return mount classifications.
+ }
+
+ return new Ec2VolumeData(deviceNameToVolumeId, nvmeDeviceToVolumeId, volumeTypeById);
+ }
+
+ private static List readMountEntries() {
+ var mountsPath = Files.isReadable(Path.of("/proc/self/mounts"))
+ ? Path.of("/proc/self/mounts")
+ : Path.of("/proc/mounts");
+
+ if (!Files.isReadable(mountsPath)) {
+ return new ArrayList<>();
+ }
+
+ var entries = new ArrayList();
+ try (Stream lines = Files.lines(mountsPath)) {
+ lines.forEach(line -> {
+ var parts = line.split(" ");
+ if (parts.length < 3) {
+ return;
+ }
+ var source = decodeMountToken(parts[0]);
+ var mountPoint = decodeMountToken(parts[1]);
+ var filesystemType = decodeMountToken(parts[2]);
+ entries.add(new MountEntry(source, mountPoint, filesystemType));
+ });
+ } catch (IOException ignored) {
+ return new ArrayList<>();
+ }
+ return entries;
+ }
+
+ private static Map mapNvmeDevicesToVolumeIds() {
+ var byIdDir = Path.of("/dev/disk/by-id");
+ if (!Files.isDirectory(byIdDir)) {
+ return Map.of();
+ }
+
+ var mapping = new LinkedHashMap();
+ try (Stream entries = Files.list(byIdDir)) {
+ entries.filter(Files::isSymbolicLink).forEach(link -> {
+ var name = link.getFileName().toString();
+ if (!name.startsWith("nvme-Amazon_Elastic_Block_Store_")) {
+ return;
+ }
+ var volumeId = extractVolumeId(name);
+ if (volumeId == null) {
+ return;
+ }
+
+ try {
+ var target = normalizeDevice(link.toRealPath().toString());
+ mapping.put(target, volumeId);
+ } catch (IOException ignored) {
+ // continue
+ }
+ });
+ } catch (IOException ignored) {
+ return Map.of();
+ }
+ return mapping;
+ }
+
+ private static String resolveVolumeId(String mountSource, Ec2VolumeData ec2Data) {
+ if (mountSource == null || !mountSource.startsWith("/dev/")) {
+ return null;
+ }
+
+ var normalized = normalizeDevice(mountSource);
+ var byNvme = ec2Data.nvmeDeviceToVolumeId().get(normalized);
+ if (byNvme != null) {
+ return byNvme;
+ }
+ return ec2Data.deviceNameToVolumeId().get(normalized);
+ }
+
+ private static StorageClass classify(MountEntry mount, String volumeId, String volumeType) {
+ var fsType = safeLower(mount.filesystemType());
+ var source = mount.source();
+ var sourceLower = safeLower(source);
+
+ if ("tmpfs".equals(fsType)) {
+ return StorageClass.MEMORY_TMPFS;
+ }
+ if (NETWORK_FILESYSTEM_TYPES.contains(fsType)) {
+ return StorageClass.NETWORK_FILESYSTEM;
+ }
+ if (isPseudoFileSystem(fsType, sourceLower)) {
+ return StorageClass.PSEUDO_FILESYSTEM;
+ }
+
+ if (volumeId != null) {
+ return mapEbsVolumeType(volumeType);
+ }
+
+ if (source != null && source.startsWith("/dev/")) {
+ if (sourceLower.contains("nvme")) {
+ return StorageClass.INSTANCE_STORE_NVME;
+ }
+ return StorageClass.INSTANCE_STORE_SSD;
+ }
+ return StorageClass.UNKNOWN;
+ }
+
+ private static StorageClass mapEbsVolumeType(String volumeType) {
+ if (volumeType == null) {
+ return StorageClass.EBS_GP3;
+ }
+
+ switch (safeLower(volumeType)) {
+ case "sc1":
+ return StorageClass.EBS_COLD_HDD;
+ case "st1":
+ return StorageClass.EBS_THROUGHPUT_HDD;
+ case "standard":
+ return StorageClass.EBS_MAGNETIC;
+ case "io1":
+ case "io2":
+ return StorageClass.EBS_PROVISIONED_IOPS_SSD;
+ case "gp2":
+ return StorageClass.EBS_GP2;
+ case "gp3":
+ return StorageClass.EBS_GP3;
+ default:
+ return StorageClass.EBS_GP3;
+ }
+ }
+
+ private static boolean isPseudoFileSystem(String fsType, String sourceLower) {
+ return fsType.equals("proc")
+ || fsType.equals("sysfs")
+ || fsType.equals("devpts")
+ || fsType.equals("devtmpfs")
+ || fsType.equals("cgroup")
+ || fsType.equals("cgroup2")
+ || fsType.equals("autofs")
+ || fsType.equals("mqueue")
+ || fsType.equals("tracefs")
+ || fsType.equals("pstore")
+ || fsType.equals("securityfs")
+ || fsType.equals("debugfs")
+ || fsType.equals("configfs")
+ || fsType.equals("fusectl")
+ || fsType.equals("binfmt_misc")
+ || fsType.equals("rpc_pipefs")
+ || sourceLower.equals("proc")
+ || sourceLower.equals("sysfs")
+ || sourceLower.equals("tmpfs");
+ }
+
+ private static String decodeMountToken(String token) {
+ return token
+ .replace("\\040", " ")
+ .replace("\\011", "\t")
+ .replace("\\012", "\n")
+ .replace("\\134", "\\");
+ }
+
+ private static String extractVolumeId(String value) {
+ var matcher = VOL_ID_PATTERN.matcher(value);
+ if (!matcher.find()) {
+ return null;
+ }
+ var raw = matcher.group();
+ if (raw.startsWith("vol-")) {
+ return raw.toLowerCase(Locale.ROOT);
+ }
+ return "vol-" + raw.substring(3).toLowerCase(Locale.ROOT);
+ }
+
+ private static String normalizeDevice(String device) {
+ if (device == null) {
+ return null;
+ }
+ if (!device.startsWith("/dev/")) {
+ return device;
+ }
+
+ if (device.startsWith("/dev/nvme")) {
+ return NVME_PARTITION_SUFFIX.matcher(device).replaceAll("");
+ }
+ return GENERIC_PARTITION_SUFFIX.matcher(device).replaceAll("");
+ }
+
+ private static String safeLower(String value) {
+ return value == null ? "" : value.toLowerCase(Locale.ROOT);
+ }
+
+ private static final class MountEntry {
+ private final String source;
+ private final String mountPoint;
+ private final String filesystemType;
+
+ private MountEntry(String source, String mountPoint, String filesystemType) {
+ this.source = source;
+ this.mountPoint = mountPoint;
+ this.filesystemType = filesystemType;
+ }
+
+ private String source() {
+ return source;
+ }
+
+ private String mountPoint() {
+ return mountPoint;
+ }
+
+ private String filesystemType() {
+ return filesystemType;
+ }
+ }
+
+ private static final class Ec2Identity {
+ private final String instanceId;
+ private final String instanceType;
+ private final String region;
+
+ private Ec2Identity(String instanceId, String instanceType, String region) {
+ this.instanceId = instanceId;
+ this.instanceType = instanceType;
+ this.region = region;
+ }
+
+ private String instanceId() {
+ return instanceId;
+ }
+
+ private String instanceType() {
+ return instanceType;
+ }
+
+ private String region() {
+ return region;
+ }
+ }
+
+ private static final class Ec2VolumeData {
+ private final Map deviceNameToVolumeId;
+ private final Map nvmeDeviceToVolumeId;
+ private final Map volumeTypeById;
+
+ private Ec2VolumeData(Map deviceNameToVolumeId,
+ Map nvmeDeviceToVolumeId,
+ Map volumeTypeById) {
+ Objects.requireNonNull(deviceNameToVolumeId, "deviceNameToVolumeId");
+ Objects.requireNonNull(nvmeDeviceToVolumeId, "nvmeDeviceToVolumeId");
+ Objects.requireNonNull(volumeTypeById, "volumeTypeById");
+ this.deviceNameToVolumeId = deviceNameToVolumeId;
+ this.nvmeDeviceToVolumeId = nvmeDeviceToVolumeId;
+ this.volumeTypeById = volumeTypeById;
+ }
+
+ private Map deviceNameToVolumeId() {
+ return deviceNameToVolumeId;
+ }
+
+ private Map nvmeDeviceToVolumeId() {
+ return nvmeDeviceToVolumeId;
+ }
+
+ private Map volumeTypeById() {
+ return volumeTypeById;
+ }
+
+ private static Ec2VolumeData empty() {
+ return new Ec2VolumeData(Map.of(), Map.of(), Map.of());
+ }
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/package-info.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/package-info.java
new file mode 100644
index 000000000..a553c8b23
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/storage/package-info.java
@@ -0,0 +1,28 @@
+/*
+ * 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.
+ */
+
+/**
+ * Provides utilities for characterizing the underlying storage hardware and layout.
+ *
+ * This package contains logic to detect and classify storage tiers (e.g., Local SSD,
+ * Persistent Disk, Network Filesystem) across different environments including
+ * AWS, GCP, and local development machines.
+ *
+ * The primary entry point is {@link io.github.jbellis.jvector.example.util.storage.CloudStorageLayoutUtil},
+ * which provides a unified view of the system's mount points and their corresponding
+ * {@link io.github.jbellis.jvector.example.util.storage.CloudStorageLayoutUtil.StorageClass}.
+ */
+package io.github.jbellis.jvector.example.util.storage;
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/TestDataPartition.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/TestDataPartition.java
new file mode 100644
index 000000000..648e4942a
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/TestDataPartition.java
@@ -0,0 +1,96 @@
+/*
+ * 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.example.yaml;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Specifically for defining how data is partitioned for testing compaction.
+ */
+public class TestDataPartition {
+ public List numSplits;
+ public List splitDistribution;
+
+ public TestDataPartition() {
+ this.numSplits = Collections.singletonList(1);
+ this.splitDistribution = Collections.singletonList(Distribution.UNIFORM);
+ }
+
+ public TestDataPartition(int numSplits) {
+ this.numSplits = Collections.singletonList(numSplits);
+ this.splitDistribution = Collections.singletonList(Distribution.UNIFORM);
+ }
+
+ public enum Distribution {
+ UNIFORM,
+ FIBONACCI,
+ LOG2N,
+ /** First partition gets 10%, last gets 90%. For N>2, middle partitions are empty. */
+ TIERED_10_90,
+ /** First partition gets 1%, last gets 99%. For N>2, middle partitions are empty. */
+ TIERED_1_99;
+
+ public List computeSplitSizes(int total, int numSplits) {
+ int[] weights = new int[numSplits];
+ switch (this) {
+ case UNIFORM:
+ for (int i = 0; i < numSplits; i++) weights[i] = 1;
+ break;
+ case FIBONACCI:
+ int a = 1, b = 2;
+ weights[0] = 1;
+ for (int i = 1; i < numSplits; i++) {
+ weights[i] = b;
+ int next = a + b;
+ a = b;
+ b = next;
+ }
+ break;
+ case LOG2N:
+ for (int i = 0; i < numSplits; i++) weights[i] = 1 << i;
+ break;
+ case TIERED_10_90:
+ weights[0] = 1;
+ weights[numSplits - 1] = 9;
+ break;
+ case TIERED_1_99:
+ weights[0] = 1;
+ weights[numSplits - 1] = 99;
+ break;
+ }
+
+ long weightSum = 0;
+ for (int w : weights) weightSum += w;
+
+ List sizes = new ArrayList<>(numSplits);
+ int assigned = 0;
+ for (int i = 0; i < numSplits; i++) {
+ int size;
+ if (i == numSplits - 1) {
+ size = total - assigned;
+ } else {
+ size = (int) (((long) weights[i] * total) / weightSum);
+ }
+ sizes.add(size);
+ assigned += size;
+ }
+ return sizes;
+ }
+ }
+}
diff --git a/jvector-examples/src/main/resources/log4j2.xml b/jvector-examples/src/main/resources/log4j2.xml
new file mode 100644
index 000000000..83c77bced
--- /dev/null
+++ b/jvector-examples/src/main/resources/log4j2.xml
@@ -0,0 +1,14 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
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
new file mode 100644
index 000000000..3517c7c40
--- /dev/null
+++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java
@@ -0,0 +1,1090 @@
+/*
+ * 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.graph.disk;
+
+import com.carrotsearch.randomizedtesting.RandomizedTest;
+import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope;
+import io.github.jbellis.jvector.TestUtil;
+import io.github.jbellis.jvector.disk.ReaderSupplier;
+import io.github.jbellis.jvector.disk.ReaderSupplierFactory;
+import io.github.jbellis.jvector.disk.SimpleMappedReader;
+import io.github.jbellis.jvector.example.util.AccuracyMetrics;
+import io.github.jbellis.jvector.graph.*;
+import io.github.jbellis.jvector.graph.disk.feature.Feature;
+import io.github.jbellis.jvector.graph.disk.feature.FeatureId;
+import io.github.jbellis.jvector.graph.disk.feature.FusedPQ;
+import io.github.jbellis.jvector.graph.disk.feature.InlineVectors;
+import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider;
+import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider;
+import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider;
+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.BoundedLongHeap;
+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 org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.*;
+import java.util.concurrent.ForkJoinPool;
+import java.util.function.IntFunction;
+
+import static io.github.jbellis.jvector.TestUtil.createRandomVectors;
+import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertFalse;
+
+@ThreadLeakScope(ThreadLeakScope.Scope.NONE)
+public class TestOnDiskGraphIndexCompactor extends RandomizedTest {
+ private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport();
+
+ private ImmutableGraphIndex golden;
+ private Path testDirectory;
+ List> allVecs = new ArrayList<>();
+ int dimension = 32;
+ int numVectorsPerGraph = 256;
+ int numSources = 3;
+ int numQueries = 20;
+ VectorSimilarityFunction similarityFunction = VectorSimilarityFunction.COSINE;
+ RandomAccessVectorValues allravv;
+ private final ForkJoinPool simdExecutor = ForkJoinPool.commonPool();
+ private final ForkJoinPool parallelExecutor = ForkJoinPool.commonPool();
+
+ @Before
+ public void setup() throws IOException {
+ testDirectory = Files.createTempDirectory("jvector_test");
+ buildFusedPQ();
+ buildGoldenPQ();
+ }
+
+ /**
+ * Builds source graphs with FusedPQ feature enabled.
+ * Uses random vectors with COSINE similarity.
+ */
+ void buildFusedPQ() throws IOException {
+ for(int i = 0; i < numSources; ++i) {
+ List> vecs = createRandomVectors(numVectorsPerGraph, dimension);
+
+ RandomAccessVectorValues ravv = new ListRandomAccessVectorValues(vecs, dimension);
+ ProductQuantization pq = ProductQuantization.compute(ravv, 8, 256, true, UNWEIGHTED, simdExecutor, parallelExecutor);
+ PQVectors pqv = (PQVectors) pq.encodeAll(ravv, simdExecutor);
+ var bsp = BuildScoreProvider.pqBuildScoreProvider(similarityFunction, pqv);
+ var builder = new GraphIndexBuilder(bsp, dimension, 16, 100, 1.2f, 1.2f, false, true, simdExecutor, parallelExecutor);
+ var graph = builder.getGraph();
+
+ var outputPath = testDirectory.resolve("test_graph_" + i);
+ Map> writeSuppliers = new EnumMap<>(FeatureId.class);
+ writeSuppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(ravv.getVector(ordinal)));
+
+ var identityMapper = new OrdinalMapper.IdentityMapper(ravv.size() - 1);
+ var writerBuilder = new OnDiskGraphIndexWriter.Builder(graph, outputPath);
+ writerBuilder.withMapper(identityMapper);
+ writerBuilder.with(new InlineVectors(dimension));
+ writerBuilder.with(new FusedPQ(graph.maxDegree(), pq));
+ var writer = writerBuilder.build();
+
+ for (var node = 0; node < ravv.size(); node++) {
+ var stateMap = new EnumMap(FeatureId.class);
+ stateMap.put(FeatureId.INLINE_VECTORS, writeSuppliers.get(FeatureId.INLINE_VECTORS).apply(node));
+ writer.writeInline(node, stateMap);
+ builder.addGraphNode(node, ravv.getVector(node));
+ }
+ builder.cleanup();
+
+ writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(graph.getView(), pqv, ordinal));
+ writer.write(writeSuppliers);
+ allVecs.addAll(vecs);
+ }
+ }
+
+ /**
+ * Builds the golden graph from all vectors combined.
+ * This represents the ideal case of building from scratch.
+ */
+ void buildGoldenPQ() throws IOException {
+ allravv = new ListRandomAccessVectorValues(allVecs, dimension);
+
+ ProductQuantization pq = ProductQuantization.compute(allravv, 8, 256, true, UNWEIGHTED, simdExecutor, parallelExecutor);
+ PQVectors pqv = (PQVectors) pq.encodeAll(allravv, simdExecutor);
+ var bsp = BuildScoreProvider.pqBuildScoreProvider(similarityFunction, pqv);
+ var builder = new GraphIndexBuilder(bsp, dimension, 16, 100, 1.2f, 1.2f, false, true, simdExecutor, parallelExecutor);
+ for (var i = 0; i < allravv.size(); i++) {
+ builder.addGraphNode(i, allravv.getVector(i));
+ }
+ builder.cleanup();
+ golden = builder.getGraph();
+ }
+ List searchFromAll(List> queries, int topK) {
+ List srs = new ArrayList<>();
+ try (GraphSearcher searcher = new GraphSearcher(golden)) {
+ for(VectorFloat q: queries) {
+ var row = new ArrayList();
+ SearchScoreProvider ssp = DefaultSearchScoreProvider.exact(q, similarityFunction, allravv);
+ SearchResult sr = searcher.search(ssp, topK, Bits.ALL);
+ srs.add(sr);
+ }
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ return srs;
+ }
+ List> buildGT(List> queries, int topK) {
+ List> rows = new ArrayList<>();
+
+ for(int i = 0; i < queries.size(); ++i) {
+ NodeQueue expected = new NodeQueue(new BoundedLongHeap(topK), NodeQueue.Order.MIN_HEAP);
+ for (int j = 0; j < allVecs.size(); j++) {
+ expected.push(j, similarityFunction.compare(queries.get(i), allVecs.get(j)));
+ }
+
+ var row = new ArrayList();
+ for(int k = 0; k < topK; ++k) {
+ row.add(expected.pop());
+ }
+ rows.add(row);
+ }
+ return rows;
+ }
+
+ @After
+ public void tearDown() {
+ TestUtil.deleteQuietly(testDirectory);
+ }
+
+ /**
+ * Builds a small source graph with InlineVectors only (no FusedPQ), using exact scoring.
+ * Returns the path to the written graph file.
+ */
+ private Path buildSimpleSourceGraph(List> vecs, int dim, VectorSimilarityFunction vsf, String name) throws IOException {
+ RandomAccessVectorValues ravv = new ListRandomAccessVectorValues(vecs, dim);
+ var bsp = BuildScoreProvider.randomAccessScoreProvider(ravv, vsf);
+ var builder = new GraphIndexBuilder(bsp, dim, 4, 20, 1.2f, 1.2f, false, true, simdExecutor, parallelExecutor);
+ for (int i = 0; i < vecs.size(); i++) {
+ builder.addGraphNode(i, vecs.get(i));
+ }
+ builder.cleanup();
+ var graph = builder.getGraph();
+
+ var outputPath = testDirectory.resolve(name);
+ var identityMapper = new OrdinalMapper.IdentityMapper(vecs.size() - 1);
+ var writerBuilder = new OnDiskGraphIndexWriter.Builder(graph, outputPath);
+ writerBuilder.withMapper(identityMapper);
+ writerBuilder.with(new InlineVectors(dim));
+ var writer = writerBuilder.build();
+
+ Map> writeSuppliers = new EnumMap<>(FeatureId.class);
+ writeSuppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(ravv.getVector(ordinal)));
+
+ for (int node = 0; node < vecs.size(); node++) {
+ var stateMap = new EnumMap(FeatureId.class);
+ stateMap.put(FeatureId.INLINE_VECTORS, writeSuppliers.get(FeatureId.INLINE_VECTORS).apply(node));
+ writer.writeInline(node, stateMap);
+ }
+ writer.write(writeSuppliers);
+ return outputPath;
+ }
+
+ /** Creates a vector of the given dimension with value at index {@code hot} set to {@code val}, rest 0. */
+ private VectorFloat makeVec(int dim, int hot, float val) {
+ VectorFloat v = vectorTypeSupport.createFloatVector(dim);
+ for (int d = 0; d < dim; d++) {
+ v.set(d, d == hot ? val : 0.0f);
+ }
+ return v;
+ }
+
+ private void assertVecEquals(VectorFloat expected, VectorFloat actual, int ordinal) {
+ int dim = expected.length();
+ assertEquals("dimension mismatch at ordinal " + ordinal, dim, actual.length());
+ for (int d = 0; d < dim; d++) {
+ assertEquals(String.format("vector[%d] dim %d mismatch", ordinal, d), expected.get(d), actual.get(d), 0.0f);
+ }
+ }
+
+ /**
+ * Tests that vectors are stored exactly at the expected global ordinals after compaction.
+ * Uses two small sources with simple, known float values and identity mapping.
+ */
+ @Test
+ public void testExactVectorValuesAfterCompaction() throws Exception {
+ int dim = 4;
+ int n = 6; // nodes per source
+ VectorSimilarityFunction vsf = VectorSimilarityFunction.EUCLIDEAN;
+
+ // Source 0: vectors with first dim varying by index
+ List> vecs0 = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ vecs0.add(makeVec(dim, 0, (float)(i + 1)));
+ }
+ // Source 1: vectors with second dim varying by index
+ List> vecs1 = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ vecs1.add(makeVec(dim, 1, (float)(i + 10)));
+ }
+
+ Path path0 = buildSimpleSourceGraph(vecs0, dim, vsf, "simple_src_0");
+ Path path1 = buildSimpleSourceGraph(vecs1, dim, vsf, "simple_src_1");
+
+ ReaderSupplier rs0 = ReaderSupplierFactory.open(path0);
+ ReaderSupplier rs1 = ReaderSupplierFactory.open(path1);
+ OnDiskGraphIndex g0 = OnDiskGraphIndex.load(rs0);
+ OnDiskGraphIndex g1 = OnDiskGraphIndex.load(rs1);
+
+ // Identity remapping: source i -> global ordinals [i*n, (i+1)*n)
+ Map map0 = new HashMap<>();
+ Map map1 = new HashMap<>();
+ for (int i = 0; i < n; i++) {
+ map0.put(i, i);
+ map1.put(i, n + i);
+ }
+
+ FixedBitSet live0 = new FixedBitSet(n);
+ live0.set(0, n);
+ FixedBitSet live1 = new FixedBitSet(n);
+ live1.set(0, n);
+
+ var compactor = new OnDiskGraphIndexCompactor(
+ List.of(g0, g1),
+ List.of(live0, live1),
+ List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)),
+ vsf, null);
+
+ Path outPath = testDirectory.resolve("simple_compact_out");
+ compactor.compact(outPath);
+
+ ReaderSupplier rsOut = ReaderSupplierFactory.open(outPath);
+ OnDiskGraphIndex compacted = OnDiskGraphIndex.load(rsOut);
+ assertEquals(2 * n, compacted.size(0));
+
+ var view = compacted.getView();
+ VectorFloat buf = vectorTypeSupport.createFloatVector(dim);
+
+ // Source 0 vectors must be at ordinals 0..n-1
+ for (int i = 0; i < n; i++) {
+ view.getVectorInto(i, buf, 0);
+ assertVecEquals(vecs0.get(i), buf, i);
+ }
+ // Source 1 vectors must be at ordinals n..2n-1
+ for (int i = 0; i < n; i++) {
+ view.getVectorInto(n + i, buf, 0);
+ assertVecEquals(vecs1.get(i), buf, n + i);
+ }
+ }
+
+ /**
+ * Tests that only live vectors appear after compaction, placed at the correct remapped ordinals.
+ * Deletes every other node from each source and verifies the compacted output exactly.
+ */
+ @Test
+ public void testExactVectorValuesWithDeletions() throws Exception {
+ int dim = 4;
+ int n = 8; // nodes per source
+ VectorSimilarityFunction vsf = VectorSimilarityFunction.EUCLIDEAN;
+
+ // Source 0: vectors [1,0,0,0] through [8,0,0,0]
+ List> vecs0 = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ vecs0.add(makeVec(dim, 0, (float)(i + 1)));
+ }
+ // Source 1: vectors [0,10,0,0] through [0,170,0,0]
+ List> vecs1 = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ vecs1.add(makeVec(dim, 1, (float)((i + 1) * 10)));
+ }
+
+ Path path0 = buildSimpleSourceGraph(vecs0, dim, vsf, "del_src_0");
+ Path path1 = buildSimpleSourceGraph(vecs1, dim, vsf, "del_src_1");
+
+ ReaderSupplier rs0 = ReaderSupplierFactory.open(path0);
+ ReaderSupplier rs1 = ReaderSupplierFactory.open(path1);
+ OnDiskGraphIndex g0 = OnDiskGraphIndex.load(rs0);
+ OnDiskGraphIndex g1 = OnDiskGraphIndex.load(rs1);
+
+ // Keep only even-indexed nodes (0, 2, 4, 6) in both sources
+ FixedBitSet live0 = new FixedBitSet(n);
+ FixedBitSet live1 = new FixedBitSet(n);
+ Map map0 = new HashMap<>();
+ Map map1 = new HashMap<>();
+ int globalOrdinal = 0;
+ for (int i = 0; i < n; i++) {
+ if (i % 2 == 0) {
+ live0.set(i);
+ map0.put(i, globalOrdinal++);
+ }
+ }
+ for (int i = 0; i < n; i++) {
+ if (i % 2 == 0) {
+ live1.set(i);
+ map1.put(i, globalOrdinal++);
+ }
+ }
+ int expectedTotal = globalOrdinal;
+
+ var compactor = new OnDiskGraphIndexCompactor(
+ List.of(g0, g1),
+ List.of(live0, live1),
+ List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)),
+ vsf, null);
+
+ Path outPath = testDirectory.resolve("del_compact_out");
+ compactor.compact(outPath);
+
+ ReaderSupplier rsOut = ReaderSupplierFactory.open(outPath);
+ OnDiskGraphIndex compacted = OnDiskGraphIndex.load(rsOut);
+ assertEquals(expectedTotal, compacted.size(0));
+
+ var view = compacted.getView();
+ VectorFloat buf = vectorTypeSupport.createFloatVector(dim);
+
+ // Verify source 0 live nodes at their mapped ordinals
+ for (int i = 0; i < n; i++) {
+ if (i % 2 == 0) {
+ int ord = map0.get(i);
+ view.getVectorInto(ord, buf, 0);
+ assertVecEquals(vecs0.get(i), buf, ord);
+ }
+ }
+ // Verify source 1 live nodes at their mapped ordinals
+ for (int i = 0; i < n; i++) {
+ if (i % 2 == 0) {
+ int ord = map1.get(i);
+ view.getVectorInto(ord, buf, 0);
+ assertVecEquals(vecs1.get(i), buf, ord);
+ }
+ }
+ }
+
+ /**
+ * Tests that vectors end up at the correct ordinals when a non-sequential remapping is used.
+ * Source 0 is mapped in reverse order; source 1 is mapped in forward order.
+ * Verifies exact vector values at every remapped position.
+ */
+ @Test
+ public void testExactVectorValuesWithCustomRemapping() throws Exception {
+ int dim = 4;
+ int n = 6;
+ VectorSimilarityFunction vsf = VectorSimilarityFunction.EUCLIDEAN;
+
+ List> vecs0 = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ vecs0.add(makeVec(dim, 2, (float)(i + 1)));
+ }
+ List> vecs1 = new ArrayList<>();
+ for (int i = 0; i < n; i++) {
+ vecs1.add(makeVec(dim, 3, (float)(i + 100)));
+ }
+
+ Path path0 = buildSimpleSourceGraph(vecs0, dim, vsf, "remap_src_0");
+ Path path1 = buildSimpleSourceGraph(vecs1, dim, vsf, "remap_src_1");
+
+ ReaderSupplier rs0 = ReaderSupplierFactory.open(path0);
+ ReaderSupplier rs1 = ReaderSupplierFactory.open(path1);
+ OnDiskGraphIndex g0 = OnDiskGraphIndex.load(rs0);
+ OnDiskGraphIndex g1 = OnDiskGraphIndex.load(rs1);
+
+ // Source 0: reverse mapping (local 0 -> global n-1, local 1 -> global n-2, ...)
+ Map map0 = new HashMap<>();
+ for (int i = 0; i < n; i++) {
+ map0.put(i, n - 1 - i);
+ }
+ // Source 1: forward mapping (local 0 -> global n, local 1 -> global n+1, ...)
+ Map map1 = new HashMap<>();
+ for (int i = 0; i < n; i++) {
+ map1.put(i, n + i);
+ }
+
+ FixedBitSet live0 = new FixedBitSet(n);
+ live0.set(0, n);
+ FixedBitSet live1 = new FixedBitSet(n);
+ live1.set(0, n);
+
+ var compactor = new OnDiskGraphIndexCompactor(
+ List.of(g0, g1),
+ List.of(live0, live1),
+ List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)),
+ vsf, null);
+
+ Path outPath = testDirectory.resolve("remap_compact_out");
+ compactor.compact(outPath);
+
+ ReaderSupplier rsOut = ReaderSupplierFactory.open(outPath);
+ OnDiskGraphIndex compacted = OnDiskGraphIndex.load(rsOut);
+ assertEquals(2 * n, compacted.size(0));
+
+ var view = compacted.getView();
+ VectorFloat buf = vectorTypeSupport.createFloatVector(dim);
+
+ for (int i = 0; i < n; i++) {
+ int ord = map0.get(i);
+ view.getVectorInto(ord, buf, 0);
+ assertVecEquals(vecs0.get(i), buf, ord);
+ }
+ for (int i = 0; i < n; i++) {
+ int ord = map1.get(i);
+ view.getVectorInto(ord, buf, 0);
+ assertVecEquals(vecs1.get(i), buf, ord);
+ }
+ }
+
+ /**
+ * Tests basic compaction: merging multiple graphs without deletions.
+ * Verifies that compacted graph recall is comparable to golden graph.
+ */
+ @Test
+ public void testCompact() throws Exception {
+ List graphs = new ArrayList<>();
+ List rss = new ArrayList<>();
+ List liveNodes = new ArrayList<>();
+ List remappers = new ArrayList<>();
+
+ // Load all source graphs
+ for(int i = 0; i < numSources; ++i) {
+ var outputPath = testDirectory.resolve("test_graph_" + i);
+ rss.add(ReaderSupplierFactory.open(outputPath.toAbsolutePath()));
+ var onDiskGraph = OnDiskGraphIndex.load(rss.get(i));
+ graphs.add(onDiskGraph);
+ }
+
+ // Create identity mapping and all nodes live
+ int globalOrdinal = 0;
+ for (int n = 0; n < numSources; n++) {
+ Map map = new HashMap<>(numVectorsPerGraph);
+ for (int i = 0; i < numVectorsPerGraph; i++) {
+ map.put(i, globalOrdinal++);
+ }
+ remappers.add(new OrdinalMapper.MapMapper(map));
+
+ var lives = new FixedBitSet(numVectorsPerGraph);
+ lives.set(0, numVectorsPerGraph);
+ liveNodes.add(lives);
+ }
+
+ var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null);
+ int topK = 10;
+
+ // Select query vectors from the dataset
+ var outputPath = testDirectory.resolve("test_compact_graph_");
+ List> queries = new ArrayList<>();
+ for(int i = 0; i < numQueries; ++i) {
+ queries.add(allVecs.get(randomIntBetween(0, allVecs.size() - 1)));
+ }
+
+ // Get golden results and ground truth
+ List goldenResults = searchFromAll(queries, topK);
+ List> groundTruth = buildGT(queries, topK);
+
+ // Compact and test
+ compactor.compact(outputPath);
+
+ ReaderSupplier rs = ReaderSupplierFactory.open(outputPath);
+ var compactGraph = OnDiskGraphIndex.load(rs);
+
+ // Verify basic properties
+ assertEquals("Compacted graph should have all nodes", numSources * numVectorsPerGraph, compactGraph.size(0));
+
+ GraphSearcher searcher = new GraphSearcher(compactGraph);
+ List compactResults = new ArrayList<>();
+ for(VectorFloat q: queries) {
+ SearchScoreProvider ssp = DefaultSearchScoreProvider.exact(q, similarityFunction, allravv);
+ compactResults.add(searcher.search(ssp, topK, Bits.ALL));
+ }
+
+ // Calculate recalls
+ double goldenRecall = AccuracyMetrics.recallFromSearchResults(groundTruth, goldenResults, topK, topK);
+ double compactRecall = AccuracyMetrics.recallFromSearchResults(groundTruth, compactResults, topK, topK);
+
+ System.out.printf("Golden (built from scratch) Recall: %.4f%n", goldenRecall);
+ System.out.printf("Compacted Recall: %.4f%n", compactRecall);
+ System.out.printf("Recall difference: %.4f%n", Math.abs(goldenRecall - compactRecall));
+
+ // For random vectors with COSINE, both golden and compact should have similar recall
+ // The key is that they're comparable to each other, showing compaction preserves graph quality
+ double recallDifference = Math.abs(goldenRecall - compactRecall);
+ assertTrue(String.format("Compacted recall (%.4f) should be comparable to golden recall (%.4f), difference: %.4f",
+ compactRecall, goldenRecall, recallDifference),
+ recallDifference < 0.2); // Allow up to 20% difference for random vectors
+
+ // Verify both are reasonable (not completely broken)
+ assertTrue(String.format("Golden recall should be at least 0.2, got %.4f", goldenRecall),
+ goldenRecall >= 0.2);
+ assertTrue(String.format("Compacted recall should be at least 0.2, got %.4f", compactRecall),
+ compactRecall >= 0.2);
+
+ searcher.close();
+ }
+
+ /**
+ * Tests compaction with deleted nodes.
+ * Verifies that deleted nodes are properly excluded from the compacted graph.
+ */
+ @Test
+ public void testCompactWithDeletions() throws Exception {
+ List graphs = new ArrayList<>();
+ List rss = new ArrayList<>();
+ List liveNodes = new ArrayList<>();
+ List remappers = new ArrayList<>();
+
+ for(int i = 0; i < numSources; ++i) {
+ var outputPath = testDirectory.resolve("test_graph_" + i);
+ rss.add(ReaderSupplierFactory.open(outputPath.toAbsolutePath()));
+ var onDiskGraph = OnDiskGraphIndex.load(rss.get(i));
+ graphs.add(onDiskGraph);
+ }
+
+ // Mark some nodes as deleted (not live)
+ int globalOrdinal = 0;
+ int totalLiveNodes = 0;
+ Set