Cooperative embedding

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);

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<String>`.
+
+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
 

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(...).
+    }
+}

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.
+ *
+ * <p>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 <b>not</b> 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);
+    }
+}

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<OrdinalMapper> 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<OrdinalMapper> remappers,
             int dimension,
             int maxOrdinal,
-            ForkJoinPool executor,
+            ExecutorService executor,
             int taskWindowSize) {
         this.sources = Collections.unmodifiableList(sources);
         this.sourceCompressed = sourceCompressed == null ? null : Collections.unmodifiableList(sourceCompressed);

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.
+ *
+ * <pre>{@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)
+ * }</pre>
+ *
+ * <p>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);
+    }
+}