Skip to content

Support batched vector index training#52

Merged
JingsongLi merged 12 commits into
apache:mainfrom
jerry-024:fix_paimon_vector_index_train_limit
Jun 30, 2026
Merged

Support batched vector index training#52
JingsongLi merged 12 commits into
apache:mainfrom
jerry-024:fix_paimon_vector_index_train_limit

Conversation

@jerry-024

@jerry-024 jerry-024 commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

What is changed

  • Add Java/JNI staged training APIs: VectorIndexWriter.addTrainingVectors(...) and finishTraining().
  • Wrap the native writer with an explicit training state machine so invalid direct-training, staged-training, add, and write transitions fail with clear errors.
  • Validate null JNI arrays and staged training batches before buffering them, then release accumulated training data after finishTraining() completes or fails.
  • Extend Java native validation tests with staged training roundtrip and state-transition coverage.

Why

This lets Java callers feed training vectors in bounded batches instead of materializing one large float[] for a single train(...) call.

@jerry-024 jerry-024 marked this pull request as draft June 25, 2026 02:42
@jerry-024 jerry-024 changed the title [feature] support train batch read data [java] Support batched vector index training Jun 25, 2026
@jerry-024 jerry-024 marked this pull request as ready for review June 29, 2026 08:16
@JingsongLi

Copy link
Copy Markdown
Contributor

Thanks for adding staged training. Before we merge this API shape, I think it is worth considering a cleaner split between training and writing, and doing that split consistently in Rust core and the Java/JNI layer.

If backward compatibility is not a blocker, I would prefer not to keep training as part of VectorIndexWriter. Conceptually, training and writing are two different phases:

  1. A trainer/builder owns the config and training data batches, then produces a trained state.
  2. A writer owns a trained state, accepts vectors, and serializes the final index.

Concretely, the Rust core could expose something like:

let trained = VectorIndexTrainer::new(config)?
    .add_training_vectors(batch1, n1)?
    .add_training_vectors(batch2, n2)?
    .finish()?;

let mut writer = VectorIndexWriter::new(trained);
writer.add_vectors(&ids, &vectors, vector_count)?;
writer.write(&mut output)?;

The Java/JNI API could mirror the same split:

VectorIndexTraining training =
        VectorIndexTrainer.create(options)
                .addTrainingVectors(batch1, n1)
                .addTrainingVectors(batch2, n2)
                .finishTraining();

try (VectorIndexWriter writer = new VectorIndexWriter(training)) {
    writer.addVectors(rowIds, vectors, vectorCount);
    writer.writeIndex(output);
}

A single-shot helper can still live on the trainer, for example VectorIndexTrainer.train(options, data, vectorCount), but VectorIndexWriter itself would no longer have a train() method or a staged-training lifecycle.

This would make the API easier to reason about and remove the need for a writer-side state machine that mixes NotTrained, CollectingTraining, Trained, and AddingOrWritten. It also prevents invalid interleavings by construction: once a writer exists, training is already complete, and the writer only has to care about adding vectors and writing the index.

@JingsongLi JingsongLi changed the title [java] Support batched vector index training Support batched vector index training Jun 30, 2026
@JingsongLi

Copy link
Copy Markdown
Contributor

+1

@JingsongLi JingsongLi merged commit 35c461e into apache:main Jun 30, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants