feat: distributed centroid training by summaryzb · Pull Request #7321 · lance-format/lance

summaryzb · 2026-06-17T10:04:36Z

Summary

This PR adds scheduler-neutral distributed IVF centroid training primitives so Spark, Ray, or custom RPC systems can train one global centroid set across many Lance fragments. It introduces a shared Arrow PartialStats wire format in Rust, dataset-aware IVF wrapp
ers, Python and Java bindings, and documentation for driving Round-0 sampling plus iterative Lloyd rounds outside Lance.

Problem

Lance already supports distributed index segment construction, but callers did not have a reusable way to train shared IVF centroids across workers. That forced distributed builders either to centralize training data on one process or to train independent per-segm
ent centroid models, which prevents a single global model from being reused across worker-built vector index segments. The missing piece was a small, stable API surface for the k-means math itself: sample local data, compute per-worker E-step statistics, merge tho
se statistics, and finalize the next centroids.

Approach

The implementation splits centroid training into two layers. lance_index::vector::kmeans::distributed provides pure Arrow-native primitives: PartialStats, compute_partial_stats, merge_partial_stats, reduce_partial_stats, finalize_centroids, reservoir s
ampling, and driver-side centroid initialization. PartialStats is represented as a fixed-schema RecordBatch with metadata for version, shape, distance type, and a centroid fingerprint so reducers reject stats from incompatible rounds.

lance::index::vector::ivf::distributed adds dataset-aware async wrappers that scan optional fragment slices, normalize cosine input when needed, filter non-finite training data, and delegate to the pure primitives. The existing streaming IVF refinement path is r
efactored to use these same partial-stats primitives, keeping the new distributed path aligned with the in-process training behavior. Python and Java expose the same operations through Arrow IPC-friendly APIs while leaving broadcast, tree reduction, convergence ch
ecks, and scheduling to the caller.

Changes

rust/lance-index/src/vector/kmeans/distributed.rs: adds the core distributed k-means primitives, PartialStats schema/metadata validation, centroid fingerprinting, E-step accumulation, merge/reduce, M-step finalization, reservoir sampling, and bootstrap helpe
rs.
rust/lance-index/src/vector/kmeans.rs, rust/lance-index/Cargo.toml, Cargo.toml, Cargo.lock: wires the distributed k-means module into lance-index and adds the sha2 dependency used for centroid fingerprints.
rust/lance/src/index/vector/ivf/distributed.rs: adds dataset-aware distributed IVF training wrappers for worker-side sampling and partial-stat computation over optional fragment id slices.
rust/lance/src/index/vector/ivf.rs: refactors streaming IVF refinement to accumulate and finalize centroids through the new PartialStats primitives and adds a regression test against the legacy sum/weight calculation.
python/src/indices.rs, python/python/lance/indices/distributed_kmeans.py, python/python/lance/indices/__init__.py: registers PyO3 bindings and a Python facade for sampling, partial-stat computation, merge/reduce, centroid finalization, and centroid initial
ization.
java/lance-jni/src/distributed_kmeans.rs, java/lance-jni/src/lib.rs, java/src/main/java/org/lance/index/vector/DistributedKMeans.java: adds JNI and Java public APIs using Arrow IPC for samples/stats and flat float arrays for centroid outputs.
docs/src/guide/distributed_indexing.md: documents distributed centroid training roles, the Python workflow, and the Arrow RecordBatch interchange contract.

Test Coverage

Rust unit tests cover PartialStats schema/version validation, centroid fingerprint determinism, basic E-step counts/loss, Hamming rejection, all-NaN/empty handling, merge identity and fingerprint mismatch rejection, centroid finalization including empty-cluste
r fallback, deterministic reservoir sampling, initial centroid selection, and bootstrap k-means.
Rust property tests cover merge commutativity and associativity across random splits.
Rust IVF tests cover dataset-aware Round-0 sampling, layer-2 partial-stat computation, a simulated four-worker end-to-end training loop compared with single-machine IVF training, and streaming refinement parity with the legacy accumulation path.
Python tests cover the public lance.indices.distributed_kmeans round trip, Arrow IPC serialization of partial stats, merge behavior after IPC restore, and initial centroid selection shape.
Java JUnit tests cover sample/bootstrap/partial-stats/reduce/finalize flow through the public Java API and JNI Arrow IPC bridge.

Squashed commits: - docs(spec): distributed IVF centroid training abstraction - docs(plan): distributed IVF centroid training implementation plan - feat(kmeans): scaffold distributed kmeans submodule - feat(kmeans): add PartialStats Arrow wrapper - feat(kmeans): centroids_fingerprint helper for distributed stats - feat(kmeans): compute_partial_stats E-step kernel - feat(kmeans): merge_partial_stats / reduce_partial_stats - feat(kmeans): finalize_centroids M-step - test(kmeans): I1/I2/I4 invariants for distributed kmeans - feat(kmeans): reservoir sample + driver-side bootstrap helpers - refactor(ivf): streaming refine uses distributed kmeans primitives - feat(ivf): scaffold layer-2 distributed module - feat(ivf): sample_round_0 + compute_partial_stats layer-2 wrappers - test(ivf): end-to-end 4-worker distributed centroid training - feat(python): pyo3 bindings for distributed kmeans primitives - feat(python): distributed kmeans facade + ipc round-trip tests - feat(java): jni shim for distributed kmeans - feat(java): DistributedKMeans public api + junit round-trip - docs: distributed centroid training user guide - chore: cargo fmt --all after distributed kmeans rollout Change-Id: Ifff14151f0c021c502ad7619de6aacc9fdd6daae

codecov · 2026-06-17T10:46:52Z

Codecov Report

❌ Patch coverage is 84.92377% with 267 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
rust/lance-index/src/vector/kmeans/distributed.rs	84.08%	170 Missing and 27 partials ⚠️
rust/lance/src/index/vector/ivf.rs	63.88%	26 Missing ⚠️
rust/lance/src/index/vector/ivf/distributed.rs	92.17%	9 Missing and 14 partials ⚠️
rust/lance/src/index/vector/utils.rs	86.75%	12 Missing and 8 partials ⚠️
rust/lance-index/src/vector/kmeans.rs	93.75%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

Change-Id: I077166ef875c2ce9cc966b48c20482e12df9b629

Xuanwo · 2026-06-20T02:03:11Z

+    let combined = concat_samples(samples)?;
+    let dim = combined.value_length() as usize;
+    let params = KMeansParams::default()
+        .with_distance_type(distance_type)


Cosine bootstrap passes the cosine metric into the existing float kmeans path; ordinary-size runs can hit the unsupported-metric branch and panic instead of returning centroids.

fixed inside bootstrap_centroids. When distance_type == Cosine we now L2-normalize the combined samples up front, dispatch the inner k-means with DistanceType::L2 (mirroring the Layer-1 contract in compute_partial_stats), and then re-normalize the returned means so callers still see unit-norm centroids

Xuanwo · 2026-06-20T02:03:11Z

+    }
+
+    /// Validate and adopt an externally-built RecordBatch.
+    pub fn from_record_batch(batch: RecordBatch) -> Result<Self> {


Externally supplied stats are accepted after only version and column-name checks, while later merge/finalize paths unwrap concrete column types; malformed Arrow IPC can panic instead of returning a validation error.

Xuanwo · 2026-06-20T02:03:11Z

+    stats: &PartialStats,
+    prev: &FixedSizeListArray,
+) -> Result<FixedSizeListArray> {
+    if stats.k() != prev.len() {


Finalization never checks that the stats fingerprint matches the previous centroids, so stale stats from another training round can be accepted and produce centroids for the wrong assignments.

The fingerprint is the same 8-byte SHA-256 prefix already used to gate merge_partial_stats, so the trust boundary is consistent on both the merge and finalize paths

Xuanwo · 2026-06-20T02:03:11Z

+) -> Result<PartialStats> {
+    // Pull all matching rows; the worker is expected to have a small enough fragment
+    // slice that scanning it whole is cheap.
+    let raw = maybe_sample_training_data(dataset, column, usize::MAX, fragment_ids).await?;


The worker E-step materializes the entire fragment slice before computing stats, so large shards can create memory usage proportional to all rows times dimension instead of the bounded partial-stats state.

Pick memory is O(k·d) instead of O(N·d). Empty input is handled by returning PartialStats::empty(k, dim, distance_type, fp) with the right fingerprint so downstream merge/finalize stays uniform

Xuanwo · 2026-06-20T02:03:12Z

+) -> Result<RecordBatch> {
+    // Round-0 oversamples to give the driver-side bootstrap enough material.
+    let raw =
+        maybe_sample_training_data(dataset, column, target.saturating_mul(2), fragment_ids).await?;


The public seed only controls the second reservoir pass; the earlier dataset sampling remains nondeterministic, so retrying the same worker with the same seed can train from different samples.

The seed is now plumbed all the way through

Xuanwo · 2026-06-20T02:03:12Z

+    env: &mut JNIEnv<'a>,
+    fsl: &FixedSizeListArray,
+) -> Result<jni::objects::JFloatArray<'a>> {
+    let values = fsl.values().as_primitive::<Float32Type>().values().to_vec();


The Java return path unconditionally reads centroid values as Float32 even though the Rust primitive can return Float16 or Float64, so supported vector dtypes can fail across JNI.

now return jbyteArray carrying an Arrow-IPC RecordBatch whose single column is FixedSizeList<original_dtype>, so callers can dispatch on it instead of silently downcasting

Xuanwo · 2026-06-20T02:03:12Z

+            &dataset_guard.inner,
+            &column,
+            frags.as_deref(),
+            target as usize,


A negative Java target is cast to a huge unsigned size before sampling, so invalid input can turn Round-0 sampling into a full-slice scan and return.

nativeSampleRound0 now validates target < 0 up front

Change-Id: I44ca6c6bacf77527caf02ccca94ca7c97d4fb4a1

Change-Id: I343fbe5a53c9d10a266f19496133a36548d6f47a

Xuanwo · 2026-07-08T18:02:39Z

+        }
+
+        let batch_idx = self.held.len();
+        self.held.push(chunk.clone());


The streaming reservoir still stores every input chunk even after selected rows from that chunk are replaced, so Round-0 sampling can retain the full streamed vector payload in memory.

Xuanwo · 2026-07-08T18:02:40Z

+            DistanceType::L2,
+        )
+    } else {
+        (combined, distance_type)


Hamming requests still flow into the float k-means bootstrap path, so a public distributed bootstrap call can panic instead of returning a validation error.

Xuanwo · 2026-07-08T18:02:40Z

+            "select/bootstrap requires at least one sample batch",
+        ));
+    }
+    let arrays: Vec<&dyn Array> = samples.iter().map(|b| b.column(0).as_ref()).collect();


Externally supplied sample batches are indexed before validating that they contain a vector column, so malformed IPC can panic the driver process instead of producing an error.

Xuanwo · 2026-07-08T18:02:40Z

+    fragment_ids: Option<&[u32]>,
+    seed: Option<u64>,
+) -> Result<Pin<Box<dyn Stream<Item = Result<RecordBatch>> + Send>>> {
+    let num_rows = count_rows(dataset, fragment_ids).await?;


Duplicate fragment filters are counted before the later deduplication path, so a repeated id can hit the ordered-fragment assertion and panic public distributed calls.

Xuanwo · 2026-07-08T18:02:40Z

+    a.values()
+        .iter()
+        .zip(b.values().iter())
+        .map(|(x, y)| x + y)


Merging externally supplied stats still adds counts with unchecked u64 arithmetic, so oversized counts can panic in debug builds or wrap in release and corrupt centroid averages.

Xuanwo · 2026-07-08T18:02:40Z

+}
+
+fn array_data_to_fsl(array: ArrayData) -> PyResult<FixedSizeListArray> {
+    Ok(FixedSizeListArray::from(array))


Python centroid input is converted with a panic-on-wrong-type Arrow constructor, so malformed user input can abort the process instead of raising a Python error.

sezruby · 2026-07-30T20:52:24Z

Thanks for this @summaryzb — the scheduler-neutral PartialStats design (Arrow wire format + centroid fingerprint so reducers reject mismatched stats) is exactly what distributed IVF training needs, and I'd like to build a Spark orchestration layer on top of it.

One suggestion that might help it land faster: would you be open to splitting it into stacked PRs? At +3959 / 20 files across four surfaces (Rust core, IVF integration, Java, Python) it's a lot to review in one pass, which may be part of why it's been quiet since early July. Splitting lets each surface get focused review and hit its own coverage bar (codecov currently flags 84.9%, 267 lines uncovered).

A natural split — core lands first, the rest fan out in parallel:

Core primitive — lance-index/src/vector/kmeans/distributed.rs + the kmeans.rs refactor: PartialStats, compute/merge/reduce_partial_stats, finalize_centroids, reservoir sampling, centroid init, Rust unit tests. Pure Arrow, no bindings — the foundation everything else depends on. Docs can ride along here.
IVF / dataset integration — lance/src/index/vector/ivf/distributed.rs, ivf.rs, utils.rs. Depends on Cpp code dump #1.
Java bindings — JNI distributed_kmeans.rs + DistributedKMeans.java + tests. Depends on Cpp code dump #1.
Python bindings — distributed_kmeans.py + indices.rs + tests. Depends on Cpp code dump #1.

Once #1 is in, #2–#4 are independent and reviewable in parallel.

Happy to help however's useful — I can review, rebase against main, or take one of the binding/integration slices off your plate.

github-actions Bot added A-python Python bindings A-index Vector index, linalg, tokenizer A-java Java bindings + JNI A-deps Dependency updates A-docs Documentation labels Jun 17, 2026

Merge branch 'main' into distributed-centroid-training

0e6ddfe

github-actions Bot added the enhancement New feature or request label Jun 17, 2026

fix ci

f14506f

Change-Id: I077166ef875c2ce9cc966b48c20482e12df9b629

Xuanwo reviewed Jun 20, 2026

View reviewed changes

summaryzb added 2 commits June 23, 2026 17:03

fix distributed kmeans

2ad6dab

Change-Id: I44ca6c6bacf77527caf02ccca94ca7c97d4fb4a1

fix doc

9a12fa5

Change-Id: I343fbe5a53c9d10a266f19496133a36548d6f47a

summaryzb requested a review from Xuanwo June 23, 2026 12:49

Xuanwo reviewed Jul 8, 2026

View reviewed changes

Uh oh!

Conversation

summaryzb commented Jun 17, 2026

Summary

Problem

Approach

Changes

Test Coverage

Uh oh!

codecov Bot commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

summaryzb Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sezruby commented Jul 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov Bot commented Jun 17, 2026 •

edited

Loading

summaryzb Jun 23, 2026 •

edited

Loading