Support batched vector index training

.github/workflows/ci.yml

@@ -254,7 +254,7 @@ jobs:
           python - <<'PY'
           import io
           import numpy as np
-          from paimon_vindex import VectorIndexReader, VectorIndexWriter
+          from paimon_vindex import VectorIndexReader, VectorIndexTrainer, VectorIndexWriter
 
           class Input:
               def __init__(self, data):
@@ -266,8 +266,9 @@ jobs:
           data = np.array([[0.0, 0.0], [0.1, 0.0], [10.0, 10.0], [10.1, 10.0]], dtype=np.float32)
           ids = np.array([1, 2, 3, 4], dtype=np.int64)
           output = io.BytesIO()
-          writer = VectorIndexWriter({"index.type": "ivf_flat", "dimension": "2", "nlist": "2", "metric": "l2"})
-          writer.train(data)
+          options = {"index.type": "ivf_flat", "dimension": "2", "nlist": "2", "metric": "l2"}
+          training = VectorIndexTrainer.train(options, data)
+          writer = VectorIndexWriter(training)
           writer.add_vectors(ids, data)
           writer.write(output)
           writer.close()
@@ -291,7 +292,7 @@ jobs:
           python -c @'
           import io
           import numpy as np
-          from paimon_vindex import VectorIndexReader, VectorIndexWriter
+          from paimon_vindex import VectorIndexReader, VectorIndexTrainer, VectorIndexWriter
 
           class Input:
               def __init__(self, data):
@@ -303,8 +304,9 @@ jobs:
           data = np.array([[0.0, 0.0], [0.1, 0.0], [10.0, 10.0], [10.1, 10.0]], dtype=np.float32)
           ids = np.array([1, 2, 3, 4], dtype=np.int64)
           output = io.BytesIO()
-          writer = VectorIndexWriter({"index.type": "ivf_flat", "dimension": "2", "nlist": "2", "metric": "l2"})
-          writer.train(data)
+          options = {"index.type": "ivf_flat", "dimension": "2", "nlist": "2", "metric": "l2"}
+          training = VectorIndexTrainer.train(options, data)
+          writer = VectorIndexWriter(training)
           writer.add_vectors(ids, data)
           writer.write(output)
           writer.close()

README.md

@@ -87,7 +87,7 @@ use std::fs::File;
 use paimon_vindex_core::distance::MetricType;
 use paimon_vindex_core::hnsw::HnswBuildParams;
 use paimon_vindex_core::index::{
-    VectorIndexConfig, VectorIndexReader, VectorIndexWriter, VectorSearchParams,
+    VectorIndexConfig, VectorIndexReader, VectorIndexTrainer, VectorIndexWriter, VectorSearchParams,
 };
 use paimon_vindex_core::io::PosWriter;
 
@@ -98,8 +98,8 @@ let config = VectorIndexConfig::IvfHnswSq {
     hnsw: HnswBuildParams::default(),
 };
 
-let mut writer = VectorIndexWriter::new(config)?;
-writer.train(&training_vectors, training_count)?;
+let training = VectorIndexTrainer::train(config, &training_vectors, training_count)?;
+let mut writer = VectorIndexWriter::new(training);
 writer.add_vectors(&row_ids, &vectors, vector_count)?;
 
 let mut file = File::create("vectors.pvindex")?;
@@ -150,9 +150,14 @@ to `include/paimon_vindex.h`.
 const char *keys[] = {"index.type", "dimension", "nlist", "metric"};
 const char *values[] = {"ivf_flat", "128", "1024", "l2"};
 
-PaimonVindexWriterHandle *writer =
-    paimon_vindex_writer_open(keys, values, 4);
-paimon_vindex_writer_train(writer, training_vectors, training_count);
+PaimonVindexTrainerHandle *trainer =
+    paimon_vindex_trainer_open(keys, values, 4);
+paimon_vindex_trainer_add_training_vectors(trainer, training_vectors, training_count);
+PaimonVindexTrainingHandle *training = paimon_vindex_trainer_finish(trainer);
+paimon_vindex_trainer_free(trainer);
+
+PaimonVindexWriterHandle *writer = paimon_vindex_writer_open(training);
+paimon_vindex_training_free(training);
 paimon_vindex_writer_add_vectors(writer, row_ids, vectors, vector_count);
 paimon_vindex_writer_write_index(writer, output_file);
 paimon_vindex_writer_free(writer);
@@ -189,8 +194,9 @@ std::vector<std::pair<std::string, std::string>> options = {
     {"metric", "l2"},
 };
 
-paimon::vindex::Writer writer(options);
-writer.train(training_vectors.data(), training_count);
+paimon::vindex::Training training =
+    paimon::vindex::Trainer::train(options, training_vectors.data(), training_count);
+paimon::vindex::Writer writer(std::move(training));
 writer.add_vectors(row_ids.data(), vectors.data(), vector_count);
 writer.write_index(output_file);
 
@@ -203,30 +209,55 @@ auto result = reader.search(query.data(), 10, 16, 80);
 ### Java/JNI
 
 ```java
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 
 import org.apache.paimon.index.vector.VectorIndexInput;
 import org.apache.paimon.index.vector.VectorIndexMetadata;
 import org.apache.paimon.index.vector.VectorIndexReader;
+import org.apache.paimon.index.vector.VectorIndexTrainer;
+import org.apache.paimon.index.vector.VectorIndexTraining;
 import org.apache.paimon.index.vector.VectorSearchResult;
 import org.apache.paimon.index.vector.VectorIndexWriter;
 
+int dimension = 128;
 Map<String, String> options = new HashMap<>();
 options.put("index.type", "ivf_hnsw_sq");
-options.put("dimension", "128");
+options.put("dimension", Integer.toString(dimension));
 options.put("nlist", "1024");
 options.put("metric", "l2");
 options.put("hnsw.m", "20");
 options.put("hnsw.ef-construction", "150");
 options.put("hnsw.max-level", "7");
 
-try (VectorIndexWriter writer = new VectorIndexWriter(options)) {
-    writer.train(trainingVectors, trainingCount);
+try (VectorIndexTraining training =
+                VectorIndexTrainer.train(options, trainingVectors, trainingCount);
+        VectorIndexWriter writer = new VectorIndexWriter(training)) {
     writer.addVectors(rowIds, vectors, vectorCount);
     writer.writeIndex(vectorIndexOutput);
 }
 
+// Large training sets can avoid one large Java float[] by staging batches in trainer-owned
+// native memory. The batches are accumulated natively and released after finishTraining()
+// returns a trained state. This reduces JVM heap pressure and avoids the Java array length
+// limit; it does not reduce native peak training memory.
+int firstBatchCount = trainingCount / 2;
+float[][] trainingBatches = {
+    Arrays.copyOfRange(trainingVectors, 0, firstBatchCount * dimension),
+    Arrays.copyOfRange(trainingVectors, firstBatchCount * dimension, trainingCount * dimension),
+};
+try (VectorIndexTrainer trainer = VectorIndexTrainer.create(options)) {
+    for (float[] batch : trainingBatches) {
+        trainer.addTrainingVectors(batch, batch.length / dimension);
+    }
+    try (VectorIndexTraining training = trainer.finishTraining();
+            VectorIndexWriter writer = new VectorIndexWriter(training)) {
+        writer.addVectors(rowIds, vectors, vectorCount);
+        writer.writeIndex(vectorIndexOutput);
+    }
+}
+
 try (VectorIndexReader reader = new VectorIndexReader(vectorIndexInput)) {
     VectorIndexMetadata metadata = reader.metadata();
     reader.optimizeForSearch();
@@ -236,12 +267,12 @@ try (VectorIndexReader reader = new VectorIndexReader(vectorIndexInput)) {
 
 The Java package is `org.apache.paimon.index.vector`, and the API surface uses
 string options so it maps directly to Paimon table/index properties. Rust parses
-and validates the options when the writer is created.
+and validates the options when the trainer is created.
 
 ### Python
 
 ```python
-from paimon_vindex import VectorIndexReader, VectorIndexWriter
+from paimon_vindex import VectorIndexReader, VectorIndexTrainer, VectorIndexWriter
 
 
 class VectorIndexInput:
@@ -261,8 +292,8 @@ options = {
     "hnsw.ef-construction": "150",
     "hnsw.max-level": "7",
 }
-writer = VectorIndexWriter(options)
-writer.train(training_vectors)
+training = VectorIndexTrainer.train(options, training_vectors)
+writer = VectorIndexWriter(training)
 writer.add_vectors(row_ids, vectors)
 writer.write(output)
 

c/test_vindex.c

