This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-vector-index.git
The following commit(s) were added to refs/heads/main by this push:
new 35c461e Convert API to batched vector index training (#52)
35c461e is described below
commit 35c461e15fab82312b709b60ce2fce3e4e55ac84
Author: jerry <[email protected]>
AuthorDate: Tue Jun 30 10:43:00 2026 +0800
Convert API to batched vector index training (#52)
---
.github/workflows/ci.yml | 14 +-
README.md | 61 ++++--
c/test_vindex.c | 48 +++--
core/benches/ivfhnswsq_filter_bench.rs | 20 +-
core/src/index.rs | 146 ++++++++++++---
cpp/test_vindex.cpp | 11 +-
ffi/src/lib.rs | 161 ++++++++++++++--
include/paimon_vindex.hpp | 144 +++++++++++---
.../paimon/index/vector/VectorIndexNative.java | 16 +-
...torIndexWriter.java => VectorIndexTrainer.java} | 69 ++++---
.../paimon/index/vector/VectorIndexTraining.java | 81 ++++++++
.../paimon/index/vector/VectorIndexWriter.java | 24 +--
.../index/vector/VectorIndexJavaApiTest.java | 58 +++++-
.../vector/VectorIndexNativeHandleSafetyTest.java | 5 +-
.../vector/VectorIndexNativePanicBoundaryTest.java | 48 +++--
.../vector/VectorIndexNativeValidationTest.java | 190 +++++++++++++++++--
jni/src/lib.rs | 207 +++++++++++++++++++--
python/paimon_vindex/__init__.py | 158 ++++++++++++++--
python/paimon_vindex/_ffi.py | 27 ++-
python/tests/test_vindex.py | 19 +-
20 files changed, 1257 insertions(+), 250 deletions(-)
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index d1e5275..4dbb1bd 100644
--- a/.github/workflows/ci.yml
+++ b/.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()
diff --git a/README.md b/README.md
index 82bd6f1..f04a838 100644
--- a/README.md
+++ b/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)
diff --git a/c/test_vindex.c b/c/test_vindex.c
index b08660e..99f3f63 100644
--- a/c/test_vindex.c
+++ b/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");
}
diff --git a/core/benches/ivfhnswsq_filter_bench.rs
b/core/benches/ivfhnswsq_filter_bench.rs
index d49a92a..ab82310 100644
--- a/core/benches/ivfhnswsq_filter_bench.rs
+++ b/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))?;
diff --git a/core/src/index.rs b/core/src/index.rs
index c0b890c..25ea689 100644
--- a/core/src/index.rs
+++ b/core/src/index.rs
@@ -333,6 +333,75 @@ pub struct VectorIndexMetadata {
pub hnsw: Option<HnswBuildParams>,
}
+pub struct VectorIndexTrainer {
+ writer: VectorIndexWriter,
+ training_data: Vec<f32>,
+ training_vector_count: usize,
+}
+
+impl VectorIndexTrainer {
+ pub fn new(config: VectorIndexConfig) -> io::Result<Self> {
+ Ok(Self {
+ writer: VectorIndexWriter::from_config(config)?,
+ training_data: Vec::new(),
+ training_vector_count: 0,
+ })
+ }
+
+ pub fn train(
+ config: VectorIndexConfig,
+ data: &[f32],
+ n: usize,
+ ) -> io::Result<VectorIndexTraining> {
+ Self::new(config)?.add_training_vectors(data, n)?.finish()
+ }
+
+ pub fn dimension(&self) -> usize {
+ self.writer.dimension()
+ }
+
+ pub fn add_training_vectors(mut self, data: &[f32], n: usize) ->
io::Result<Self> {
+ self.add_training_vectors_mut(data, n)?;
+ Ok(self)
+ }
+
+ pub fn add_training_vectors_mut(&mut self, data: &[f32], n: usize) ->
io::Result<&mut Self> {
+ validate_vectors(data, n, self.dimension(), "training data")?;
+ let training_vector_count =
self.training_vector_count.checked_add(n).ok_or_else(|| {
+ io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "training vector count overflows usize",
+ )
+ })?;
+ self.training_data.extend_from_slice(data);
+ self.training_vector_count = training_vector_count;
+ Ok(self)
+ }
+
+ pub fn finish(mut self) -> io::Result<VectorIndexTraining> {
+ if self.training_vector_count == 0 || self.training_data.is_empty() {
+ return Err(invalid_input("no training vectors added"));
+ }
+ self.writer
+ .train_internal(&self.training_data, self.training_vector_count)?;
+ Ok(VectorIndexTraining { inner: self.writer })
+ }
+}
+
+pub struct VectorIndexTraining {
+ inner: VectorIndexWriter,
+}
+
+impl VectorIndexTraining {
+ pub fn index_type(&self) -> IndexType {
+ self.inner.index_type()
+ }
+
+ pub fn dimension(&self) -> usize {
+ self.inner.dimension()
+ }
+}
+
pub enum VectorIndexWriter {
IvfFlat(IVFFlatIndex),
IvfPq(IVFPQIndex),
@@ -341,7 +410,11 @@ pub enum VectorIndexWriter {
}
impl VectorIndexWriter {
- pub fn new(config: VectorIndexConfig) -> io::Result<Self> {
+ pub fn new(training: VectorIndexTraining) -> Self {
+ training.inner
+ }
+
+ fn from_config(config: VectorIndexConfig) -> io::Result<Self> {
validate_config(&config)?;
Ok(match config {
VectorIndexConfig::IvfFlat {
@@ -399,8 +472,8 @@ impl VectorIndexWriter {
}
}
- pub fn train(&mut self, data: &[f32], n: usize) -> io::Result<()> {
- validate_vectors(data, n, self.dimension(), "training data")?;
+ fn train_internal(&mut self, data: &[f32], n: usize) -> io::Result<()> {
+ debug_assert_eq!(Some(data.len()), n.checked_mul(self.dimension()));
match self {
Self::IvfFlat(index) => index.train(data, n),
Self::IvfPq(index) => index.train(data, n),
@@ -813,9 +886,8 @@ mod tests {
let data = generate_clustered_data(n, d, nlist);
let ids = (0..n as i64).collect::<Vec<_>>();
- let mut writer = VectorIndexWriter::new(config.clone()).unwrap();
+ let mut writer = build_writer(config.clone(), &data, n);
assert_eq!(writer.index_type(), config.index_type());
- writer.train(&data, n).unwrap();
writer.add_vectors(&ids, &data, n).unwrap();
let mut buf = Vec::new();
@@ -842,8 +914,7 @@ mod tests {
let data = generate_clustered_data(n, d, nlist);
let ids = (0..n as i64).collect::<Vec<_>>();
- let mut writer = VectorIndexWriter::new(config).unwrap();
- writer.train(&data, n).unwrap();
+ let mut writer = build_writer(config, &data, n);
writer.add_vectors(&ids, &data, n).unwrap();
let mut buf = Vec::new();
@@ -852,13 +923,15 @@ mod tests {
}
fn build_ivfflat_reader() -> VectorIndexReader<Cursor<Vec<u8>>> {
- let mut writer = VectorIndexWriter::new(VectorIndexConfig::IvfFlat {
- dimension: 1,
- nlist: 1,
- metric: MetricType::L2,
- })
- .unwrap();
- writer.train(&[0.0, 1.0], 2).unwrap();
+ let mut writer = build_writer(
+ VectorIndexConfig::IvfFlat {
+ dimension: 1,
+ nlist: 1,
+ metric: MetricType::L2,
+ },
+ &[0.0, 1.0],
+ 2,
+ );
writer.add_vectors(&[1, 2], &[0.0, 1.0], 2).unwrap();
let mut bytes = Vec::new();
@@ -866,6 +939,11 @@ mod tests {
VectorIndexReader::open(Cursor::new(bytes)).unwrap()
}
+ fn build_writer(config: VectorIndexConfig, data: &[f32], n: usize) ->
VectorIndexWriter {
+ let training = VectorIndexTrainer::train(config, data, n).unwrap();
+ VectorIndexWriter::new(training)
+ }
+
fn assert_invalid_input_contains(result: io::Result<()>, expected: &str) {
let err = result.expect_err("invalid input should be rejected");
assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
@@ -968,7 +1046,7 @@ mod tests {
#[test]
fn unified_config_rejects_invalid_pq_m() {
- let err = match VectorIndexWriter::new(VectorIndexConfig::IvfPq {
+ let err = match VectorIndexTrainer::new(VectorIndexConfig::IvfPq {
dimension: 10,
nlist: 4,
m: 3,
@@ -1038,7 +1116,7 @@ mod tests {
}
#[test]
- fn unified_writer_rejects_non_finite_training_data() {
+ fn unified_trainer_rejects_non_finite_training_data() {
for (value, expected) in [
(
f32::NAN,
@@ -1053,13 +1131,19 @@ mod tests {
"training data contains non-finite value at offset 0: -inf",
),
] {
- let mut writer = VectorIndexWriter::new(VectorIndexConfig::IvfFlat
{
- dimension: 1,
- nlist: 1,
- metric: MetricType::L2,
- })
- .unwrap();
- assert_invalid_input_contains(writer.train(&[value, 1.0], 2),
expected);
+ assert_invalid_input_contains(
+ VectorIndexTrainer::train(
+ VectorIndexConfig::IvfFlat {
+ dimension: 1,
+ nlist: 1,
+ metric: MetricType::L2,
+ },
+ &[value, 1.0],
+ 2,
+ )
+ .map(|_| ()),
+ expected,
+ );
}
}
@@ -1130,13 +1214,15 @@ mod tests {
"vector data contains non-finite value at offset 0: -inf",
),
] {
- let mut writer = VectorIndexWriter::new(VectorIndexConfig::IvfFlat
{
- dimension: 1,
- nlist: 1,
- metric: MetricType::L2,
- })
- .unwrap();
- writer.train(&[0.0, 1.0], 2).unwrap();
+ let mut writer = build_writer(
+ VectorIndexConfig::IvfFlat {
+ dimension: 1,
+ nlist: 1,
+ metric: MetricType::L2,
+ },
+ &[0.0, 1.0],
+ 2,
+ );
assert_invalid_input_contains(writer.add_vectors(&[1, 2], &[value,
1.0], 2), expected);
}
}
diff --git a/cpp/test_vindex.cpp b/cpp/test_vindex.cpp
index 96bd2dd..ec1b5e3 100644
--- a/cpp/test_vindex.cpp
+++ b/cpp/test_vindex.cpp
@@ -110,12 +110,15 @@ static void run_roundtrip(
uint32_t expected_index_type,
size_t expected_pq_m,
size_t expected_hnsw_m) {
- paimon::vindex::Writer writer(options);
- ASSERT_EQ(writer.dimension(), 2);
-
std::vector<float> data = roundtrip_data();
std::vector<int64_t> ids = roundtrip_ids();
- writer.train(data.data(), kRoundtripVectorCount);
+ paimon::vindex::Trainer trainer(options);
+ ASSERT_EQ(trainer.dimension(), 2);
+ paimon::vindex::Training training =
+ trainer.add_training_vectors(data.data(),
kRoundtripVectorCount).finish_training();
+
+ paimon::vindex::Writer writer(std::move(training));
+ ASSERT_EQ(writer.dimension(), 2);
writer.add_vectors(ids.data(), data.data(), kRoundtripVectorCount);
MemBuffer buf;
diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs
index 5e58f92..1e719e3 100644
--- a/ffi/src/lib.rs
+++ b/ffi/src/lib.rs
@@ -19,8 +19,8 @@
use paimon_vindex_core::distance::MetricType;
use paimon_vindex_core::index::{
- VectorIndexConfig, VectorIndexMetadata, VectorIndexReader,
VectorIndexWriter,
- VectorSearchParams,
+ VectorIndexConfig, VectorIndexMetadata, VectorIndexReader,
VectorIndexTrainer,
+ VectorIndexTraining, VectorIndexWriter, VectorSearchParams,
};
use paimon_vindex_core::io::{ReadRequest, SeekRead, SeekWrite};
use std::cell::RefCell;
@@ -214,6 +214,14 @@ pub struct PaimonVindexMetadata {
pub hnsw_max_level: usize,
}
+pub struct PaimonVindexTrainerHandle {
+ inner: Option<VectorIndexTrainer>,
+}
+
+pub struct PaimonVindexTrainingHandle {
+ inner: Option<VectorIndexTraining>,
+}
+
pub struct PaimonVindexWriterHandle {
inner: VectorIndexWriter,
}
@@ -301,6 +309,36 @@ unsafe fn writer_mut<'a>(
}
}
+unsafe fn trainer_mut<'a>(
+ handle: *mut PaimonVindexTrainerHandle,
+) -> Result<&'a mut PaimonVindexTrainerHandle, String> {
+ if handle.is_null() {
+ Err("null trainer handle".to_string())
+ } else {
+ Ok(unsafe { &mut *handle })
+ }
+}
+
+unsafe fn trainer_ref<'a>(
+ handle: *const PaimonVindexTrainerHandle,
+) -> Result<&'a PaimonVindexTrainerHandle, String> {
+ if handle.is_null() {
+ Err("null trainer handle".to_string())
+ } else {
+ Ok(unsafe { &*handle })
+ }
+}
+
+unsafe fn training_mut<'a>(
+ handle: *mut PaimonVindexTrainingHandle,
+) -> Result<&'a mut PaimonVindexTrainingHandle, String> {
+ if handle.is_null() {
+ Err("null training handle".to_string())
+ } else {
+ Ok(unsafe { &mut *handle })
+ }
+}
+
unsafe fn reader_mut<'a>(
handle: *mut PaimonVindexReaderHandle,
) -> Result<&'a mut PaimonVindexReaderHandle, String> {
@@ -385,27 +423,28 @@ fn copy_search_result(
Ok(())
}
-// ======================== Writer ========================
+// ======================== Trainer / Writer ========================
#[no_mangle]
-pub unsafe extern "C" fn paimon_vindex_writer_open(
+pub unsafe extern "C" fn paimon_vindex_trainer_open(
keys: *const *const c_char,
values: *const *const c_char,
num_options: usize,
-) -> *mut PaimonVindexWriterHandle {
+) -> *mut PaimonVindexTrainerHandle {
ffi_ptr(|| {
let options = unsafe { options_from_raw(keys, values, num_options) }?;
let config = VectorIndexConfig::from_options(&options)
.map_err(|e| format!("invalid vector index options: {}", e))?;
- let writer = VectorIndexWriter::new(config).map_err(|e|
format!("create writer: {}", e))?;
- Ok(Box::into_raw(Box::new(PaimonVindexWriterHandle {
- inner: writer,
+ let trainer =
+ VectorIndexTrainer::new(config).map_err(|e| format!("create
trainer: {}", e))?;
+ Ok(Box::into_raw(Box::new(PaimonVindexTrainerHandle {
+ inner: Some(trainer),
})))
})
}
#[no_mangle]
-pub unsafe extern "C" fn paimon_vindex_writer_free(handle: *mut
PaimonVindexWriterHandle) {
+pub unsafe extern "C" fn paimon_vindex_trainer_free(handle: *mut
PaimonVindexTrainerHandle) {
if !handle.is_null() {
unsafe {
drop(Box::from_raw(handle));
@@ -414,36 +453,118 @@ pub unsafe extern "C" fn
paimon_vindex_writer_free(handle: *mut PaimonVindexWrit
}
#[no_mangle]
-pub unsafe extern "C" fn paimon_vindex_writer_dimension(
- handle: *const PaimonVindexWriterHandle,
+pub unsafe extern "C" fn paimon_vindex_trainer_dimension(
+ handle: *const PaimonVindexTrainerHandle,
out: *mut usize,
) -> c_int {
ffi_status(|| {
if out.is_null() {
return Err("out pointer is null".to_string());
}
- let handle = unsafe { writer_ref(handle) }?;
+ let handle = unsafe { trainer_ref(handle) }?;
+ let trainer = handle
+ .inner
+ .as_ref()
+ .ok_or_else(|| "trainer has already finished".to_string())?;
unsafe {
- *out = handle.inner.dimension();
+ *out = trainer.dimension();
}
Ok(())
})
}
#[no_mangle]
-pub unsafe extern "C" fn paimon_vindex_writer_train(
- handle: *mut PaimonVindexWriterHandle,
+pub unsafe extern "C" fn paimon_vindex_trainer_add_training_vectors(
+ handle: *mut PaimonVindexTrainerHandle,
data: *const f32,
vector_count: usize,
) -> c_int {
ffi_status(|| {
- let handle = unsafe { writer_mut(handle) }?;
- let len = checked_len(vector_count, handle.inner.dimension(),
"training data")?;
+ let handle = unsafe { trainer_mut(handle) }?;
+ let trainer = handle
+ .inner
+ .as_mut()
+ .ok_or_else(|| "trainer has already finished".to_string())?;
+ let len = checked_len(vector_count, trainer.dimension(), "training
data")?;
let data = unsafe { const_slice(data, len, "data") }?;
- handle
+ trainer
+ .add_training_vectors_mut(data, vector_count)
+ .map(|_| ())
+ .map_err(|e| format!("add training vectors: {}", e))
+ })
+}
+
+/// Finishes training and consumes the trainer's internal state, but does not
free `handle`.
+/// Callers must still call `paimon_vindex_trainer_free(handle)` after this
returns.
+#[no_mangle]
+pub unsafe extern "C" fn paimon_vindex_trainer_finish(
+ handle: *mut PaimonVindexTrainerHandle,
+) -> *mut PaimonVindexTrainingHandle {
+ ffi_ptr(|| {
+ let handle = unsafe { trainer_mut(handle) }?;
+ let trainer = handle
.inner
- .train(data, vector_count)
- .map_err(|e| format!("train: {}", e))
+ .take()
+ .ok_or_else(|| "trainer has already finished".to_string())?;
+ let training = trainer
+ .finish()
+ .map_err(|e| format!("finish training: {}", e))?;
+ Ok(Box::into_raw(Box::new(PaimonVindexTrainingHandle {
+ inner: Some(training),
+ })))
+ })
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn paimon_vindex_training_free(handle: *mut
PaimonVindexTrainingHandle) {
+ if !handle.is_null() {
+ unsafe {
+ drop(Box::from_raw(handle));
+ }
+ }
+}
+
+/// Opens a writer by consuming the training state inside `training`, but does
not free the handle.
+/// Callers must still call `paimon_vindex_training_free(training)` after this
returns.
+#[no_mangle]
+pub unsafe extern "C" fn paimon_vindex_writer_open(
+ training: *mut PaimonVindexTrainingHandle,
+) -> *mut PaimonVindexWriterHandle {
+ ffi_ptr(|| {
+ let training = unsafe { training_mut(training) }?;
+ let training = training
+ .inner
+ .take()
+ .ok_or_else(|| "training has already been consumed".to_string())?;
+ Ok(Box::into_raw(Box::new(PaimonVindexWriterHandle {
+ inner: VectorIndexWriter::new(training),
+ })))
+ })
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn paimon_vindex_writer_free(handle: *mut
PaimonVindexWriterHandle) {
+ if !handle.is_null() {
+ unsafe {
+ drop(Box::from_raw(handle));
+ }
+ }
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn paimon_vindex_writer_dimension(
+ handle: *const PaimonVindexWriterHandle,
+ out: *mut usize,
+) -> c_int {
+ ffi_status(|| {
+ if out.is_null() {
+ return Err("out pointer is null".to_string());
+ }
+ let handle = unsafe { writer_ref(handle) }?;
+ unsafe {
+ *out = handle.inner.dimension();
+ }
+ Ok(())
})
}
diff --git a/include/paimon_vindex.hpp b/include/paimon_vindex.hpp
index b0ab679..1216835 100644
--- a/include/paimon_vindex.hpp
+++ b/include/paimon_vindex.hpp
@@ -114,14 +114,44 @@ struct SearchResult {
std::vector<float> distances;
};
-class Writer {
+class Training {
+public:
+ explicit Training(PaimonVindexTrainingHandle* handle = nullptr) :
handle_(handle) {}
+
+ Training(const Training&) = delete;
+ Training& operator=(const Training&) = delete;
+
+ Training(Training&& other) noexcept : handle_(other.handle_) {
+ other.handle_ = nullptr;
+ }
+
+ Training& operator=(Training&& other) noexcept {
+ if (this != &other) {
+ if (handle_) paimon_vindex_training_free(handle_);
+ handle_ = other.handle_;
+ other.handle_ = nullptr;
+ }
+ return *this;
+ }
+
+ ~Training() {
+ if (handle_) paimon_vindex_training_free(handle_);
+ }
+
+private:
+ friend class Writer;
+
+ PaimonVindexTrainingHandle* handle_ = nullptr;
+};
+
+class Trainer {
public:
- Writer(const char* const* keys, const char* const* values, size_t
num_options) {
- handle_ = paimon_vindex_writer_open(keys, values, num_options);
- if (!handle_) throw Error("failed to open vector index writer");
+ Trainer(const char* const* keys, const char* const* values, size_t
num_options) {
+ handle_ = paimon_vindex_trainer_open(keys, values, num_options);
+ if (!handle_) throw Error("failed to open vector index trainer");
}
- explicit Writer(const std::vector<std::pair<std::string, std::string>>&
options) {
+ explicit Trainer(const std::vector<std::pair<std::string, std::string>>&
options) {
option_keys_.reserve(options.size());
option_values_.reserve(options.size());
key_ptrs_.reserve(options.size());
@@ -134,16 +164,15 @@ public:
key_ptrs_.push_back(option_keys_[i].c_str());
value_ptrs_.push_back(option_values_[i].c_str());
}
- handle_ = paimon_vindex_writer_open(key_ptrs_.data(),
value_ptrs_.data(), options.size());
- if (!handle_) throw Error("failed to open vector index writer");
+ handle_ = paimon_vindex_trainer_open(key_ptrs_.data(),
value_ptrs_.data(), options.size());
+ if (!handle_) throw Error("failed to open vector index trainer");
}
- Writer(const Writer&) = delete;
- Writer& operator=(const Writer&) = delete;
+ Trainer(const Trainer&) = delete;
+ Trainer& operator=(const Trainer&) = delete;
- Writer(Writer&& other) noexcept
+ Trainer(Trainer&& other) noexcept
: handle_(other.handle_),
- output_(std::move(other.output_)),
option_keys_(std::move(other.option_keys_)),
option_values_(std::move(other.option_values_)),
key_ptrs_(std::move(other.key_ptrs_)),
@@ -151,11 +180,10 @@ public:
other.handle_ = nullptr;
}
- Writer& operator=(Writer&& other) noexcept {
+ Trainer& operator=(Trainer&& other) noexcept {
if (this != &other) {
- if (handle_) paimon_vindex_writer_free(handle_);
+ if (handle_) paimon_vindex_trainer_free(handle_);
handle_ = other.handle_;
- output_ = std::move(other.output_);
option_keys_ = std::move(other.option_keys_);
option_values_ = std::move(other.option_values_);
key_ptrs_ = std::move(other.key_ptrs_);
@@ -165,6 +193,86 @@ public:
return *this;
}
+ ~Trainer() {
+ if (handle_) paimon_vindex_trainer_free(handle_);
+ }
+
+ size_t dimension() const {
+ size_t out = 0;
+ check(paimon_vindex_trainer_dimension(handle_, &out));
+ return out;
+ }
+
+ Trainer& add_training_vectors(const float* data, size_t vector_count) {
+ check(paimon_vindex_trainer_add_training_vectors(handle_, data,
vector_count));
+ return *this;
+ }
+
+ // C ABI note: finish consumes the trainer state but leaves the trainer
handle owned by caller.
+ // This RAII wrapper frees the trainer handle after a successful finish.
+ Training finish_training() {
+ PaimonVindexTrainingHandle* training =
paimon_vindex_trainer_finish(handle_);
+ if (!training) {
+ const char* err = paimon_vindex_last_error();
+ throw Error(err ? err : "failed to finish vector index training");
+ }
+ paimon_vindex_trainer_free(handle_);
+ handle_ = nullptr;
+ return Training(training);
+ }
+
+ static Training train(
+ const std::vector<std::pair<std::string, std::string>>& options,
+ const float* data,
+ size_t vector_count) {
+ Trainer trainer(options);
+ trainer.add_training_vectors(data, vector_count);
+ return trainer.finish_training();
+ }
+
+private:
+ PaimonVindexTrainerHandle* handle_ = nullptr;
+ std::vector<std::string> option_keys_;
+ std::vector<std::string> option_values_;
+ std::vector<const char*> key_ptrs_;
+ std::vector<const char*> value_ptrs_;
+};
+
+class Writer {
+public:
+ explicit Writer(Training&& training) {
+ if (!training.handle_) throw Error("training has already been
consumed");
+ PaimonVindexTrainingHandle* training_handle = training.handle_;
+ training.handle_ = nullptr;
+ // C ABI note: writer_open consumes the training state but leaves the
handle owned by caller.
+ // This RAII wrapper frees the consumed training handle after opening
the writer.
+ handle_ = paimon_vindex_writer_open(training_handle);
+ paimon_vindex_training_free(training_handle);
+ if (!handle_) {
+ const char* err = paimon_vindex_last_error();
+ throw Error(err ? err : "failed to open vector index writer");
+ }
+ }
+
+ Writer(const Writer&) = delete;
+ Writer& operator=(const Writer&) = delete;
+
+ Writer(Writer&& other) noexcept
+ : handle_(other.handle_),
+ output_(std::move(other.output_)) {
+ other.handle_ = nullptr;
+ }
+
+ Writer& operator=(Writer&& other) noexcept {
+ if (this != &other) {
+ if (handle_) paimon_vindex_writer_free(handle_);
+ handle_ = other.handle_;
+ output_ = std::move(other.output_);
+ other.handle_ = nullptr;
+ }
+ return *this;
+ }
+
~Writer() {
if (handle_) paimon_vindex_writer_free(handle_);
}
@@ -175,10 +283,6 @@ public:
return out;
}
- void train(const float* data, size_t vector_count) {
- check(paimon_vindex_writer_train(handle_, data, vector_count));
- }
-
void add_vectors(const int64_t* ids, const float* data, size_t
vector_count) {
check(paimon_vindex_writer_add_vectors(handle_, ids, data,
vector_count));
}
@@ -196,10 +300,6 @@ public:
private:
PaimonVindexWriterHandle* handle_ = nullptr;
std::shared_ptr<OutputFile> output_;
- std::vector<std::string> option_keys_;
- std::vector<std::string> option_values_;
- std::vector<const char*> key_ptrs_;
- std::vector<const char*> value_ptrs_;
};
class Reader {
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 4adf3c9..e7cf220 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
@@ -21,11 +21,19 @@ final class VectorIndexNative {
private VectorIndexNative() {}
- static native long createWriter(String[] optionKeys, String[]
optionValues);
+ static native long createTrainer(String[] optionKeys, String[]
optionValues);
- static native int writerDimension(long ptr);
+ static native int trainerDimension(long ptr);
+
+ static native void trainerAddTrainingVectors(long ptr, float[] data, int
n);
+
+ static native long trainerFinishTraining(long ptr);
+
+ static native void freeTrainer(long ptr);
- static native void train(long ptr, float[] data, int n);
+ static native long createWriter(long trainingPtr);
+
+ static native int writerDimension(long ptr);
static native void addVectors(long ptr, long[] ids, float[] data, int n);
@@ -33,6 +41,8 @@ final class VectorIndexNative {
static native void freeWriter(long ptr);
+ static native void freeTraining(long ptr);
+
static native long openReader(Object streamInput);
static native VectorIndexMetadata metadata(long ptr);
diff --git
a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexWriter.java
b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexTrainer.java
similarity index 56%
copy from
java/src/main/java/org/apache/paimon/index/vector/VectorIndexWriter.java
copy to
java/src/main/java/org/apache/paimon/index/vector/VectorIndexTrainer.java
index 0dda2a3..1eec2e3 100644
--- a/java/src/main/java/org/apache/paimon/index/vector/VectorIndexWriter.java
+++ b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexTrainer.java
@@ -19,13 +19,27 @@ package org.apache.paimon.index.vector;
import java.util.Map;
-public final class VectorIndexWriter implements AutoCloseable {
+/**
+ * Builds a trained vector-index state from one or more training batches.
+ *
+ * <p>Staging batches avoids requiring one large Java {@code float[]} and its
array length limit,
+ * but all batches are accumulated in native memory until {@link
#finishTraining()}; this does not
+ * reduce native peak training memory.
+ */
+public final class VectorIndexTrainer implements AutoCloseable {
private final Object nativeHandleLock = new Object();
private long nativePtr;
private Thread nativeHandleOwner;
- public VectorIndexWriter(Map<String, String> options) {
+ private VectorIndexTrainer(long nativePtr) {
+ this.nativePtr = nativePtr;
+ }
+
+ public static VectorIndexTrainer create(Map<String, String> options) {
+ if (options == null) {
+ throw new NullPointerException("options");
+ }
String[] keys = new String[options.size()];
String[] values = new String[options.size()];
int index = 0;
@@ -34,60 +48,59 @@ public final class VectorIndexWriter implements
AutoCloseable {
values[index] = entry.getValue();
index++;
}
- this.nativePtr = VectorIndexNative.createWriter(keys, values);
+ return new VectorIndexTrainer(VectorIndexNative.createTrainer(keys,
values));
}
- private VectorIndexWriter(long nativePtr) {
- this.nativePtr = nativePtr;
+ public static VectorIndexTraining train(
+ Map<String, String> options, float[] data, int vectorCount) {
+ try (VectorIndexTrainer trainer = create(options)) {
+ return trainer.addTrainingVectors(data,
vectorCount).finishTraining();
+ }
}
- static VectorIndexWriter fromNativePointerForTesting(long nativePtr) {
- return new VectorIndexWriter(nativePtr);
+ static VectorIndexTrainer fromNativePointerForTesting(long nativePtr) {
+ return new VectorIndexTrainer(nativePtr);
}
public int dimension() {
- return VectorIndexNative.writerDimension(requireOpen());
- }
-
- public void train(float[] data, int vectorCount) {
- if (data == null) {
- throw new NullPointerException("data");
- }
synchronized (nativeHandleLock) {
enterNativeHandle();
try {
- VectorIndexNative.train(requireOpen(), data, vectorCount);
+ return VectorIndexNative.trainerDimension(requireOpen());
} finally {
exitNativeHandle();
}
}
}
- public void addVectors(long[] ids, float[] data, int vectorCount) {
- if (ids == null) {
- throw new NullPointerException("ids");
- }
+ public VectorIndexTrainer addTrainingVectors(float[] data, int
vectorCount) {
if (data == null) {
throw new NullPointerException("data");
}
synchronized (nativeHandleLock) {
enterNativeHandle();
try {
- VectorIndexNative.addVectors(requireOpen(), ids, data,
vectorCount);
+ VectorIndexNative.trainerAddTrainingVectors(requireOpen(),
data, vectorCount);
+ return this;
} finally {
exitNativeHandle();
}
}
}
- public void writeIndex(Object output) {
- if (output == null) {
- throw new NullPointerException("output");
- }
+ /**
+ * Finishes training and returns a trained state for creating a {@link
VectorIndexWriter}.
+ *
+ * <p>This method consumes this trainer on both success and failure. After
it returns or throws,
+ * this trainer is closed and cannot accept more training batches; create
a new trainer to retry.
+ */
+ public VectorIndexTraining finishTraining() {
synchronized (nativeHandleLock) {
enterNativeHandle();
try {
- VectorIndexNative.writeIndex(requireOpen(), output);
+ long ptr = requireOpen();
+ nativePtr = 0L;
+ return new
VectorIndexTraining(VectorIndexNative.trainerFinishTraining(ptr));
} finally {
exitNativeHandle();
}
@@ -102,7 +115,7 @@ public final class VectorIndexWriter implements
AutoCloseable {
long ptr = nativePtr;
nativePtr = 0L;
if (ptr != 0L) {
- VectorIndexNative.freeWriter(ptr);
+ VectorIndexNative.freeTrainer(ptr);
}
} finally {
exitNativeHandle();
@@ -112,7 +125,7 @@ public final class VectorIndexWriter implements
AutoCloseable {
private long requireOpen() {
if (nativePtr == 0L) {
- throw new IllegalStateException("VectorIndexWriter is closed");
+ throw new IllegalStateException("VectorIndexTrainer is closed");
}
return nativePtr;
}
@@ -120,7 +133,7 @@ public final class VectorIndexWriter implements
AutoCloseable {
private void enterNativeHandle() {
Thread current = Thread.currentThread();
if (nativeHandleOwner == current) {
- throw new IllegalStateException("VectorIndexWriter native handle
is already in use");
+ throw new IllegalStateException("VectorIndexTrainer native handle
is already in use");
}
nativeHandleOwner = current;
}
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
new file mode 100644
index 0000000..65780e7
--- /dev/null
+++ b/java/src/main/java/org/apache/paimon/index/vector/VectorIndexTraining.java
@@ -0,0 +1,81 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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 org.apache.paimon.index.vector;
+
+public final class VectorIndexTraining implements AutoCloseable {
+
+ private final Object nativeHandleLock = new Object();
+ private long nativePtr;
+ private Thread nativeHandleOwner;
+
+ VectorIndexTraining(long nativePtr) {
+ this.nativePtr = nativePtr;
+ }
+
+ static VectorIndexTraining fromNativePointerForTesting(long nativePtr) {
+ return new VectorIndexTraining(nativePtr);
+ }
+
+ long takeNativePointer() {
+ synchronized (nativeHandleLock) {
+ enterNativeHandle();
+ try {
+ long ptr = requireOpen();
+ nativePtr = 0L;
+ return ptr;
+ } finally {
+ exitNativeHandle();
+ }
+ }
+ }
+
+ @Override
+ public void close() {
+ synchronized (nativeHandleLock) {
+ enterNativeHandle();
+ try {
+ long ptr = nativePtr;
+ nativePtr = 0L;
+ if (ptr != 0L) {
+ VectorIndexNative.freeTraining(ptr);
+ }
+ } finally {
+ exitNativeHandle();
+ }
+ }
+ }
+
+ private long requireOpen() {
+ if (nativePtr == 0L) {
+ throw new IllegalStateException("VectorIndexTraining is closed");
+ }
+ return nativePtr;
+ }
+
+ private void enterNativeHandle() {
+ Thread current = Thread.currentThread();
+ if (nativeHandleOwner == current) {
+ throw new IllegalStateException("VectorIndexTraining native handle
is already in use");
+ }
+ nativeHandleOwner = current;
+ }
+
+ private void exitNativeHandle() {
+ nativeHandleOwner = null;
+ }
+}
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 0dda2a3..1f1863b 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
@@ -17,24 +17,17 @@
package org.apache.paimon.index.vector;
-import java.util.Map;
-
public final class VectorIndexWriter implements AutoCloseable {
private final Object nativeHandleLock = new Object();
private long nativePtr;
private Thread nativeHandleOwner;
- public VectorIndexWriter(Map<String, String> options) {
- String[] keys = new String[options.size()];
- String[] values = new String[options.size()];
- int index = 0;
- for (Map.Entry<String, String> entry : options.entrySet()) {
- keys[index] = entry.getKey();
- values[index] = entry.getValue();
- index++;
+ public VectorIndexWriter(VectorIndexTraining training) {
+ if (training == null) {
+ throw new NullPointerException("training");
}
- this.nativePtr = VectorIndexNative.createWriter(keys, values);
+ this.nativePtr =
VectorIndexNative.createWriter(training.takeNativePointer());
}
private VectorIndexWriter(long nativePtr) {
@@ -46,17 +39,10 @@ public final class VectorIndexWriter implements
AutoCloseable {
}
public int dimension() {
- return VectorIndexNative.writerDimension(requireOpen());
- }
-
- public void train(float[] data, int vectorCount) {
- if (data == null) {
- throw new NullPointerException("data");
- }
synchronized (nativeHandleLock) {
enterNativeHandle();
try {
- VectorIndexNative.train(requireOpen(), data, vectorCount);
+ return VectorIndexNative.writerDimension(requireOpen());
} finally {
exitNativeHandle();
}
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 6fd37a6..5d75836 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
@@ -28,6 +28,8 @@ public class VectorIndexJavaApiTest {
testBatchResultCopiesArraysAndSlicesRows();
testMetadata();
testClosedReaderRejectsOperations();
+ testClosedTrainerRejectsOperations();
+ testClosedTrainingRejectsOperations();
testClosedWriterRejectsOperations();
testReaderAndWriterApiCompile();
}
@@ -156,7 +158,7 @@ public class VectorIndexJavaApiTest {
assertThrows(IllegalStateException.class, new ThrowingRunnable() {
@Override
public void run() {
- writer.train(new float[] {0.0f, 1.0f}, 1);
+ writer.dimension();
}
});
assertThrows(IllegalStateException.class, new ThrowingRunnable() {
@@ -173,6 +175,44 @@ public class VectorIndexJavaApiTest {
});
}
+ private static void testClosedTrainerRejectsOperations() {
+ final VectorIndexTrainer trainer =
VectorIndexTrainer.fromNativePointerForTesting(0L);
+ trainer.close();
+ trainer.close();
+
+ assertThrows(IllegalStateException.class, new ThrowingRunnable() {
+ @Override
+ public void run() {
+ trainer.dimension();
+ }
+ });
+ assertThrows(IllegalStateException.class, new ThrowingRunnable() {
+ @Override
+ public void run() {
+ trainer.addTrainingVectors(new float[] {0.0f, 1.0f}, 1);
+ }
+ });
+ assertThrows(IllegalStateException.class, new ThrowingRunnable() {
+ @Override
+ public void run() {
+ trainer.finishTraining();
+ }
+ });
+ }
+
+ private static void testClosedTrainingRejectsOperations() {
+ final VectorIndexTraining training =
VectorIndexTraining.fromNativePointerForTesting(0L);
+ training.close();
+ training.close();
+
+ assertThrows(IllegalStateException.class, new ThrowingRunnable() {
+ @Override
+ public void run() {
+ new VectorIndexWriter(training);
+ }
+ });
+ }
+
private static void testReaderAndWriterApiCompile() {
Map<String, String> options = ivfPqOptions(2, 4, 1);
VectorIndexReader closedReader =
VectorIndexReader.fromNativePointerForTesting(0L);
@@ -200,10 +240,22 @@ public class VectorIndexJavaApiTest {
reader.searchBatch(
new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, 10, 4, 32, new
byte[] {1, 2});
- VectorIndexWriter writer = new VectorIndexWriter(options);
- writer.train(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2);
+ VectorIndexTraining training =
+ VectorIndexTrainer.train(options, new float[] {0.0f, 1.0f,
2.0f, 3.0f}, 2);
+ VectorIndexWriter writer = new VectorIndexWriter(training);
+ writer.dimension();
writer.addVectors(new long[] {1L, 2L}, new float[] {0.0f, 1.0f,
2.0f, 3.0f}, 2);
writer.writeIndex(new Object());
+
+ VectorIndexTrainer trainer = VectorIndexTrainer.create(options);
+ trainer.dimension();
+ VectorIndexTraining stagedTraining =
+ trainer.addTrainingVectors(new float[] {0.0f, 1.0f}, 1)
+ .addTrainingVectors(new float[] {2.0f, 3.0f}, 1)
+ .finishTraining();
+ VectorIndexWriter stagedWriter = new
VectorIndexWriter(stagedTraining);
+ stagedWriter.addVectors(new long[] {1L}, new float[] {0.0f, 1.0f},
1);
+ stagedWriter.writeIndex(new Object());
}
}
diff --git
a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeHandleSafetyTest.java
b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeHandleSafetyTest.java
index efcfdcb..3609cdd 100644
---
a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeHandleSafetyTest.java
+++
b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativeHandleSafetyTest.java
@@ -75,8 +75,9 @@ public class VectorIndexNativeHandleSafetyTest {
}
private static VectorIndexWriter newPopulatedWriter() {
- VectorIndexWriter writer = new VectorIndexWriter(ivfFlatOptions());
- writer.train(new float[] {0.0f, 1.0f}, 2);
+ VectorIndexWriter writer =
+ new VectorIndexWriter(
+ VectorIndexTrainer.train(ivfFlatOptions(), new float[]
{0.0f, 1.0f}, 2));
writer.addVectors(new long[] {1L, 2L}, new float[] {0.0f, 1.0f}, 2);
return writer;
}
diff --git
a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativePanicBoundaryTest.java
b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativePanicBoundaryTest.java
index 739defe..963b26f 100644
---
a/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativePanicBoundaryTest.java
+++
b/java/src/test/java/org/apache/paimon/index/vector/VectorIndexNativePanicBoundaryTest.java
@@ -30,32 +30,36 @@ public class VectorIndexNativePanicBoundaryTest {
System.load(args[0]);
- testVoidEntrypointPanicBecomesRuntimeException();
+ testVoidEntrypointErrorBecomesRuntimeException();
testObjectEntrypointPanicBecomesRuntimeException();
- VectorIndexWriter survivor = new VectorIndexWriter(ivfFlatOptions());
+ VectorIndexTrainer survivor =
VectorIndexTrainer.create(ivfFlatOptions());
survivor.close();
}
- private static void testVoidEntrypointPanicBecomesRuntimeException() {
- final VectorIndexWriter writer = new
VectorIndexWriter(ivfFlatOptions());
+ private static void testVoidEntrypointErrorBecomesRuntimeException() {
+ final VectorIndexTrainer trainer =
VectorIndexTrainer.create(ivfFlatOptions());
try {
- assertThrows(RuntimeException.class, new ThrowingRunnable() {
- @Override
- public void run() {
- writer.addVectors(new long[] {1L}, new float[] {1.0f}, 1);
- }
- });
+ assertThrowsMessage(
+ RuntimeException.class,
+ "training data length 1 does not match vector count *
dimension 2",
+ new ThrowingRunnable() {
+ @Override
+ public void run() {
+ trainer.addTrainingVectors(new float[] {1.0f}, 2);
+ }
+ });
} finally {
- writer.close();
+ trainer.close();
}
}
private static void testObjectEntrypointPanicBecomesRuntimeException() {
ByteArrayPositionOutputStream output = new
ByteArrayPositionOutputStream();
- VectorIndexWriter writer = new VectorIndexWriter(ivfFlatOptions());
+ VectorIndexWriter writer =
+ new VectorIndexWriter(
+ VectorIndexTrainer.train(ivfFlatOptions(), new float[]
{0.0f, 1.0f}, 2));
try {
- writer.train(new float[] {0.0f, 1.0f}, 2);
writer.addVectors(new long[] {1L, 2L}, new float[] {0.0f, 1.0f},
2);
writer.writeIndex(output);
} finally {
@@ -117,6 +121,24 @@ public class VectorIndexNativePanicBoundaryTest {
throw new AssertionError("expected " + expected.getName());
}
+ private static void assertThrowsMessage(
+ Class<? extends Throwable> expected, String expectedMessage,
ThrowingRunnable runnable) {
+ try {
+ runnable.run();
+ } catch (Throwable t) {
+ if (!expected.isInstance(t)) {
+ throw new AssertionError(
+ "expected " + expected.getName() + " but got " +
t.getClass().getName(), t);
+ }
+ String message = t.getMessage();
+ if (message == null || !message.contains(expectedMessage)) {
+ throw new AssertionError("unexpected exception message: " +
message, t);
+ }
+ return;
+ }
+ throw new AssertionError("expected " + expected.getName());
+ }
+
private static void corruptFirstIvfFlatVector(byte[] indexBytes, float
value) {
int dimension = readIntLe(indexBytes, 8);
int nlist = readIntLe(indexBytes, 12);
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 320f2ec..07db0ca 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
@@ -37,13 +37,15 @@ public class VectorIndexNativeValidationTest {
testWriterValidationComesFromCore();
testWriterRejectsNonFiniteValues();
+ testStagedTrainingRoundtrip();
+ testStagedTrainingStateValidation();
testReaderValidationComesFromCore();
testReaderRejectsNonFiniteQueries();
testSupportedIndexRoundtrips();
}
private static void testWriterValidationComesFromCore() {
- final VectorIndexWriter writer = new
VectorIndexWriter(ivfFlatOptions());
+ final VectorIndexTrainer trainingDataTrainer =
VectorIndexTrainer.create(ivfFlatOptions());
try {
assertThrowsMessage(
RuntimeException.class,
@@ -51,29 +53,42 @@ public class VectorIndexNativeValidationTest {
new ThrowingRunnable() {
@Override
public void run() {
- writer.train(new float[] {0.0f, 1.0f}, 1);
+ trainingDataTrainer.addTrainingVectors(new float[]
{0.0f, 1.0f}, 1);
}
});
+ } finally {
+ trainingDataTrainer.close();
+ }
+
+ final VectorIndexWriter addVectorWriter =
+ newWriter(ivfFlatOptions(), new float[] {0.0f, 1.0f}, 2);
+ try {
assertThrowsMessage(
RuntimeException.class,
"ids length 2 does not match vector count 1",
new ThrowingRunnable() {
@Override
public void run() {
- writer.addVectors(new long[] {1L, 2L}, new float[]
{0.0f}, 1);
+ addVectorWriter.addVectors(new long[] {1L, 2L},
new float[] {0.0f}, 1);
}
});
+ } finally {
+ addVectorWriter.close();
+ }
+
+ final VectorIndexTrainer vectorCountTrainer =
VectorIndexTrainer.create(ivfFlatOptions());
+ try {
assertThrowsMessage(
RuntimeException.class,
"vector count must be greater than 0",
new ThrowingRunnable() {
@Override
public void run() {
- writer.train(new float[0], 0);
+ vectorCountTrainer.addTrainingVectors(new
float[0], 0);
}
});
} finally {
- writer.close();
+ vectorCountTrainer.close();
}
}
@@ -114,7 +129,7 @@ public class VectorIndexNativeValidationTest {
}
private static void testWriterRejectsNonFiniteValues() {
- final VectorIndexWriter trainingWriter = new
VectorIndexWriter(ivfFlatOptions());
+ final VectorIndexTrainer trainingTrainer =
VectorIndexTrainer.create(ivfFlatOptions());
try {
assertThrowsMessage(
RuntimeException.class,
@@ -122,16 +137,16 @@ public class VectorIndexNativeValidationTest {
new ThrowingRunnable() {
@Override
public void run() {
- trainingWriter.train(new float[] {Float.NaN,
1.0f}, 2);
+ trainingTrainer.addTrainingVectors(new float[]
{Float.NaN, 1.0f}, 2);
}
});
} finally {
- trainingWriter.close();
+ trainingTrainer.close();
}
- final VectorIndexWriter vectorWriter = new
VectorIndexWriter(ivfFlatOptions());
+ final VectorIndexWriter vectorWriter =
+ newWriter(ivfFlatOptions(), new float[] {0.0f, 1.0f}, 2);
try {
- vectorWriter.train(new float[] {0.0f, 1.0f}, 2);
assertThrowsMessage(
RuntimeException.class,
"vector data contains non-finite value at offset 0: inf",
@@ -149,6 +164,108 @@ public class VectorIndexNativeValidationTest {
}
}
+ private static void testStagedTrainingRoundtrip() {
+ runStagedTrainingRoundtrip(
+ "ivf_flat", ivfFlatOptions(ROUNDTRIP_DIMENSION,
ROUNDTRIP_NLIST), 0, 0);
+ runStagedTrainingRoundtrip(
+ "ivf_pq", ivfPqOptions(ROUNDTRIP_DIMENSION, ROUNDTRIP_NLIST,
1), 1, 0);
+ runStagedTrainingRoundtrip(
+ "ivf_hnsw_flat",
+ ivfHnswOptions("ivf_hnsw_flat", ROUNDTRIP_DIMENSION,
ROUNDTRIP_NLIST),
+ 0,
+ 4);
+ runStagedTrainingRoundtrip(
+ "ivf_hnsw_sq",
+ ivfHnswOptions("ivf_hnsw_sq", ROUNDTRIP_DIMENSION,
ROUNDTRIP_NLIST),
+ 0,
+ 4);
+ }
+
+ private static void testStagedTrainingStateValidation() {
+ final VectorIndexTrainer emptyTrainer =
VectorIndexTrainer.create(ivfFlatOptions());
+ try {
+ assertThrowsMessage(
+ RuntimeException.class,
+ "no training vectors added",
+ new ThrowingRunnable() {
+ @Override
+ public void run() {
+ emptyTrainer.finishTraining();
+ }
+ });
+ assertThrowsMessage(
+ IllegalStateException.class,
+ "VectorIndexTrainer is closed",
+ new ThrowingRunnable() {
+ @Override
+ public void run() {
+ emptyTrainer.addTrainingVectors(new float[]
{0.0f}, 1);
+ }
+ });
+ } finally {
+ emptyTrainer.close();
+ }
+
+ final VectorIndexTrainer stagedTrainer =
VectorIndexTrainer.create(ivfFlatOptions());
+ final VectorIndexTraining stagedTraining;
+ try {
+ stagedTraining =
+ stagedTrainer.addTrainingVectors(new float[] {0.0f}, 1)
+ .addTrainingVectors(new float[] {1.0f}, 1)
+ .finishTraining();
+ assertThrowsMessage(
+ IllegalStateException.class,
+ "VectorIndexTrainer is closed",
+ new ThrowingRunnable() {
+ @Override
+ public void run() {
+ stagedTrainer.addTrainingVectors(new float[]
{2.0f}, 1);
+ }
+ });
+ } finally {
+ stagedTrainer.close();
+ }
+
+ final VectorIndexWriter writer = new VectorIndexWriter(stagedTraining);
+ try {
+ assertThrowsMessage(
+ IllegalStateException.class,
+ "VectorIndexTraining is closed",
+ new ThrowingRunnable() {
+ @Override
+ public void run() {
+ new VectorIndexWriter(stagedTraining);
+ }
+ });
+ } finally {
+ writer.close();
+ }
+
+ final VectorIndexTrainer invalidBatchTrainer =
VectorIndexTrainer.create(ivfFlatOptions());
+ try {
+ assertThrowsMessage(
+ RuntimeException.class,
+ "training data length 2 does not match vector count *
dimension 1",
+ new ThrowingRunnable() {
+ @Override
+ public void run() {
+ invalidBatchTrainer.addTrainingVectors(new float[]
{0.0f, 1.0f}, 1);
+ }
+ });
+ assertThrowsMessage(
+ RuntimeException.class,
+ "vector count must be greater than 0",
+ new ThrowingRunnable() {
+ @Override
+ public void run() {
+ invalidBatchTrainer.addTrainingVectors(new
float[0], 0);
+ }
+ });
+ } finally {
+ invalidBatchTrainer.close();
+ }
+ }
+
private static void testReaderRejectsNonFiniteQueries() {
VectorIndexReader reader =
new VectorIndexReader(new
ByteArraySeekableInputStream(buildIndexBytes()));
@@ -202,6 +319,19 @@ public class VectorIndexNativeValidationTest {
byte[] indexBytes =
buildIndexBytes(
options, roundtripData(), roundtripIds(),
ROUNDTRIP_VECTOR_COUNT);
+ assertRoundtrip(indexType, indexBytes, expectedPqM, expectedHnswM);
+ }
+
+ private static void runStagedTrainingRoundtrip(
+ String indexType, Map<String, String> options, int expectedPqM,
int expectedHnswM) {
+ byte[] indexBytes =
+ buildStagedIndexBytes(
+ options, roundtripData(), roundtripIds(),
ROUNDTRIP_VECTOR_COUNT);
+ assertRoundtrip(indexType, indexBytes, expectedPqM, expectedHnswM);
+ }
+
+ private static void assertRoundtrip(
+ String indexType, byte[] indexBytes, int expectedPqM, int
expectedHnswM) {
VectorIndexReader reader =
new VectorIndexReader(new
ByteArraySeekableInputStream(indexBytes));
try {
@@ -237,10 +367,9 @@ public class VectorIndexNativeValidationTest {
private static byte[] buildIndexBytes(
Map<String, String> options, float[] data, long[] ids, int
vectorCount) {
- VectorIndexWriter writer = new VectorIndexWriter(options);
+ VectorIndexWriter writer = newWriter(options, data, vectorCount);
ByteArrayPositionOutputStream output = new
ByteArrayPositionOutputStream();
try {
- writer.train(data, vectorCount);
writer.addVectors(ids, data, vectorCount);
writer.writeIndex(output);
return output.toByteArray();
@@ -249,6 +378,43 @@ public class VectorIndexNativeValidationTest {
}
}
+ private static byte[] buildStagedIndexBytes(
+ Map<String, String> options, float[] data, long[] ids, int
vectorCount) {
+ VectorIndexTrainer trainer = VectorIndexTrainer.create(options);
+ VectorIndexWriter writer = null;
+ ByteArrayPositionOutputStream output = new
ByteArrayPositionOutputStream();
+ int dimension = data.length / vectorCount;
+ try {
+ int offset = 0;
+ while (offset < vectorCount) {
+ int batchCount = Math.min(ROUNDTRIP_PER_LIST / 2, vectorCount
- offset);
+ trainer.addTrainingVectors(
+ copyVectors(data, dimension, offset, batchCount),
batchCount);
+ offset += batchCount;
+ }
+ writer = new VectorIndexWriter(trainer.finishTraining());
+ writer.addVectors(ids, data, vectorCount);
+ writer.writeIndex(output);
+ return output.toByteArray();
+ } finally {
+ trainer.close();
+ if (writer != null) {
+ writer.close();
+ }
+ }
+ }
+
+ private static VectorIndexWriter newWriter(
+ Map<String, String> options, float[] data, int vectorCount) {
+ return new VectorIndexWriter(VectorIndexTrainer.train(options, data,
vectorCount));
+ }
+
+ private static float[] copyVectors(float[] data, int dimension, int
offset, int count) {
+ float[] copy = new float[count * dimension];
+ System.arraycopy(data, offset * dimension, copy, 0, copy.length);
+ return copy;
+ }
+
private static Map<String, String> ivfFlatOptions() {
return ivfFlatOptions(1, 1);
}
diff --git a/jni/src/lib.rs b/jni/src/lib.rs
index 186fa84..82f92c7 100644
--- a/jni/src/lib.rs
+++ b/jni/src/lib.rs
@@ -21,8 +21,8 @@ use jni::objects::{JByteArray, JClass, JFloatArray,
JLongArray, JObject, JValue}
use jni::sys::{jint, jlong, jobject, jobjectArray};
use jni::JNIEnv;
use paimon_vindex_core::index::{
- VectorIndexConfig, VectorIndexMetadata, VectorIndexReader,
VectorIndexWriter,
- VectorSearchParams,
+ VectorIndexConfig, VectorIndexMetadata, VectorIndexReader,
VectorIndexTrainer,
+ VectorIndexTraining, VectorIndexWriter, VectorSearchParams,
};
use std::any::Any;
use std::collections::HashMap;
@@ -63,11 +63,75 @@ fn throw_panic_and_return<T: Default>(env: &mut JNIEnv,
payload: &(dyn Any + Sen
throw_and_return(env, &format!("Rust panic in JNI call: {}", payload))
}
-fn deref_writer(ptr: jlong) -> Option<&'static mut VectorIndexWriter> {
+struct JniVectorIndexTrainer {
+ trainer: Option<VectorIndexTrainer>,
+}
+
+impl JniVectorIndexTrainer {
+ fn new(trainer: VectorIndexTrainer) -> Self {
+ Self {
+ trainer: Some(trainer),
+ }
+ }
+
+ fn trainer_mut(&mut self) -> Result<&mut VectorIndexTrainer, String> {
+ self.trainer
+ .as_mut()
+ .ok_or_else(|| "trainer has already finished".to_string())
+ }
+
+ fn take(&mut self) -> Result<VectorIndexTrainer, String> {
+ self.trainer
+ .take()
+ .ok_or_else(|| "trainer has already finished".to_string())
+ }
+}
+
+struct JniVectorIndexTraining {
+ training: Option<VectorIndexTraining>,
+}
+
+impl JniVectorIndexTraining {
+ fn new(training: VectorIndexTraining) -> Self {
+ Self {
+ training: Some(training),
+ }
+ }
+
+ fn take(&mut self) -> Result<VectorIndexTraining, String> {
+ self.training
+ .take()
+ .ok_or_else(|| "training has already been consumed".to_string())
+ }
+}
+
+struct JniVectorIndexWriter {
+ writer: VectorIndexWriter,
+}
+
+impl JniVectorIndexWriter {
+ fn new(writer: VectorIndexWriter) -> Self {
+ Self { writer }
+ }
+
+ fn dimension(&self) -> usize {
+ self.writer.dimension()
+ }
+}
+
+fn deref_trainer(ptr: jlong) -> Option<&'static mut JniVectorIndexTrainer> {
if ptr == 0 {
None
} else {
- Some(unsafe { &mut *(ptr as *mut VectorIndexWriter) })
+ Some(unsafe { &mut *(ptr as *mut JniVectorIndexTrainer) })
+ }
+}
+
+fn deref_writer(ptr: jlong) -> Option<&'static mut JniVectorIndexWriter> {
+ if ptr == 0 {
+ None
+ } else {
+ Some(unsafe { &mut *(ptr as *mut JniVectorIndexWriter) })
}
}
@@ -170,6 +234,9 @@ fn read_byte_array(env: &mut JNIEnv, array: JByteArray) ->
Result<Vec<u8>, Strin
}
fn read_float_array(env: &mut JNIEnv, array: &JFloatArray, name: &str) ->
Result<Vec<f32>, String> {
+ if array.as_raw().is_null() {
+ return Err(format!("{} float array is null", name));
+ }
let len = env
.get_array_length(array)
.map_err(|e| format!("get_array_length({}): {}", name, e))? as usize;
@@ -180,6 +247,9 @@ fn read_float_array(env: &mut JNIEnv, array: &JFloatArray,
name: &str) -> Result
}
fn read_long_array(env: &mut JNIEnv, array: &JLongArray, name: &str) ->
Result<Vec<i64>, String> {
+ if array.as_raw().is_null() {
+ return Err(format!("{} long array is null", name));
+ }
let len = env
.get_array_length(array)
.map_err(|e| format!("get_array_length({}): {}", name, e))? as usize;
@@ -311,10 +381,10 @@ fn search_params(k: jint, nprobe: jint, ef_search: jint)
-> Option<VectorSearchP
}
}
-// --- Unified Writer API ---
+// --- Unified Trainer / Writer API ---
#[no_mangle]
-pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_createWriter(
+pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_createTrainer(
env: JNIEnv,
_class: JClass,
keys: jobjectArray,
@@ -326,16 +396,35 @@ pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_cre
None => return 0,
};
- let writer = match VectorIndexWriter::new(config) {
- Ok(writer) => writer,
- Err(e) => return throw_and_return(env, &format!("create writer:
{}", e)),
+ let trainer = match VectorIndexTrainer::new(config) {
+ Ok(trainer) => trainer,
+ Err(e) => return throw_and_return(env, &format!("create trainer:
{}", e)),
+ };
+ Box::into_raw(Box::new(JniVectorIndexTrainer::new(trainer))) as jlong
+ })
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_trainerDimension(
+ env: JNIEnv,
+ _class: JClass,
+ ptr: jlong,
+) -> jint {
+ jni_call(env, |env| {
+ let trainer = match deref_trainer(ptr) {
+ Some(trainer) => trainer,
+ None => return throw_and_return(env, "null native pointer (trainer
already freed?)"),
};
- Box::into_raw(Box::new(writer)) as jlong
+ let trainer = match trainer.trainer_mut() {
+ Ok(trainer) => trainer,
+ Err(e) => return throw_and_return(env, &e),
+ };
+ trainer.dimension() as jint
})
}
#[no_mangle]
-pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_train(
+pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_trainerAddTrainingVectors(
env: JNIEnv,
_class: JClass,
ptr: jlong,
@@ -343,10 +432,6 @@ pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_tra
n: jint,
) {
jni_call_void(env, |env| {
- let writer = match deref_writer(ptr) {
- Some(writer) => writer,
- None => return throw_and_return(env, "null native pointer (writer
already freed?)"),
- };
if n < 0 {
return throw_and_return(env, &format!("invalid vector count: {}",
n));
}
@@ -355,8 +440,90 @@ pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_tra
Ok(buf) => buf,
Err(e) => return throw_and_return(env, &e),
};
- if let Err(e) = writer.train(&data_buf, n) {
- throw_and_return::<()>(env, &format!("train: {}", e));
+ let trainer = match deref_trainer(ptr) {
+ Some(trainer) => trainer,
+ None => return throw_and_return(env, "null native pointer (trainer
already freed?)"),
+ };
+ let trainer = match trainer.trainer_mut() {
+ Ok(trainer) => trainer,
+ Err(e) => return throw_and_return(env, &e),
+ };
+ if let Err(e) = trainer.add_training_vectors_mut(&data_buf, n) {
+ throw_and_return::<()>(env, &e.to_string());
+ }
+ })
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_trainerFinishTraining(
+ env: JNIEnv,
+ _class: JClass,
+ ptr: jlong,
+) -> jlong {
+ jni_call(env, |env| {
+ if ptr == 0 {
+ return throw_and_return(env, "null native pointer (trainer already
freed?)");
+ }
+ let mut trainer_handle = unsafe { Box::from_raw(ptr as *mut
JniVectorIndexTrainer) };
+ let trainer = match trainer_handle.take() {
+ Ok(trainer) => trainer,
+ Err(e) => return throw_and_return(env, &e),
+ };
+ let training = match trainer.finish() {
+ Ok(training) => training,
+ Err(e) => return throw_and_return(env, &format!("finishTraining:
{}", e)),
+ };
+ Box::into_raw(Box::new(JniVectorIndexTraining::new(training))) as jlong
+ })
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_freeTrainer(
+ env: JNIEnv,
+ _class: JClass,
+ ptr: jlong,
+) {
+ jni_call_void(env, |_env| {
+ if ptr != 0 {
+ unsafe {
+ drop(Box::from_raw(ptr as *mut JniVectorIndexTrainer));
+ }
+ }
+ })
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_createWriter(
+ 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 mut training_handle =
+ unsafe { Box::from_raw(training_ptr as *mut
JniVectorIndexTraining) };
+ let training = match training_handle.take() {
+ Ok(training) => training,
+ Err(e) => return throw_and_return(env, &e),
+ };
+ let writer = VectorIndexWriter::new(training);
+ Box::into_raw(Box::new(JniVectorIndexWriter::new(writer))) as jlong
+ })
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_freeTraining(
+ env: JNIEnv,
+ _class: JClass,
+ ptr: jlong,
+) {
+ jni_call_void(env, |_env| {
+ if ptr != 0 {
+ unsafe {
+ drop(Box::from_raw(ptr as *mut JniVectorIndexTraining));
+ }
}
})
}
@@ -402,7 +569,7 @@ pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_add
Ok(buf) => buf,
Err(e) => return throw_and_return(env, &e),
};
- if let Err(e) = writer.add_vectors(&id_buf, &data_buf, n) {
+ if let Err(e) = writer.writer.add_vectors(&id_buf, &data_buf, n) {
throw_and_return::<()>(env, &format!("add_vectors: {}", e));
}
})
@@ -431,7 +598,7 @@ pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_wri
};
let mut output = JniOutputStream::new(jvm, global_ref);
- if let Err(e) = writer.write(&mut output) {
+ if let Err(e) = writer.writer.write(&mut output) {
throw_and_return::<()>(env, &format!("write index: {}", e));
}
})
@@ -446,7 +613,7 @@ pub extern "system" fn
Java_org_apache_paimon_index_vector_VectorIndexNative_fre
jni_call_void(env, |_env| {
if ptr != 0 {
unsafe {
- drop(Box::from_raw(ptr as *mut VectorIndexWriter));
+ drop(Box::from_raw(ptr as *mut JniVectorIndexWriter));
}
}
})
diff --git a/python/paimon_vindex/__init__.py b/python/paimon_vindex/__init__.py
index ddaedfa..f66c47c 100644
--- a/python/paimon_vindex/__init__.py
+++ b/python/paimon_vindex/__init__.py
@@ -105,37 +105,94 @@ def _bytes_buffer(value, name):
return buf, len(data), data
-class VectorIndexWriter:
+def _option_arrays(options: Mapping[str, str]):
+ option_items = list(options.items())
+ key_bytes = []
+ value_bytes = []
+ for key, value in option_items:
+ if not isinstance(key, str) or not isinstance(value, str):
+ raise ValueError("options must be a mapping of str to str")
+ key_bytes.append(key.encode("utf-8"))
+ value_bytes.append(value.encode("utf-8"))
+ keys = (ctypes.c_char_p * len(key_bytes))(*key_bytes)
+ values = (ctypes.c_char_p * len(value_bytes))(*value_bytes)
+ return option_items, key_bytes, value_bytes, keys, values
+
+
+class VectorIndexTraining:
+ def __init__(self, handle):
+ self._closed = False
+ self._handle = handle
+
+ def _require_open(self):
+ if self._closed or not self._handle:
+ raise RuntimeError("VectorIndexTraining is closed")
+
+ def _take_handle(self):
+ self._require_open()
+ handle = self._handle
+ self._handle = None
+ self._closed = True
+ return handle
+
+ def close(self):
+ if self._handle:
+ lib.paimon_vindex_training_free(self._handle)
+ self._handle = None
+ self._closed = True
+
+ def __enter__(self):
+ self._require_open()
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ self.close()
+ return False
+
+ def __del__(self):
+ try:
+ self.close()
+ except Exception:
+ pass
+
+
+class VectorIndexTrainer:
def __init__(self, options: Mapping[str, str]):
self._closed = False
- option_items = list(options.items())
- self._key_bytes = []
- self._value_bytes = []
- for key, value in option_items:
- if not isinstance(key, str) or not isinstance(value, str):
- raise ValueError("options must be a mapping of str to str")
- self._key_bytes.append(key.encode("utf-8"))
- self._value_bytes.append(value.encode("utf-8"))
- self._keys = (ctypes.c_char_p * len(self._key_bytes))(*self._key_bytes)
- self._values = (ctypes.c_char_p *
len(self._value_bytes))(*self._value_bytes)
- self._handle = lib.paimon_vindex_writer_open(
+ (
+ option_items,
+ self._key_bytes,
+ self._value_bytes,
+ self._keys,
+ self._values,
+ ) = _option_arrays(options)
+ self._handle = lib.paimon_vindex_trainer_open(
self._keys,
self._values,
len(option_items),
)
if not self._handle:
- _check_error("failed to open writer")
+ _check_error("failed to open trainer")
self._dimension = self._read_dimension()
+ @classmethod
+ def create(cls, options: Mapping[str, str]):
+ return cls(options)
+
+ @classmethod
+ def train(cls, options: Mapping[str, str], data):
+ with cls(options) as trainer:
+ return trainer.add_training_vectors(data).finish_training()
+
def _require_open(self):
if self._closed or not self._handle:
- raise RuntimeError("VectorIndexWriter is closed")
+ raise RuntimeError("VectorIndexTrainer is closed")
def _read_dimension(self):
out = ctypes.c_size_t(0)
- rc = lib.paimon_vindex_writer_dimension(self._handle,
ctypes.byref(out))
+ rc = lib.paimon_vindex_trainer_dimension(self._handle,
ctypes.byref(out))
if rc != 0:
- _check_error("writer dimension failed")
+ _check_error("trainer dimension failed")
return out.value
@property
@@ -143,7 +200,7 @@ class VectorIndexWriter:
self._require_open()
return self._dimension
- def train(self, data):
+ def add_training_vectors(self, data):
self._require_open()
data = _float32_matrix(data, "data")
if data.shape[1] != self._dimension:
@@ -151,13 +208,74 @@ class VectorIndexWriter:
f"training data length {data.size} does not match vector count
"
f"* dimension {data.shape[0] * self._dimension}"
)
- rc = lib.paimon_vindex_writer_train(
+ rc = lib.paimon_vindex_trainer_add_training_vectors(
self._handle,
data.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
data.shape[0],
)
if rc != 0:
- _check_error("train failed")
+ _check_error("add training vectors failed")
+ return self
+
+ def finish_training(self):
+ self._require_open()
+ handle = self._handle
+ training = lib.paimon_vindex_trainer_finish(handle)
+ lib.paimon_vindex_trainer_free(handle)
+ self._handle = None
+ self._closed = True
+ if not training:
+ _check_error("finish training failed")
+ return VectorIndexTraining(training)
+
+ def close(self):
+ if self._handle:
+ lib.paimon_vindex_trainer_free(self._handle)
+ self._handle = None
+ self._closed = True
+
+ def __enter__(self):
+ self._require_open()
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ self.close()
+ return False
+
+ def __del__(self):
+ try:
+ self.close()
+ except Exception:
+ pass
+
+
+class VectorIndexWriter:
+ def __init__(self, training: VectorIndexTraining):
+ if not isinstance(training, VectorIndexTraining):
+ raise TypeError("training must be a VectorIndexTraining")
+ self._closed = False
+ training_handle = training._take_handle()
+ self._handle = lib.paimon_vindex_writer_open(training_handle)
+ lib.paimon_vindex_training_free(training_handle)
+ if not self._handle:
+ _check_error("failed to open writer")
+ self._dimension = self._read_dimension()
+
+ def _require_open(self):
+ if self._closed or not self._handle:
+ raise RuntimeError("VectorIndexWriter is closed")
+
+ def _read_dimension(self):
+ out = ctypes.c_size_t(0)
+ rc = lib.paimon_vindex_writer_dimension(self._handle,
ctypes.byref(out))
+ if rc != 0:
+ _check_error("writer dimension failed")
+ return out.value
+
+ @property
+ def dimension(self):
+ self._require_open()
+ return self._dimension
def add_vectors(self, ids, data):
self._require_open()
@@ -415,5 +533,7 @@ class VectorIndexReader:
__all__ = [
"VectorIndexMetadata",
"VectorIndexReader",
+ "VectorIndexTrainer",
+ "VectorIndexTraining",
"VectorIndexWriter",
]
diff --git a/python/paimon_vindex/_ffi.py b/python/paimon_vindex/_ffi.py
index 72cc25c..38c0ea7 100644
--- a/python/paimon_vindex/_ffi.py
+++ b/python/paimon_vindex/_ffi.py
@@ -121,11 +121,33 @@ class PaimonVindexMetadata(Structure):
lib.paimon_vindex_last_error.argtypes = []
lib.paimon_vindex_last_error.restype = c_char_p
-lib.paimon_vindex_writer_open.argtypes = [
+lib.paimon_vindex_trainer_open.argtypes = [
POINTER(c_char_p),
POINTER(c_char_p),
c_size_t,
]
+lib.paimon_vindex_trainer_open.restype = c_void_p
+
+lib.paimon_vindex_trainer_free.argtypes = [c_void_p]
+lib.paimon_vindex_trainer_free.restype = None
+
+lib.paimon_vindex_trainer_dimension.argtypes = [c_void_p, POINTER(c_size_t)]
+lib.paimon_vindex_trainer_dimension.restype = c_int
+
+lib.paimon_vindex_trainer_add_training_vectors.argtypes = [
+ c_void_p,
+ POINTER(c_float),
+ c_size_t,
+]
+lib.paimon_vindex_trainer_add_training_vectors.restype = c_int
+
+lib.paimon_vindex_trainer_finish.argtypes = [c_void_p]
+lib.paimon_vindex_trainer_finish.restype = c_void_p
+
+lib.paimon_vindex_training_free.argtypes = [c_void_p]
+lib.paimon_vindex_training_free.restype = None
+
+lib.paimon_vindex_writer_open.argtypes = [c_void_p]
lib.paimon_vindex_writer_open.restype = c_void_p
lib.paimon_vindex_writer_free.argtypes = [c_void_p]
@@ -134,9 +156,6 @@ lib.paimon_vindex_writer_free.restype = None
lib.paimon_vindex_writer_dimension.argtypes = [c_void_p, POINTER(c_size_t)]
lib.paimon_vindex_writer_dimension.restype = c_int
-lib.paimon_vindex_writer_train.argtypes = [c_void_p, POINTER(c_float),
c_size_t]
-lib.paimon_vindex_writer_train.restype = c_int
-
lib.paimon_vindex_writer_add_vectors.argtypes = [
c_void_p,
POINTER(c_int64),
diff --git a/python/tests/test_vindex.py b/python/tests/test_vindex.py
index 8380c84..a6e79bc 100644
--- a/python/tests/test_vindex.py
+++ b/python/tests/test_vindex.py
@@ -20,7 +20,7 @@ import io
import numpy as np
import pytest
-from paimon_vindex import VectorIndexReader, VectorIndexWriter
+from paimon_vindex import VectorIndexReader, VectorIndexTrainer,
VectorIndexWriter
class VectorIndexInput:
@@ -44,8 +44,8 @@ def build_index(options, d, n=512):
data = clustered_data(n, d, int(options.get("nlist", "4")))
ids = np.arange(n, dtype=np.int64)
output = io.BytesIO()
- with VectorIndexWriter(options) as writer:
- writer.train(data)
+ training = VectorIndexTrainer.train(options, data)
+ with VectorIndexWriter(training) as writer:
writer.add_vectors(ids, data)
writer.write(output)
return output.getvalue(), data
@@ -151,15 +151,16 @@ def test_python_ffi_delegates_validation():
"pq.m": "4",
"metric": "l2",
}
- writer = VectorIndexWriter(options)
- with pytest.raises(RuntimeError, match="training data length 17"):
- writer.train(np.zeros((1, 17), dtype=np.float32))
+ with VectorIndexTrainer.create(options) as trainer:
+ with pytest.raises(RuntimeError, match="training data length 17"):
+ trainer.add_training_vectors(np.zeros((1, 17), dtype=np.float32))
data = np.zeros((1, 16), dtype=np.float32)
ids = np.array([1, 2], dtype=np.int64)
- with pytest.raises(RuntimeError, match="ids length 2 does not match vector
count 1"):
- writer.add_vectors(ids, data)
- writer.close()
+ training = VectorIndexTrainer.train(options, data)
+ with VectorIndexWriter(training) as writer:
+ with pytest.raises(RuntimeError, match="ids length 2 does not match
vector count 1"):
+ writer.add_vectors(ids, data)
index_bytes, data = build_index(options, 16)
with reader_from_bytes(index_bytes) as reader:
