Skip to content

feat: allow training results to create multiple writers - #87

Open
wangzhigang1999 wants to merge 3 commits into
apache:mainfrom
wangzhigang1999:feature/reusable-training-multi-writer
Open

feat: allow training results to create multiple writers#87
wangzhigang1999 wants to merge 3 commits into
apache:mainfrom
wangzhigang1999:feature/reusable-training-multi-writer

Conversation

@wangzhigang1999

@wangzhigang1999 wangzhigang1999 commented Sep 2, 2026

Copy link
Copy Markdown

Summary

Allow one completed VectorIndexTraining to create multiple independent writers. Callers building several segment or file indexes from the same data distribution can train once and reuse the trained centroids, quantizers, and rotations instead of repeating training for every writer.

The existing consuming APIs remain unchanged for one-shot callers. Reusable writer creation is explicit and opt-in.

Closes #86.

Changes

  • Add VectorIndexTraining::create_writer() for IVF-Flat, IVF-SQ, IVF-PQ, IVF-RQ, and DiskANN.
  • Initialize each reusable writer with the trained model and empty segment-local IDs, vectors, codes, and graph state.
  • Expose equivalent non-consuming entry points through C, C++, Java/JNI, and Python while retaining the existing consuming forms.
  • Document both lifecycle choices without changing the primary one-shot examples.
  • Cover independent writer payloads with write/open/search round trips and keep binding tests focused on their ownership and ABI boundaries.

Benchmark

Release-mode synthetic benchmark on an Intel Xeon Platinum 8575C host with 32 logical CPUs and 60 GiB RAM. The workload used 128-dimensional vectors, 64 IVF lists, 8,192 training rows, eight segments, and 2,048 vectors per segment. retrain trains separately for every segment; reuse trains once and creates eight writers. The table reports the median of three alternating process runs.

Index Prepare: retrain -> reuse Speedup Full build: retrain -> reuse Speedup
IVF-Flat 112 -> 18 ms 6.2x 138 -> 33 ms 4.2x
IVF-SQ 123 -> 20 ms 6.2x 140 -> 27 ms 5.2x
IVF-PQ 1,355 -> 173 ms 7.8x 1,384 -> 194 ms 7.1x
IVF-RQ 109 -> 16 ms 6.8x 132 -> 28 ms 4.7x
DiskANN 1,240 -> 161 ms 7.7x 1,485 -> 407 ms 3.6x

The preparation measurement isolates training and writer creation. The full-build measurement also adds and serializes all segment vectors. Serialized output sizes were unchanged for each index type. This benchmark measures the removal of repeated training work; it does not claim an index-algorithm speedup.

Testing

  • cargo fmt --all -- --check
  • cargo test --workspace --locked (487 passed, 2 ignored)
  • cargo clippy --all-targets --workspace --locked -- -D warnings -A clippy::chunks-exact-to-as-chunks
  • C tests (10 passed)
  • C++ tests (8 passed)
  • Java Maven tests and JNI native validation
  • Python 3.12 tests (29 passed)

The Clippy allow is limited to the chunks_exact_to_as_chunks lint newly enabled by Rust 1.98; strict Clippy currently reports it at 28 unchanged locations on main. No call site introduced by this PR requires the allow.

Compatibility

  • No default behavior change; reuse requires an explicit new API call.
  • Existing consuming APIs remain source compatible.
  • No on-disk format change.
  • No cross-process training-model serialization or index merging.

Copilot AI lite review requested due to automatic review settings September 2, 2026 08:21

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

The new Python _from_handle path can leak a native writer handle on initialization failure (exception path) and should eagerly free the handle.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR adds an opt-in lifecycle for reusing a completed VectorIndexTraining to create multiple independent writers (across Rust core plus C/C++/Java/Python bindings), enabling segment-by-segment builds to share one training pass while preserving the existing consuming/one-shot APIs.

Changes:

  • Add VectorIndexTraining::create_writer() in core and per-index from_trained constructors to initialize empty writers with the trained model state.
  • Expose non-consuming “open writer from training” entry points through FFI (C), C++ RAII wrapper, Java/JNI, and Python ctypes wrapper, while retaining consuming constructors/APIs.
  • Add cross-language tests + documentation updates covering independent writer payloads and lifecycle guidance.
File summaries
File Description
python/tests/test_vindex.py Adds Python coverage for creating multiple writers from one training and validating isolation via round trips.
python/paimon_vindex/_ffi.py Declares the new ctypes symbol for opening a writer from a training handle.
python/paimon_vindex/init.py Adds VectorIndexTraining.create_writer() and a writer constructor path from an existing native handle.
jni/src/lib.rs Exposes createWriterFromTraining JNI entry point without consuming the training state.
java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java Adds native validation that reusable training produces independent writer outputs.
java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java Adds API test asserting createWriter() fails after the training has been consumed.
java/src/main/java/org/apache/paimon/index/vector/VectorIndexWriter.java Adds internal constructor path for creating writers from a training pointer (non-consuming).
java/src/main/java/org/apache/paimon/index/vector/VectorIndexTraining.java Adds createWriter() Java API for reusable writer creation.
java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java Declares the new native method createWriterFromTraining.
include/paimon_vindex.hpp Adds C++ Writer(const Training&) to open a writer without consuming training.
ffi/src/lib.rs Adds paimon_vindex_writer_open_from_training and supporting handle access for non-consuming writer creation.
docs/api.html Documents both lifecycle choices and lists reusable vs consuming APIs per language.
cpp/test_vindex.cpp Adds C++ test covering multiple writer opens before consuming training.
core/src/ivfsq.rs Adds IVFSQIndex::from_trained to reuse trained quantizers with empty payloads.
core/src/ivfrq.rs Adds IVFRQIndex::from_trained to reuse trained coarse quantizer/rotation with empty payloads.
core/src/ivfflat.rs Adds IVFFlatIndex::from_trained to reuse trained centroids with empty payloads.
core/src/index.rs Adds VectorIndexTraining::create_writer() and a core test ensuring writer independence across index types.
core/src/diskann.rs Adds DiskAnnIndex::from_trained to reuse the trained PQ state with empty payloads.
c/test_vindex.c Adds C test covering multiple writer opens before consuming training.
Review details
  • Files reviewed: 19/19 changed files
  • Comments generated: 1
  • Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread python/paimon_vindex/__init__.py Outdated
…ning-multi-writer

# Conflicts:
#	core/src/ivfflat.rs
#	core/src/ivfrq.rs
#	core/src/ivfsq.rs
#	docs/api.html
Comment thread python/paimon_vindex/__init__.py Outdated

@classmethod
def _from_handle(cls, handle):
writer = cls.__new__(cls)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Free the native handle if wrapper allocation fails

At this point create_writer() already owns a newly allocated native writer, but cls.__new__() and _NativeHandleLock() can raise before writer._handle is installed. Under memory pressure, no object owns the raw pointer, so __del__ cannot recover it and the cloned model is leaked. Please make this method an explicit ownership-transfer boundary: catch BaseException and call paimon_vindex_writer_free(handle) exactly once when construction fails. Fault-injection coverage should exercise failures both before and after _handle assignment. The existing _read_dimension() comment covers a later failure, but this earlier window is the unrecoverable one.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @JingsongLi, fixed in ddf6b4a. I added fault-injection tests before and after handle assignment, and all CI checks pass. Could you take another look when you have time?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support reusing a training result across multiple vector index writers

3 participants