diff --git a/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java b/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java index 3a3b9c8bf..bf3acb695 100644 --- a/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java +++ b/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java @@ -704,7 +704,7 @@ private long compactPartitions() throws Exception { liveNodes.add(randomLiveNodes(size, liveNodesRate, n)); globalOrdinal += size; } - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null, -1); long startNanos = System.nanoTime(); compactor.compact(compactOutputPath); diff --git a/docs/compaction.md b/docs/compaction.md index af0c12317..a24c4485a 100644 --- a/docs/compaction.md +++ b/docs/compaction.md @@ -33,7 +33,8 @@ for (var src : sources) { var compactor = new OnDiskGraphIndexCompactor( sources, liveNodes, remappers, VectorSimilarityFunction.COSINE, - /* executor= */ null // null = use by default shared threadpool in compactor + /* executor= */ null, // null = jvector's shared physical-core pool + /* taskWindowSize= */ -1 // <= 0 derives the window from the pool's parallelism ); compactor.compact(Path.of("compacted.index")); @@ -57,6 +58,99 @@ for (int i = 0; i < source.size(); i++) { remappers.add(new OrdinalMapper.MapMapper(oldToNew)); ``` +## Embedding API + +For embedding the compactor into a host system (for example, a database's compaction pipeline), this branch adds `@Experimental` extension points that let the host supply its own threads, observe and throttle the merge, and place the output inside its own container file. All are additive: with none of them used, behavior and output are unchanged (a jvector-owned pool, no throttling, a standalone output file). + +### Supplying an executor + +The compactor dispatches its batch work through an internal `ExecutorCompletionService`, so it runs on any `Executor` the host supplies, with an explicit in-flight window: + +```java +new OnDiskGraphIndexCompactor(sources, liveNodes, remappers, sim, executor, taskWindowSize); +``` + +- Pass a **`ForkJoinPool`** (work-stealing makes submit-and-block safe), or a **caller-runs** executor (`Runnable::run`) to run the entire merge on the *calling* thread — no jvector-owned pool. A host then gets parallelism by running multiple compactions concurrently rather than from a per-compaction pool. +- **Do not** pass a bounded `ThreadPoolExecutor` that is *also* running the calling thread: the caller blocks in `take()` while its sub-tasks queue behind it, which can thread-starvation-deadlock. +- `taskWindowSize` bounds in-flight batches (concurrency and peak write-side memory); `<= 0` derives it from a `ForkJoinPool`'s parallelism, else defaults to `1`. Read it back with `getTaskWindowSize()`. + +A `ForkJoinPool` is passed as the `executor` (with `taskWindowSize <= 0` to derive its window) — there is no separate `ForkJoinPool` constructor. + +### Progress and throttling + +`setProgressLimiter(ProgressLimiter)` installs a control surface (package `io.github.jbellis.jvector.util.work`) that melds two independently-optional facets: + +- **`onProgress(WorkStage stage, long completed, long total)`** — progress observation. Stages are `OnDiskGraphIndexCompactor.Phase.{MERGE_LEVELS, REFINE}`; `total` may be `-1` until known. +- **`acquire(long amount) → Grant`** — blocking work admission. For the compactor the unit is **bytes about to be written**. `acquire` is called on the orchestrating thread (never a pool worker), so a blocking limiter back-pressures dispatch without a `ForkJoinPool.ManagedBlocker` — exactly like ordinary code blocking on a rate limiter. + +The default is `ProgressLimiter.UNLIMITED` (both facets no-op), so output and timing are unchanged when unset. + +Two ready-made implementations compose as decorators: + +```java +// Cap merge write bandwidth to 100 MB/s, logging throttled writes and progress. +compactor.setProgressLimiter( + ProgressLimiter.logging( + ProgressLimiter.rateLimited(100.0 * 1024 * 1024), // leaky-bucket meter, bytes/sec + msg -> log.info("compaction: {}", msg))); +``` + +- `rateLimited(unitsPerSecond)` — a leaky-bucket rate meter (drains when idle, interruptible; grant is a no-op). +- `logging(delegate, sink)` / `logging(sink)` — logs each `onProgress` and each *blocked* `acquire`, delegating both facets; a `sink` is any `Consumer`. + +A host that draws from its own shared budget implements `acquire` directly: + +```java +compactor.setProgressLimiter(new ProgressLimiter() { + @Override public void onProgress(WorkStage stage, long completed, long total) { + metrics.report(stage, completed, total); + } + @Override public Grant acquire(long bytes) throws InterruptedException { + hostIoBudget.acquire(bytes); // block against a host-wide IO limiter + return Grant.NOOP; // rate-limiter model: nothing to release + } +}); +``` + +The returned `Grant` is closed when the admitted work completes. A **rate-limiter** realization pays its cost at `acquire` and returns `Grant.NOOP`; a **semaphore** realization (permits = in-flight bytes) releases on `Grant.close()`. + +> **Caveat (semaphore realization).** During `REFINE`, batches are submitted and drained through a sliding window of `taskWindowSize`, so a permit-releasing semaphore must admit at least `taskWindowSize` batches' worth of bytes or the window can't fill. `MERGE_LEVELS` has no such constraint, and the rate-limiter realization has none in either phase. + +### Output destination (no temp-file copy) + +By default `compact(Path)` writes a standalone file. To place the graph body directly inside a host container after a reserved header — eliminating a temp-file-and-copy — use either overload: + +```java +// (1) write the body into an existing file at a reserved offset; bytes in [0, startOffset) are preserved. +compactor.compact(componentPath, /* startOffset= */ headerSize); + +// (2) resource-scoped destination with a commit/abort lifecycle; returns the body length. +long bodyLength = compactor.compact(CompactionDestination.toFile(outputPath)); +``` + +`compact(CompactionDestination)` opens one `Target` per call and drives its lifecycle: + +```java +CompactionDestination dest = () -> { + FileChannel ch = FileChannel.open(componentPath, CREATE, WRITE, READ); + writeHostHeader(ch); // reserve [0, headerSize) + return new CompactionDestination.Target() { + public Path file() { return componentPath; } + public long startOffset() { return headerSize; } + public void commit(long bodyLength) throws IOException { // success: finalize the container + writeHostFooter(ch, bodyLength); + ch.force(true); + } + public void close() throws IOException { ch.close(); } // always runs; no commit ⇒ host discards the file + }; +}; +long bodyLength = compactor.compact(dest); +``` + +- `commit(bodyLength)` fires exactly once, only on success, after the body is written and forced. `close()` always runs (try-with-resources); reaching it without a prior `commit` is an unambiguous abort — discard the partial output. +- The compactor writes into `file()` at `startOffset()` using its own random-access and memory-mapped IO — the destination expresses *where*, not *how*. +- To read the committed body back (e.g. for a checksum), address the region with the `SeekableSink` primitive (package `io.github.jbellis.jvector.disk`): `SeekableSink.over(channel, target.startOffset())` gives region-relative `writeAt`/`readAt`. + ## Algorithm ### Ordinal Remapping @@ -104,7 +198,7 @@ for alpha in [1.0, 1.2]: Level 0 (base layer) stores inline vectors, FusedPQ codes, and the neighbor list. Upper levels store only the neighbor list (plus PQ codes at level 1 for cross-level searching). -Processing is batched per source and run in parallel across sources using a `ForkJoinPool`. A backpressure window keeps at most `taskWindowSize` batches in-flight at once, bounding memory use. +Processing is batched per source and run in parallel across sources on the supplied executor (a `ForkJoinPool` by default; see [Embedding API](#embedding-api)). A backpressure window keeps at most `taskWindowSize` batches in-flight at once, bounding memory use. ### Entry Node diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/FileChannelSeekableSink.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/FileChannelSeekableSink.java new file mode 100644 index 000000000..bcc3bd412 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/FileChannelSeekableSink.java @@ -0,0 +1,70 @@ +/* + * 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.disk; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; + +/** + * {@link SeekableSink} over a {@link FileChannel}, translating region-relative positions by a fixed + * base offset. The channel is owned by the caller; {@link #close()} does not close it. + */ +final class FileChannelSeekableSink implements SeekableSink { + private final FileChannel channel; + private final long baseOffset; + + FileChannelSeekableSink(FileChannel channel, long baseOffset) { + if (channel == null) { + throw new NullPointerException("channel"); + } + if (baseOffset < 0) { + throw new IllegalArgumentException("baseOffset must be >= 0, got " + baseOffset); + } + this.channel = channel; + this.baseOffset = baseOffset; + } + + @Override + public void writeAt(long position, ByteBuffer src) throws IOException { + if (position < 0) { + throw new IllegalArgumentException("position must be >= 0, got " + position); + } + long abs = baseOffset + position; + while (src.hasRemaining()) { + abs += channel.write(src, abs); + } + } + + @Override + public int readAt(long position, ByteBuffer dst) throws IOException { + if (position < 0) { + throw new IllegalArgumentException("position must be >= 0, got " + position); + } + return channel.read(dst, baseOffset + position); + } + + @Override + public void force() throws IOException { + channel.force(false); + } + + @Override + public void close() { + // The channel is owned by the caller, per SeekableSink.over(...). + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/SeekableSink.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/SeekableSink.java new file mode 100644 index 000000000..85833062f --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/SeekableSink.java @@ -0,0 +1,65 @@ +/* + * 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.disk; + +import io.github.jbellis.jvector.annotations.Experimental; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; + +/** + * A seekable region that supports positional reads and writes, addressed in coordinates relative + * to the region's start (0-based). An embedder uses it to hand a compactor (or other writer) a + * bounded window inside a larger container file: positions are region-relative and the + * implementation adds the container's base offset, so the writer never needs to know the absolute + * offset. + * + *

Implementations must support concurrent positional writes and reads to disjoint ranges (a + * {@link FileChannel} does). This is a generic IO primitive; the compaction extension point that + * hands one out is {@code io.github.jbellis.jvector.graph.disk.CompactionDestination}. + */ +@Experimental +public interface SeekableSink extends AutoCloseable { + + /** Write {@code src} fully at region-relative {@code position} (must be {@code >= 0}). */ + void writeAt(long position, ByteBuffer src) throws IOException; + + /** + * Read up to {@code dst.remaining()} bytes at region-relative {@code position} (must be + * {@code >= 0}); returns the number of bytes read, or {@code -1} at end of region. + */ + int readAt(long position, ByteBuffer dst) throws IOException; + + /** Force written bytes to durable storage. */ + void force() throws IOException; + + @Override + void close() throws IOException; + + /** + * Reference implementation over a {@link FileChannel} region. Every region-relative position is + * translated by {@code baseOffset}. The channel's lifecycle is owned by the caller — this + * {@link #close()} does not close the channel. + * + * @param channel the backing channel, opened for read and write + * @param baseOffset the absolute offset of the region's start within {@code channel} ({@code >= 0}) + */ + static SeekableSink over(FileChannel channel, long baseOffset) { + return new FileChannelSeekableSink(channel, baseOffset); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java index b01901832..b0deca442 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java @@ -21,7 +21,7 @@ import java.util.Collections; import java.util.List; -import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.ExecutorService; /** * Bundle of inputs that {@link QuantizationCompactionStrategy} implementations need to do their work @@ -38,7 +38,7 @@ public final class CompactionContext { public final List remappers; public final int dimension; public final int maxOrdinal; - public final ForkJoinPool executor; + public final ExecutorService executor; public final int taskWindowSize; public CompactionContext( @@ -48,7 +48,7 @@ public CompactionContext( List remappers, int dimension, int maxOrdinal, - ForkJoinPool executor, + ExecutorService executor, int taskWindowSize) { this.sources = Collections.unmodifiableList(sources); this.sourceCompressed = sourceCompressed == null ? null : Collections.unmodifiableList(sourceCompressed); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionDestination.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionDestination.java new file mode 100644 index 000000000..5049333ad --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionDestination.java @@ -0,0 +1,83 @@ +/* + * 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.annotations.Experimental; +import io.github.jbellis.jvector.disk.SeekableSink; + +import java.io.IOException; +import java.nio.file.Path; + +/** + * Embedding extension point: tells {@link OnDiskGraphIndexCompactor} WHERE to write its compacted + * graph, so the body lands directly inside the embedder's container (after a header the embedder + * reserves) — eliminating the temp-file-and-copy. Resource-scoped: the compactor uses one + * {@link Target} per {@code compact(...)} call, commits on success, and always closes. + * + * 

{@code
+ *   try (CompactionDestination.Target t = destination.open()) {
+ *       // ...compactor writes the graph into t.file() at t.startOffset()...
+ *       t.commit(bodyLength);   // success: body written & durable; embedder finalizes its footer
+ *   }                           // close() always runs; no commit() => aborted (discard partial output)
+ * }
+ * + *

The compactor needs a real file (it uses a memory-mapped read-back during refinement and a + * random-access writer), so a {@link Target} is expressed as a container {@link Path} plus a base + * offset rather than an opaque stream. The generic {@link SeekableSink} primitive addresses the same + * window in region-relative coordinates and is what an embedder uses to read the committed body back + * for its checksum, e.g. {@code SeekableSink.over(channel, target.startOffset())}. + */ +@FunctionalInterface +@Experimental +public interface CompactionDestination { + + /** Open a fresh target for one compaction. */ + Target open() throws IOException; + + /** One compaction's output region plus its commit/abort lifecycle. */ + interface Target extends AutoCloseable { + + /** The container file the graph body is written into. */ + Path file(); + + /** The byte offset within {@link #file()} at which the graph body begins ({@code >= 0}). */ + long startOffset(); + + /** + * Signalled exactly once, after the body has been fully written and forced, reporting its + * length ({@code file() size - startOffset()}). The embedder finalizes its container here + * (e.g. writes a footer/checksum). MUST be called before {@link #close()} on the success path. + */ + void commit(long bodyLength) throws IOException; + + /** + * Always runs (try-with-resources). If reached without a prior {@link #commit}, the + * compaction failed and the embedder discards the partial output; releases embedder resources. + */ + @Override + void close() throws IOException; + } + + /** + * Default standalone destination: writes to its own file at offset {@code 0} (today's + * {@code compact(Path)} behaviour). {@code commit} is a no-op marker; a {@code close} without a + * prior commit deletes the partial file. + */ + static CompactionDestination toFile(Path path) { + return new FileCompactionDestination(path); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FileCompactionDestination.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FileCompactionDestination.java new file mode 100644 index 000000000..24d610d0e --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FileCompactionDestination.java @@ -0,0 +1,66 @@ +/* + * 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 java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; + +/** + * {@link CompactionDestination} that writes a standalone graph file at offset 0. Backs + * {@link CompactionDestination#toFile(Path)}. + */ +final class FileCompactionDestination implements CompactionDestination { + private final Path path; + + FileCompactionDestination(Path path) { + if (path == null) { + throw new NullPointerException("path"); + } + this.path = path; + } + + @Override + public Target open() { + return new Target() { + private boolean committed; + + @Override + public Path file() { + return path; + } + + @Override + public long startOffset() { + return 0L; + } + + @Override + public void commit(long bodyLength) { + // Standalone file: the graph IS the whole file; compact() already wrote and flushed it. + committed = true; + } + + @Override + public void close() throws IOException { + if (!committed) { + Files.deleteIfExists(path); // abort: discard the partial file + } + } + }; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index deffb719f..8e2219ec4 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -35,6 +35,9 @@ import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider; import io.github.jbellis.jvector.util.*; +import io.github.jbellis.jvector.util.work.ProgressLimiter; +import io.github.jbellis.jvector.util.work.WorkLimiter; +import io.github.jbellis.jvector.util.work.WorkStage; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.graph.similarity.ScoreFunction; import io.github.jbellis.jvector.vector.types.VectorFloat; @@ -81,33 +84,82 @@ public final class OnDiskGraphIndexCompactor implements Accountable { private final int dimension; private int maxOrdinal = -1; private int numTotalNodes = 0; - private final ForkJoinPool executor; + private final Executor executor; private final int taskWindowSize; private final VectorSimilarityFunction similarityFunction; + // Embedder progress + work-admission control surface (see io.github.jbellis.jvector.util.work). + // Default UNLIMITED = no observation and no throttling → byte-identical output and equivalent + // timing to no SPI installed. + private volatile ProgressLimiter limiter = ProgressLimiter.UNLIMITED; + /** - * Constructs a new OnDiskGraphIndexCompactor for graphs without a non-fused compressed sidecar. - * Equivalent to calling the 6-arg constructor with {@code sourceCompressed = null}. + * The stages a compaction reports progress and acquires work admission against. Names are + * stable so an embedder can distinguish them. The write side of {@code MERGE_LEVELS} carries the + * graph-body IO — there is no separate flush phase. */ @Experimental - public OnDiskGraphIndexCompactor( - List sources, - List liveNodes, - List remappers, - VectorSimilarityFunction similarityFunction, - ForkJoinPool executor) { - this(sources, null, liveNodes, remappers, similarityFunction, executor); + public enum Phase implements WorkStage { + /** Merging source graphs level-by-level and writing the compacted graph body. */ + MERGE_LEVELS, + /** Second pass refining neighborhoods in the compacted graph. */ + REFINE } /** - * Constructs a new OnDiskGraphIndexCompactor to merge multiple graph indexes. - * Initializes thread pool, validates inputs, and prepares metadata for compaction. + * Installs an embedder control surface for progress observation and work admission. Both facets + * are optional (see {@link ProgressLimiter}); passing {@code null} restores + * {@link ProgressLimiter#UNLIMITED}. Returns {@code this} for chaining. + * + *

Admission ({@code acquire}) is invoked only on the orchestrating thread — the caller of + * {@code compact} — never on a pool worker, so a blocking limiter back-pressures batch dispatch + * without a {@code ForkJoinPool.ManagedBlocker}. The unit of {@code acquire} for this consumer + * is bytes about to be written. + */ + @Experimental + public OnDiskGraphIndexCompactor setProgressLimiter(ProgressLimiter limiter) { + this.limiter = (limiter == null) ? ProgressLimiter.UNLIMITED : limiter; + return this; + } + + /** + * The compactor's task window: the number of batches kept in flight, equal to the injected + * pool's {@code getParallelism()} (or the {@link PhysicalCoreExecutor#pool()} default's). It + * bounds both compaction concurrency and the in-flight memory window; embedders can log or + * assert it to confirm the parallelism they injected. + */ + @Experimental + public int getTaskWindowSize() { + return taskWindowSize; + } + + /** + * Primary constructor: merges multiple graph indexes using any {@link Executor} with an + * explicit in-flight window. * * @param sourceCompressed parallel to {@code sources}, supplying the non-fused compressed * vectors (e.g. {@link io.github.jbellis.jvector.quantization.PQVectors}) * that ship alongside each graph. Pass {@code null} when sources carry * quantization inline (FUSED_PQ) or have none. Must not be combined * with sources that carry the FUSED_PQ feature. + * @param executor runs compaction batches submitted via an internal + * {@link java.util.concurrent.ExecutorCompletionService} while the calling + * thread blocks for their completion. Safe to pass a {@link ForkJoinPool} + * (work-stealing + managed blocking make submit-and-block safe) or a + * caller-runs / same-thread executor (e.g. {@code Runnable::run}), which + * runs each batch synchronously on the calling thread — no worker threads, no + * separate pool. It is not safe to pass a bounded + * {@code ThreadPoolExecutor} that is also running the calling thread: the + * caller blocks in {@code take()} while its sub-tasks queue behind it, which can + * thread-starvation-deadlock (worst case a single-thread pool). Embedders that + * want their own bounded pool without a dedicated fan-out pool should pass a + * caller-runs executor and derive parallelism from the number of concurrent + * compactions instead. Pass {@code null} to use the shared + * {@link PhysicalCoreExecutor#pool()} default. The compactor never owns or shuts + * down the executor. + * @param taskWindowSize bounds the number of in-flight batches (both concurrency and peak + * write-side memory). {@code <= 0} derives it from a {@link ForkJoinPool}'s + * {@code getParallelism()}, else defaults to 1 (serial). */ @Experimental public OnDiskGraphIndexCompactor( @@ -116,22 +168,16 @@ public OnDiskGraphIndexCompactor( List liveNodes, List remappers, VectorSimilarityFunction similarityFunction, - ForkJoinPool executor) { + Executor executor, + int taskWindowSize) { checkBeforeCompact(sources, sourceCompressed, liveNodes, remappers); - if (executor != null) { - this.executor = executor; - } else { - // Default to the shared physical-core pool. Compaction (PQ encode + parallel record - // flush + refinement) is compute- and memory-bandwidth-bound, so sizing to logical - // cores oversubscribes hyperthreaded hosts and costs throughput. This pool is - // process-wide and shared with index construction and quantization; the compactor - // never owns or shuts it down. - this.executor = PhysicalCoreExecutor.pool(); - } - // Track the pool's real parallelism so task-window / backpressure sizing stays correct - // whether the executor is the shared default or a caller-injected pool. - this.taskWindowSize = this.executor.getParallelism(); + // Default to the shared physical-core pool. Compaction (PQ encode + parallel record flush + + // refinement) is compute- and memory-bandwidth-bound, so sizing to logical cores + // oversubscribes hyperthreaded hosts. This pool is process-wide and shared with index + // construction and quantization; the compactor never owns or shuts it down. + this.executor = (executor != null) ? executor : PhysicalCoreExecutor.pool(); + this.taskWindowSize = resolveWindow(this.executor, taskWindowSize); this.sources = sources; this.sourceCompressed = (sourceCompressed == null || sourceCompressed.isEmpty()) ? null : sourceCompressed; @@ -155,6 +201,49 @@ public OnDiskGraphIndexCompactor( this.similarityFunction = similarityFunction; } + /** + * Convenience {@link Executor} constructor without a non-fused compressed sidecar. Equivalent + * to the 7-arg form with {@code sourceCompressed = null}. + */ + @Experimental + public OnDiskGraphIndexCompactor( + List sources, + List liveNodes, + List remappers, + VectorSimilarityFunction similarityFunction, + Executor executor, + int taskWindowSize) { + this(sources, null, liveNodes, remappers, similarityFunction, executor, taskWindowSize); + } + + private static int resolveWindow(Executor executor, int requested) { + if (requested > 0) return requested; + if (executor instanceof ForkJoinPool) return ((ForkJoinPool) executor).getParallelism(); + return 1; // arbitrary Executor with no declared parallelism → serial window + } + + /** + * Adapts an arbitrary {@link Executor} to the {@link ExecutorService} the quantization + * strategies need for {@code invokeAll} during PQ pre-encode. A real {@code ExecutorService} + * (including a {@link ForkJoinPool}) is used directly; a plain {@code Executor} — notably a + * caller-runs {@code Runnable::run} — is wrapped so its tasks run on whatever thread the + * executor dispatches to (the calling thread, for caller-runs). No pool is created, and the + * lifecycle methods are inert since the compactor never owns the executor. + */ + private static ExecutorService asExecutorService(Executor executor) { + if (executor instanceof ExecutorService) { + return (ExecutorService) executor; + } + return new AbstractExecutorService() { + @Override public void execute(Runnable command) { executor.execute(command); } + @Override public void shutdown() { } + @Override public List shutdownNow() { return Collections.emptyList(); } + @Override public boolean isShutdown() { return false; } + @Override public boolean isTerminated() { return false; } + @Override public boolean awaitTermination(long timeout, TimeUnit unit) { return false; } + }; + } + /** * Validates that all source indexes have compatible configurations and required features * before attempting compaction. Ensures consistent dimensions, max degrees, hierarchical @@ -294,11 +383,32 @@ private void validateFeatures(List sources) { */ @Experimental public void compact(Path outputPath) throws FileNotFoundException { + compact(outputPath, 0L); + } + + /** + * No-copy compaction entry point: writes the compacted graph into {@code outputPath} at + * {@code startOffset}, leaving any bytes in {@code [0, startOffset)} untouched. An embedder + * that wraps the graph in its own container reserves its header by passing the header size as + * {@code startOffset}, so jvector's body lands directly inside the container — removing the + * temp-file-and-copy. jvector writes only {@code [startOffset, projectedSize)} and never reads + * or clobbers the reserved prefix. The file is opened read/write and not truncated, so a + * prefix the embedder pre-wrote survives. {@code compact(path, 0)} equals {@link #compact(Path)}. + * + * @param outputPath the file to write into + * @param startOffset the byte offset at which jvector's output begins; {@code 0} for a + * standalone file + */ + @Experimental + public void compact(Path outputPath, long startOffset) throws FileNotFoundException { + if (startOffset < 0) { + throw new IllegalArgumentException("startOffset must be >= 0, got " + startOffset); + } QuantizationCompactionStrategy strategy = detectInlineStrategy(); try { - compactGraphImpl(outputPath, strategy); + compactGraphImpl(outputPath, startOffset, strategy); releaseSourcesBeforeRefine(strategy); - refineCompactedGraph(outputPath, strategy); + refineCompactedGraph(outputPath, startOffset, 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. @@ -306,6 +416,28 @@ public void compact(Path outputPath) throws FileNotFoundException { } } + /** + * No-copy compaction into an embedder-supplied {@link CompactionDestination}: writes the graph + * body into the destination's container at its reserved offset, then {@code commit}s the body + * length on success (so the embedder can finalize its footer/checksum), or — on any failure — + * closes the target without committing (so the embedder discards the partial output). The + * compactor writes into {@code target.file()} at {@code target.startOffset()} using its own + * random-access + memory-mapped IO; the destination expresses where, not how. + * + * @return the graph body length written, i.e. {@code size(file) - startOffset} + */ + @Experimental + public long compact(CompactionDestination destination) throws IOException { + try (CompactionDestination.Target target = destination.open()) { + Path file = target.file(); + long base = target.startOffset(); + compact(file, base); + long bodyLength = java.nio.file.Files.size(file) - base; + target.commit(bodyLength); + return bodyLength; + } + } + /** * 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 @@ -329,8 +461,8 @@ public void compact(Path graphPath, Path compressedPath) throws FileNotFoundExce QuantizationCompactionStrategy sidecarStrategy = detectSidecarStrategy(); try { sidecarStrategy.retrain(similarityFunction); - compactGraphImpl(graphPath, inlineStrategy); - refineCompactedGraph(graphPath, inlineStrategy); + compactGraphImpl(graphPath, 0L, inlineStrategy); + refineCompactedGraph(graphPath, 0L, inlineStrategy); sidecarStrategy.writeSidecar(compressedPath); } catch (IOException e) { throw new RuntimeException("Sidecar compaction failed", e); @@ -388,7 +520,7 @@ private QuantizationCompactionStrategy detectSidecarStrategy() { /** Snapshot the compactor's state into a {@link CompactionContext} for strategies to consume. */ private CompactionContext buildContext() { return new CompactionContext(sources, sourceCompressed, liveNodes, remappers, - dimension, maxOrdinal, executor, taskWindowSize); + dimension, maxOrdinal, asExecutorService(executor), taskWindowSize); } /** @@ -401,7 +533,7 @@ private CompactionContext buildContext() { * 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 { + private void compactGraphImpl(Path outputPath, long startOffset, QuantizationCompactionStrategy strategy) throws FileNotFoundException { strategy.retrain(similarityFunction); boolean fusedPQEnabled = strategy.writesCodesInline(); @@ -417,7 +549,7 @@ private void compactGraphImpl(Path outputPath, QuantizationCompactionStrategy st log.info("Writing compacted graph : {} total nodes, maxOrdinal={}, dimension={}, degree={}", numTotalNodes, maxOrdinal, dimension, maxDegrees.get(0)); - try (CompactWriter writer = new CompactWriter(outputPath, maxOrdinal, numTotalNodes, 0, layerInfo, entryNode, dimension, maxDegrees, outputFusedFeature)) { + try (CompactWriter writer = new CompactWriter(outputPath, maxOrdinal, numTotalNodes, startOffset, 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(); @@ -457,7 +589,7 @@ private void compactGraphImpl(Path outputPath, QuantizationCompactionStrategy st * 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) { + private void refineCompactedGraph(Path outputPath, long startOffset, QuantizationCompactionStrategy strategy) { log.info("Refining compacted graph: {}", outputPath); long t0 = System.nanoTime(); @@ -484,7 +616,7 @@ private void refineCompactedGraph(Path outputPath, QuantizationCompactionStrateg // 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); + OnDiskGraphIndex mergedGraph = OnDiskGraphIndex.load(supplier, startOffset, 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 @@ -520,33 +652,51 @@ private void refineCompactedGraph(Path outputPath, QuantizationCompactionStrateg 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; + // Windowed submit/drain (mirrors runBatchesWithBackpressure) so a blocking WorkLimiter — + // a rate limiter, or a semaphore that releases permits on Grant.close() — is + // deadlock-free: admission is on the orchestrator before each submit, and exactly one + // grant is released per completed batch. This also bounds in-flight refine batches (and + // thus peak memory) to taskWindowSize, which the prior submit-all loop did not. Amount is + // an estimate of the per-node record bytes rewritten. UNLIMITED makes it all no-ops. + final long refinedRecordSize = (long) dimension * Float.BYTES + + (long) baseDegree * Integer.BYTES + codeSize; + final int nBatches = (total + batchSize - 1) / batchSize; + java.util.ArrayDeque grants = new java.util.ArrayDeque<>(); + + int nextStart = 0, inFlight = 0, completed = 0, 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; + try { + while (completed < nBatches) { + while (inFlight < taskWindowSize && nextStart < total) { + final int s = nextStart; + final int e = Math.min(nextStart + batchSize, total); + nextStart = e; + grants.add(limiter.acquire((long) (e - s) * refinedRecordSize)); // may park the orchestrator + 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; + }); + inFlight++; + } + nodesDone += ecs.take().get(); + completed++; + inFlight--; + WorkLimiter.Grant g = grants.poll(); + if (g != null) g.close(); + limiter.onProgress(Phase.REFINE, nodesDone, total); + if (nodesDone >= nextProgress) { + log.info("Refinement progress: {}/{} nodes", nodesDone, total); + nextProgress += progressStep; + } } + } finally { + for (WorkLimiter.Grant g : grants) g.close(); } // Per-thread scratches live in worker-thread ThreadLocals; closing the supplier in @@ -868,6 +1018,11 @@ private void compactLevels(CompactWriter writer, new Scratch(maxCandidateSize, scratchDegree, dimension, sources, pq) ); + // MERGE_LEVELS progress denominator: base-layer live nodes. Level 0 (the bulk) runs first, + // so progress[0] climbs 0 -> numTotalNodes across it, then holds while the small upper tail + // runs (completed is clamped to the total). progress persists across all level calls. + long[] mergeProgress = { 0L, numTotalNodes }; + 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); @@ -908,7 +1063,14 @@ private void compactLevels(CompactWriter writer, } catch (IOException e) { throw new RuntimeException(e); } - } + }, + Phase.MERGE_LEVELS, + (results) -> { // exact bytes: read before the write consumes the buffers + long s = 0; + for (WriteResult r : results) s += r.data.remaining(); + return s; + }, + mergeProgress ); } @@ -945,7 +1107,15 @@ private void compactLevels(CompactWriter writer, } catch (IOException e) { throw new RuntimeException(e); } - } + }, + Phase.MERGE_LEVELS, + (results) -> { // estimate: neighbor ids + optional PQ code per node + long s = 0; + for (UpperLayerWriteResult r : results) + s += (long) r.neighbors.length * Integer.BYTES + (r.pqCode == null ? 0 : r.pqCode.length()); + return s; + }, + mergeProgress ); } } @@ -1291,7 +1461,10 @@ private void runBatchesWithBackpressure( List batches, ExecutorCompletionService> ecs, java.util.function.Consumer submitOne, - java.util.function.Consumer> onComplete + java.util.function.Consumer> onComplete, + WorkStage stage, + java.util.function.ToLongFunction> batchBytes, + long[] progress ) throws InterruptedException, ExecutionException { final int total = batches.size(); @@ -1307,11 +1480,22 @@ private void runBatchesWithBackpressure( int completed = 0; while (completed < total) { List results = ecs.take().get(); - onComplete.accept(results); + + // Admission runs on this (orchestrating) thread, before the write. A blocking limiter + // back-pressures dispatch/consume while in-flight workers keep computing — no + // ManagedBlocker. The amount is the bytes this batch is about to write, read here before + // onComplete consumes the buffers. For the default UNLIMITED limiter both calls are no-ops. + long amount = batchBytes.applyAsLong(results); + try (WorkLimiter.Grant g = limiter.acquire(amount)) { + onComplete.accept(results); + } completed++; inFlight--; + progress[0] = Math.min(progress[0] + results.size(), progress[1]); + limiter.onProgress(stage, progress[0], progress[1]); + if (nextToSubmit < total) { submitOne.accept(batches.get(nextToSubmit++)); inFlight++; @@ -1354,11 +1538,18 @@ private List computeLayerInfoFromSources() { int count = 0; for (int s = 0; s < sources.size(); s++) { if (level > sources.get(s).getMaxLevel()) continue; - NodesIterator it = sources.get(s).getNodes(level); - FixedBitSet alive = liveNodes.get(s); - while (it.hasNext()) { - int node = it.next(); - if (alive.get(node)) count++; + if (level == 0) { + // Every live node is present at level 0 (HNSW base layer invariant), + // so count directly from the in-memory bitset instead of scanning node + // records on disk (which touches gigabytes of source data on a cold cache). + count += liveNodes.get(s).cardinality(); + } else { + NodesIterator it = sources.get(s).getNodes(level); + FixedBitSet alive = liveNodes.get(s); + while (it.hasNext()) { + int node = it.next(); + if (alive.get(node)) count++; + } } } layerInfo.add(new CommonHeader.LayerInfo(count, maxDegrees.get(level))); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java index 280d75c9c..c45120f70 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java @@ -22,6 +22,7 @@ import io.github.jbellis.jvector.quantization.ProductQuantization; import io.github.jbellis.jvector.util.DocIdSetIterator; import io.github.jbellis.jvector.util.FixedBitSet; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.types.VectorFloat; @@ -32,6 +33,7 @@ import java.util.ArrayList; import java.util.Comparator; import java.util.List; +import java.util.concurrent.ForkJoinPool; import java.util.concurrent.ThreadLocalRandom; /** @@ -96,22 +98,27 @@ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, log.info("Collected {} training samples", samples.size()); - // Extract vectors sequentially in sorted (source, node) order so disk reads are - // purely sequential and the OS read-ahead can cover them efficiently. We do this - // here rather than letting ProductQuantization.compute() drive the reads via its - // parallel stream, which would scatter page faults across a potentially very large - // file and cause I/O that scales with dataset size rather than sample count. List> trainingVectors = extractVectorsSequential(samples); - var ravv = new ListRandomAccessVectorValues(trainingVectors, dimension); - boolean center = similarityFunction == VectorSimilarityFunction.EUCLIDEAN; + long t0 = System.nanoTime(); + log.info("Extracted {} vectors in {}ms; starting PQ refinement", + trainingVectors.size(), (System.nanoTime() - t0) / 1_000_000L); + + var ravv = new ListRandomAccessVectorValues(trainingVectors, dimension); - return ProductQuantization.compute( - ravv, - basePQ.getSubspaceCount(), - basePQ.getClusterCount(), - center - ); + // Warm-start from the existing codebook via Lloyd's-only refinement rather than + // re-running k-means++ from scratch. k-means++ initialization visits every point + // once per centroid (256 passes for k=256), which dominates training time. + // Since the source codebooks are already trained on data from the same underlying + // distribution, this warm-start converges in far fewer passes with no recall loss. + long t1 = System.nanoTime(); + ProductQuantization result = basePQ.refine(ravv, + ProductQuantization.K_MEANS_ITERATIONS, + -1.0f, // UNWEIGHTED / isotropic + PhysicalCoreExecutor.pool(), + ForkJoinPool.commonPool()); + log.info("PQ refinement complete in {}ms", (System.nanoTime() - t1) / 1_000_000L); + return result; } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java index c84b7b955..6d9d23879 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java @@ -60,7 +60,7 @@ public class ProductQuantization implements VectorCompressor>, A private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); static final int DEFAULT_CLUSTERS = 256; // number of clusters per subspace = one byte's worth - static final int K_MEANS_ITERATIONS = 6; + public static final int K_MEANS_ITERATIONS = 6; public static final int MAX_PQ_TRAINING_SET_SIZE = 128000; final VectorFloat[] codebooks; // array of codebooks, where each codebook is a VectorFloat consisting of k contiguous subvectors each of length M diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/LeakyBucketLimiter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/LeakyBucketLimiter.java new file mode 100644 index 000000000..0ebba8f3d --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/LeakyBucketLimiter.java @@ -0,0 +1,67 @@ +/* + * 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.util.work; + +import java.util.concurrent.TimeUnit; + +/** + * A leaky-bucket rate meter realizing the {@link WorkLimiter} facet: {@link #acquire} paces the + * aggregate admitted amount to a fixed {@code unitsPerSecond}, blocking the caller when the rate + * would be exceeded. The bucket drains during idle gaps (a burst after a quiet period is not + * charged for the idle time), and the first request after an idle period is admitted without + * delay — the cost of each request is paid by the next one, which is the standard smooth + * shaping behaviour. {@link #onProgress} is inherited as a no-op: this limiter only throttles. + * + *

Thread-safe and reentrant: the emission clock is advanced under a short lock, then the caller + * sleeps outside the lock, so concurrent callers serialize their reservations but wait + * independently. The returned grant is a no-op — the cost is paid entirely at {@code acquire}. + * + *

Obtain instances via {@link ProgressLimiter#rateLimited(double)}. + */ +final class LeakyBucketLimiter implements ProgressLimiter { + private final double nanosPerUnit; + private final Object lock = new Object(); + // Earliest nanoTime at which the next reservation may start. Long.MIN_VALUE until the first + // acquire, so Math.max(now, nextFreeNanos) == now (a fully drained bucket) on the first call. + private long nextFreeNanos = Long.MIN_VALUE; + + LeakyBucketLimiter(double unitsPerSecond) { + if (!(unitsPerSecond > 0) || Double.isInfinite(unitsPerSecond)) { + throw new IllegalArgumentException("unitsPerSecond must be finite and > 0, got " + unitsPerSecond); + } + this.nanosPerUnit = 1_000_000_000.0 / unitsPerSecond; + } + + @Override + public Grant acquire(long amount) throws InterruptedException { + if (amount <= 0) { + return Grant.NOOP; + } + long startAt; + synchronized (lock) { + long now = System.nanoTime(); + startAt = Math.max(now, nextFreeNanos); // drain if idle, else queue behind backlog + long cost = (long) Math.min((double) Long.MAX_VALUE, amount * nanosPerUnit); + nextFreeNanos = startAt + cost; + } + // Sleep (interruptibly, so cancellation aborts) until this request's slot opens. + for (long remaining = startAt - System.nanoTime(); remaining > 0; remaining = startAt - System.nanoTime()) { + TimeUnit.NANOSECONDS.sleep(remaining); + } + return Grant.NOOP; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressLimiter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressLimiter.java new file mode 100644 index 000000000..34a23d9dd --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressLimiter.java @@ -0,0 +1,97 @@ +/* + * 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.util.work; + +import io.github.jbellis.jvector.annotations.Experimental; + +import java.util.Objects; +import java.util.function.Consumer; + +/** + * The {@link ProgressTracker tracker} and the {@link WorkLimiter throttle} melded into one control + * surface. A long-running jvector operation accepts a single {@code ProgressLimiter} and uses both + * facets; an embedder may override only the facet it needs — the other defaults to a no-op, so + * {@link #UNLIMITED} behaves exactly as if no SPI were installed. + * + *

Both methods default to no-ops here. A consumer that wants only one facet can still accept a + * lambda via the single-method parents ({@link ProgressTracker}, {@link WorkLimiter}); a consumer + * that wants both accepts a {@code ProgressLimiter}. + */ +@Experimental +public interface ProgressLimiter extends ProgressTracker, WorkLimiter { + + @Override + default void onProgress(WorkStage stage, long completed, long total) { } + + @Override + default Grant acquire(long amount) throws InterruptedException { return Grant.NOOP; } + + /** Observes nothing and limits nothing — behaviour identical to no SPI installed. */ + ProgressLimiter UNLIMITED = new ProgressLimiter() { }; + + /** + * A leaky-bucket rate meter realizing the throttle facet: {@link #acquire} paces the aggregate + * admitted amount to {@code unitsPerSecond} (bytes/sec for the compaction consumer), blocking + * the caller when the rate would be exceeded and draining during idle gaps. {@link #onProgress} + * is a no-op and the returned grant is a no-op (cost is paid at {@code acquire}). Compose with + * {@link #logging(ProgressLimiter, Consumer)} to also log. + * + * @param unitsPerSecond the sustained admission rate; must be finite and {@code > 0} + * @throws IllegalArgumentException if {@code unitsPerSecond} is not finite and positive + */ + static ProgressLimiter rateLimited(double unitsPerSecond) { + return new LeakyBucketLimiter(unitsPerSecond); + } + + /** + * Wraps {@code delegate}, emitting a one-line message to {@code sink} on each + * {@link #onProgress} and on each {@link #acquire} that actually blocked, then delegating both + * facets to {@code delegate}. Composes over any limiter — e.g. + * {@code logging(rateLimited(bytesPerSecond), log::info)} logs a rate-limited operation. The + * delegate's grant is returned unchanged, so a semaphore delegate still releases on close. + * + * @param delegate the limiter to observe and delegate to; {@code null} means {@link #UNLIMITED} + * @param sink receives formatted log lines (e.g. {@code msg -> logger.info(msg)}) + */ + static ProgressLimiter logging(ProgressLimiter delegate, Consumer sink) { + Objects.requireNonNull(sink, "sink"); + final ProgressLimiter d = (delegate == null) ? UNLIMITED : delegate; + return new ProgressLimiter() { + @Override + public void onProgress(WorkStage stage, long completed, long total) { + sink.accept("progress[" + stage.name() + "] " + completed + "/" + (total < 0 ? "?" : Long.toString(total))); + d.onProgress(stage, completed, total); + } + + @Override + public Grant acquire(long amount) throws InterruptedException { + long startNanos = System.nanoTime(); + Grant g = d.acquire(amount); + long waitedMs = (System.nanoTime() - startNanos) / 1_000_000L; + if (waitedMs > 0) { + sink.accept("acquire " + amount + " units - throttled " + waitedMs + "ms"); + } + return g; + } + }; + } + + /** Logging over no throttle: equivalent to {@code logging(UNLIMITED, sink)}. */ + static ProgressLimiter logging(Consumer sink) { + return logging(UNLIMITED, sink); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressTracker.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressTracker.java new file mode 100644 index 000000000..d7d117269 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressTracker.java @@ -0,0 +1,43 @@ +/* + * 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.util.work; + +import io.github.jbellis.jvector.annotations.Experimental; + +/** + * Observation contract: receives progress updates for a stage of a long-running operation. + * + *

Best-effort and cheap: implementations must not throw (the caller invokes this on its + * orchestrating thread and treats it as fire-and-forget). See {@link ProgressLimiter} for the + * melded progress + throttle surface that most consumers accept. + */ +@Experimental +@FunctionalInterface +public interface ProgressTracker { + /** + * Reports progress for {@code stage}. + * + * @param stage the stage reporting progress + * @param completed work done so far in this stage, in stage-defined units; monotonically + * non-decreasing within a stage + * @param total total work for this stage, or {@code -1} if not yet known + */ + void onProgress(WorkStage stage, long completed, long total); + + /** A tracker that discards every update. */ + ProgressTracker NOOP = (stage, completed, total) -> { }; +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkLimiter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkLimiter.java new file mode 100644 index 000000000..5355f0757 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkLimiter.java @@ -0,0 +1,59 @@ +/* + * 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.util.work; + +import io.github.jbellis.jvector.annotations.Experimental; + +/** + * Admission contract: blocks until an amount of work may proceed, returning a {@link Grant} that + * the caller closes once the admitted work has completed. + * + *

The unit of {@code amount} is defined by the consumer (e.g. bytes for IO, or rows, + * nodes, items); jvector fixes only the blocking-grant mechanism, never the meaning of the + * quantity. Implementations must be thread-safe and reentrant. {@code acquire} may block but must + * not throw for ordinary back-pressure. + */ +@Experimental +@FunctionalInterface +public interface WorkLimiter { + /** + * Blocks until {@code amount} units of work may proceed. + * + * @param amount the amount of work about to be performed, in consumer-defined units + * @return a non-null grant to {@link Grant#close() close} once that work has completed + * @throws InterruptedException if the calling thread is interrupted while blocked, which + * aborts the operation + */ + Grant acquire(long amount) throws InterruptedException; + + /** + * A handle released by the consumer once the admitted work has completed. For a rate-limiter + * realization (cost paid at {@link WorkLimiter#acquire}) {@link #close()} is a no-op; for a + * semaphore-style in-flight-amount realization it releases the permits taken by {@code acquire}. + */ + interface Grant extends AutoCloseable { + /** Releases the grant. Never throws. */ + @Override + void close(); + + /** A grant that holds nothing and releases nothing. */ + Grant NOOP = () -> { }; + } + + /** A limiter that admits everything immediately. */ + WorkLimiter UNLIMITED = amount -> Grant.NOOP; +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkStage.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkStage.java new file mode 100644 index 000000000..001197a00 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkStage.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.util.work; + +import io.github.jbellis.jvector.annotations.Experimental; + +/** + * Identifies a stage of a long-running operation. The consumer defines its own stages; an + * {@code enum} satisfies this for free via {@link Enum#name()}. + * + *

Part of the generic progress + work-admission SPI ({@link ProgressTracker}, + * {@link WorkLimiter}, {@link ProgressLimiter}). Neither the stage identity nor the unit of work + * is fixed by jvector; both are supplied by the consumer. + */ +@Experimental +public interface WorkStage { + /** The stage's name, stable within a single operation. */ + String name(); +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java index 15543ebbc..a60cb2ed2 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java @@ -172,7 +172,7 @@ private static BenchResult compactAndMeasure(DataSet ds, PartitionConfig cfg, // Compact Path compactPath = tempDir.resolve("compacted"); - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, vsf, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, vsf, null, -1); long t0 = System.currentTimeMillis(); compactor.compact(compactPath); long compactionMs = System.currentTimeMillis() - t0; diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestSeekableSink.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestSeekableSink.java new file mode 100644 index 000000000..9fde6b162 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestSeekableSink.java @@ -0,0 +1,85 @@ +/* + * 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.disk; + +import org.junit.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +public class TestSeekableSink { + + @Test + public void writesAndReadsInRegionRelativeCoordinates() throws IOException { + Path f = Files.createTempFile("sink", ".bin"); + try (FileChannel ch = FileChannel.open(f, StandardOpenOption.WRITE, StandardOpenOption.READ)) { + long base = 100; + SeekableSink sink = SeekableSink.over(ch, base); + sink.writeAt(0, ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))); + sink.writeAt(5, ByteBuffer.wrap("WORLD".getBytes(StandardCharsets.UTF_8))); + sink.force(); + + // Region-relative read returns what was written. + ByteBuffer dst = ByteBuffer.allocate(10); + assertEquals(10, sink.readAt(0, dst)); + assertEquals("helloWORLD", new String(dst.array(), StandardCharsets.UTF_8)); + + // The bytes actually land at the absolute base offset (region-relative -> absolute). + ByteBuffer raw = ByteBuffer.allocate(10); + ch.read(raw, base); + assertEquals("helloWORLD", new String(raw.array(), StandardCharsets.UTF_8)); + + // Nothing was written before the region. + ByteBuffer before = ByteBuffer.allocate((int) base); + ch.read(before, 0); + for (byte b : before.array()) { + assertEquals("region must not write before its base", 0, b); + } + + // close() must NOT close the caller-owned channel. + sink.close(); + assertTrue("sink.close() must not close the caller's channel", ch.isOpen()); + } + Files.deleteIfExists(f); + } + + @Test + public void rejectsNegativeBaseAndPosition() throws IOException { + Path f = Files.createTempFile("sink", ".bin"); + try (FileChannel ch = FileChannel.open(f, StandardOpenOption.WRITE, StandardOpenOption.READ)) { + try { SeekableSink.over(ch, -1); fail("negative base"); } catch (IllegalArgumentException expected) { } + SeekableSink sink = SeekableSink.over(ch, 0); + try { sink.writeAt(-1, ByteBuffer.allocate(1)); fail("negative write pos"); } catch (IllegalArgumentException expected) { } + try { sink.readAt(-1, ByteBuffer.allocate(1)); fail("negative read pos"); } catch (IllegalArgumentException expected) { } + } + Files.deleteIfExists(f); + } + + @Test(expected = NullPointerException.class) + public void rejectsNullChannel() { + SeekableSink.over(null, 0); + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java index 3517c7c40..ce32b8d43 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java @@ -36,6 +36,9 @@ 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.util.work.ProgressLimiter; +import io.github.jbellis.jvector.util.work.WorkLimiter; +import io.github.jbellis.jvector.util.work.WorkStage; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorFloat; @@ -48,14 +51,21 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.*; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executor; import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.Semaphore; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; 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.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.fail; @ThreadLeakScope(ThreadLeakScope.Scope.NONE) public class TestOnDiskGraphIndexCompactor extends RandomizedTest { @@ -270,7 +280,7 @@ public void testExactVectorValuesAfterCompaction() throws Exception { List.of(g0, g1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path outPath = testDirectory.resolve("simple_compact_out"); compactor.compact(outPath); @@ -347,7 +357,7 @@ public void testExactVectorValuesWithDeletions() throws Exception { List.of(g0, g1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path outPath = testDirectory.resolve("del_compact_out"); compactor.compact(outPath); @@ -425,7 +435,7 @@ public void testExactVectorValuesWithCustomRemapping() throws Exception { List.of(g0, g1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path outPath = testDirectory.resolve("remap_compact_out"); compactor.compact(outPath); @@ -482,7 +492,7 @@ public void testCompact() throws Exception { liveNodes.add(lives); } - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null, -1); int topK = 10; // Select query vectors from the dataset @@ -536,6 +546,326 @@ public void testCompact() throws Exception { searcher.close(); } + // ---- ProgressLimiter SPI (io.github.jbellis.jvector.util.work) ---- + + /** A ProgressLimiter that records observations and hands out no-op grants (rate-limiter style). */ + private static final class RecordingLimiter implements ProgressLimiter { + final Set stages = ConcurrentHashMap.newKeySet(); + final Map lastCompleted = new ConcurrentHashMap<>(); + final Map lastTotal = new ConcurrentHashMap<>(); + final Map monotonic = new ConcurrentHashMap<>(); + final AtomicInteger acquires = new AtomicInteger(); + final AtomicInteger closes = new AtomicInteger(); + final AtomicLong bytesAcquired = new AtomicLong(); + + @Override + public void onProgress(WorkStage stage, long completed, long total) { + String s = stage.name(); + stages.add(s); + lastTotal.put(s, total); + Long prev = lastCompleted.put(s, completed); + if (prev != null && completed < prev) monotonic.put(s, false); + else monotonic.putIfAbsent(s, true); + } + + @Override + public WorkLimiter.Grant acquire(long amount) { + acquires.incrementAndGet(); + bytesAcquired.addAndGet(amount); + return () -> { closes.incrementAndGet(); }; + } + } + + /** A ProgressLimiter whose grants hold real semaphore permits, released on close. */ + private static final class SemaphoreBytesLimiter implements ProgressLimiter { + final Semaphore permits; + final int cap; + final AtomicInteger open = new AtomicInteger(); + + SemaphoreBytesLimiter(int totalPermits) { + this.permits = new Semaphore(totalPermits); + this.cap = totalPermits; + } + + @Override + public WorkLimiter.Grant acquire(long amount) throws InterruptedException { + final int n = (int) Math.max(1, Math.min(amount, cap)); // never exceed total -> no single-acquire deadlock + permits.acquire(n); + open.incrementAndGet(); + return () -> { permits.release(n); open.decrementAndGet(); }; + } + } + + /** Loads the fixture's source graphs with every node live and identity remapping. */ + private OnDiskGraphIndexCompactor newAllLiveCompactor(List rss) throws IOException { + return newAllLiveCompactor(rss, (Executor) null, -1); + } + + private OnDiskGraphIndexCompactor newAllLiveCompactor(List rss, Executor executor, int taskWindowSize) throws IOException { + List graphs = new ArrayList<>(); + List liveNodes = new ArrayList<>(); + List remappers = new ArrayList<>(); + for (int i = 0; i < numSources; ++i) { + var rs = ReaderSupplierFactory.open(testDirectory.resolve("test_graph_" + i).toAbsolutePath()); + rss.add(rs); + graphs.add(OnDiskGraphIndex.load(rs)); + } + 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); + } + return new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, executor, taskWindowSize); + } + + @Test + public void testTaskWindowSizeResolution() throws Exception { + List rss1 = new ArrayList<>(); + try { + assertEquals("explicit window must be honored", + 5, newAllLiveCompactor(rss1, (Executor) Runnable::run, 5).getTaskWindowSize()); + } finally { for (var rs : rss1) rs.close(); } + + List rss2 = new ArrayList<>(); + try { + assertEquals("non-FJP executor with window<=0 must default to 1 (serial)", + 1, newAllLiveCompactor(rss2, (Executor) Runnable::run, 0).getTaskWindowSize()); + } finally { for (var rs : rss2) rs.close(); } + + List rss3 = new ArrayList<>(); + ForkJoinPool pool = new ForkJoinPool(3); + try { + assertEquals("window<=0 with a ForkJoinPool must derive getParallelism()", + 3, newAllLiveCompactor(rss3, pool, -1).getTaskWindowSize()); + } finally { pool.shutdown(); for (var rs : rss3) rs.close(); } + } + + @Test + public void testCallerRunsProducesValidGraph() throws Exception { + // Caller-runs + serial must produce a correct, searchable graph of equivalent quality to + // the ForkJoinPool path. (Not asserted byte-identical: fused-PQ retraining uses parallel + // floating-point reduction, which is order-sensitive, so serialization can shift low bits.) + List rss = new ArrayList<>(); + try { + var path = testDirectory.resolve("cr_caller"); + newAllLiveCompactor(rss, (Executor) Runnable::run, 1).compact(path); + + try (ReaderSupplier out = ReaderSupplierFactory.open(path)) { + var g = OnDiskGraphIndex.load(out); + assertEquals(numSources * numVectorsPerGraph, g.size(0)); + + List> queries = new ArrayList<>(); + for (int i = 0; i < numQueries; ++i) queries.add(allVecs.get(randomIntBetween(0, allVecs.size() - 1))); + List> gt = buildGT(queries, 10); + GraphSearcher searcher = new GraphSearcher(g); + List results = new ArrayList<>(); + for (VectorFloat q : queries) { + SearchScoreProvider ssp = DefaultSearchScoreProvider.exact(q, similarityFunction, allravv); + results.add(searcher.search(ssp, 10, Bits.ALL)); + } + double recall = AccuracyMetrics.recallFromSearchResults(gt, results, 10, 10); + assertTrue("caller-runs graph recall should be reasonable, got " + recall, recall >= 0.2); + searcher.close(); + } + } finally { + for (var rs : rss) rs.close(); + } + } + + @Test + public void testCallerRunsRunsOnCallingThreadOnly() throws Exception { + List rss = new ArrayList<>(); + Set taskThreads = Collections.synchronizedSet(new HashSet<>()); + Executor recordingCallerRuns = command -> { + taskThreads.add(Thread.currentThread()); + command.run(); + }; + try { + newAllLiveCompactor(rss, recordingCallerRuns, 1).compact(testDirectory.resolve("cr_thread")); + assertEquals("all batch + pre-encode work must run on the calling thread only", + Collections.singleton(Thread.currentThread()), taskThreads); + } finally { + for (var rs : rss) rs.close(); + } + } + + @Test + public void testProgressLimiterObservesAndPairsGrants() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + var rec = new RecordingLimiter(); + compactor.setProgressLimiter(rec); + + var outputPath = testDirectory.resolve("test_compact_progress"); + compactor.compact(outputPath); + + // Both stages observed, progress monotonic and finishing at 100% of a known total. + assertTrue("MERGE_LEVELS progress not reported", rec.stages.contains("MERGE_LEVELS")); + assertTrue("REFINE progress not reported", rec.stages.contains("REFINE")); + assertTrue("MERGE_LEVELS progress not monotonic", rec.monotonic.getOrDefault("MERGE_LEVELS", true)); + assertTrue("REFINE progress not monotonic", rec.monotonic.getOrDefault("REFINE", true)); + assertEquals("MERGE_LEVELS did not reach total", + rec.lastTotal.get("MERGE_LEVELS"), rec.lastCompleted.get("MERGE_LEVELS")); + assertEquals("REFINE did not reach total", + rec.lastTotal.get("REFINE"), rec.lastCompleted.get("REFINE")); + assertEquals("MERGE_LEVELS total should be the live-node count", + (long) numSources * numVectorsPerGraph, (long) rec.lastTotal.get("MERGE_LEVELS")); + + // Throttle invoked with real byte amounts, and every acquire paired with exactly one close. + assertTrue("acquire never called", rec.acquires.get() > 0); + assertTrue("acquire amounts were all zero", rec.bytesAcquired.get() > 0); + assertEquals("every grant must be closed exactly once", rec.acquires.get(), rec.closes.get()); + + // Output is a valid, complete graph. + try (ReaderSupplier rs = ReaderSupplierFactory.open(outputPath)) { + var compactGraph = OnDiskGraphIndex.load(rs); + assertEquals(numSources * numVectorsPerGraph, compactGraph.size(0)); + } + for (var rs : rss) rs.close(); + } + + @Test + public void testSemaphoreLimiterReleasesAllGrantsWithoutDeadlock() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + var sem = new SemaphoreBytesLimiter(64 * 1024 * 1024); // ample: exercises release-on-close, not undersized + compactor.setProgressLimiter(sem); + + var outputPath = testDirectory.resolve("test_compact_semaphore"); + final Throwable[] thrown = new Throwable[1]; + Thread t = new Thread(() -> { + try { compactor.compact(outputPath); } + catch (Throwable e) { thrown[0] = e; } + }, "compact-semaphore"); + t.start(); + t.join(90_000); + + assertFalse("compaction did not finish within 90s — possible throttle deadlock", t.isAlive()); + if (thrown[0] != null) throw new AssertionError("compaction failed under semaphore limiter", thrown[0]); + assertEquals("all semaphore grants must be released (acquire/close paired)", 0, sem.open.get()); + + try (ReaderSupplier rs = ReaderSupplierFactory.open(outputPath)) { + var compactGraph = OnDiskGraphIndex.load(rs); + assertEquals(numSources * numVectorsPerGraph, compactGraph.size(0)); + } + for (var rs : rss) rs.close(); + } + + @Test + public void testCompactAtNonZeroStartOffset() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + + // Reserve a prefix (as an embedder would for its container header) and record its bytes. + var outputPath = testDirectory.resolve("test_compact_offset"); + byte[] prefix = new byte[64]; + for (int i = 0; i < prefix.length; i++) prefix[i] = (byte) (0xA0 + (i % 16)); + Files.write(outputPath, prefix); + long startOffset = prefix.length; + + compactor.compact(outputPath, startOffset); + + // The reserved prefix must be untouched by the no-copy write. + byte[] afterPrefix = Arrays.copyOf(Files.readAllBytes(outputPath), prefix.length); + assertArrayEquals("compaction clobbered the reserved prefix", prefix, afterPrefix); + + // The graph loads from startOffset and is complete + searchable (proves refine wrote valid + // records at the base-shifted offsets). + try (ReaderSupplier rs = ReaderSupplierFactory.open(outputPath)) { + var compactGraph = OnDiskGraphIndex.load(rs, startOffset); + assertEquals(numSources * numVectorsPerGraph, compactGraph.size(0)); + + GraphSearcher searcher = new GraphSearcher(compactGraph); + for (int i = 0; i < 5; i++) { + VectorFloat q = allVecs.get(randomIntBetween(0, allVecs.size() - 1)); + SearchScoreProvider ssp = DefaultSearchScoreProvider.exact(q, similarityFunction, allravv); + SearchResult r = searcher.search(ssp, 10, Bits.ALL); + assertTrue("search from offset-loaded graph returned nothing", r.getNodes().length > 0); + } + searcher.close(); + } + for (var rs : rss) rs.close(); + } + + @Test + public void testCompactToFileDestination() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + var path = testDirectory.resolve("dest_tofile"); + + long bodyLength = compactor.compact(CompactionDestination.toFile(path)); + assertEquals("returned body length must equal the standalone file size", + Files.size(path), bodyLength); + try (ReaderSupplier out = ReaderSupplierFactory.open(path)) { + assertEquals(numSources * numVectorsPerGraph, OnDiskGraphIndex.load(out).size(0)); + } + for (var rs : rss) rs.close(); + } + + @Test + public void testCompactToContainerDestinationCommitsAndPreservesPrefix() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + var container = testDirectory.resolve("dest_container"); + byte[] prefix = new byte[32]; + for (int i = 0; i < prefix.length; i++) prefix[i] = (byte) (0xC0 + (i % 16)); + Files.write(container, prefix); + final long base = prefix.length; + + final long[] committed = { -1 }; + final boolean[] closed = { false }; + CompactionDestination dest = () -> new CompactionDestination.Target() { + public Path file() { return container; } + public long startOffset() { return base; } + public void commit(long bodyLength) { committed[0] = bodyLength; } + public void close() { closed[0] = true; } + }; + + long returned = compactor.compact(dest); + assertTrue("commit must have been called", committed[0] >= 0); + assertEquals("commit bodyLength must equal the returned value", committed[0], returned); + assertTrue("target must be closed", closed[0]); + assertEquals("body length must be file size minus base", Files.size(container) - base, returned); + + byte[] afterPrefix = Arrays.copyOf(Files.readAllBytes(container), prefix.length); + assertArrayEquals("the reserved prefix must be preserved", prefix, afterPrefix); + try (ReaderSupplier out = ReaderSupplierFactory.open(container)) { + assertEquals(numSources * numVectorsPerGraph, OnDiskGraphIndex.load(out, base).size(0)); + } + for (var rs : rss) rs.close(); + } + + @Test + public void testCompactToDestinationAbortsWithoutCommitOnFailure() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + // Parent directory does not exist, so the graph writer fails partway. + final var badFile = testDirectory.resolve("no_such_dir").resolve("graph"); + + final boolean[] committed = { false }; + final boolean[] closed = { false }; + CompactionDestination dest = () -> new CompactionDestination.Target() { + public Path file() { return badFile; } + public long startOffset() { return 0; } + public void commit(long bodyLength) { committed[0] = true; } + public void close() { closed[0] = true; } + }; + + try { + compactor.compact(dest); + fail("expected compaction to fail for an unwritable destination"); + } catch (Exception expected) { + // ok + } + assertFalse("commit must NOT be called on failure", committed[0]); + assertTrue("close (abort) must run on failure", closed[0]); + for (var rs : rss) rs.close(); + } + /** * Tests compaction with deleted nodes. * Verifies that deleted nodes are properly excluded from the compacted graph. @@ -579,7 +909,7 @@ public void testCompactWithDeletions() throws Exception { liveNodes.add(lives); } - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null, -1); var outputPath = testDirectory.resolve("test_compact_with_deletions"); compactor.compact(outputPath); @@ -646,7 +976,7 @@ public void testOrdinalMapping() throws Exception { liveNodes.add(lives); } - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null, -1); var outputPath = testDirectory.resolve("test_compact_with_ordinal_mapping"); compactor.compact(outputPath); @@ -724,7 +1054,7 @@ public void testDeletionsAndOrdinalMapping() throws Exception { liveNodes.add(lives); } - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null, -1); var outputPath = testDirectory.resolve("test_compact_deletions_and_mapping"); compactor.compact(outputPath); @@ -830,7 +1160,7 @@ public void testCompactWithCompressedSidecar() throws Exception { List.of(pqv0, pqv1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path graphOut = testDirectory.resolve("sidecar_graph_out"); Path compressedOut = testDirectory.resolve("sidecar_pq_out"); @@ -913,7 +1243,7 @@ public void testCompactCompressedSidecarRejectsFusedPQ() throws Exception { List.of(pqv0, pqv1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - similarityFunction, null); + similarityFunction, null, -1); org.junit.Assert.fail("expected IllegalArgumentException for FUSED_PQ + sourceCompressed"); } catch (IllegalArgumentException expected) { assertTrue("error message mentions FUSED_PQ", @@ -956,7 +1286,7 @@ public void testCompactCompressedSidecarRejectsSizeMismatch() throws Exception { List.of(pqv0), // size 1 vs sources size 2 List.of(live, live), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); org.junit.Assert.fail("expected IllegalArgumentException for size mismatch"); } catch (IllegalArgumentException expected) { assertTrue("error message mentions size", @@ -992,7 +1322,7 @@ public void testCompactTwoArgRequiresSourceCompressed() throws Exception { List.of(g0, g1), List.of(live, live), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path graphOut = testDirectory.resolve("noarg_graph_out"); Path compressedOut = testDirectory.resolve("noarg_pq_out"); @@ -1052,7 +1382,7 @@ public void testCompactCompressedSidecarWithDeletions() throws Exception { List.of(pqv0, pqv1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path graphOut = testDirectory.resolve("delsidecar_graph_out"); Path compressedOut = testDirectory.resolve("delsidecar_pq_out"); diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/util/work/TestProgressLimiter.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/util/work/TestProgressLimiter.java new file mode 100644 index 000000000..16d02f31c --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/util/work/TestProgressLimiter.java @@ -0,0 +1,289 @@ +/* + * 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.util.work; + +import io.github.jbellis.jvector.util.work.WorkLimiter.Grant; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +public class TestProgressLimiter { + + private static final WorkStage STAGE = () -> "TEST"; + + private static long millisFor(ThrowingRunnable r) throws Exception { + long t0 = System.nanoTime(); + r.run(); + return (System.nanoTime() - t0) / 1_000_000L; + } + + private interface ThrowingRunnable { void run() throws Exception; } + + // ---- rateLimited (leaky bucket) ---- + + @Test + public void rateLimitedRejectsNonPositiveOrNonFiniteRate() { + for (double bad : new double[]{0.0, -1.0, -0.0, Double.NaN, Double.POSITIVE_INFINITY}) { + try { + ProgressLimiter.rateLimited(bad); + fail("expected IllegalArgumentException for rate " + bad); + } catch (IllegalArgumentException expected) { + // ok + } + } + } + + @Test + public void rateLimitedAdmitsZeroOrNegativeAmountImmediately() throws Exception { + ProgressLimiter limiter = ProgressLimiter.rateLimited(1.0); // 1 unit/sec: any real wait would be seconds + long ms = millisFor(() -> { + try (Grant g = limiter.acquire(0)) { assertNotNull(g); } + try (Grant g = limiter.acquire(-100)) { assertNotNull(g); } + }); + assertTrue("zero/negative amount must not block, waited " + ms + "ms", ms < 500); + } + + @Test + public void rateLimitedPacesSubsequentAcquire() throws Exception { + ProgressLimiter limiter = ProgressLimiter.rateLimited(1000.0); // 1 unit/ms + limiter.acquire(200).close(); // warmup: drained bucket admits the first request immediately + + long ms = millisFor(() -> limiter.acquire(200).close()); // must wait ~200ms behind the warmup reservation + assertTrue("expected pacing >= ~100ms at 1000 units/s after a 200-unit warmup, got " + ms + "ms", ms >= 100); + } + + @Test + public void rateLimitedFirstAcquireIsNotDelayed() throws Exception { + ProgressLimiter limiter = ProgressLimiter.rateLimited(10.0); // slow: a delayed first call would be seconds + long ms = millisFor(() -> limiter.acquire(1000).close()); + assertTrue("first acquire on a drained bucket must not block, waited " + ms + "ms", ms < 500); + } + + @Test + public void rateLimitedIsInterruptible() throws Exception { + ProgressLimiter limiter = ProgressLimiter.rateLimited(100.0); // 100 units/sec + limiter.acquire(100).close(); // warmup reserves ~1s of future emission time + + AtomicReference caught = new AtomicReference<>(); + AtomicInteger returnedNormally = new AtomicInteger(); + Thread t = new Thread(() -> { + try { + limiter.acquire(1).close(); // blocks ~1s behind the warmup reservation + returnedNormally.incrementAndGet(); + } catch (Throwable e) { + caught.set(e); + } + }, "rate-limited-blocked"); + t.start(); + Thread.sleep(150); // let it reach the interruptible sleep + t.interrupt(); + t.join(5_000); + + assertFalse("interrupted acquire should not hang", t.isAlive()); + assertEquals("acquire should not have returned normally", 0, returnedNormally.get()); + assertTrue("expected InterruptedException, got " + caught.get(), + caught.get() instanceof InterruptedException); + } + + @Test + public void rateLimitedGrantIsNoopAndProgressIsNoop() throws Exception { + ProgressLimiter limiter = ProgressLimiter.rateLimited(1_000_000.0); + Grant g = limiter.acquire(10); + assertNotNull(g); + g.close(); + g.close(); // idempotent no-op + limiter.onProgress(STAGE, 1, 2); // rate limiter does not track progress; must not throw + } + + // ---- logging wrapper ---- + + @Test(expected = NullPointerException.class) + public void loggingRejectsNullSinkWithDelegate() { + ProgressLimiter.logging(ProgressLimiter.UNLIMITED, null); + } + + @Test(expected = NullPointerException.class) + public void loggingRejectsNullSink() { + ProgressLimiter.logging((java.util.function.Consumer) null); + } + + @Test + public void loggingNullDelegateBehavesAsUnlimited() throws Exception { + List log = Collections.synchronizedList(new ArrayList<>()); + ProgressLimiter limiter = ProgressLimiter.logging(null, log::add); + long ms = millisFor(() -> limiter.acquire(Long.MAX_VALUE).close()); // UNLIMITED: instant + assertTrue("null delegate should not throttle, waited " + ms + "ms", ms < 500); + } + + @Test + public void loggingDelegatesBothFacets() { + RecordingLimiter delegate = new RecordingLimiter(); + List log = Collections.synchronizedList(new ArrayList<>()); + ProgressLimiter limiter = ProgressLimiter.logging(delegate, log::add); + + limiter.onProgress(STAGE, 3, 10); + assertEquals("onProgress must be delegated", 1, delegate.progressCalls.get()); + assertEquals(3, delegate.lastCompleted); + assertEquals(10, delegate.lastTotal); + assertTrue("onProgress should have been logged", + log.stream().anyMatch(s -> s.contains("TEST") && s.contains("3/10"))); + } + + @Test + public void loggingPreservesDelegateGrant() throws Exception { + RecordingLimiter delegate = new RecordingLimiter(); + ProgressLimiter limiter = ProgressLimiter.logging(delegate, s -> { }); + + Grant g = limiter.acquire(1234); + assertEquals("acquire must be delegated", 1, delegate.acquireCalls.get()); + assertEquals(1234, delegate.lastAmount); + assertEquals("grant must not be closed yet", 0, delegate.grantCloses.get()); + g.close(); + assertEquals("closing the wrapper grant must close the delegate's grant", 1, delegate.grantCloses.get()); + } + + @Test + public void loggingLogsAcquireOnlyWhenItBlocks() throws Exception { + List log = Collections.synchronizedList(new ArrayList<>()); + + // Instant delegate (UNLIMITED): no throttled line expected. + ProgressLimiter fast = ProgressLimiter.logging(ProgressLimiter.UNLIMITED, log::add); + fast.acquire(500).close(); + assertTrue("unblocked acquire should not log a throttle line", + log.stream().noneMatch(s -> s.contains("throttled"))); + + // Blocking delegate: a throttled line is expected. + log.clear(); + ProgressLimiter slow = ProgressLimiter.logging(new SleepingLimiter(60), log::add); + slow.acquire(500).close(); + assertTrue("blocked acquire should log a throttle line", + log.stream().anyMatch(s -> s.contains("throttled") && s.contains("500"))); + } + + // ---- composition ---- + + @Test + public void loggingComposesWithRateLimited() throws Exception { + List log = Collections.synchronizedList(new ArrayList<>()); + ProgressLimiter limiter = ProgressLimiter.logging(ProgressLimiter.rateLimited(1000.0), log::add); + + limiter.acquire(200).close(); // warmup + long ms = millisFor(() -> limiter.acquire(200).close()); + + assertTrue("composed limiter should still pace, got " + ms + "ms", ms >= 100); + assertTrue("composed limiter should log the throttled acquire", + log.stream().anyMatch(s -> s.contains("throttled"))); + limiter.onProgress(STAGE, 5, 5); + assertTrue("composed limiter should log progress", + log.stream().anyMatch(s -> s.contains("TEST") && s.contains("5/5"))); + } + + // ---- melded SPI defaults ---- + + @Test + public void unlimitedIsFullyNoop() throws Exception { + long ms = millisFor(() -> { + try (Grant g = ProgressLimiter.UNLIMITED.acquire(Long.MAX_VALUE)) { + assertNotNull(g); + } + }); + assertTrue("UNLIMITED.acquire must not block, waited " + ms + "ms", ms < 500); + ProgressLimiter.UNLIMITED.onProgress(STAGE, 7, -1); // no-op, must not throw + + // Facet no-op constants exist and are safe. + WorkLimiter.Grant.NOOP.close(); + ProgressTracker.NOOP.onProgress(STAGE, 1, 1); + try (Grant g = WorkLimiter.UNLIMITED.acquire(99)) { + assertNotNull(g); + } + } + + @Test + public void facetsAreIndependentlyOverridable() throws Exception { + // Tracker-only: overrides onProgress, inherits no-op acquire. + AtomicInteger progressSeen = new AtomicInteger(); + ProgressLimiter trackerOnly = new ProgressLimiter() { + @Override public void onProgress(WorkStage stage, long completed, long total) { + progressSeen.incrementAndGet(); + } + }; + try (Grant g = trackerOnly.acquire(1_000_000)) { // inherited no-op: must not block + assertNotNull(g); + } + trackerOnly.onProgress(STAGE, 1, 1); + assertEquals(1, progressSeen.get()); + + // Throttle-only: overrides acquire, inherits no-op onProgress. + AtomicInteger acquireSeen = new AtomicInteger(); + ProgressLimiter throttleOnly = new ProgressLimiter() { + @Override public Grant acquire(long amount) { + acquireSeen.incrementAndGet(); + return Grant.NOOP; + } + }; + throttleOnly.onProgress(STAGE, 1, 1); // inherited no-op: must not throw + throttleOnly.acquire(5).close(); + assertEquals(1, acquireSeen.get()); + } + + // ---- test doubles ---- + + /** Records both facets and hands out a grant whose close is counted. */ + private static final class RecordingLimiter implements ProgressLimiter { + final AtomicInteger progressCalls = new AtomicInteger(); + final AtomicInteger acquireCalls = new AtomicInteger(); + final AtomicInteger grantCloses = new AtomicInteger(); + volatile long lastCompleted, lastTotal, lastAmount; + + @Override + public void onProgress(WorkStage stage, long completed, long total) { + progressCalls.incrementAndGet(); + lastCompleted = completed; + lastTotal = total; + } + + @Override + public Grant acquire(long amount) { + acquireCalls.incrementAndGet(); + lastAmount = amount; + return grantCloses::incrementAndGet; + } + } + + /** A throttle that always blocks for a fixed number of milliseconds. */ + private static final class SleepingLimiter implements ProgressLimiter { + private final long sleepMillis; + + SleepingLimiter(long sleepMillis) { this.sleepMillis = sleepMillis; } + + @Override + public Grant acquire(long amount) throws InterruptedException { + Thread.sleep(sleepMillis); + return Grant.NOOP; + } + } +}