@@ -176,15 +176,15 @@ static void run_roundtrip(
         uint32_t expected_index_type,
         uintptr_t expected_pq_m,
         uintptr_t expected_hnsw_m) {
-    PaimonVindexWriterHandle *writer =
-        paimon_vindex_writer_open(keys, values, num_options);
-    if (writer == NULL) {
-        fail_ffi("writer open failed");
+    PaimonVindexTrainerHandle *trainer =
+        paimon_vindex_trainer_open(keys, values, num_options);
+    if (trainer == NULL) {
+        fail_ffi("trainer open failed");
     }
 
     uintptr_t dimension = 0;
-    if (paimon_vindex_writer_dimension(writer, &dimension) != 0) {
-        fail_ffi("writer dimension failed");
+    if (paimon_vindex_trainer_dimension(trainer, &dimension) != 0) {
+        fail_ffi("trainer dimension failed");
     }
     ASSERT_EQ_I64(dimension, 2);
 
@@ -194,9 +194,20 @@ static void run_roundtrip(
     ASSERT_TRUE(ids != NULL);
     fill_roundtrip_data(data, ids);
 
-    if (paimon_vindex_writer_train(writer, data, ROUNDTRIP_VECTOR_COUNT) != 0) {
-        fail_ffi("writer train failed");
+    if (paimon_vindex_trainer_add_training_vectors(trainer, data, ROUNDTRIP_VECTOR_COUNT) != 0) {
+        fail_ffi("trainer add training vectors failed");
+    }
+    PaimonVindexTrainingHandle *training = paimon_vindex_trainer_finish(trainer);
+    if (training == NULL) {
+        fail_ffi("trainer finish failed");
+    }
+    paimon_vindex_trainer_free(trainer);
+
+    PaimonVindexWriterHandle *writer = paimon_vindex_writer_open(training);
+    if (writer == NULL) {
+        fail_ffi("writer open failed");
     }
+    paimon_vindex_training_free(training);
     if (paimon_vindex_writer_add_vectors(writer, ids, data, ROUNDTRIP_VECTOR_COUNT) != 0) {
         fail_ffi("writer add failed");
     }
@@ -269,16 +280,27 @@ static void run_roundtrip(
 static PaimonVindexWriterHandle *new_trained_flat_writer(void) {
     const char *keys[] = {"index.type", "dimension", "nlist", "metric"};
     const char *values[] = {"ivf_flat", "1", "1", "l2"};
-    PaimonVindexWriterHandle *writer = paimon_vindex_writer_open(keys, values, 4);
-    if (writer == NULL) {
-        fail_ffi("writer open failed");
+    PaimonVindexTrainerHandle *trainer = paimon_vindex_trainer_open(keys, values, 4);
+    if (trainer == NULL) {
+        fail_ffi("trainer open failed");
     }
 
     const float data[] = {0.0f, 1.0f};
     const int64_t ids[] = {1, 2};
-    if (paimon_vindex_writer_train(writer, data, 2) != 0) {
-        fail_ffi("writer train failed");
+    if (paimon_vindex_trainer_add_training_vectors(trainer, data, 2) != 0) {
+        fail_ffi("trainer add training vectors failed");
+    }
+    PaimonVindexTrainingHandle *training = paimon_vindex_trainer_finish(trainer);
+    if (training == NULL) {
+        fail_ffi("trainer finish failed");
+    }
+    paimon_vindex_trainer_free(trainer);
+
+    PaimonVindexWriterHandle *writer = paimon_vindex_writer_open(training);
+    if (writer == NULL) {
+        fail_ffi("writer open failed");
     }
+    paimon_vindex_training_free(training);
     if (paimon_vindex_writer_add_vectors(writer, ids, data, 2) != 0) {
         fail_ffi("writer add failed");
     }

core/benches/ivfhnswsq_filter_bench.rs

@@ -18,7 +18,7 @@
 use paimon_vindex_core::distance::MetricType;
 use paimon_vindex_core::hnsw::HnswBuildParams;
 use paimon_vindex_core::index::{
-    VectorIndexConfig, VectorIndexReader, VectorIndexWriter, VectorSearchParams,
+    VectorIndexConfig, VectorIndexReader, VectorIndexTrainer, VectorIndexWriter, VectorSearchParams,
 };
 use paimon_vindex_core::io::PosWriter;
 use roaring::RoaringTreemap;
@@ -34,13 +34,17 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let ids: Vec<i64> = (0..cfg.n as i64).collect();
 
     let start = Instant::now();
-    let mut writer = VectorIndexWriter::new(VectorIndexConfig::IvfHnswSq {
-        dimension: cfg.d,
-        nlist: cfg.nlist,
-        metric: MetricType::L2,
-        hnsw: cfg.hnsw_params(),
-    })?;
-    writer.train(&data, cfg.n)?;
+    let training = VectorIndexTrainer::train(
+        VectorIndexConfig::IvfHnswSq {
+            dimension: cfg.d,
+            nlist: cfg.nlist,
+            metric: MetricType::L2,
+            hnsw: cfg.hnsw_params(),
+        },
+        &data,
+        cfg.n,
+    )?;
+    let mut writer = VectorIndexWriter::new(training);
     writer.add_vectors(&ids, &data, cfg.n)?;
     let mut index_bytes = Vec::new();
     writer.write(&mut PosWriter::new(&mut index_bytes))?;