feat: allow training results to create multiple writers - #87
feat: allow training results to create multiple writers#87wangzhigang1999 wants to merge 3 commits into
Conversation
There was a problem hiding this comment.
🟡 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-indexfrom_trainedconstructors 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.
…ning-multi-writer # Conflicts: # core/src/ivfflat.rs # core/src/ivfrq.rs # core/src/ivfsq.rs # docs/api.html
|
|
||
| @classmethod | ||
| def _from_handle(cls, handle): | ||
| writer = cls.__new__(cls) |
There was a problem hiding this comment.
[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.
There was a problem hiding this comment.
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?
Summary
Allow one completed
VectorIndexTrainingto 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
VectorIndexTraining::create_writer()for IVF-Flat, IVF-SQ, IVF-PQ, IVF-RQ, and DiskANN.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.
retraintrains separately for every segment;reusetrains once and creates eight writers. The table reports the median of three alternating process runs.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 -- --checkcargo test --workspace --locked(487 passed, 2 ignored)cargo clippy --all-targets --workspace --locked -- -D warnings -A clippy::chunks-exact-to-as-chunksThe Clippy allow is limited to the
chunks_exact_to_as_chunkslint newly enabled by Rust 1.98; strict Clippy currently reports it at 28 unchanged locations onmain. No call site introduced by this PR requires the allow.Compatibility