diff --git a/c/test_vindex.c b/c/test_vindex.c index 474aba7..5fbdb2a 100644 --- a/c/test_vindex.c +++ b/c/test_vindex.c @@ -536,6 +536,40 @@ static void test_supported_index_roundtrips(void) { 4); } +static void test_training_opens_multiple_writers_before_consumption(void) { + const char *keys[] = {"index.type", "dimension", "nlist", "metric"}; + const char *values[] = {"ivf_flat", "1", "1", "l2"}; + const float data[] = {0.0f, 1.0f}; + + PaimonVindexTrainerHandle *trainer = paimon_vindex_trainer_open(keys, values, 4); + ASSERT_TRUE(trainer != NULL); + ASSERT_EQ_I64( + paimon_vindex_trainer_add_training_vectors(trainer, data, 2), + 0); + PaimonVindexTrainingHandle *training = paimon_vindex_trainer_finish(trainer); + ASSERT_TRUE(training != NULL); + paimon_vindex_trainer_free(trainer); + + PaimonVindexWriterHandle *first = + paimon_vindex_writer_open_from_training(training); + PaimonVindexWriterHandle *second = + paimon_vindex_writer_open_from_training(training); + PaimonVindexWriterHandle *consuming = paimon_vindex_writer_open(training); + ASSERT_TRUE(first != NULL); + ASSERT_TRUE(second != NULL); + ASSERT_TRUE(consuming != NULL); + + uintptr_t dimension = 0; + ASSERT_EQ_I64(paimon_vindex_writer_dimension(first, &dimension), 0); + ASSERT_EQ_I64(dimension, 1); + + paimon_vindex_writer_free(first); + paimon_vindex_writer_free(second); + paimon_vindex_writer_free(consuming); + paimon_vindex_training_free(training); + printf("PASS training_opens_multiple_writers_before_consumption\n"); +} + static void test_extensible_search_params_defaults(void) { PaimonVindexSearchParamsEx params = paimon_vindex_search_params_ex_default(); @@ -554,6 +588,7 @@ static void test_extensible_search_params_defaults(void) { int main(void) { test_extensible_search_params_defaults(); test_supported_index_roundtrips(); + test_training_opens_multiple_writers_before_consumption(); test_output_write_callback_error_propagates(); test_output_flush_callback_error_propagates(); test_input_read_ranges_callback_error_propagates(); diff --git a/core/src/diskann.rs b/core/src/diskann.rs index 4b2cb95..cb61000 100644 --- a/core/src/diskann.rs +++ b/core/src/diskann.rs @@ -281,6 +281,27 @@ impl DiskAnnIndex { } } + /// Creates an empty index that reuses the trained product quantizer. + pub(crate) fn from_trained(trained: &DiskAnnIndex) -> Self { + Self { + d: trained.d, + metric: trained.metric, + pq: ProductQuantizer { + d: trained.pq.d, + m: trained.pq.m, + nbits: trained.pq.nbits, + dsub: trained.pq.dsub, + ksub: trained.pq.ksub, + chunk_offsets: trained.pq.chunk_offsets.clone(), + centroids: trained.pq.centroids.clone(), + centroid_norms_cache: trained.pq.centroid_norms_cache.clone(), + }, + build_params: trained.build_params, + ids: Vec::new(), + vectors: Vec::new(), + } + } + pub fn train(&mut self, data: &[f32], n: usize) -> io::Result<()> { if n == 0 { return Err(invalid_input( diff --git a/core/src/index.rs b/core/src/index.rs index 198eaed..dad8f52 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -1294,6 +1294,31 @@ impl VectorIndexTraining { pub fn dimension(&self) -> usize { self.inner.dimension() } + + /// Creates an empty writer that reuses this training result. + /// + /// The training result remains available, so callers can create independent writers for + /// multiple segments without repeating the training step. Use [`VectorIndexWriter::new`] when + /// the training result is only needed once and may be consumed. + pub fn create_writer(&self) -> VectorIndexWriter { + match &self.inner { + VectorIndexWriter::IvfFlat(index) => { + VectorIndexWriter::IvfFlat(IVFFlatIndex::from_trained(index)) + } + VectorIndexWriter::IvfSq(index) => { + VectorIndexWriter::IvfSq(IVFSQIndex::from_trained(index)) + } + VectorIndexWriter::IvfPq(index) => { + VectorIndexWriter::IvfPq(IVFPQIndex::from_trained(index)) + } + VectorIndexWriter::IvfRq(index) => { + VectorIndexWriter::IvfRq(IVFRQIndex::from_trained(index)) + } + VectorIndexWriter::DiskAnn(index) => { + VectorIndexWriter::DiskAnn(DiskAnnIndex::from_trained(index)) + } + } + } } pub enum VectorIndexWriter { @@ -2632,6 +2657,60 @@ mod tests { } } + fn assert_reusable_training_creates_independent_writers(config: VectorIndexConfig) { + let dimension = config.dimension(); + let nlist = config.nlist(); + let training_count = 256; + let segment_count = 48; + let data = generate_clustered_data(training_count, dimension, nlist); + let training = VectorIndexTrainer::train(config, &data, training_count).unwrap(); + + let mut first = training.create_writer(); + let first_ids = (1_000..1_000 + segment_count as i64).collect::>(); + first + .add_vectors( + &first_ids, + &data[..segment_count * dimension], + segment_count, + ) + .unwrap(); + + let mut second = training.create_writer(); + let second_ids = (2_000..2_000 + segment_count as i64).collect::>(); + second + .add_vectors( + &second_ids, + &data[segment_count * dimension..2 * segment_count * dimension], + segment_count, + ) + .unwrap(); + + for (mut writer, expected_ids, query) in [ + ( + first, + 1_000..1_000 + segment_count as i64, + &data[..dimension], + ), + ( + second, + 2_000..2_000 + segment_count as i64, + &data[segment_count * dimension..(segment_count + 1) * dimension], + ), + ] { + let mut bytes = Vec::new(); + writer.write(&mut PosWriter::new(&mut bytes)).unwrap(); + let mut reader = VectorIndexReader::open(Cursor::new(bytes)).unwrap(); + assert_eq!(reader.metadata().total_vectors, segment_count as i64); + let params = if reader.metadata().index_type == IndexType::DiskAnn { + VectorSearchParams::with_l_search(5, 32) + } else { + VectorSearchParams::new(5, nlist) + }; + let (ids, _) = reader.search(query, params).unwrap(); + assert!(ids.into_iter().all(|id| expected_ids.contains(&id))); + } + } + fn build_reader(config: VectorIndexConfig) -> (VectorIndexReader>>, Vec) { let d = config.dimension(); let nlist = config.nlist(); @@ -3037,6 +3116,43 @@ mod tests { ); } + #[test] + fn reusable_training_creates_independent_writers_for_all_index_types() { + assert_reusable_training_creates_independent_writers(VectorIndexConfig::IvfFlat { + dimension: 8, + nlist: 4, + metric: MetricType::L2, + use_approximate_coarse_assignment: true, + }); + assert_reusable_training_creates_independent_writers( + VectorIndexConfig::ivf_pq(16, 4, MetricType::L2, false).unwrap(), + ); + assert_reusable_training_creates_independent_writers(VectorIndexConfig::IvfRq { + dimension: 8, + nlist: 4, + bits: DEFAULT_RQ_BITS, + metric: MetricType::L2, + use_approximate_coarse_assignment: true, + }); + assert_reusable_training_creates_independent_writers(VectorIndexConfig::IvfSq { + dimension: 8, + nlist: 4, + metric: MetricType::L2, + use_approximate_coarse_assignment: true, + }); + assert_reusable_training_creates_independent_writers(VectorIndexConfig::DiskAnn { + dimension: 8, + metric: MetricType::L2, + pq_m: 4, + pq_bits: 4, + build: DiskAnnBuildParams { + max_degree: 8, + build_search_list_size: 16, + ..DiskAnnBuildParams::default() + }, + }); + } + #[test] fn unified_config_rejects_invalid_pq_m() { let err = match VectorIndexTrainer::new(VectorIndexConfig::IvfPq { diff --git a/core/src/ivfflat.rs b/core/src/ivfflat.rs index ab1d622..950e573 100644 --- a/core/src/ivfflat.rs +++ b/core/src/ivfflat.rs @@ -44,6 +44,21 @@ impl IVFFlatIndex { } } + /// Creates an empty index that reuses the trained coarse quantizer. + pub(crate) fn from_trained(trained: &IVFFlatIndex) -> Self { + let mut index = Self { + d: trained.d, + nlist: trained.nlist, + metric: trained.metric, + quantizer_centroids: trained.quantizer_centroids.clone(), + ids: vec![Vec::new(); trained.nlist], + vectors: vec![Vec::new(); trained.nlist], + coarse_assignment: CoarseAssignment::default(), + }; + index.set_approximate_coarse_assignment(trained.coarse_assignment.approximate_enabled()); + index + } + pub fn quantizer_centroids(&self) -> &[f32] { &self.quantizer_centroids } diff --git a/core/src/ivfrq.rs b/core/src/ivfrq.rs index 86596ae..f834dcd 100644 --- a/core/src/ivfrq.rs +++ b/core/src/ivfrq.rs @@ -121,6 +121,30 @@ impl IVFRQIndex { } } + /// Creates an empty index that reuses the trained coarse quantizer and rotation. + pub(crate) fn from_trained(trained: &IVFRQIndex) -> Self { + let mut index = Self { + d: trained.d, + padded_d: trained.padded_d, + nlist: trained.nlist, + bits: trained.bits, + metric: trained.metric, + quantizer_centroids: trained.quantizer_centroids.clone(), + quantizer_centroid_norms: trained.quantizer_centroid_norms.clone(), + rotated_centroids: trained.rotated_centroids.clone(), + rotation_seed: trained.rotation_seed, + rotation_rounds: trained.rotation_rounds, + ids: vec![Vec::new(); trained.nlist], + codes: vec![Vec::new(); trained.nlist], + factors: vec![Vec::new(); trained.nlist], + quantizer: trained.quantizer.clone(), + rotation: trained.rotation.clone(), + coarse_assignment: CoarseAssignment::default(), + }; + index.set_approximate_coarse_assignment(trained.coarse_assignment.approximate_enabled()); + index + } + pub fn quantizer_centroids(&self) -> &[f32] { &self.quantizer_centroids } diff --git a/core/src/ivfsq.rs b/core/src/ivfsq.rs index c39a62e..87959f4 100644 --- a/core/src/ivfsq.rs +++ b/core/src/ivfsq.rs @@ -53,6 +53,23 @@ impl IVFSQIndex { } } + /// Creates an empty index that reuses the trained coarse and scalar quantizers. + pub(crate) fn from_trained(trained: &IVFSQIndex) -> Self { + let mut index = Self { + d: trained.d, + nlist: trained.nlist, + metric: trained.metric, + quantizer_centroids: trained.quantizer_centroids.clone(), + sq: trained.sq.clone(), + list_sqs: trained.list_sqs.clone(), + ids: vec![Vec::new(); trained.nlist], + codes: vec![Vec::new(); trained.nlist], + coarse_assignment: CoarseAssignment::default(), + }; + index.set_approximate_coarse_assignment(trained.coarse_assignment.approximate_enabled()); + index + } + pub fn quantizer_centroids(&self) -> &[f32] { &self.quantizer_centroids } diff --git a/cpp/test_vindex.cpp b/cpp/test_vindex.cpp index af7104a..fec0acc 100644 --- a/cpp/test_vindex.cpp +++ b/cpp/test_vindex.cpp @@ -302,6 +302,23 @@ static void test_supported_index_roundtrips() { 4); } +static void test_training_opens_multiple_writers_before_consumption() { + const float data[] = {0.0f, 1.0f}; + paimon::vindex::Trainer trainer({ + {"index.type", "ivf_flat"}, + {"dimension", "1"}, + {"nlist", "1"}, + {"metric", "l2"}, + }); + auto training = trainer.add_training_vectors(data, 2).finish_training(); + + paimon::vindex::Writer first(training); + paimon::vindex::Writer second(training); + paimon::vindex::Writer consuming(std::move(training)); + ASSERT_EQ(first.dimension(), 1); + printf("PASS training_opens_multiple_writers_before_consumption\n"); +} + static void test_worker_callback_reentry_is_rejected() { int callback_context = 0; paimon::vindex::detail::NativeHandleMutex mutex; @@ -341,6 +358,7 @@ static void test_extensible_search_params_forward_query_tuning() { int main() { test_supported_index_roundtrips(); + test_training_opens_multiple_writers_before_consumption(); test_worker_callback_reentry_is_rejected(); test_extensible_search_params_forward_query_tuning(); return 0; diff --git a/docs/api.html b/docs/api.html index 0be8fbb..2d31a5a 100644 --- a/docs/api.html +++ b/docs/api.html @@ -47,8 +47,16 @@

Public integration layers

Shared lifecycle

-
01Create a Trainer
Parse and validate options
02Submit one or more
training batches
03Finish training and
create a one-shot Writer
04Add row IDs / vectors
and write the file
05Detect file magic
and execute searches
-
  • Vectors are contiguous f32 values; length must equal vector_count × dimension.
  • Training data may arrive in batches. Every IVF trainer keeps a deterministic reservoir of at most max(65,536, 64 × resolved nlist) vectors. DiskANN starts from a 50,000-row cap and lowers it when necessary so the retained sample, optional cosine-normalized copy, codebook, and parallel PQ-training scratch fit diskann.memory-budget-bytes. Sampling is independent of batch boundaries.
  • The Python and Java one-shot train helpers infer dimension from the matrix and use its row count for automatic nlist. When the matrix is only a sample, pass the final corpus size as expected-vector-count. Streaming Trainer APIs require a concrete dimension before their first batch.
  • A Writer may receive production vectors in multiple batches. Row-ID count must equal vector count.
  • Readers expose metadata, single-query search, batch search, and Roaring64-filtered variants.
  • Files carry their type and resolved model sections. Callers do not pass index options again when opening a Reader.
+
01Create a Trainer
Parse and validate options
02Submit one or more
training batches
03Finish training and
create one or more Writers
04Add row IDs / vectors
and write each file
05Detect file magic
and execute searches
+
  • Vectors are contiguous f32 values; length must equal vector_count × dimension.
  • Training data may arrive in batches. Every IVF trainer keeps a deterministic reservoir of at most max(65,536, 64 × resolved nlist) vectors. DiskANN starts from a 50,000-row cap and lowers it when necessary so the retained sample, optional cosine-normalized copy, codebook, and parallel PQ-training scratch fit diskann.memory-budget-bytes. Sampling is independent of batch boundaries.
  • The Python and Java one-shot train helpers infer dimension from the matrix and use its row count for automatic nlist. When the matrix is only a sample, pass the final corpus size as expected-vector-count. Streaming Trainer APIs require a concrete dimension before their first batch.
  • A completed training result can create multiple independent Writers. Each Writer starts with the same trained centroids, quantizers, and rotations but empty vector payloads, so separate segments can share one training pass.
  • A Writer may receive production vectors in multiple batches. Row-ID count must equal vector count.
  • Readers expose metadata, single-query search, batch search, and Roaring64-filtered variants.
  • Files carry their type and resolved model sections. Callers do not pass index options again when opening a Reader.
+
+ + + + + +
LanguageReusable trainingExisting consuming form
Rusttraining.create_writer()VectorIndexWriter::new(training)
Cpaimon_vindex_writer_open_from_training(training)paimon_vindex_writer_open(training)
C++Writer writer(training)Writer writer(std::move(training))
Javatraining.createWriter()new VectorIndexWriter(training)
Pythontraining.create_writer()VectorIndexWriter(training)
+

Reusable Writer creation is opt-in; the existing consuming forms and their behavior remain unchanged. Do not consume or close a training result until every reusable Writer has been created.

IVF coarse assignment is approximate by default for large centroid matricesWhen dimension × nlist ≥ 1,000,000, ivf.coarse-assignment=auto uses a Vamana graph while training and adding vectors. Search still selects lists by exact centroid distance, so graph assignment can lower recall at small nprobe and does not guarantee that a vector is found by a self-query with nprobe=1. Set ivf.coarse-assignment=exact to disable the graph, preserve exact nearest-centroid assignment, and avoid graph startup cost for small non-empty batches. Empty batches never build the graph.
diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs index 90b9916..f75d3b3 100644 --- a/ffi/src/lib.rs +++ b/ffi/src/lib.rs @@ -463,6 +463,16 @@ unsafe fn training_mut<'a>( } } +unsafe fn training_ref<'a>( + handle: *const PaimonVindexTrainingHandle, +) -> Result<&'a PaimonVindexTrainingHandle, String> { + if handle.is_null() { + Err("null training handle".to_string()) + } else { + Ok(unsafe { &*handle }) + } +} + unsafe fn reader_mut<'a>( handle: *mut PaimonVindexReaderHandle, ) -> Result<&'a mut PaimonVindexReaderHandle, String> { @@ -828,6 +838,26 @@ pub unsafe extern "C" fn paimon_vindex_writer_open( }) } +/// Opens an independent writer without consuming `training`. +/// +/// The same training handle may be used to open additional writers until it is consumed by +/// `paimon_vindex_writer_open` or freed by `paimon_vindex_training_free`. +#[no_mangle] +pub unsafe extern "C" fn paimon_vindex_writer_open_from_training( + training: *const PaimonVindexTrainingHandle, +) -> *mut PaimonVindexWriterHandle { + ffi_ptr(|| { + let training = unsafe { training_ref(training) }?; + let training = training + .inner + .as_ref() + .ok_or_else(|| "training has already been consumed".to_string())?; + Ok(Box::into_raw(Box::new(PaimonVindexWriterHandle { + inner: training.create_writer(), + }))) + }) +} + #[no_mangle] pub unsafe extern "C" fn paimon_vindex_writer_free(handle: *mut PaimonVindexWriterHandle) { if !handle.is_null() { diff --git a/include/paimon_vindex.hpp b/include/paimon_vindex.hpp index f3638bb..aa83b4e 100644 --- a/include/paimon_vindex.hpp +++ b/include/paimon_vindex.hpp @@ -412,6 +412,15 @@ class Trainer { class Writer { public: + explicit Writer(const Training& training) { + if (!training.handle_) throw Error("training has already been consumed"); + handle_ = paimon_vindex_writer_open_from_training(training.handle_); + if (!handle_) { + const char* err = paimon_vindex_last_error(); + throw Error(err ? err : "failed to open vector index writer from training"); + } + } + explicit Writer(Training&& training) { if (!training.handle_) throw Error("training has already been consumed"); PaimonVindexTrainingHandle* training_handle = training.handle_; diff --git a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java index 7d48be3..cd8b879 100644 --- a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java +++ b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexNative.java @@ -37,6 +37,8 @@ private VectorIndexNative() {} static native long createWriter(long trainingPtr); + static native long createWriterFromTraining(long trainingPtr); + static native int writerDimension(long ptr); static native void addVectors(long ptr, long[] ids, float[] data, int n); diff --git a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexTraining.java b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexTraining.java index 65780e7..4a65f33 100644 --- a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexTraining.java +++ b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexTraining.java @@ -44,6 +44,18 @@ long takeNativePointer() { } } + /** Creates an independent writer while keeping this training result reusable. */ + public VectorIndexWriter createWriter() { + synchronized (nativeHandleLock) { + enterNativeHandle(); + try { + return VectorIndexWriter.fromTrainingPointer(requireOpen()); + } finally { + exitNativeHandle(); + } + } + } + @Override public void close() { synchronized (nativeHandleLock) { diff --git a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexWriter.java b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexWriter.java index 1f1863b..e38f04b 100644 --- a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexWriter.java +++ b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexWriter.java @@ -23,6 +23,7 @@ public final class VectorIndexWriter implements AutoCloseable { private long nativePtr; private Thread nativeHandleOwner; + /** Creates a writer by consuming the training result. */ public VectorIndexWriter(VectorIndexTraining training) { if (training == null) { throw new NullPointerException("training"); @@ -38,6 +39,10 @@ static VectorIndexWriter fromNativePointerForTesting(long nativePtr) { return new VectorIndexWriter(nativePtr); } + static VectorIndexWriter fromTrainingPointer(long trainingPtr) { + return new VectorIndexWriter(VectorIndexNative.createWriterFromTraining(trainingPtr)); + } + public int dimension() { synchronized (nativeHandleLock) { enterNativeHandle(); diff --git a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java index 761762e..40c9c10 100644 --- a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java +++ b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexJavaApiTest.java @@ -342,6 +342,12 @@ public void run() { new VectorIndexWriter(training); } }); + assertThrows(IllegalStateException.class, new ThrowingRunnable() { + @Override + public void run() { + training.createWriter(); + } + }); } private static void testReaderAndWriterApiCompile() { diff --git a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java index a7fd445..8045aae 100644 --- a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java +++ b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeValidationTest.java @@ -35,6 +35,7 @@ public static void main(String[] args) { testWriterRejectsNonFiniteValues(); testStagedTrainingRoundtrip(); testStagedTrainingStateValidation(); + testReusableTrainingCreatesIndependentWriters(); testReaderValidationComesFromCore(); testReaderRejectsNonFiniteQueries(); testReaderCapabilityFailuresArePropagatedBeforeOpen(); @@ -44,6 +45,48 @@ public static void main(String[] args) { testDiskAnnInnerProductAndCosine(); } + private static void testReusableTrainingCreatesIndependentWriters() { + VectorIndexTraining training = + VectorIndexTrainer.train(ivfFlatOptions(), new float[] {0.0f, 1.0f}, 2); + byte[] first; + byte[] second; + try { + first = writeReusableSegment(training, new long[] {10L}, new float[] {0.0f}); + second = writeReusableSegment(training, new long[] {20L}, new float[] {1.0f}); + } finally { + training.close(); + } + + assertSingleRowIndex(first, 10L, 0.0f); + assertSingleRowIndex(second, 20L, 1.0f); + } + + private static byte[] writeReusableSegment( + VectorIndexTraining training, long[] ids, float[] data) { + VectorIndexWriter writer = training.createWriter(); + ByteArrayPositionOutputStream output = new ByteArrayPositionOutputStream(); + try { + writer.addVectors(ids, data, ids.length); + writer.writeIndex(output); + return output.toByteArray(); + } finally { + writer.close(); + } + } + + private static void assertSingleRowIndex(byte[] bytes, long expectedId, float query) { + VectorIndexReader reader = + new VectorIndexReader(new ByteArraySeekableInputStream(bytes)); + try { + assertEquals(1L, reader.totalVectors()); + VectorSearchResult result = + reader.search(new float[] {query}, new VectorSearchParams(1, 1)); + assertEquals(expectedId, result.ids()[0]); + } finally { + reader.close(); + } + } + private static void testReaderCapabilityFailuresArePropagatedBeforeOpen() { final VectorIndexInput throwingInput = new VectorIndexInput() { diff --git a/jni/src/lib.rs b/jni/src/lib.rs index e025cbd..0e1c640 100644 --- a/jni/src/lib.rs +++ b/jni/src/lib.rs @@ -105,6 +105,12 @@ impl JniVectorIndexTraining { .take() .ok_or_else(|| "training has already been consumed".to_string()) } + + fn training(&self) -> Result<&VectorIndexTraining, String> { + self.training + .as_ref() + .ok_or_else(|| "training has already been consumed".to_string()) + } } struct JniVectorIndexWriter { @@ -618,6 +624,27 @@ pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_cre }) } +#[no_mangle] +pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_createWriterFromTraining( + env: JNIEnv, + _class: JClass, + training_ptr: jlong, +) -> jlong { + jni_call(env, |env| { + if training_ptr == 0 { + return throw_and_return(env, "null native pointer (training already freed?)"); + } + let training_handle = unsafe { &*(training_ptr as *const JniVectorIndexTraining) }; + let training = match training_handle.training() { + Ok(training) => training, + Err(e) => return throw_and_return(env, &e), + }; + Box::into_raw(Box::new(JniVectorIndexWriter::new( + training.create_writer(), + ))) as jlong + }) +} + #[no_mangle] pub extern "system" fn Java_org_apache_paimon_index_vector_VectorIndexNative_freeTraining( env: JNIEnv, diff --git a/python/paimon_vindex/__init__.py b/python/paimon_vindex/__init__.py index bba1ef0..989760c 100644 --- a/python/paimon_vindex/__init__.py +++ b/python/paimon_vindex/__init__.py @@ -363,6 +363,15 @@ def _take_handle(self): self._closed = True return handle + def create_writer(self): + """Create an independent writer while keeping this training reusable.""" + with self._native_handle_lock: + self._require_open() + handle = lib.paimon_vindex_writer_open_from_training(self._handle) + if not handle: + _check_error("failed to open writer from training") + return VectorIndexWriter._from_handle(handle) + def close(self): with self._native_handle_lock: if self._handle: @@ -512,6 +521,22 @@ def __init__(self, training: VectorIndexTraining): _check_error("failed to open writer") self._dimension = self._read_dimension() + @classmethod + def _from_handle(cls, handle): + writer = None + try: + writer = cls.__new__(cls) + writer._native_handle_lock = _NativeHandleLock() + writer._closed = False + writer._handle = handle + writer._dimension = writer._read_dimension() + return writer + except BaseException: + if writer is not None: + writer._handle = None + lib.paimon_vindex_writer_free(handle) + raise + def _require_open(self): if self._closed or not self._handle: raise RuntimeError("VectorIndexWriter is closed") diff --git a/python/paimon_vindex/_ffi.py b/python/paimon_vindex/_ffi.py index 02cd900..ea0daee 100644 --- a/python/paimon_vindex/_ffi.py +++ b/python/paimon_vindex/_ffi.py @@ -208,6 +208,9 @@ class PaimonVindexReadPlan(Structure): lib.paimon_vindex_writer_open.argtypes = [c_void_p] lib.paimon_vindex_writer_open.restype = c_void_p +lib.paimon_vindex_writer_open_from_training.argtypes = [c_void_p] +lib.paimon_vindex_writer_open_from_training.restype = c_void_p + lib.paimon_vindex_writer_free.argtypes = [c_void_p] lib.paimon_vindex_writer_free.restype = None diff --git a/python/tests/test_vindex.py b/python/tests/test_vindex.py index 950f714..0f47423 100644 --- a/python/tests/test_vindex.py +++ b/python/tests/test_vindex.py @@ -22,6 +22,7 @@ import numpy as np import pytest +import paimon_vindex from paimon_vindex import ( IvfPqBatchTableReuseMode, SearchParams, @@ -129,6 +130,65 @@ def test_python_high_level_training_infers_dimension_and_ivf_shape(): assert result_ids[0] == 0 +def test_python_training_creates_independent_writers(): + data = clustered_data(128, 8, 4) + options = { + "index.type": "ivf_flat", + "dimension": "8", + "nlist": "4", + "metric": "l2", + } + training = VectorIndexTrainer.train(options, data) + outputs = [] + try: + for offset, id_base in [(0, 1_000), (32, 2_000)]: + output = io.BytesIO() + ids = np.arange(id_base, id_base + 32, dtype=np.int64) + with training.create_writer() as writer: + writer.add_vectors(ids, data[offset : offset + 32]) + writer.write(output) + outputs.append((output.getvalue(), id_base, data[offset])) + finally: + training.close() + + for index_bytes, id_base, query in outputs: + with reader_from_bytes(index_bytes) as reader: + assert reader.metadata().total_vectors == 32 + ids, _ = reader.search(query, SearchParams.ivf(top_k=5, nprobe=4)) + assert np.all((ids >= id_base) & (ids < id_base + 32)) + + +@pytest.mark.parametrize("failure_point", ["before_handle", "after_handle"]) +def test_python_writer_from_handle_frees_handle_on_failure( + monkeypatch, failure_point +): + handle = 1234 + freed = [] + partial_writers = [] + + def fail_construction(*args): + partial_writers.extend(args) + raise MemoryError("injected construction failure") + + target = ( + (paimon_vindex, "_NativeHandleLock") + if failure_point == "before_handle" + else (VectorIndexWriter, "_read_dimension") + ) + monkeypatch.setattr(*target, fail_construction) + monkeypatch.setattr( + paimon_vindex.lib, "paimon_vindex_writer_free", freed.append + ) + + with pytest.raises(MemoryError, match="injected construction failure"): + VectorIndexWriter._from_handle(handle) + + for writer in partial_writers: + assert writer._handle is None + writer.close() + assert freed == [handle] + + def test_python_high_level_training_preserves_explicit_expected_count(): data = clustered_data(512, 16, 8) options = {