diff --git a/Cargo.lock b/Cargo.lock index 63819e6f678..de84cb2dd14 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4827,6 +4827,7 @@ dependencies = [ "ndarray", "num-traits", "object_store", + "proptest", "prost", "prost-build", "prost-types", @@ -4840,6 +4841,7 @@ dependencies = [ "rstest", "serde", "serde_json", + "sha2 0.10.9", "smallvec", "tempfile", "test-log", diff --git a/Cargo.toml b/Cargo.toml index 6e79a26e69f..639eb41907d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -187,6 +187,7 @@ serde = { version = "^1" } serde_json = { version = "1" } semver = "1.0" serial_test = "3" +sha2 = "0.10" snafu = "0.9" strum = "0.26" lindera = { version = "3.0.7" } diff --git a/docs/src/guide/distributed_indexing.md b/docs/src/guide/distributed_indexing.md index 389e5a1bc09..34fd644e30a 100644 --- a/docs/src/guide/distributed_indexing.md +++ b/docs/src/guide/distributed_indexing.md @@ -193,3 +193,60 @@ unreferenced index files. This split keeps distributed scheduling outside the storage engine while still letting Lance own the on-disk index format. + +## Distributed centroid training + +Lance exposes primitive functions that let an external scheduler (Spark, Ray, +custom RPC) train a single global IVF centroid set across N workers. The +caller controls broadcast, tree-reduce, and convergence; Lance provides only +the math: one E-step, one merge, one M-step, plus a Round-0 reservoir-sample +init. + +The primitives live under `lance.indices.distributed_kmeans` (Python) and +`org.lance.index.vector.DistributedKMeans` (Java). Internally they delegate to +the Rust modules: + +- `lance_index::vector::kmeans::distributed` — pure-function Arrow primitives. +- `lance::index::vector::ivf::distributed` — async dataset-aware wrappers + (`sample_round_0`, `compute_partial_stats`). + +### Python example + +```python +import lance +import lance.indices.distributed_kmeans as dk + +ds = lance.dataset("s3://bucket/vec.lance") +fragment_groups = chunked( + [f.fragment_id for f in ds.get_fragments()], + num_workers=200, +) + +# Round 0 — each worker reservoir-samples its slice; driver bootstraps centroids. +samples = spark.parallelize(fragment_groups).map( + lambda fids: dk.sample_round_0( + ds, "vec", target=256 * K, fragment_ids=fids, + distance_type="l2", rng_seed=42, + ) +).collect() +centroids = dk.bootstrap_centroids(samples, k=K, distance_type="l2", rng_seed=42) + +# Rounds r = 1..max_iter — broadcast centroids, treeReduce partial stats. +for r in range(50): + centroids_b = spark.broadcast(centroids) + merged = ( + spark.parallelize(fragment_groups) + .map(lambda fids: dk.compute_partial_stats( + ds, "vec", centroids_b.value, fragment_ids=fids, + )) + .treeReduce(dk.merge_partial_stats) + ) + next_centroids = dk.finalize_centroids(merged, centroids) + if converged(merged, centroids): + break + centroids = next_centroids +``` + +`PartialStats` batches are plain Arrow `RecordBatch`es and round-trip cleanly +through Arrow IPC, so they can flow through any RPC layer that already speaks +Arrow. diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 5dba72718b3..6d4a60a847f 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -4031,6 +4031,7 @@ dependencies = [ "roaring", "serde", "serde_json", + "sha2 0.10.9", "smallvec", "tempfile", "tokio", diff --git a/java/lance-jni/src/distributed_kmeans.rs b/java/lance-jni/src/distributed_kmeans.rs new file mode 100644 index 00000000000..80123e3c849 --- /dev/null +++ b/java/lance-jni/src/distributed_kmeans.rs @@ -0,0 +1,353 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! JNI shim for the distributed IVF centroid-training primitives. +//! +//! Mirrors `python/src/indices.rs`. Callers (Spark, custom RPC) own broadcast, +//! tree-reduce, and convergence; this module exposes only the math. Every +//! native that crosses the JNI boundary moves data as Arrow-IPC `byte[]`, +//! including the three centroid-returning helpers (`finalizeCentroids`, +//! `selectInitialCentroids`, `bootstrapCentroids`). Float16/Float32/Float64 +//! centroids all round-trip without dtype collapse — Java callers reconstruct +//! a `VectorSchemaRoot` whose child vector preserves the original element +//! type. The legacy `float[]` interface was removed to avoid silently +//! downcasting Float16/Float64 outputs to Float32. + +use crate::RT; +use crate::blocking_dataset::{BlockingDataset, NATIVE_DATASET}; +use crate::error::{Error, Result}; + +use arrow::ipc::reader::StreamReader; +use arrow::ipc::writer::StreamWriter; +use arrow_array::{Array, FixedSizeListArray, RecordBatch}; +use arrow_schema::{Field, Schema}; +use jni::JNIEnv; +use jni::objects::{JByteArray, JClass, JIntArray, JObject, JObjectArray, JString}; +use jni::sys::jbyteArray; +use std::sync::Arc; + +use lance::index::vector::ivf::distributed as l2; +use lance_index::vector::kmeans::distributed as l1; +use lance_linalg::distance::DistanceType; + +fn parse_distance_type(env: &mut JNIEnv, s: &JString) -> Result { + let raw: String = env.get_string(s)?.into(); + DistanceType::try_from(raw.as_str()).map_err(|e| Error::input_error(e.to_string())) +} + +fn arrow_err(e: arrow::error::ArrowError) -> Error { + Error::input_error(e.to_string()) +} + +fn record_batch_to_ipc(batch: &RecordBatch) -> Result> { + let mut buf = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut buf, &batch.schema()).map_err(arrow_err)?; + writer.write(batch).map_err(arrow_err)?; + writer.finish().map_err(arrow_err)?; + } + Ok(buf) +} + +fn ipc_to_record_batch(env: &mut JNIEnv, jba: &JByteArray) -> Result { + let bytes = env.convert_byte_array(jba)?; + let mut reader = StreamReader::try_new(std::io::Cursor::new(bytes), None).map_err(arrow_err)?; + reader + .next() + .ok_or_else(|| Error::input_error("empty IPC stream".to_string()))? + .map_err(arrow_err) +} + +fn ipc_to_centroids_fsl(env: &mut JNIEnv, jba: &JByteArray) -> Result { + let batch = ipc_to_record_batch(env, jba)?; + if batch.num_columns() != 1 { + return Err(Error::input_error(format!( + "centroids IPC must have a single column, got {}", + batch.num_columns() + ))); + } + let fsl = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::input_error("centroids column must be FixedSizeList".to_string()))? + .clone(); + if !matches!( + fsl.value_type(), + arrow_schema::DataType::Float16 + | arrow_schema::DataType::Float32 + | arrow_schema::DataType::Float64 + ) { + return Err(Error::input_error(format!( + "centroids inner dtype must be Float16/Float32/Float64, got {}", + fsl.value_type() + ))); + } + Ok(fsl) +} + +/// Wrap a centroid FSL into a single-column Arrow-IPC payload. The schema +/// preserves the original inner dtype so Float16/Float32/Float64 round-trip +/// across the JNI boundary unchanged. Java callers reconstruct a +/// `VectorSchemaRoot` and dispatch on the child vector's type. +fn fsl_to_centroids_ipc(fsl: &FixedSizeListArray) -> Result> { + let schema = Arc::new(Schema::new(vec![Field::new( + "vec", + fsl.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(fsl.clone())]).map_err(arrow_err)?; + record_batch_to_ipc(&batch) +} + +fn read_optional_fragment_ids(env: &mut JNIEnv, arr: &JIntArray) -> Result>> { + if arr.is_null() { + return Ok(None); + } + let len = env.get_array_length(arr)? as usize; + let mut buf = vec![0i32; len]; + env.get_int_array_region(arr, 0, &mut buf)?; + Ok(Some(buf.into_iter().map(|x| x as u32).collect())) +} + +fn read_byte_array_2d(env: &mut JNIEnv, arr: &JObjectArray) -> Result> { + let len = env.get_array_length(arr)? as usize; + let mut out = Vec::with_capacity(len); + for i in 0..len { + let element = env.get_object_array_element(arr, i as i32)?; + let jba: JByteArray = JByteArray::from(element); + out.push(ipc_to_record_batch(env, &jba)?); + } + Ok(out) +} + +trait Pipe: Sized { + fn pipe R>(self, f: F) -> R { + f(self) + } +} +impl Pipe for T {} + +#[unsafe(no_mangle)] +pub extern "system" fn Java_org_lance_index_vector_DistributedKMeans_nativeSampleRound0<'local>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + dataset_obj: JObject<'local>, + column_jstr: JString<'local>, + target: i64, + distance_type_jstr: JString<'local>, + rng_seed: i64, + fragment_ids_arr: JIntArray<'local>, +) -> jbyteArray { + let mut inner = || -> Result> { + let column: String = env.get_string(&column_jstr)?.into(); + let dt = parse_distance_type(&mut env, &distance_type_jstr)?; + let frags = read_optional_fragment_ids(&mut env, &fragment_ids_arr)?; + if target < 0 { + return Err(Error::input_error(format!( + "target must be >= 0, got {}", + target + ))); + } + let dataset_guard = + unsafe { env.get_rust_field::<_, _, BlockingDataset>(&dataset_obj, NATIVE_DATASET) }?; + let batch = RT.block_on(l2::sample_round_0( + &dataset_guard.inner, + &column, + frags.as_deref(), + target as usize, + dt, + rng_seed as u64, + ))?; + record_batch_to_ipc(&batch) + }; + crate::ok_or_throw_with_return!(env, inner(), JByteArray::default().into_raw()).pipe(|bytes| { + match env.byte_array_from_slice(&bytes) { + Ok(arr) => arr.into_raw(), + Err(e) => { + let _ = env.throw_new("java/lang/RuntimeException", e.to_string()); + JByteArray::default().into_raw() + } + } + }) +} + +#[unsafe(no_mangle)] +pub extern "system" fn Java_org_lance_index_vector_DistributedKMeans_nativeComputePartialStats< + 'local, +>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + dataset_obj: JObject<'local>, + column_jstr: JString<'local>, + centroids_ipc: JByteArray<'local>, + distance_type_jstr: JString<'local>, + fragment_ids_arr: JIntArray<'local>, +) -> jbyteArray { + let mut inner = || -> Result> { + let column: String = env.get_string(&column_jstr)?.into(); + let dt = parse_distance_type(&mut env, &distance_type_jstr)?; + let frags = read_optional_fragment_ids(&mut env, &fragment_ids_arr)?; + let centroids = ipc_to_centroids_fsl(&mut env, ¢roids_ipc)?; + let dataset_guard = + unsafe { env.get_rust_field::<_, _, BlockingDataset>(&dataset_obj, NATIVE_DATASET) }?; + let stats = RT.block_on(l2::compute_partial_stats( + &dataset_guard.inner, + &column, + frags.as_deref(), + ¢roids, + dt, + ))?; + record_batch_to_ipc(stats.record_batch()) + }; + crate::ok_or_throw_with_return!(env, inner(), JByteArray::default().into_raw()).pipe(|bytes| { + match env.byte_array_from_slice(&bytes) { + Ok(arr) => arr.into_raw(), + Err(e) => { + let _ = env.throw_new("java/lang/RuntimeException", e.to_string()); + JByteArray::default().into_raw() + } + } + }) +} + +#[unsafe(no_mangle)] +pub extern "system" fn Java_org_lance_index_vector_DistributedKMeans_nativeMergePartialStats< + 'local, +>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + a_ipc: JByteArray<'local>, + b_ipc: JByteArray<'local>, +) -> jbyteArray { + let mut inner = || -> Result> { + let a = l1::PartialStats::from_record_batch(ipc_to_record_batch(&mut env, &a_ipc)?)?; + let b = l1::PartialStats::from_record_batch(ipc_to_record_batch(&mut env, &b_ipc)?)?; + let merged = l1::merge_partial_stats(a, b)?; + record_batch_to_ipc(merged.record_batch()) + }; + crate::ok_or_throw_with_return!(env, inner(), JByteArray::default().into_raw()).pipe(|bytes| { + match env.byte_array_from_slice(&bytes) { + Ok(arr) => arr.into_raw(), + Err(e) => { + let _ = env.throw_new("java/lang/RuntimeException", e.to_string()); + JByteArray::default().into_raw() + } + } + }) +} + +#[unsafe(no_mangle)] +pub extern "system" fn Java_org_lance_index_vector_DistributedKMeans_nativeReducePartialStats< + 'local, +>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + stats_arr: JObjectArray<'local>, +) -> jbyteArray { + let mut inner = || -> Result> { + let batches = read_byte_array_2d(&mut env, &stats_arr)?; + let mut parsed = Vec::with_capacity(batches.len()); + for b in batches { + parsed.push(l1::PartialStats::from_record_batch(b)?); + } + let merged = l1::reduce_partial_stats(parsed)?; + record_batch_to_ipc(merged.record_batch()) + }; + crate::ok_or_throw_with_return!(env, inner(), JByteArray::default().into_raw()).pipe(|bytes| { + match env.byte_array_from_slice(&bytes) { + Ok(arr) => arr.into_raw(), + Err(e) => { + let _ = env.throw_new("java/lang/RuntimeException", e.to_string()); + JByteArray::default().into_raw() + } + } + }) +} + +#[unsafe(no_mangle)] +pub extern "system" fn Java_org_lance_index_vector_DistributedKMeans_nativeFinalizeCentroids< + 'local, +>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + stats_ipc: JByteArray<'local>, + prev_ipc: JByteArray<'local>, +) -> jbyteArray { + let mut inner = || -> Result> { + let stats = + l1::PartialStats::from_record_batch(ipc_to_record_batch(&mut env, &stats_ipc)?)?; + let prev = ipc_to_centroids_fsl(&mut env, &prev_ipc)?; + let fsl = l1::finalize_centroids(&stats, &prev)?; + fsl_to_centroids_ipc(&fsl) + }; + crate::ok_or_throw_with_return!(env, inner(), JByteArray::default().into_raw()).pipe(|bytes| { + match env.byte_array_from_slice(&bytes) { + Ok(arr) => arr.into_raw(), + Err(e) => { + let _ = env.throw_new("java/lang/RuntimeException", e.to_string()); + JByteArray::default().into_raw() + } + } + }) +} + +#[unsafe(no_mangle)] +pub extern "system" fn Java_org_lance_index_vector_DistributedKMeans_nativeSelectInitialCentroids< + 'local, +>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + samples_arr: JObjectArray<'local>, + k: i32, + rng_seed: i64, +) -> jbyteArray { + let mut inner = || -> Result> { + if k < 0 { + return Err(Error::input_error(format!("k must be >= 0, got {}", k))); + } + let batches = read_byte_array_2d(&mut env, &samples_arr)?; + let fsl = l1::select_initial_centroids(batches, k as usize, rng_seed as u64)?; + fsl_to_centroids_ipc(&fsl) + }; + crate::ok_or_throw_with_return!(env, inner(), JByteArray::default().into_raw()).pipe(|bytes| { + match env.byte_array_from_slice(&bytes) { + Ok(arr) => arr.into_raw(), + Err(e) => { + let _ = env.throw_new("java/lang/RuntimeException", e.to_string()); + JByteArray::default().into_raw() + } + } + }) +} + +#[unsafe(no_mangle)] +pub extern "system" fn Java_org_lance_index_vector_DistributedKMeans_nativeBootstrapCentroids< + 'local, +>( + mut env: JNIEnv<'local>, + _class: JClass<'local>, + samples_arr: JObjectArray<'local>, + k: i32, + distance_type_jstr: JString<'local>, + rng_seed: i64, +) -> jbyteArray { + let mut inner = || -> Result> { + if k < 0 { + return Err(Error::input_error(format!("k must be >= 0, got {}", k))); + } + let dt = parse_distance_type(&mut env, &distance_type_jstr)?; + let batches = read_byte_array_2d(&mut env, &samples_arr)?; + let fsl = l1::bootstrap_centroids(batches, k as usize, dt, rng_seed as u64)?; + fsl_to_centroids_ipc(&fsl) + }; + crate::ok_or_throw_with_return!(env, inner(), JByteArray::default().into_raw()).pipe(|bytes| { + match env.byte_array_from_slice(&bytes) { + Ok(arr) => arr.into_raw(), + Err(e) => { + let _ = env.throw_new("java/lang/RuntimeException", e.to_string()); + JByteArray::default().into_raw() + } + } + }) +} diff --git a/java/lance-jni/src/lib.rs b/java/lance-jni/src/lib.rs index 37eeff66693..69599e0e085 100644 --- a/java/lance-jni/src/lib.rs +++ b/java/lance-jni/src/lib.rs @@ -45,6 +45,7 @@ mod blocking_dataset; mod blocking_scanner; mod delta; mod dispatcher; +mod distributed_kmeans; pub mod error; pub mod ffi; mod file_reader; diff --git a/java/src/main/java/org/lance/index/vector/DistributedKMeans.java b/java/src/main/java/org/lance/index/vector/DistributedKMeans.java new file mode 100644 index 00000000000..5cb24d24870 --- /dev/null +++ b/java/src/main/java/org/lance/index/vector/DistributedKMeans.java @@ -0,0 +1,218 @@ +/* + * Licensed 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.lance.index.vector; + +import org.lance.Dataset; +import org.lance.JniLoader; +import org.lance.index.DistanceType; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.VectorLoader; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.VectorUnloader; +import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.arrow.vector.types.pojo.Schema; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.nio.channels.Channels; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; + +/** + * Distributed IVF centroid-training primitives. + * + *

Mirrors {@code lance::index::vector::ivf::distributed}. Callers (Spark, custom RPC) own + * broadcast, tree-reduce, and convergence; this class exposes only the math. + * + *

Every payload that crosses the JNI boundary — sampled rows, partial stats, and centroid arrays + * — moves as Arrow IPC byte arrays. The centroid-returning helpers ({@link #finalizeCentroids}, + * {@link #selectInitialCentroids}, {@link #bootstrapCentroids}) return a {@link VectorSchemaRoot} + * whose child vector preserves the original Float16/Float32/Float64 element dtype. + */ +public final class DistributedKMeans { + + static { + JniLoader.ensureLoaded(); + } + + private DistributedKMeans() {} + + /** Round-0 reservoir-sample on the worker's fragment slice. */ + public static VectorSchemaRoot sampleRound0( + Dataset dataset, + String column, + long target, + DistanceType distanceType, + long rngSeed, + int[] fragmentIds, + BufferAllocator allocator) { + Objects.requireNonNull(dataset, "dataset"); + Objects.requireNonNull(column, "column"); + Objects.requireNonNull(distanceType, "distanceType"); + Objects.requireNonNull(allocator, "allocator"); + byte[] ipc = + nativeSampleRound0(dataset, column, target, distanceType.toString(), rngSeed, fragmentIds); + return readIpc(ipc, allocator); + } + + /** Round-r E-step on the worker's fragment slice. */ + public static VectorSchemaRoot computePartialStats( + Dataset dataset, + String column, + VectorSchemaRoot centroids, + DistanceType distanceType, + int[] fragmentIds, + BufferAllocator allocator) { + Objects.requireNonNull(dataset, "dataset"); + Objects.requireNonNull(column, "column"); + Objects.requireNonNull(centroids, "centroids"); + Objects.requireNonNull(distanceType, "distanceType"); + Objects.requireNonNull(allocator, "allocator"); + byte[] centroidsIpc = writeIpc(centroids); + byte[] statsIpc = + nativeComputePartialStats( + dataset, column, centroidsIpc, distanceType.toString(), fragmentIds); + return readIpc(statsIpc, allocator); + } + + /** Combine two partial stats. */ + public static VectorSchemaRoot mergePartialStats( + VectorSchemaRoot a, VectorSchemaRoot b, BufferAllocator allocator) { + Objects.requireNonNull(a, "a"); + Objects.requireNonNull(b, "b"); + return readIpc(nativeMergePartialStats(writeIpc(a), writeIpc(b)), allocator); + } + + /** Fold a list of partial stats. */ + public static VectorSchemaRoot reducePartialStats( + List stats, BufferAllocator allocator) { + Objects.requireNonNull(stats, "stats"); + List serialized = new ArrayList<>(stats.size()); + for (VectorSchemaRoot s : stats) { + serialized.add(writeIpc(s)); + } + return readIpc(nativeReducePartialStats(serialized.toArray(new byte[0][])), allocator); + } + + /** + * Compute new centroids; the returned VectorSchemaRoot has a single FixedSizeList column whose + * inner dtype matches {@code prev} (Float16/Float32/Float64). + */ + public static VectorSchemaRoot finalizeCentroids( + VectorSchemaRoot stats, VectorSchemaRoot prev, BufferAllocator allocator) { + Objects.requireNonNull(stats, "stats"); + Objects.requireNonNull(prev, "prev"); + Objects.requireNonNull(allocator, "allocator"); + return readIpc(nativeFinalizeCentroids(writeIpc(stats), writeIpc(prev)), allocator); + } + + /** + * Driver-side: pick {@code k} rows uniformly at random from worker samples. The returned + * VectorSchemaRoot has a single FixedSizeList column whose inner dtype matches the samples. + */ + public static VectorSchemaRoot selectInitialCentroids( + List samples, int k, long rngSeed, BufferAllocator allocator) { + Objects.requireNonNull(samples, "samples"); + Objects.requireNonNull(allocator, "allocator"); + List serialized = new ArrayList<>(samples.size()); + for (VectorSchemaRoot s : samples) { + serialized.add(writeIpc(s)); + } + return readIpc( + nativeSelectInitialCentroids(serialized.toArray(new byte[0][]), k, rngSeed), allocator); + } + + /** + * Driver-side: bootstrap centroids by running single-machine kmeans on worker samples. The + * returned VectorSchemaRoot has a single FixedSizeList column whose inner dtype matches the + * samples. + */ + public static VectorSchemaRoot bootstrapCentroids( + List samples, + int k, + DistanceType distanceType, + long rngSeed, + BufferAllocator allocator) { + Objects.requireNonNull(samples, "samples"); + Objects.requireNonNull(distanceType, "distanceType"); + Objects.requireNonNull(allocator, "allocator"); + List serialized = new ArrayList<>(samples.size()); + for (VectorSchemaRoot s : samples) { + serialized.add(writeIpc(s)); + } + return readIpc( + nativeBootstrapCentroids( + serialized.toArray(new byte[0][]), k, distanceType.toString(), rngSeed), + allocator); + } + + // -- helpers ------------------------------------------------------------- + + private static byte[] writeIpc(VectorSchemaRoot root) { + try (ByteArrayOutputStream out = new ByteArrayOutputStream(); + ArrowStreamWriter writer = new ArrowStreamWriter(root, null, Channels.newChannel(out))) { + writer.start(); + writer.writeBatch(); + writer.end(); + return out.toByteArray(); + } catch (Exception e) { + throw new RuntimeException("failed to serialize Arrow IPC", e); + } + } + + private static VectorSchemaRoot readIpc(byte[] bytes, BufferAllocator allocator) { + try (ArrowStreamReader reader = + new ArrowStreamReader(new ByteArrayInputStream(bytes), allocator)) { + Schema schema = reader.getVectorSchemaRoot().getSchema(); + VectorSchemaRoot dest = VectorSchemaRoot.create(schema, allocator); + VectorLoader loader = new VectorLoader(dest); + VectorUnloader unloader = new VectorUnloader(reader.getVectorSchemaRoot()); + reader.loadNextBatch(); + try (ArrowRecordBatch batch = unloader.getRecordBatch()) { + loader.load(batch); + } + return dest; + } catch (Exception e) { + throw new RuntimeException("failed to deserialize Arrow IPC", e); + } + } + + // -- native --------------------------------------------------------------- + + private static native byte[] nativeSampleRound0( + Dataset dataset, + String column, + long target, + String distanceType, + long rngSeed, + int[] fragmentIds); + + private static native byte[] nativeComputePartialStats( + Dataset dataset, String column, byte[] centroidsIpc, String distanceType, int[] fragmentIds); + + private static native byte[] nativeMergePartialStats(byte[] a, byte[] b); + + private static native byte[] nativeReducePartialStats(byte[][] stats); + + private static native byte[] nativeFinalizeCentroids(byte[] stats, byte[] prev); + + private static native byte[] nativeSelectInitialCentroids(byte[][] samples, int k, long rngSeed); + + private static native byte[] nativeBootstrapCentroids( + byte[][] samples, int k, String distanceType, long rngSeed); +} diff --git a/java/src/test/java/org/lance/index/vector/DistributedKMeansTest.java b/java/src/test/java/org/lance/index/vector/DistributedKMeansTest.java new file mode 100644 index 00000000000..025ab324b36 --- /dev/null +++ b/java/src/test/java/org/lance/index/vector/DistributedKMeansTest.java @@ -0,0 +1,339 @@ +/* + * Licensed 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.lance.index.vector; + +import org.lance.Dataset; +import org.lance.WriteParams; +import org.lance.index.DistanceType; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.Float2Vector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.FixedSizeListVector; +import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.apache.arrow.vector.ipc.ArrowStreamWriter; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.ByteArrayOutputStream; +import java.nio.file.Path; +import java.util.Collections; +import java.util.Random; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class DistributedKMeansTest { + + private static final int DIM = 8; + private static final int N = 1_000; + + private enum Precision { + HALF, + SINGLE, + DOUBLE + } + + /** Build a small in-memory Lance dataset of FixedSizeList<Float?, DIM> vectors. */ + private Dataset writeVectorDataset(String uri, BufferAllocator allocator, Precision precision) + throws Exception { + ArrowType.FloatingPoint floatType; + switch (precision) { + case HALF: + floatType = new ArrowType.FloatingPoint(FloatingPointPrecision.HALF); + break; + case SINGLE: + floatType = new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); + break; + case DOUBLE: + floatType = new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); + break; + default: + throw new IllegalStateException("unreachable"); + } + Field child = new Field("item", new FieldType(true, floatType, null), Collections.emptyList()); + Field vec = + new Field( + "vec", + new FieldType(true, new ArrowType.FixedSizeList(DIM), null), + Collections.singletonList(child)); + Schema schema = new Schema(Collections.singletonList(vec)); + + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) { + FixedSizeListVector list = (FixedSizeListVector) root.getVector("vec"); + list.allocateNew(); + + Random rng = new Random(42); + switch (precision) { + case HALF: + { + Float2Vector inner = (Float2Vector) list.getDataVector(); + inner.allocateNew(N * DIM); + for (int i = 0; i < N; i++) { + list.setNotNull(i); + for (int d = 0; d < DIM; d++) { + inner.setWithPossibleTruncate(i * DIM + d, (float) rng.nextGaussian()); + } + } + inner.setValueCount(N * DIM); + break; + } + case SINGLE: + { + Float4Vector inner = (Float4Vector) list.getDataVector(); + inner.allocateNew(N * DIM); + for (int i = 0; i < N; i++) { + list.setNotNull(i); + for (int d = 0; d < DIM; d++) { + inner.set(i * DIM + d, (float) rng.nextGaussian()); + } + } + inner.setValueCount(N * DIM); + break; + } + case DOUBLE: + { + Float8Vector inner = (Float8Vector) list.getDataVector(); + inner.allocateNew(N * DIM); + for (int i = 0; i < N; i++) { + list.setNotNull(i); + for (int d = 0; d < DIM; d++) { + inner.set(i * DIM + d, rng.nextGaussian()); + } + } + inner.setValueCount(N * DIM); + break; + } + } + list.setValueCount(N); + root.setRowCount(N); + + ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (ArrowStreamWriter writer = new ArrowStreamWriter(root, null, out)) { + writer.start(); + writer.writeBatch(); + writer.end(); + } + try (ArrowStreamReader reader = + new ArrowStreamReader( + new ByteArrayReadableSeekableByteChannel(out.toByteArray()), allocator)) { + return Dataset.write() + .allocator(allocator) + .reader(reader) + .uri(uri) + .mode(WriteParams.WriteMode.OVERWRITE) + .execute(); + } + } + } + + private Dataset writeVectorDataset(String uri, BufferAllocator allocator) throws Exception { + return writeVectorDataset(uri, allocator, Precision.SINGLE); + } + + private static FloatingPointPrecision innerPrecision(VectorSchemaRoot root) { + FixedSizeListVector list = (FixedSizeListVector) root.getVector("vec"); + ArrowType inner = list.getDataVector().getField().getType(); + return ((ArrowType.FloatingPoint) inner).getPrecision(); + } + + @Test + void roundTripFourPrimitives(@TempDir Path tmp) throws Exception { + String datasetUri = tmp.resolve("vec.lance").toString(); + try (BufferAllocator allocator = new RootAllocator(); + Dataset dataset = writeVectorDataset(datasetUri, allocator)) { + + VectorSchemaRoot samples = + DistributedKMeans.sampleRound0( + dataset, "vec", 256, DistanceType.L2, 42L, null, allocator); + try { + assertEquals(256, samples.getRowCount()); + + VectorSchemaRoot bootstrap = + DistributedKMeans.bootstrapCentroids( + Collections.singletonList(samples), 16, DistanceType.L2, 7L, allocator); + try { + assertEquals(16, bootstrap.getRowCount()); + assertEquals(FloatingPointPrecision.SINGLE, innerPrecision(bootstrap)); + FixedSizeListVector bootstrapList = (FixedSizeListVector) bootstrap.getVector("vec"); + Float4Vector bootstrapInner = (Float4Vector) bootstrapList.getDataVector(); + assertEquals(16 * DIM, bootstrapInner.getValueCount()); + for (int i = 0; i < bootstrapInner.getValueCount(); i++) { + assertTrue(Float.isFinite(bootstrapInner.get(i)), "non-finite centroid value"); + } + + VectorSchemaRoot partial = + DistributedKMeans.computePartialStats( + dataset, "vec", bootstrap, DistanceType.L2, null, allocator); + try { + assertEquals(16, partial.getRowCount()); + + VectorSchemaRoot merged = + DistributedKMeans.reducePartialStats(Collections.singletonList(partial), allocator); + try { + assertEquals(16, merged.getRowCount()); + VectorSchemaRoot next = + DistributedKMeans.finalizeCentroids(merged, bootstrap, allocator); + try { + assertEquals(16, next.getRowCount()); + assertEquals(FloatingPointPrecision.SINGLE, innerPrecision(next)); + Float4Vector nextInner = + (Float4Vector) ((FixedSizeListVector) next.getVector("vec")).getDataVector(); + assertEquals(16 * DIM, nextInner.getValueCount()); + for (int i = 0; i < nextInner.getValueCount(); i++) { + assertTrue(Float.isFinite(nextInner.get(i)), "non-finite centroid value"); + } + } finally { + next.close(); + } + } finally { + merged.close(); + } + } finally { + partial.close(); + } + } finally { + bootstrap.close(); + } + } finally { + samples.close(); + } + } + } + + @Test + void roundTripFloat16(@TempDir Path tmp) throws Exception { + String datasetUri = tmp.resolve("vec16.lance").toString(); + try (BufferAllocator allocator = new RootAllocator(); + Dataset dataset = writeVectorDataset(datasetUri, allocator, Precision.HALF)) { + runEndToEnd(dataset, allocator, FloatingPointPrecision.HALF); + } + } + + @Test + void roundTripFloat64(@TempDir Path tmp) throws Exception { + String datasetUri = tmp.resolve("vec64.lance").toString(); + try (BufferAllocator allocator = new RootAllocator(); + Dataset dataset = writeVectorDataset(datasetUri, allocator, Precision.DOUBLE)) { + runEndToEnd(dataset, allocator, FloatingPointPrecision.DOUBLE); + } + } + + private void runEndToEnd( + Dataset dataset, BufferAllocator allocator, FloatingPointPrecision expected) + throws Exception { + VectorSchemaRoot samples = + DistributedKMeans.sampleRound0(dataset, "vec", 256, DistanceType.L2, 42L, null, allocator); + try { + assertEquals(expected, innerPrecision(samples)); + VectorSchemaRoot bootstrap = + DistributedKMeans.bootstrapCentroids( + Collections.singletonList(samples), 16, DistanceType.L2, 7L, allocator); + try { + assertEquals(16, bootstrap.getRowCount()); + assertEquals(expected, innerPrecision(bootstrap)); + + VectorSchemaRoot partial = + DistributedKMeans.computePartialStats( + dataset, "vec", bootstrap, DistanceType.L2, null, allocator); + try { + assertEquals(16, partial.getRowCount()); + + VectorSchemaRoot merged = + DistributedKMeans.reducePartialStats(Collections.singletonList(partial), allocator); + try { + assertEquals(16, merged.getRowCount()); + VectorSchemaRoot next = + DistributedKMeans.finalizeCentroids(merged, bootstrap, allocator); + try { + assertEquals(16, next.getRowCount()); + assertEquals(expected, innerPrecision(next)); + } finally { + next.close(); + } + } finally { + merged.close(); + } + } finally { + partial.close(); + } + } finally { + bootstrap.close(); + } + } finally { + samples.close(); + } + } + + @Test + void sampleRound0NegativeTargetThrows(@TempDir Path tmp) throws Exception { + String datasetUri = tmp.resolve("vec.lance").toString(); + try (BufferAllocator allocator = new RootAllocator(); + Dataset dataset = writeVectorDataset(datasetUri, allocator)) { + assertThrows( + IllegalArgumentException.class, + () -> + DistributedKMeans.sampleRound0( + dataset, "vec", -1L, DistanceType.L2, 42L, null, allocator)); + } + } + + @Test + void selectInitialCentroidsNegativeKThrows(@TempDir Path tmp) throws Exception { + String datasetUri = tmp.resolve("vec.lance").toString(); + try (BufferAllocator allocator = new RootAllocator(); + Dataset dataset = writeVectorDataset(datasetUri, allocator)) { + VectorSchemaRoot samples = + DistributedKMeans.sampleRound0(dataset, "vec", 64, DistanceType.L2, 1L, null, allocator); + try { + assertThrows( + IllegalArgumentException.class, + () -> + DistributedKMeans.selectInitialCentroids( + Collections.singletonList(samples), -1, 1L, allocator)); + } finally { + samples.close(); + } + } + } + + @Test + void bootstrapCentroidsNegativeKThrows(@TempDir Path tmp) throws Exception { + String datasetUri = tmp.resolve("vec.lance").toString(); + try (BufferAllocator allocator = new RootAllocator(); + Dataset dataset = writeVectorDataset(datasetUri, allocator)) { + VectorSchemaRoot samples = + DistributedKMeans.sampleRound0(dataset, "vec", 64, DistanceType.L2, 1L, null, allocator); + try { + assertThrows( + IllegalArgumentException.class, + () -> + DistributedKMeans.bootstrapCentroids( + Collections.singletonList(samples), -1, DistanceType.L2, 1L, allocator)); + } finally { + samples.close(); + } + } + } +} diff --git a/python/Cargo.lock b/python/Cargo.lock index f4e52846476..980d18918ca 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -4435,6 +4435,7 @@ dependencies = [ "roaring", "serde", "serde_json", + "sha2 0.10.9", "smallvec", "tempfile", "tokio", diff --git a/python/python/lance/indices/__init__.py b/python/python/lance/indices/__init__.py index 675754cc2d0..c06e1b5f880 100644 --- a/python/python/lance/indices/__init__.py +++ b/python/python/lance/indices/__init__.py @@ -4,6 +4,7 @@ from enum import Enum from .. import lance as _lance +from . import distributed_kmeans # noqa: F401 from .builder import IndexConfig, IndicesBuilder from .ivf import IvfModel from .pq import PqModel diff --git a/python/python/lance/indices/distributed_kmeans.py b/python/python/lance/indices/distributed_kmeans.py new file mode 100644 index 00000000000..636d8a200f1 --- /dev/null +++ b/python/python/lance/indices/distributed_kmeans.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Distributed IVF centroid training primitives. + +Mirrors :mod:`lance::index::vector::ivf::distributed`. The caller (Spark / Ray / +custom RPC) is responsible for fragment partitioning, broadcast, treeReduce, +and convergence. Lance only provides the math: one E-step, one merge, one +M-step, plus a Round-0 reservoir-sample initializer. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Iterable, Optional, Sequence, Union + +import numpy as np +import pyarrow as pa + +from .. import lance as _lance + +if TYPE_CHECKING: + from ..dataset import LanceDataset + +_indices = _lance.indices + + +def _to_array_data( + centroids: Union[np.ndarray, pa.FixedSizeListArray], +) -> pa.FixedSizeListArray: + """Coerce centroids to a ``pa.FixedSizeListArray`` for the FFI boundary.""" + if isinstance(centroids, np.ndarray): + if centroids.ndim != 2: + raise ValueError(f"expected 2-D centroids, got shape {centroids.shape}") + flat = pa.array(centroids.reshape(-1)) + return pa.FixedSizeListArray.from_arrays(flat, centroids.shape[1]) + return centroids + + +def _fsl_to_ndarray(arr: pa.FixedSizeListArray) -> np.ndarray: + """Reshape a 1-D ``FixedSizeListArray`` flat values buffer into ``(k, dim)``.""" + return np.asarray(arr.values).reshape(-1, arr.type.list_size) + + +def sample_round_0( + dataset: "LanceDataset", + column: str, + target: int, + *, + fragment_ids: Optional[Sequence[int]] = None, + distance_type: str = "l2", + rng_seed: int = 0, +) -> pa.RecordBatch: + """Round-0 reservoir-sample on the worker's fragment slice.""" + return _indices.distributed_sample_round_0( + dataset._ds, + column, + target, + distance_type, + rng_seed, + list(fragment_ids) if fragment_ids is not None else None, + ) + + +def compute_partial_stats( + dataset: "LanceDataset", + column: str, + centroids: Union[np.ndarray, pa.FixedSizeListArray], + *, + distance_type: str = "l2", + fragment_ids: Optional[Sequence[int]] = None, +) -> pa.RecordBatch: + """Round-r E-step on the worker's fragment slice.""" + return _indices.distributed_compute_partial_stats( + dataset._ds, + column, + _to_array_data(centroids), + distance_type, + list(fragment_ids) if fragment_ids is not None else None, + ) + + +def merge_partial_stats(a: pa.RecordBatch, b: pa.RecordBatch) -> pa.RecordBatch: + """Combine two partial stats produced against the same centroids.""" + return _indices.distributed_merge_partial_stats(a, b) + + +def reduce_partial_stats( + stats: Iterable[pa.RecordBatch], +) -> pa.RecordBatch: + """Fold an iterable of partial stats.""" + return _indices.distributed_reduce_partial_stats(list(stats)) + + +def finalize_centroids( + stats: pa.RecordBatch, + prev_centroids: Union[np.ndarray, pa.FixedSizeListArray], +) -> np.ndarray: + """Compute the new centroids as a ``(k, dim)`` ndarray.""" + arr = _indices.distributed_finalize_centroids(stats, _to_array_data(prev_centroids)) + return _fsl_to_ndarray(arr) + + +def select_initial_centroids( + samples: Sequence[pa.RecordBatch], + k: int, + *, + rng_seed: int = 0, +) -> np.ndarray: + """Driver-side: pick ``k`` rows uniformly at random from worker samples.""" + arr = _indices.distributed_select_initial_centroids(list(samples), k, rng_seed) + return _fsl_to_ndarray(arr) + + +def bootstrap_centroids( + samples: Sequence[pa.RecordBatch], + k: int, + *, + distance_type: str = "l2", + rng_seed: int = 0, +) -> np.ndarray: + """Driver-side: run single-machine kmeans over the union of worker samples.""" + arr = _indices.distributed_bootstrap_centroids( + list(samples), k, distance_type, rng_seed + ) + return _fsl_to_ndarray(arr) diff --git a/python/python/tests/test_distributed_kmeans.py b/python/python/tests/test_distributed_kmeans.py new file mode 100644 index 00000000000..6b74fe406b9 --- /dev/null +++ b/python/python/tests/test_distributed_kmeans.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Tests for :mod:`lance.indices.distributed_kmeans`.""" + +from __future__ import annotations + +import io +from typing import TYPE_CHECKING + +import lance +import lance.indices.distributed_kmeans as dk +import numpy as np +import pyarrow as pa +import pytest + +if TYPE_CHECKING: + from pathlib import Path + + +@pytest.fixture +def vector_dataset(tmp_path: Path): + rng = np.random.default_rng(0) + data = rng.normal(size=(2_000, 8)).astype(np.float32) + schema = pa.schema([("vec", pa.list_(pa.float32(), 8))]) + arr = pa.FixedSizeListArray.from_arrays(pa.array(data.reshape(-1)), 8) + table = pa.table([arr], schema=schema) + uri = str(tmp_path / "vec.lance") + return lance.write_dataset(table, uri, max_rows_per_file=500) + + +def test_round_trip_end_to_end(vector_dataset): + ds = vector_dataset + samples = dk.sample_round_0(ds, "vec", target=512, distance_type="l2", rng_seed=1) + assert samples.num_rows == 512 + + centroids = dk.bootstrap_centroids([samples], k=16, distance_type="l2", rng_seed=2) + assert centroids.shape == (16, 8) + + partial = dk.compute_partial_stats(ds, "vec", centroids, distance_type="l2") + assert partial.num_rows == 16 + assert partial.schema.field("count").type == pa.uint64() + + merged = dk.reduce_partial_stats([partial]) + new_centroids = dk.finalize_centroids(merged, centroids) + assert new_centroids.shape == centroids.shape + assert np.all(np.isfinite(new_centroids)) + + +def test_partial_stats_arrow_ipc_round_trip(vector_dataset): + ds = vector_dataset + centroids = np.random.RandomState(3).normal(size=(8, 8)).astype(np.float32) + partial = dk.compute_partial_stats(ds, "vec", centroids) + + sink = io.BytesIO() + with pa.ipc.new_stream(sink, partial.schema) as writer: + writer.write_batch(partial) + sink.seek(0) + reader = pa.ipc.open_stream(sink) + restored = next(reader) + + merged = dk.merge_partial_stats(partial, restored) + assert merged.column("count").to_pylist() == [ + 2 * c for c in partial.column("count").to_pylist() + ] + + +def test_select_initial_centroids_picks_k(vector_dataset): + ds = vector_dataset + samples = dk.sample_round_0(ds, "vec", target=256, rng_seed=4) + centroids = dk.select_initial_centroids([samples], k=32, rng_seed=5) + assert centroids.shape == (32, 8) diff --git a/python/src/indices.rs b/python/src/indices.rs index 7ce7a297924..7cc8c682a4c 100644 --- a/python/src/indices.rs +++ b/python/src/indices.rs @@ -748,6 +748,143 @@ impl PyIndexDescription { } } +// --------------------------------------------------------------------------- +// Distributed kmeans primitives +// --------------------------------------------------------------------------- + +use lance::index::vector::ivf::distributed as l2; +use lance_index::vector::kmeans::distributed as l1; + +fn parse_distance_type(s: &str) -> PyResult { + DistanceType::try_from(s).map_err(|e| PyValueError::new_err(e.to_string())) +} + +fn array_data_to_fsl(array: ArrayData) -> PyResult { + Ok(FixedSizeListArray::from(array)) +} + +#[pyfunction] +#[pyo3(signature = (dataset, column, target, distance_type, rng_seed, fragment_ids=None))] +fn distributed_sample_round_0<'py>( + py: Python<'py>, + dataset: &Dataset, + column: &str, + target: usize, + distance_type: &str, + rng_seed: u64, + fragment_ids: Option>, +) -> PyResult> { + let dt = parse_distance_type(distance_type)?; + let column = column.to_string(); + let dataset = dataset.ds.clone(); + let batch = rt().block_on(Some(py), async move { + l2::sample_round_0( + dataset.as_ref(), + &column, + fragment_ids.as_deref(), + target, + dt, + rng_seed, + ) + .await + })?; + let batch = batch.infer_error()?; + batch.to_pyarrow(py) +} + +#[pyfunction] +#[pyo3(signature = (dataset, column, centroids, distance_type, fragment_ids=None))] +fn distributed_compute_partial_stats<'py>( + py: Python<'py>, + dataset: &Dataset, + column: &str, + centroids: PyArrowType, + distance_type: &str, + fragment_ids: Option>, +) -> PyResult> { + let dt = parse_distance_type(distance_type)?; + let centroids_fsl = array_data_to_fsl(centroids.0)?; + let column = column.to_string(); + let dataset = dataset.ds.clone(); + let stats = rt().block_on(Some(py), async move { + l2::compute_partial_stats( + dataset.as_ref(), + &column, + fragment_ids.as_deref(), + ¢roids_fsl, + dt, + ) + .await + })?; + let stats = stats.infer_error()?; + stats.into_record_batch().to_pyarrow(py) +} + +#[pyfunction] +fn distributed_merge_partial_stats<'py>( + py: Python<'py>, + a: PyArrowType, + b: PyArrowType, +) -> PyResult> { + let sa = l1::PartialStats::from_record_batch(a.0).infer_error()?; + let sb = l1::PartialStats::from_record_batch(b.0).infer_error()?; + let merged = l1::merge_partial_stats(sa, sb).infer_error()?; + merged.into_record_batch().to_pyarrow(py) +} + +#[pyfunction] +fn distributed_reduce_partial_stats<'py>( + py: Python<'py>, + stats: Vec>, +) -> PyResult> { + let mut parsed = Vec::with_capacity(stats.len()); + for batch in stats { + parsed.push(l1::PartialStats::from_record_batch(batch.0).infer_error()?); + } + let merged = l1::reduce_partial_stats(parsed).infer_error()?; + merged.into_record_batch().to_pyarrow(py) +} + +#[pyfunction] +fn distributed_finalize_centroids<'py>( + py: Python<'py>, + stats: PyArrowType, + prev: PyArrowType, +) -> PyResult> { + let s = l1::PartialStats::from_record_batch(stats.0).infer_error()?; + let prev_fsl = array_data_to_fsl(prev.0)?; + let new = l1::finalize_centroids(&s, &prev_fsl).infer_error()?; + new.to_data().to_pyarrow(py) +} + +#[pyfunction] +#[pyo3(signature = (samples, k, rng_seed))] +fn distributed_select_initial_centroids<'py>( + py: Python<'py>, + samples: Vec>, + k: usize, + rng_seed: u64, +) -> PyResult> { + let batches: Vec = samples.into_iter().map(|s| s.0).collect(); + let centroids = l1::select_initial_centroids(batches, k, rng_seed).infer_error()?; + centroids.to_data().to_pyarrow(py) +} + +#[pyfunction] +#[pyo3(signature = (samples, k, distance_type, rng_seed))] +fn distributed_bootstrap_centroids<'py>( + py: Python<'py>, + samples: Vec>, + k: usize, + distance_type: &str, + rng_seed: u64, +) -> PyResult> { + let dt = parse_distance_type(distance_type)?; + let batches: Vec = samples.into_iter().map(|s| s.0).collect(); + let centroids = l1::bootstrap_centroids(batches, k, dt, rng_seed).infer_error()?; + centroids.to_data().to_pyarrow(py) +} + pub fn register_indices(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { let indices = PyModule::new(py, "indices")?; indices.add_wrapped(wrap_pyfunction!(train_ivf_model))?; @@ -756,6 +893,13 @@ pub fn register_indices(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { indices.add_wrapped(wrap_pyfunction!(transform_vectors))?; indices.add_wrapped(wrap_pyfunction!(shuffle_transformed_vectors))?; indices.add_wrapped(wrap_pyfunction!(load_shuffled_vectors))?; + indices.add_wrapped(wrap_pyfunction!(distributed_sample_round_0))?; + indices.add_wrapped(wrap_pyfunction!(distributed_compute_partial_stats))?; + indices.add_wrapped(wrap_pyfunction!(distributed_merge_partial_stats))?; + indices.add_wrapped(wrap_pyfunction!(distributed_reduce_partial_stats))?; + indices.add_wrapped(wrap_pyfunction!(distributed_finalize_centroids))?; + indices.add_wrapped(wrap_pyfunction!(distributed_select_initial_centroids))?; + indices.add_wrapped(wrap_pyfunction!(distributed_bootstrap_centroids))?; indices.add_class::()?; indices.add_class::()?; indices.add_class::()?; diff --git a/python/src/utils.rs b/python/src/utils.rs index 4f7d6d7dde2..cab9b168e87 100644 --- a/python/src/utils.rs +++ b/python/src/utils.rs @@ -89,11 +89,7 @@ impl KMeans { return Err(PyValueError::new_err("Must be a FixedSizeList")); } let fixed_size_arr = FixedSizeListArray::from(data); - let params = KMeansParams { - distance_type: metric_type.try_into().unwrap(), - max_iters, - ..Default::default() - }; + let params = KMeansParams::new(None, max_iters, 1, metric_type.try_into().unwrap()); let kmeans = LanceKMeans::new_with_params(&fixed_size_arr, k, ¶ms).map_err(|e| { PyRuntimeError::new_err(format!( @@ -120,11 +116,7 @@ impl KMeans { return Err(PyValueError::new_err("Must be a FixedSizeList")); } let fixed_size_arr = FixedSizeListArray::from(data); - let params = KMeansParams { - distance_type: self.metric_type, - max_iters: self.max_iters, - ..Default::default() - }; + let params = KMeansParams::new(None, self.max_iters, 1, self.metric_type); let kmeans = LanceKMeans::new_with_params(&fixed_size_arr, self.k, ¶ms) .map_err(|e| PyRuntimeError::new_err(format!("Error training KMeans: {}", e)))?; self.trained_kmeans = Some(kmeans); diff --git a/rust/lance-index/Cargo.toml b/rust/lance-index/Cargo.toml index 85de43c0f9b..8505846637e 100644 --- a/rust/lance-index/Cargo.toml +++ b/rust/lance-index/Cargo.toml @@ -74,6 +74,7 @@ bitpacking = { version = "0.9.2", features = ["bitpacker4x"] } rand_distr.workspace = true lance-datagen.workspace = true rangemap.workspace = true +sha2.workspace = true [dev-dependencies] approx.workspace = true @@ -82,6 +83,7 @@ env_logger = "0.11.6" geo-traits.workspace = true lance-datagen.workspace = true lance-testing.workspace = true +proptest.workspace = true test-log.workspace = true rstest.workspace = true chrono.workspace = true diff --git a/rust/lance-index/src/vector/kmeans.rs b/rust/lance-index/src/vector/kmeans.rs index b11fb70bed0..597b52c8014 100644 --- a/rust/lance-index/src/vector/kmeans.rs +++ b/rust/lance-index/src/vector/kmeans.rs @@ -47,10 +47,17 @@ use { use crate::vector::utils::SimpleIndex; use crate::{Error, Result}; +pub mod distributed; + /// KMean initialization method. #[derive(Debug, PartialEq)] pub enum KMeanInit { - Random, + /// Random init. When the payload is `Some(seed)`, the seed is forwarded to + /// `kmeans_random_init` via `SmallRng::seed_from_u64(seed)` so the choice + /// of starting rows is reproducible; `None` falls back to + /// `SmallRng::from_os_rng()`. Set via [`KMeansParams::with_seed`]; the + /// distributed `bootstrap_centroids` primitive relies on it. + Random(Option), Incremental(Arc), } @@ -111,7 +118,7 @@ impl Default for KMeansParams { max_iters: 50, tolerance: 1e-4, redos: 1, - init: KMeanInit::Random, + init: KMeanInit::Random(None), distance_type: DistanceType::L2, balance_factor: 0.0, hierarchical_k: 16, @@ -129,7 +136,7 @@ impl KMeansParams { ) -> Self { let init = match centroids { Some(centroids) => KMeanInit::Incremental(centroids), - None => KMeanInit::Random, + None => KMeanInit::Random(None), }; Self { max_iters, @@ -163,6 +170,31 @@ impl KMeansParams { self.hierarchical_k = hierarchical_k; self } + + /// Set the distance type for kmeans clustering. + pub fn with_distance_type(mut self, distance_type: DistanceType) -> Self { + self.distance_type = distance_type; + self + } + + /// Set a deterministic RNG seed for random initialization. + /// + /// The seed is stored as the payload of [`KMeanInit::Random`] and is only + /// consulted when the init mode is `Random`; calling this on a params + /// configured with `KMeanInit::Incremental` is a no-op because that mode + /// supplies centroids directly. Distributed kmeans relies on this so the + /// same `(samples, k, seed)` input produces the same bootstrap centroids + /// on every worker. + /// + /// Note: `train_hierarchical_kmeans` re-uses the same `params` for every + /// recursion level, so the same seed is applied at every level (matching + /// the prior single-seed-field behavior). + pub fn with_seed(mut self, seed: u64) -> Self { + if let KMeanInit::Random(slot) = &mut self.init { + *slot = Some(seed); + } + self + } } /// Randomly initialize kmeans centroids. @@ -799,11 +831,13 @@ impl KMeans { let mut cluster_sizes = vec![0; k]; let mut adjusted_balance_factor = f32::MAX; - // TODO: use seed for Rng. - let mut rng = SmallRng::from_os_rng(); + let mut rng = match ¶ms.init { + KMeanInit::Random(Some(seed)) => SmallRng::seed_from_u64(*seed), + _ => SmallRng::from_os_rng(), + }; for redo in 1..=params.redos { let mut kmeans: Self = match ¶ms.init { - KMeanInit::Random => Self::init_random::( + KMeanInit::Random(_) => Self::init_random::( data.values(), dimension, k, diff --git a/rust/lance-index/src/vector/kmeans/distributed.rs b/rust/lance-index/src/vector/kmeans/distributed.rs new file mode 100644 index 00000000000..9a930b30cfe --- /dev/null +++ b/rust/lance-index/src/vector/kmeans/distributed.rs @@ -0,0 +1,1808 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Distributed KMeans primitives. +//! +//! Pure-function, Arrow-native primitives that let an external scheduler +//! (Spark, Ray, custom RPC) drive distributed IVF centroid training. See +//! `docs/superpowers/specs/2026-06-10-distributed-centroid-training-abstraction-design.md`. + +use std::collections::HashMap; +use std::fmt::Write as _; +use std::sync::Arc; + +use arrow_array::{ + Array, ArrowPrimitiveType, FixedSizeListArray, Float16Array, Float32Array, Float64Array, + PrimitiveArray, RecordBatch, UInt32Array, UInt64Array, + builder::{FixedSizeListBuilder, Float64Builder}, + cast::AsArray, + types::{Float16Type, Float32Type, Float64Type, Int8Type}, +}; +use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; +use half::f16; +use lance_arrow::FixedSizeListArrayExt; +use lance_linalg::distance::DistanceType; +use lance_linalg::kernels::normalize_fsl_owned; + +use crate::vector::kmeans::{KMeans, KMeansParams, train_kmeans}; +use crate::{Error, Result}; + +pub const PARTIAL_STATS_VERSION: &str = "1"; +pub const META_VERSION: &str = "lance.partial_stats.version"; +pub const META_K: &str = "lance.partial_stats.k"; +pub const META_DIM: &str = "lance.partial_stats.dim"; +pub const META_DT: &str = "lance.partial_stats.distance_type"; +pub const META_FP: &str = "lance.partial_stats.centroids_fingerprint"; + +pub const COL_CLUSTER_ID: &str = "cluster_id"; +pub const COL_COUNT: &str = "count"; +pub const COL_SUM: &str = "sum"; +pub const COL_SQ_NORM_SUM: &str = "sq_norm_sum"; +pub const COL_LOSS: &str = "loss"; +pub const COL_RADIUS: &str = "radius"; + +/// One worker's contribution to a kmeans round. +/// +/// Wraps a fixed-schema [`RecordBatch`] of `k` rows so the wire format is just +/// Arrow IPC. See module docs for the schema contract. +#[derive(Debug, Clone)] +pub struct PartialStats { + pub(crate) batch: RecordBatch, +} + +fn dt_metadata_value(dt: DistanceType) -> &'static str { + match dt { + DistanceType::L2 => "l2", + DistanceType::Dot => "dot", + DistanceType::Cosine => "cosine", + DistanceType::Hamming => "hamming", + } +} + +fn fingerprint_to_hex(fp: &[u8; 8]) -> String { + let mut s = String::with_capacity(16); + for byte in fp { + // Write into a String never errors. + let _ = write!(s, "{:02x}", byte); + } + s +} + +fn fingerprint_from_hex(s: &str) -> Result<[u8; 8]> { + if s.len() != 16 { + return Err(Error::index(format!( + "invalid fingerprint hex length: {}", + s.len() + ))); + } + let mut out = [0u8; 8]; + for i in 0..8 { + out[i] = u8::from_str_radix(&s[i * 2..i * 2 + 2], 16) + .map_err(|e| Error::index(format!("invalid fingerprint hex: {}", e)))?; + } + Ok(out) +} + +pub(crate) fn build_schema(k: usize, dim: usize, dt: DistanceType, fp: [u8; 8]) -> SchemaRef { + // The inner Float64 field is nullable so it round-trips through the default + // [`FixedSizeListBuilder`] (whose primitive builder produces a nullable + // child by default). + let sum_field = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float64, true)), + dim as i32, + ); + let mut metadata = HashMap::new(); + metadata.insert(META_VERSION.into(), PARTIAL_STATS_VERSION.into()); + metadata.insert(META_K.into(), k.to_string()); + metadata.insert(META_DIM.into(), dim.to_string()); + metadata.insert(META_DT.into(), dt_metadata_value(dt).into()); + metadata.insert(META_FP.into(), fingerprint_to_hex(&fp)); + Arc::new( + Schema::new(vec![ + Field::new(COL_CLUSTER_ID, DataType::UInt32, false), + Field::new(COL_COUNT, DataType::UInt64, false), + Field::new(COL_SUM, sum_field, false), + Field::new(COL_SQ_NORM_SUM, DataType::Float64, false), + Field::new(COL_LOSS, DataType::Float64, false), + Field::new(COL_RADIUS, DataType::Float32, false), + ]) + .with_metadata(metadata), + ) +} + +impl PartialStats { + /// Build a zero-filled stats buffer for `k` clusters of `dim` dimension. + /// `fingerprint` identifies the centroids these stats are computed against. + pub fn empty(k: usize, dim: usize, dt: DistanceType, fingerprint: [u8; 8]) -> Self { + let schema = build_schema(k, dim, dt, fingerprint); + + let cluster_id: UInt32Array = (0..k as u32).collect(); + let count = UInt64Array::from(vec![0u64; k]); + let sq_norm = Float64Array::from(vec![0.0f64; k]); + let loss = Float64Array::from(vec![0.0f64; k]); + let radius = Float32Array::from(vec![0.0f32; k]); + + let mut sum_builder = FixedSizeListBuilder::new(Float64Builder::new(), dim as i32); + for _ in 0..k { + for _ in 0..dim { + sum_builder.values().append_value(0.0); + } + sum_builder.append(true); + } + let sum = sum_builder.finish(); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(cluster_id), + Arc::new(count), + Arc::new(sum), + Arc::new(sq_norm), + Arc::new(loss), + Arc::new(radius), + ], + ) + .expect("partial stats schema is internally consistent"); + Self { batch } + } + + /// Validate and adopt an externally-built RecordBatch. + /// + /// This is the trust boundary for partial stats arriving over the wire + /// (Arrow IPC, file, etc.). Validation is strict and exhaustive: + /// + /// * required metadata: `version`, `k`, `dim`, `distance_type`, `centroids_fingerprint` + /// * row count equals `k` + /// * column names match the canonical schema + /// * column dtypes match: `UInt32`, `UInt64`, `FixedSizeList`, + /// `Float64`, `Float64`, `Float32` + /// * every column has `null_count() == 0` (downstream `merge_*` / + /// `finalize_*` read raw `.values()` slices and ignore null bitmaps, + /// so a stray null bit would silently corrupt aggregates) + /// * the inner Float64 buffer of `sum` also has `null_count() == 0` + /// * `cluster_id` column equals `0..k` (a partial-stats batch is sorted by + /// cluster id, so the first column is fully redundant — but a mismatch + /// indicates a malformed sender, so we reject rather than silently drop) + pub fn from_record_batch(batch: RecordBatch) -> Result { + let schema = batch.schema(); + let md = schema.metadata(); + + // --- Metadata --- + let version = md.get(META_VERSION).map(String::as_str).unwrap_or(""); + if version != PARTIAL_STATS_VERSION { + return Err(Error::index(format!( + "PartialStats version mismatch: got {:?}, expected {}", + version, PARTIAL_STATS_VERSION + ))); + } + let k: usize = md + .get(META_K) + .ok_or_else(|| Error::index(format!("PartialStats missing metadata `{}`", META_K)))? + .parse() + .map_err(|e| { + Error::index(format!("PartialStats metadata `{}` invalid: {}", META_K, e)) + })?; + let dim: usize = md + .get(META_DIM) + .ok_or_else(|| Error::index(format!("PartialStats missing metadata `{}`", META_DIM)))? + .parse() + .map_err(|e| { + Error::index(format!( + "PartialStats metadata `{}` invalid: {}", + META_DIM, e + )) + })?; + let dt_str = md + .get(META_DT) + .map(String::as_str) + .ok_or_else(|| Error::index(format!("PartialStats missing metadata `{}`", META_DT)))?; + if !matches!(dt_str, "l2" | "dot" | "cosine" | "hamming") { + return Err(Error::index(format!( + "PartialStats metadata `{}` has unknown value {:?}", + META_DT, dt_str + ))); + } + let fp_str = md + .get(META_FP) + .ok_or_else(|| Error::index(format!("PartialStats missing metadata `{}`", META_FP)))?; + fingerprint_from_hex(fp_str).map_err(|e| { + Error::index(format!( + "PartialStats metadata `{}` invalid: {}", + META_FP, e + )) + })?; + + // --- Row count --- + if batch.num_rows() != k { + return Err(Error::index(format!( + "PartialStats row count {} does not match metadata k={}", + batch.num_rows(), + k + ))); + } + + // --- Column names + dtypes --- + let expected: [(&str, DataType); 6] = [ + (COL_CLUSTER_ID, DataType::UInt32), + (COL_COUNT, DataType::UInt64), + ( + COL_SUM, + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float64, true)), + dim as i32, + ), + ), + (COL_SQ_NORM_SUM, DataType::Float64), + (COL_LOSS, DataType::Float64), + (COL_RADIUS, DataType::Float32), + ]; + if schema.fields().len() != expected.len() { + return Err(Error::index(format!( + "PartialStats schema has {} columns, expected {}", + schema.fields().len(), + expected.len() + ))); + } + for (idx, (name, want_dt)) in expected.iter().enumerate() { + let field = schema.field(idx); + if field.name() != name { + return Err(Error::index(format!( + "PartialStats schema mismatch at col {}: got {}, expected {}", + idx, + field.name(), + name + ))); + } + // For FixedSizeList we compare the inner length but allow the inner + // field name (e.g. "item" vs "element") to differ — that is purely + // cosmetic and varies by Arrow producer. + let ok = match (field.data_type(), want_dt) { + ( + DataType::FixedSizeList(got_inner, got_n), + DataType::FixedSizeList(want_inner, want_n), + ) => got_n == want_n && got_inner.data_type() == want_inner.data_type(), + (got, want) => got == want, + }; + if !ok { + return Err(Error::index(format!( + "PartialStats column `{}` has dtype {}, expected {}", + name, + field.data_type(), + want_dt + ))); + } + } + + // --- Null counts --- + // + // The writer emits dense buffers for every column (see `build_schema` + // and `compute_partial_stats`). The downstream `merge_partial_stats`, + // `pairwise_*` helpers, and `finalize_centroids` all read `.values()` + // directly and ignore null bitmaps, so a null bit in any of these + // columns would silently corrupt the aggregate. Reject up front + // instead of letting raw buffer slots leak through. + for (idx, (name, _)) in expected.iter().enumerate() { + let col = batch.column(idx); + if col.null_count() != 0 { + return Err(Error::index(format!( + "PartialStats column `{}` has {} nulls; expected dense buffer", + name, + col.null_count() + ))); + } + } + // The `sum` column is a FixedSizeList. Its top-level validity + // is already covered above; also ensure the inner Float64 buffer is + // dense, since `pairwise_sum_fsl_f64` reads it via `.values()` slice. + let sum_inner_nulls = batch + .column(2) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::index("PartialStats column `sum` is not FixedSizeList"))? + .values() + .null_count(); + if sum_inner_nulls != 0 { + return Err(Error::index(format!( + "PartialStats column `sum` has {} nulls in its inner Float64 buffer", + sum_inner_nulls + ))); + } + + // --- cluster_id values must equal 0..k --- + let cluster_id = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::index("PartialStats column 0 is not UInt32 after dtype check"))?; + for (i, v) in cluster_id.values().iter().enumerate() { + if (*v as usize) != i { + return Err(Error::index(format!( + "PartialStats `cluster_id` row {} = {}, expected sequential 0..k", + i, v + ))); + } + } + + Ok(Self { batch }) + } + + pub fn into_record_batch(self) -> RecordBatch { + self.batch + } + + pub fn record_batch(&self) -> &RecordBatch { + &self.batch + } + + pub fn k(&self) -> usize { + self.batch.num_rows() + } + + pub fn dim(&self) -> usize { + match self.batch.schema().field(2).data_type() { + DataType::FixedSizeList(_, n) => *n as usize, + _ => 0, + } + } + + pub fn distance_type(&self) -> DistanceType { + match self + .batch + .schema() + .metadata() + .get(META_DT) + .map(String::as_str) + .unwrap_or("l2") + { + "dot" => DistanceType::Dot, + "cosine" => DistanceType::Cosine, + "hamming" => DistanceType::Hamming, + _ => DistanceType::L2, + } + } + + pub fn centroids_fingerprint(&self) -> [u8; 8] { + self.batch + .schema() + .metadata() + .get(META_FP) + .and_then(|s| fingerprint_from_hex(s).ok()) + .unwrap_or([0u8; 8]) + } + + pub fn total_count(&self) -> u64 { + let counts = self + .batch + .column(1) + .as_any() + .downcast_ref::() + .expect("count column is UInt64"); + counts.values().iter().sum() + } + + pub fn total_loss(&self) -> f64 { + let losses = self + .batch + .column(4) + .as_any() + .downcast_ref::() + .expect("loss column is Float64"); + losses.values().iter().sum() + } +} + +/// 8-byte SHA-256 prefix over the raw bytes of a centroids buffer. +/// +/// Used as `lance.partial_stats.centroids_fingerprint` so reducers can detect +/// mixing partials from different training rounds. +pub fn compute_centroids_fingerprint(centroids: &FixedSizeListArray) -> [u8; 8] { + use sha2::{Digest, Sha256}; + let buffers = centroids.values().to_data().buffers().to_vec(); + let mut hasher = Sha256::new(); + for buf in buffers { + hasher.update(buf.as_slice()); + } + let digest = hasher.finalize(); + let mut out = [0u8; 8]; + out.copy_from_slice(&digest[..8]); + out +} + +fn arrow_error_to_lance(e: ArrowError) -> Error { + Error::index(e.to_string()) +} + +/// E-step: assign every row of `data` to its closest centroid and accumulate +/// `count`/`sum`/`sq_norm_sum`/`loss`/`radius` per cluster. +/// +/// `data` and `centroids` must share dtype, except that `Int8` data is up-cast +/// to `Float32` before assignment, matching `train_ivf_kmeans_step` in +/// `rust/lance/src/index/vector/ivf.rs`. +/// +/// Hamming / `UInt8` is intentionally not supported here. +pub fn compute_partial_stats( + centroids: &FixedSizeListArray, + data: &FixedSizeListArray, + distance_type: DistanceType, +) -> Result { + if matches!(distance_type, DistanceType::Hamming) { + return Err(Error::index( + "distributed Hamming kmeans is not supported in v1", + )); + } + if centroids.value_length() != data.value_length() { + return Err(Error::index(format!( + "centroids dim {} does not match data dim {}", + centroids.value_length(), + data.value_length() + ))); + } + + let dim = centroids.value_length() as usize; + let k = centroids.len(); + let fingerprint = compute_centroids_fingerprint(centroids); + let stats_empty = PartialStats::empty(k, dim, distance_type, fingerprint); + if data.len() == 0 { + return Ok(stats_empty); + } + + // Up-cast Int8 -> Float32 (matches `train_ivf_kmeans_step`). + let data_for_assignment: FixedSizeListArray = match data.value_type() { + DataType::Int8 => convert_int8_to_f32(data)?, + DataType::Float16 | DataType::Float32 | DataType::Float64 => data.clone(), + other => { + return Err(Error::index(format!( + "unsupported data dtype for distributed kmeans: {}", + other + ))); + } + }; + let centroids_for_assignment: FixedSizeListArray = + if centroids.value_type() == data_for_assignment.value_type() { + centroids.clone() + } else { + return Err(Error::index(format!( + "centroids dtype {} does not match data dtype {}", + centroids.value_type(), + data_for_assignment.value_type() + ))); + }; + + // Cosine: caller (Layer 2) is responsible for normalizing both centroids + // and data before reaching here, so we run the assignment as L2. + let kmeans = KMeans::with_centroids( + centroids_for_assignment.values().clone(), + dim, + match distance_type { + DistanceType::Cosine => DistanceType::L2, + other => other, + }, + f64::MAX, + ); + let (membership, distances) = kmeans + .compute_membership_and_distances(&data_for_assignment) + .map_err(arrow_error_to_lance)?; + + let mut counts_vec = vec![0u64; k]; + let mut sum_vec = vec![0.0f64; k * dim]; + let mut sq_norm = vec![0.0f64; k]; + let mut loss = vec![0.0f64; k]; + let mut radius = vec![0.0f32; k]; + + let value_array = data_for_assignment.values(); + match value_array.data_type() { + DataType::Float32 => accumulate::( + value_array.as_primitive::().values(), + dim, + &membership, + &distances, + &mut counts_vec, + &mut sum_vec, + &mut sq_norm, + &mut loss, + &mut radius, + |v| v as f64, + ), + DataType::Float16 => accumulate::( + value_array.as_primitive::().values(), + dim, + &membership, + &distances, + &mut counts_vec, + &mut sum_vec, + &mut sq_norm, + &mut loss, + &mut radius, + |v| v.to_f64(), + ), + DataType::Float64 => accumulate::( + value_array.as_primitive::().values(), + dim, + &membership, + &distances, + &mut counts_vec, + &mut sum_vec, + &mut sq_norm, + &mut loss, + &mut radius, + |v| v, + ), + other => { + return Err(Error::index(format!( + "unsupported assignment dtype: {}", + other + ))); + } + } + + let new_count = Arc::new(UInt64Array::from(counts_vec)); + let mut sum_builder = FixedSizeListBuilder::new(Float64Builder::new(), dim as i32); + for ci in 0..k { + for d in 0..dim { + sum_builder.values().append_value(sum_vec[ci * dim + d]); + } + sum_builder.append(true); + } + let new_sum = Arc::new(sum_builder.finish()); + let new_sq = Arc::new(Float64Array::from(sq_norm)); + let new_loss = Arc::new(Float64Array::from(loss)); + let new_radius = Arc::new(Float32Array::from(radius)); + + let schema = stats_empty.batch.schema(); + let cluster_id_col = stats_empty.batch.column(0).clone(); + let batch = RecordBatch::try_new( + schema, + vec![ + cluster_id_col, + new_count, + new_sum, + new_sq, + new_loss, + new_radius, + ], + ) + .map_err(arrow_error_to_lance)?; + Ok(PartialStats { batch }) +} + +#[allow(clippy::too_many_arguments)] +fn accumulate( + values: &[T::Native], + dim: usize, + membership: &[Option], + distances: &[Option], + counts: &mut [u64], + sums: &mut [f64], + sq_norm: &mut [f64], + loss: &mut [f64], + radius: &mut [f32], + to_f64: impl Fn(T::Native) -> f64, +) where + T::Native: Copy, +{ + for (row_idx, (&m, &d)) in membership.iter().zip(distances.iter()).enumerate() { + let (Some(c), Some(dist)) = (m, d) else { + continue; + }; + let ci = c as usize; + counts[ci] += 1; + loss[ci] += dist as f64; + if dist > radius[ci] { + radius[ci] = dist; + } + let row = &values[row_idx * dim..(row_idx + 1) * dim]; + let mut row_sq = 0.0f64; + for (offset, &v) in row.iter().enumerate() { + let v64 = to_f64(v); + sums[ci * dim + offset] += v64; + row_sq += v64 * v64; + } + sq_norm[ci] += row_sq; + } +} + +fn convert_int8_to_f32(data: &FixedSizeListArray) -> Result { + let values = data + .values() + .as_any() + .downcast_ref::>() + .ok_or_else(|| Error::index("expected Int8 values"))?; + let f32_values: Float32Array = values.iter().map(|v| v.map(|x| x as f32)).collect(); + FixedSizeListArray::try_new_from_values(f32_values, data.value_length()) + .map_err(arrow_error_to_lance) +} + +/// Combine two partial stats produced against the same centroids. +pub fn merge_partial_stats(a: PartialStats, b: PartialStats) -> Result { + if a.batch.schema().metadata().get(META_VERSION) + != b.batch.schema().metadata().get(META_VERSION) + { + return Err(Error::index("PartialStats version mismatch")); + } + if a.k() != b.k() || a.dim() != b.dim() { + return Err(Error::index(format!( + "PartialStats shape mismatch: ({},{}) vs ({},{})", + a.k(), + a.dim(), + b.k(), + b.dim() + ))); + } + if a.distance_type() != b.distance_type() { + return Err(Error::index("PartialStats distance_type mismatch")); + } + if a.centroids_fingerprint() != b.centroids_fingerprint() { + return Err(Error::index("PartialStats centroids_fingerprint mismatch")); + } + + let k = a.k(); + let dim = a.dim(); + let counts = pairwise_sum_u64(a.batch.column(1), b.batch.column(1)); + let sums = pairwise_sum_fsl_f64(a.batch.column(2), b.batch.column(2), k, dim); + let sq_norm = pairwise_sum_f64(a.batch.column(3), b.batch.column(3)); + let loss = pairwise_sum_f64(a.batch.column(4), b.batch.column(4)); + let radius = pairwise_max_f32(a.batch.column(5), b.batch.column(5)); + + let schema = a.batch.schema(); + let batch = RecordBatch::try_new( + schema, + vec![ + a.batch.column(0).clone(), + Arc::new(UInt64Array::from(counts)), + Arc::new(sums), + Arc::new(Float64Array::from(sq_norm)), + Arc::new(Float64Array::from(loss)), + Arc::new(Float32Array::from(radius)), + ], + ) + .map_err(arrow_error_to_lance)?; + Ok(PartialStats { batch }) +} + +/// Fold an iterator of partial stats. Returns `Err` if the iterator is empty. +pub fn reduce_partial_stats>(iter: I) -> Result { + let mut iter = iter.into_iter(); + let mut acc = iter + .next() + .ok_or_else(|| Error::index("reduce_partial_stats: empty iterator"))?; + for next in iter { + acc = merge_partial_stats(acc, next)?; + } + Ok(acc) +} + +fn pairwise_sum_u64(a: &dyn Array, b: &dyn Array) -> Vec { + let a = a.as_any().downcast_ref::().unwrap(); + let b = b.as_any().downcast_ref::().unwrap(); + a.values() + .iter() + .zip(b.values().iter()) + .map(|(x, y)| x + y) + .collect() +} + +fn pairwise_sum_f64(a: &dyn Array, b: &dyn Array) -> Vec { + let a = a.as_any().downcast_ref::().unwrap(); + let b = b.as_any().downcast_ref::().unwrap(); + a.values() + .iter() + .zip(b.values().iter()) + .map(|(x, y)| x + y) + .collect() +} + +fn pairwise_max_f32(a: &dyn Array, b: &dyn Array) -> Vec { + let a = a.as_any().downcast_ref::().unwrap(); + let b = b.as_any().downcast_ref::().unwrap(); + a.values() + .iter() + .zip(b.values().iter()) + .map(|(x, y)| x.max(*y)) + .collect() +} + +fn pairwise_sum_fsl_f64(a: &dyn Array, b: &dyn Array, k: usize, dim: usize) -> FixedSizeListArray { + let a = a.as_any().downcast_ref::().unwrap(); + let b = b.as_any().downcast_ref::().unwrap(); + let av = a.values().as_primitive::().values(); + let bv = b.values().as_primitive::().values(); + let mut builder = FixedSizeListBuilder::new(Float64Builder::new(), dim as i32); + for ci in 0..k { + for d in 0..dim { + builder + .values() + .append_value(av[ci * dim + d] + bv[ci * dim + d]); + } + builder.append(true); + } + builder.finish() +} + +/// Compute new centroids from accumulated stats. +/// +/// `prev` provides the dtype and the fallback for empty clusters. +pub fn finalize_centroids( + stats: &PartialStats, + prev: &FixedSizeListArray, +) -> Result { + if stats.k() != prev.len() { + return Err(Error::index(format!( + "stats.k {} != prev.len {}", + stats.k(), + prev.len() + ))); + } + if stats.dim() != prev.value_length() as usize { + return Err(Error::index(format!( + "stats.dim {} != prev.dim {}", + stats.dim(), + prev.value_length() + ))); + } + if stats.total_count() == 0 { + return Err(Error::index("no training data assigned")); + } + + // Reject stats from a different round. Without this guard, accidentally + // finalizing against the wrong `prev` (e.g. a stale cached batch) would + // silently produce incorrect centroids. The fingerprint is short (8 bytes + // of SHA-256) but already used to gate `merge_partial_stats`; reusing it + // here closes the same trust boundary on the finalize path. + let expected_fp = compute_centroids_fingerprint(prev); + if stats.centroids_fingerprint() != expected_fp { + return Err(Error::index( + "finalize_centroids: PartialStats fingerprint does not match prev centroids \ + (stale stats from a different training round?)", + )); + } + + let k = stats.k(); + let dim = stats.dim(); + let counts = stats + .batch + .column(1) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::index("finalize_centroids: count column is not UInt64"))?; + let sums_arr = stats + .batch + .column(2) + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::index("finalize_centroids: sum column is not FixedSizeList"))?; + let sums = sums_arr.values().as_primitive::().values(); + + match prev.value_type() { + DataType::Float32 => { + let prev_vals = prev.values().as_primitive::().values(); + let mut out = prev_vals.to_vec(); + for ci in 0..k { + let n = counts.value(ci); + if n == 0 { + continue; + } + for d in 0..dim { + let s = sums[ci * dim + d]; + if !s.is_finite() { + return Err(Error::index(format!( + "non-finite sum at cluster {}, dim {}", + ci, d + ))); + } + out[ci * dim + d] = (s / n as f64) as f32; + } + } + FixedSizeListArray::try_new_from_values(Float32Array::from(out), dim as i32) + .map_err(arrow_error_to_lance) + } + DataType::Float16 => { + let prev_vals = prev.values().as_primitive::().values(); + let mut out: Vec = prev_vals.to_vec(); + for ci in 0..k { + let n = counts.value(ci); + if n == 0 { + continue; + } + for d in 0..dim { + let s = sums[ci * dim + d]; + if !s.is_finite() { + return Err(Error::index(format!( + "non-finite sum at cluster {}, dim {}", + ci, d + ))); + } + out[ci * dim + d] = f16::from_f64(s / n as f64); + } + } + FixedSizeListArray::try_new_from_values(Float16Array::from_iter_values(out), dim as i32) + .map_err(arrow_error_to_lance) + } + DataType::Float64 => { + let prev_vals = prev.values().as_primitive::().values(); + let mut out = prev_vals.to_vec(); + for ci in 0..k { + let n = counts.value(ci); + if n == 0 { + continue; + } + for d in 0..dim { + let s = sums[ci * dim + d]; + if !s.is_finite() { + return Err(Error::index(format!( + "non-finite sum at cluster {}, dim {}", + ci, d + ))); + } + out[ci * dim + d] = s / n as f64; + } + } + FixedSizeListArray::try_new_from_values(Float64Array::from(out), dim as i32) + .map_err(arrow_error_to_lance) + } + other => Err(Error::index(format!( + "finalize_centroids: unsupported prev dtype {}", + other + ))), + } +} + +const RESERVOIR_VEC_COL: &str = "vec"; + +fn samples_schema(value_type: DataType, dim: usize) -> SchemaRef { + let item = DataType::FixedSizeList(Arc::new(Field::new("item", value_type, true)), dim as i32); + Arc::new(Schema::new(vec![Field::new( + RESERVOIR_VEC_COL, + item, + false, + )])) +} + +/// Algorithm-R reservoir sample of `target` rows from `data`. +/// +/// Output schema: `vec: FixedSizeList`. +/// If `data.len() <= target`, returns all rows verbatim. Internally a thin +/// wrapper over [`StreamingReservoir`] so single-batch and streaming callers +/// share the same RNG consumption. +pub fn local_reservoir_sample( + data: &FixedSizeListArray, + target: usize, + rng_seed: u64, +) -> Result { + let mut reservoir = StreamingReservoir::new(target, rng_seed); + reservoir.feed(data)?; + reservoir.into_record_batch() +} + +/// Streaming Algorithm-R reservoir sampler over a sequence of FSL chunks. +/// +/// Layer-1 (`local_reservoir_sample`) and Layer-2 (`sample_round_0`) share +/// this implementation so a same-seed run consumes the RNG identically +/// regardless of how the input was chunked. +pub struct StreamingReservoir { + rng: rand::rngs::StdRng, + target: usize, + seen: usize, + /// Materialized rows currently held by the reservoir, all sliced from + /// previously-fed batches. We re-arrow them at the end. + held: Vec, + /// Parallel index into `held`: `(batch_idx, row_idx_within_batch)` pairs + /// for each row currently in the reservoir. + indices: Vec<(usize, usize)>, + value_type: Option, + dim: Option, +} + +impl StreamingReservoir { + pub fn new(target: usize, rng_seed: u64) -> Self { + use rand::SeedableRng; + Self { + rng: rand::rngs::StdRng::seed_from_u64(rng_seed), + target, + seen: 0, + held: Vec::new(), + indices: Vec::new(), + value_type: None, + dim: None, + } + } + + /// Feed a chunk of vectors. Updates the reservoir per Algorithm-R; rows + /// not selected are dropped immediately. Schema-consistency across chunks + /// is enforced (every chunk must agree on `(value_type, dim)`). + pub fn feed(&mut self, chunk: &FixedSizeListArray) -> Result<()> { + if chunk.is_empty() { + return Ok(()); + } + + let chunk_value_type = chunk.value_type(); + let chunk_dim = chunk.value_length() as usize; + match (&self.value_type, self.dim) { + (None, _) => { + self.value_type = Some(chunk_value_type); + self.dim = Some(chunk_dim); + } + (Some(prev_type), Some(prev_dim)) => { + if prev_type != &chunk_value_type || prev_dim != chunk_dim { + return Err(Error::index(format!( + "StreamingReservoir: schema mismatch across chunks (prev=FSL<{:?},{}> got=FSL<{:?},{}>)", + prev_type, prev_dim, chunk_value_type, chunk_dim + ))); + } + } + _ => unreachable!(), + } + + let batch_idx = self.held.len(); + self.held.push(chunk.clone()); + + let n = chunk.len(); + for row_idx in 0..n { + if self.indices.len() < self.target { + self.indices.push((batch_idx, row_idx)); + } else { + use rand::Rng; + // i = self.seen + row_idx is the global 0-based index of this row. + let global_idx = self.seen + row_idx; + let j = self.rng.random_range(0..=global_idx); + if j < self.target { + self.indices[j] = (batch_idx, row_idx); + } + } + } + self.seen += n; + Ok(()) + } + + /// Finalize the reservoir, producing one [`RecordBatch`] of selected rows. + /// + /// If the reservoir never saw any rows, returns an empty batch with a + /// `Float32` placeholder schema (callers that care about value type + /// should feed at least one chunk first). + pub fn into_record_batch(self) -> Result { + let value_type = self.value_type.unwrap_or(DataType::Float32); + let dim = self.dim.unwrap_or(0); + let schema = samples_schema(value_type, dim); + + if self.indices.is_empty() { + let empty = arrow_array::array::new_empty_array(schema.field(0).data_type()); + return RecordBatch::try_new(schema, vec![empty]).map_err(arrow_error_to_lance); + } + + // Take per source batch to keep concat cheap, then concatenate. + let mut per_batch_indices: HashMap> = HashMap::new(); + for (b, r) in &self.indices { + per_batch_indices.entry(*b).or_default().push(*r as u32); + } + + let mut taken: Vec = Vec::with_capacity(self.held.len()); + for (b_idx, batch) in self.held.iter().enumerate() { + let Some(rows) = per_batch_indices.get(&b_idx) else { + continue; + }; + let take_array = UInt32Array::from(rows.clone()); + let arr = + arrow::compute::take(batch, &take_array, None).map_err(arrow_error_to_lance)?; + let fsl = arr + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::index("StreamingReservoir: take produced non-FSL output".to_string()) + })? + .clone(); + taken.push(fsl); + } + + let arrays: Vec<&dyn Array> = taken.iter().map(|a| a as &dyn Array).collect(); + let combined = arrow::compute::concat(&arrays).map_err(arrow_error_to_lance)?; + let fsl = combined + .as_any() + .downcast_ref::() + .cloned() + .ok_or_else(|| Error::index("StreamingReservoir: concat result is not FSL"))?; + + RecordBatch::try_new(schema, vec![Arc::new(fsl)]).map_err(arrow_error_to_lance) + } +} + +fn concat_samples(samples: Vec) -> Result { + if samples.is_empty() { + return Err(Error::index( + "select/bootstrap requires at least one sample batch", + )); + } + let arrays: Vec<&dyn Array> = samples.iter().map(|b| b.column(0).as_ref()).collect(); + let combined = arrow::compute::concat(&arrays).map_err(arrow_error_to_lance)?; + combined + .as_any() + .downcast_ref::() + .cloned() + .ok_or_else(|| Error::index("concatenated samples are not FixedSizeList")) +} + +/// Driver-side: pick `k` rows uniformly at random from the union of worker samples. +pub fn select_initial_centroids( + samples: Vec, + k: usize, + rng_seed: u64, +) -> Result { + use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom}; + + let combined = concat_samples(samples)?; + if combined.len() < k { + return Err(Error::index(format!( + "not enough samples ({}) to select {} centroids", + combined.len(), + k + ))); + } + let mut rng = StdRng::seed_from_u64(rng_seed); + let mut idx: Vec = (0..combined.len() as u32).collect(); + idx.shuffle(&mut rng); + let take = UInt32Array::from(idx[..k].to_vec()); + let chosen = arrow::compute::take(&combined, &take, None).map_err(arrow_error_to_lance)?; + chosen + .as_any() + .downcast_ref::() + .cloned() + .ok_or_else(|| Error::index("select_initial_centroids: take returned non-FSL")) +} + +/// Driver-side: run single-machine `train_kmeans` over the union of worker samples +/// to obtain a high-quality initial set of centroids. For `k > 256`, the existing +/// kmeans engine automatically falls back to hierarchical kmeans (see +/// `train_kmeans` in this crate). +/// +/// Cosine handling: `train_kmeans` does not implement Cosine in its inner kernels +/// (`argmin_value_float_with_bias` panics). To match the Layer-1 contract in +/// `compute_partial_stats` (caller normalizes, kernel runs as L2), this function +/// L2-normalizes the combined samples up front and dispatches the inner k-means +/// with `DistanceType::L2` whenever the requested `distance_type` is Cosine. The +/// resulting centroids — means in normalized space — are then re-normalized so +/// the returned array satisfies the Cosine invariant that subsequent rounds +/// expect (unit-norm centroids, distance computed as L2). +pub fn bootstrap_centroids( + samples: Vec, + k: usize, + distance_type: DistanceType, + rng_seed: u64, +) -> Result { + let combined = concat_samples(samples)?; + let dim = combined.value_length() as usize; + + // Cosine: normalize inputs once, then run inner kmeans as L2 (mirrors the + // assignment dispatch in `compute_partial_stats`). + let is_cosine = matches!(distance_type, DistanceType::Cosine); + let (combined, inner_dt) = if is_cosine { + ( + normalize_fsl_owned(combined).map_err(arrow_error_to_lance)?, + DistanceType::L2, + ) + } else { + (combined, distance_type) + }; + let params = KMeansParams::default() + .with_distance_type(inner_dt) + .with_seed(rng_seed); + let centroids = match combined.value_type() { + DataType::Float32 => { + let arr = combined.values().as_primitive::().clone(); + let model = train_kmeans::(&arr, params, dim, k, 256)?; + FixedSizeListArray::try_new_from_values( + model.centroids.as_primitive::().clone(), + dim as i32, + ) + .map_err(arrow_error_to_lance) + } + DataType::Float16 => { + let arr = combined.values().as_primitive::().clone(); + let model = train_kmeans::(&arr, params, dim, k, 256)?; + FixedSizeListArray::try_new_from_values( + model.centroids.as_primitive::().clone(), + dim as i32, + ) + .map_err(arrow_error_to_lance) + } + DataType::Float64 => { + let arr = combined.values().as_primitive::().clone(); + let model = train_kmeans::(&arr, params, dim, k, 256)?; + FixedSizeListArray::try_new_from_values( + model.centroids.as_primitive::().clone(), + dim as i32, + ) + .map_err(arrow_error_to_lance) + } + other => Err(Error::index(format!( + "bootstrap_centroids: unsupported dtype {}", + other + ))), + }?; + + // Cosine: the inner kmeans operates on normalized samples and returns + // arithmetic means, which are not generally unit-norm. Re-normalize so + // callers (and Layer-1 `compute_partial_stats`) get unit-norm centroids. + if is_cosine { + normalize_fsl_owned(centroids).map_err(arrow_error_to_lance) + } else { + Ok(centroids) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{ArrayRef, FixedSizeListArray}; + use arrow_schema::DataType; + use lance_arrow::FixedSizeListArrayExt; + use lance_linalg::distance::DistanceType; + + #[test] + fn test_partial_stats_empty_schema() { + let stats = PartialStats::empty(64, 16, DistanceType::L2, [0u8; 8]); + let batch = stats.into_record_batch(); + + assert_eq!(batch.num_rows(), 64); + let schema = batch.schema(); + assert_eq!(schema.field(0).name(), "cluster_id"); + assert_eq!(schema.field(0).data_type(), &DataType::UInt32); + assert_eq!(schema.field(1).name(), "count"); + assert_eq!(schema.field(1).data_type(), &DataType::UInt64); + assert!(matches!( + schema.field(2).data_type(), + DataType::FixedSizeList(_, 16) + )); + let md = schema.metadata(); + assert_eq!( + md.get("lance.partial_stats.version").map(String::as_str), + Some("1") + ); + assert_eq!( + md.get("lance.partial_stats.k").map(String::as_str), + Some("64") + ); + assert_eq!( + md.get("lance.partial_stats.dim").map(String::as_str), + Some("16") + ); + assert_eq!( + md.get("lance.partial_stats.distance_type") + .map(String::as_str), + Some("l2") + ); + } + + #[test] + fn test_from_record_batch_rejects_wrong_version() { + let stats = PartialStats::empty(4, 8, DistanceType::L2, [0u8; 8]); + let batch = stats.into_record_batch(); + let mut md = batch.schema().metadata().clone(); + md.insert(META_VERSION.into(), "999".into()); + let new_schema = Arc::new((*batch.schema()).clone().with_metadata(md)); + let bad = RecordBatch::try_new(new_schema, batch.columns().to_vec()).unwrap(); + assert!(PartialStats::from_record_batch(bad).is_err()); + } + + #[test] + fn test_from_record_batch_rejects_malformed_inputs() { + // baseline: a valid empty PartialStats batch + let valid = PartialStats::empty(4, 8, DistanceType::L2, [0u8; 8]).into_record_batch(); + assert!(PartialStats::from_record_batch(valid.clone()).is_ok()); + + // helper: clone with mutated metadata + let mutate_meta = |key: &str, value: Option<&str>| -> RecordBatch { + let mut md = valid.schema().metadata().clone(); + match value { + Some(v) => { + md.insert(key.into(), v.into()); + } + None => { + md.remove(key); + } + } + let s = Arc::new((*valid.schema()).clone().with_metadata(md)); + RecordBatch::try_new(s, valid.columns().to_vec()).unwrap() + }; + + // missing each required metadata key + for key in [META_K, META_DIM, META_DT, META_FP] { + assert!( + PartialStats::from_record_batch(mutate_meta(key, None)).is_err(), + "missing `{}` should be rejected", + key + ); + } + + // unparsable k / dim / fingerprint + assert!(PartialStats::from_record_batch(mutate_meta(META_K, Some("nope"))).is_err()); + assert!(PartialStats::from_record_batch(mutate_meta(META_DIM, Some("-1"))).is_err()); + assert!(PartialStats::from_record_batch(mutate_meta(META_FP, Some("zz"))).is_err()); + + // unknown distance_type + assert!(PartialStats::from_record_batch(mutate_meta(META_DT, Some("manhattan"))).is_err()); + + // dim mismatch: metadata says dim=4 but the FSL column is dim=8 + let bad_dim = mutate_meta(META_DIM, Some("4")); + assert!(PartialStats::from_record_batch(bad_dim).is_err()); + + // k mismatch: metadata says k=999 but the batch has 4 rows + let bad_k = mutate_meta(META_K, Some("999")); + assert!(PartialStats::from_record_batch(bad_k).is_err()); + + // wrong dtype on the `count` column (UInt32 instead of UInt64) + let mut cols = valid.columns().to_vec(); + cols[1] = Arc::new(UInt32Array::from(vec![0u32; 4])); + let mut fields: Vec<_> = valid + .schema() + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + fields[1] = Field::new(COL_COUNT, DataType::UInt32, false); + let bad_count_schema = + Arc::new(Schema::new(fields).with_metadata(valid.schema().metadata().clone())); + let bad_count = RecordBatch::try_new(bad_count_schema, cols).unwrap(); + assert!(PartialStats::from_record_batch(bad_count).is_err()); + + // non-sequential cluster_id values + let cols = valid.columns().to_vec(); + let mut bad_ids = vec![0u32; 4]; + bad_ids[2] = 99; + let mut cols = cols; + cols[0] = Arc::new(UInt32Array::from(bad_ids)); + let bad_ids_batch = RecordBatch::try_new(valid.schema(), cols).unwrap(); + assert!(PartialStats::from_record_batch(bad_ids_batch).is_err()); + } + + #[test] + fn test_from_record_batch_rejects_nulls_in_dense_columns() { + use arrow::buffer::NullBuffer; + let k = 4usize; + let dim = 8usize; + let valid = PartialStats::empty(k, dim, DistanceType::L2, [0u8; 8]).into_record_batch(); + + // RecordBatch::try_new itself rejects nulls in fields declared + // nullable=false (which the canonical schema uses for every top-level + // column). To smuggle a null past that check and let + // `from_record_batch` see it, rebuild the schema so just the mutated + // column's field is nullable=true. The validator still must reject. + let make_schema_for_replaced = |idx: usize| -> SchemaRef { + let mut fields: Vec = valid + .schema() + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + let field = &fields[idx]; + fields[idx] = Field::new(field.name(), field.data_type().clone(), true); + Arc::new(Schema::new(fields).with_metadata(valid.schema().metadata().clone())) + }; + + let replace_col = |idx: usize, replacement: ArrayRef| -> RecordBatch { + let mut cols = valid.columns().to_vec(); + cols[idx] = replacement; + RecordBatch::try_new(make_schema_for_replaced(idx), cols).unwrap() + }; + + let bad_count: ArrayRef = + Arc::new(UInt64Array::from(vec![None, Some(0), Some(0), Some(0)])); + assert!( + PartialStats::from_record_batch(replace_col(1, bad_count)).is_err(), + "null in `count` must be rejected" + ); + + let bad_sq_norm: ArrayRef = Arc::new(Float64Array::from(vec![ + None, + Some(0.0), + Some(0.0), + Some(0.0), + ])); + assert!( + PartialStats::from_record_batch(replace_col(3, bad_sq_norm)).is_err(), + "null in `sq_norm_sum` must be rejected" + ); + + let bad_loss: ArrayRef = Arc::new(Float64Array::from(vec![ + None, + Some(0.0), + Some(0.0), + Some(0.0), + ])); + assert!( + PartialStats::from_record_batch(replace_col(4, bad_loss)).is_err(), + "null in `loss` must be rejected" + ); + + let bad_radius: ArrayRef = Arc::new(Float32Array::from(vec![ + None, + Some(0.0_f32), + Some(0.0), + Some(0.0), + ])); + assert!( + PartialStats::from_record_batch(replace_col(5, bad_radius)).is_err(), + "null in `radius` must be rejected" + ); + + // sum: top-level FSL row null. Build with a full inner buffer (k*dim + // values) and a top-level NullBuffer that flips row 0. + let dense_inner = Float64Array::from(vec![0.0_f64; k * dim]); + let mut top_validity = vec![true; k]; + top_validity[0] = false; + let nulls = NullBuffer::from(top_validity); + let sum_field = Arc::new(Field::new("item", DataType::Float64, true)); + let bad_sum_top: ArrayRef = Arc::new(FixedSizeListArray::new( + sum_field.clone(), + dim as i32, + Arc::new(dense_inner), + Some(nulls), + )); + assert!( + PartialStats::from_record_batch(replace_col(2, bad_sum_top)).is_err(), + "top-level null in `sum` must be rejected" + ); + + // sum: a null in the inner Float64 buffer. Build a Float64 array of + // length k*dim with a single null in the first cell, then wrap into + // a fully-valid FSL (top-level all valid). Inner field already + // declares nullable=true so the canonical schema accepts the buffer. + let mut inner_builder = Float64Builder::new(); + inner_builder.append_null(); + for _ in 1..(k * dim) { + inner_builder.append_value(0.0); + } + let inner_with_null = inner_builder.finish(); + let bad_sum_inner: ArrayRef = Arc::new(FixedSizeListArray::new( + sum_field, + dim as i32, + Arc::new(inner_with_null), + None, + )); + let bad_inner_batch = RecordBatch::try_new(valid.schema(), { + let mut cols = valid.columns().to_vec(); + cols[2] = bad_sum_inner; + cols + }) + .unwrap(); + assert!( + PartialStats::from_record_batch(bad_inner_batch).is_err(), + "inner Float64 null in `sum` must be rejected" + ); + } + + #[test] + fn test_finalize_centroids_rejects_fingerprint_mismatch() { + let prev1 = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 1.0, 1.0]), + 2, + ) + .unwrap(); + let prev2 = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![10.0_f32, 10.0, 20.0, 20.0]), + 2, + ) + .unwrap(); + // build stats against prev1 with at least one assigned point so the + // total_count check doesn't short-circuit before fingerprint check + let data = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 1.0, 1.0]), + 2, + ) + .unwrap(); + let stats = super::compute_partial_stats(&prev1, &data, DistanceType::L2).unwrap(); + // applying the stats to a different `prev` must error + let err = finalize_centroids(&stats, &prev2) + .expect_err("stats from prev1 must not finalize against prev2"); + let msg = format!("{}", err); + assert!( + msg.contains("fingerprint"), + "error should mention fingerprint mismatch, got: {}", + msg + ); + // sanity: applying to the same prev still works + finalize_centroids(&stats, &prev1).unwrap(); + } + + #[test] + fn test_fingerprint_is_stable_and_dtype_independent() { + let f32_centroids = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![1.0_f32, 2.0, 3.0, 4.0]), + 2, + ) + .unwrap(); + let fp1 = compute_centroids_fingerprint(&f32_centroids); + let fp2 = compute_centroids_fingerprint(&f32_centroids); + assert_eq!(fp1, fp2, "fingerprint must be deterministic"); + + let mutated = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![1.0_f32, 2.0, 3.0, 5.0]), + 2, + ) + .unwrap(); + assert_ne!( + compute_centroids_fingerprint(&mutated), + fp1, + "different bytes must produce different fingerprint" + ); + } + + #[test] + fn test_compute_partial_stats_l2_basic() { + // 4 centroids in 2-D, 6 vectors; manually-computed expected counts. + let centroids = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 10.0, 0.0, 0.0, 10.0, 10.0, 10.0]), + 2, + ) + .unwrap(); + let data = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![ + 0.1_f32, 0.1, // -> cluster 0 + 10.0, 0.1, // -> cluster 1 + 0.0, 9.9, // -> cluster 2 + 9.9, 9.9, // -> cluster 3 + 10.1, 0.0, // -> cluster 1 + 0.0, 0.0, // -> cluster 0 + ]), + 2, + ) + .unwrap(); + + let stats = compute_partial_stats(¢roids, &data, DistanceType::L2).unwrap(); + let counts = stats + .record_batch() + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!(counts, vec![2, 2, 1, 1]); + assert_eq!(stats.total_count(), 6); + assert!(stats.total_loss() >= 0.0); + } + + #[test] + fn test_compute_partial_stats_rejects_hamming() { + let centroids = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 1.0, 1.0]), + 2, + ) + .unwrap(); + let data = + FixedSizeListArray::try_new_from_values(Float32Array::from(vec![0.0_f32, 0.0]), 2) + .unwrap(); + assert!(compute_partial_stats(¢roids, &data, DistanceType::Hamming).is_err()); + } + + #[test] + fn test_compute_partial_stats_all_nan_returns_empty() { + let centroids = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 10.0, 10.0]), + 2, + ) + .unwrap(); + let data = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![f32::NAN, f32::NAN, f32::NAN, f32::NAN]), + 2, + ) + .unwrap(); + let stats = compute_partial_stats(¢roids, &data, DistanceType::L2).unwrap(); + assert_eq!(stats.total_count(), 0, "all NaN -> no assignments"); + } + + #[test] + fn test_merge_partial_stats_simple_sum() { + let centroids = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 10.0, 10.0]), + 2, + ) + .unwrap(); + let d1 = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 1.0, 1.0]), + 2, + ) + .unwrap(); + let d2 = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![10.0_f32, 10.0, 9.0, 9.0]), + 2, + ) + .unwrap(); + let s1 = compute_partial_stats(¢roids, &d1, DistanceType::L2).unwrap(); + let s2 = compute_partial_stats(¢roids, &d2, DistanceType::L2).unwrap(); + + let merged = merge_partial_stats(s1, s2).unwrap(); + let counts = merged + .record_batch() + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!(counts, vec![2, 2]); + assert_eq!(merged.total_count(), 4); + } + + #[test] + fn test_merge_rejects_fingerprint_mismatch() { + let c1 = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 1.0, 1.0]), + 2, + ) + .unwrap(); + let c2 = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 5.0, 5.0]), + 2, + ) + .unwrap(); + let d = FixedSizeListArray::try_new_from_values(Float32Array::from(vec![0.5_f32, 0.5]), 2) + .unwrap(); + let s1 = compute_partial_stats(&c1, &d, DistanceType::L2).unwrap(); + let s2 = compute_partial_stats(&c2, &d, DistanceType::L2).unwrap(); + assert!(merge_partial_stats(s1, s2).is_err()); + } + + #[test] + fn test_merge_empty_is_identity() { + let centroids = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 1.0, 1.0]), + 2, + ) + .unwrap(); + let d = FixedSizeListArray::try_new_from_values(Float32Array::from(vec![0.1_f32, 0.1]), 2) + .unwrap(); + let s = compute_partial_stats(¢roids, &d, DistanceType::L2).unwrap(); + let fp = compute_centroids_fingerprint(¢roids); + let empty = PartialStats::empty(2, 2, DistanceType::L2, fp); + let merged = merge_partial_stats(s.clone(), empty).unwrap(); + assert_eq!(merged.total_count(), s.total_count()); + assert_eq!(merged.total_loss(), s.total_loss()); + } + + #[test] + fn test_finalize_centroids_basic_f32() { + let prev = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 100.0, 100.0]), + 2, + ) + .unwrap(); + let data = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 2.0, 2.0, 99.0, 99.0, 101.0, 101.0]), + 2, + ) + .unwrap(); + let stats = compute_partial_stats(&prev, &data, DistanceType::L2).unwrap(); + let new = finalize_centroids(&stats, &prev).unwrap(); + let v = new.values().as_primitive::().values().to_vec(); + assert!((v[0] - 1.0).abs() < 1e-5); + assert!((v[1] - 1.0).abs() < 1e-5); + assert!((v[2] - 100.0).abs() < 1e-5); + assert!((v[3] - 100.0).abs() < 1e-5); + } + + #[test] + fn test_finalize_centroids_empty_cluster_keeps_prev() { + let prev = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 100.0, 100.0]), + 2, + ) + .unwrap(); + let data = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 0.5, 0.5]), + 2, + ) + .unwrap(); + let stats = compute_partial_stats(&prev, &data, DistanceType::L2).unwrap(); + let new = finalize_centroids(&stats, &prev).unwrap(); + let v = new.values().as_primitive::().values().to_vec(); + // cluster 1 has 0 assignments and must retain prev[1] + assert!((v[2] - 100.0).abs() < 1e-6); + assert!((v[3] - 100.0).abs() < 1e-6); + } + + #[test] + fn test_finalize_centroids_all_empty_errors() { + let prev = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 100.0, 100.0]), + 2, + ) + .unwrap(); + let fp = compute_centroids_fingerprint(&prev); + let empty = PartialStats::empty(2, 2, DistanceType::L2, fp); + assert!(finalize_centroids(&empty, &prev).is_err()); + } + + fn random_fsl_f32(seed: u64, rows: usize, dim: usize) -> FixedSizeListArray { + use rand::{Rng, SeedableRng, rngs::StdRng}; + let mut rng = StdRng::seed_from_u64(seed); + let total = rows * dim; + let v: Vec = (0..total) + .map(|_| rng.random_range(-10.0..10.0_f32)) + .collect(); + FixedSizeListArray::try_new_from_values(Float32Array::from(v), dim as i32).unwrap() + } + + fn assert_partial_stats_close(a: &PartialStats, b: &PartialStats, eps: f64) { + assert_eq!(a.k(), b.k()); + let ca = a + .record_batch() + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + let cb = b + .record_batch() + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!(ca, cb, "counts must match exactly"); + + let sa = a + .record_batch() + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_primitive::() + .values() + .to_vec(); + let sb = b + .record_batch() + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_primitive::() + .values() + .to_vec(); + for (x, y) in sa.iter().zip(sb.iter()) { + assert!((x - y).abs() < eps, "sum mismatch: {} vs {}", x, y); + } + } + + #[test] + fn test_distributed_equals_single_three_way_split() { + let centroids = random_fsl_f32(1, 16, 8); + let data = random_fsl_f32(2, 600, 8); + + let single = compute_partial_stats(¢roids, &data, DistanceType::L2).unwrap(); + + let dim = 8; + let values = data.values().as_primitive::().values(); + let split_a = FixedSizeListArray::try_new_from_values( + Float32Array::from(values[..200 * dim].to_vec()), + dim as i32, + ) + .unwrap(); + let split_b = FixedSizeListArray::try_new_from_values( + Float32Array::from(values[200 * dim..400 * dim].to_vec()), + dim as i32, + ) + .unwrap(); + let split_c = FixedSizeListArray::try_new_from_values( + Float32Array::from(values[400 * dim..].to_vec()), + dim as i32, + ) + .unwrap(); + + let distributed = reduce_partial_stats(vec![ + compute_partial_stats(¢roids, &split_a, DistanceType::L2).unwrap(), + compute_partial_stats(¢roids, &split_b, DistanceType::L2).unwrap(), + compute_partial_stats(¢roids, &split_c, DistanceType::L2).unwrap(), + ]) + .unwrap(); + + assert_partial_stats_close(&single, &distributed, 1e-6); + } + + use proptest::prelude::*; + + proptest! { + #![proptest_config(ProptestConfig::with_cases(32))] + + #[test] + fn merge_is_commutative(seed in 0u64..1024, k in 4usize..16, dim in 2usize..8) { + let centroids = random_fsl_f32(seed, k, dim); + let d1 = random_fsl_f32(seed.wrapping_add(1), 50, dim); + let d2 = random_fsl_f32(seed.wrapping_add(2), 50, dim); + let s1 = compute_partial_stats(¢roids, &d1, DistanceType::L2).unwrap(); + let s2 = compute_partial_stats(¢roids, &d2, DistanceType::L2).unwrap(); + let merged_forward = merge_partial_stats(s1.clone(), s2.clone()).unwrap(); + let merged_reverse = merge_partial_stats(s2, s1).unwrap(); + assert_partial_stats_close(&merged_forward, &merged_reverse, 1e-9); + } + + #[test] + fn merge_is_associative(seed in 0u64..1024, k in 4usize..16, dim in 2usize..8) { + let centroids = random_fsl_f32(seed, k, dim); + let d1 = random_fsl_f32(seed.wrapping_add(11), 30, dim); + let d2 = random_fsl_f32(seed.wrapping_add(22), 30, dim); + let d3 = random_fsl_f32(seed.wrapping_add(33), 30, dim); + let s1 = compute_partial_stats(¢roids, &d1, DistanceType::L2).unwrap(); + let s2 = compute_partial_stats(¢roids, &d2, DistanceType::L2).unwrap(); + let s3 = compute_partial_stats(¢roids, &d3, DistanceType::L2).unwrap(); + let lhs = merge_partial_stats( + merge_partial_stats(s1.clone(), s2.clone()).unwrap(), + s3.clone(), + ) + .unwrap(); + let rhs = merge_partial_stats(s1, merge_partial_stats(s2, s3).unwrap()).unwrap(); + assert_partial_stats_close(&lhs, &rhs, 1e-9); + } + } + + #[test] + fn test_local_reservoir_sample_size_and_seed() { + let data = random_fsl_f32(7, 1000, 4); + let s1 = local_reservoir_sample(&data, 64, 42).unwrap(); + let s2 = local_reservoir_sample(&data, 64, 42).unwrap(); + let s3 = local_reservoir_sample(&data, 64, 43).unwrap(); + + assert_eq!(s1.num_rows(), 64); + assert_eq!(s1.schema().field(0).name(), "vec"); + assert_eq!(s1, s2, "same seed -> same output (I7)"); + assert_ne!(s1, s3, "different seed -> different output"); + } + + #[test] + fn test_local_reservoir_sample_smaller_than_target() { + let data = random_fsl_f32(8, 10, 4); + let s = local_reservoir_sample(&data, 64, 1).unwrap(); + assert_eq!(s.num_rows(), 10, "less data than target -> return all rows"); + } + + #[test] + fn test_select_initial_centroids_picks_k_rows() { + let s1 = local_reservoir_sample(&random_fsl_f32(1, 200, 6), 100, 9).unwrap(); + let s2 = local_reservoir_sample(&random_fsl_f32(2, 200, 6), 100, 10).unwrap(); + let centroids = select_initial_centroids(vec![s1, s2], 32, 7).unwrap(); + assert_eq!(centroids.len(), 32); + assert_eq!(centroids.value_length(), 6); + } + + #[test] + fn test_bootstrap_centroids_runs_kmeans() { + let s = local_reservoir_sample(&random_fsl_f32(3, 5_000, 8), 4_000, 11).unwrap(); + let centroids = bootstrap_centroids(vec![s], 32, DistanceType::L2, 13).unwrap(); + assert_eq!(centroids.len(), 32); + assert_eq!(centroids.value_length(), 8); + } + + #[test] + fn test_bootstrap_centroids_is_deterministic_with_seed() { + let s = local_reservoir_sample(&random_fsl_f32(7, 5_000, 8), 4_000, 11).unwrap(); + let a = bootstrap_centroids(vec![s.clone()], 32, DistanceType::L2, 7).unwrap(); + let b = bootstrap_centroids(vec![s], 32, DistanceType::L2, 7).unwrap(); + assert_eq!(a.len(), b.len()); + assert_eq!(a.value_length(), b.value_length()); + let av = a.values().as_primitive::().values().to_vec(); + let bv = b.values().as_primitive::().values().to_vec(); + assert_eq!(av, bv, "same-seed bootstrap_centroids must be byte-equal"); + } + + #[test] + fn test_bootstrap_centroids_cosine_does_not_panic() { + let s = local_reservoir_sample(&random_fsl_f32(7, 2_000, 8), 1_500, 21).unwrap(); + let centroids = bootstrap_centroids(vec![s], 16, DistanceType::Cosine, 23).unwrap(); + assert_eq!(centroids.len(), 16); + assert_eq!(centroids.value_length(), 8); + + // Every centroid row must have unit L2 norm (within ε), because the + // inner kmeans was run on normalized data and didn't denormalize. + let values = centroids.values().as_primitive::().values(); + for (i, row) in values.chunks_exact(8).enumerate() { + let norm = row.iter().map(|x| (*x as f64).powi(2)).sum::().sqrt(); + assert!( + (norm - 1.0).abs() < 1e-3, + "centroid {} has L2 norm {} (expected ≈ 1.0)", + i, + norm + ); + } + } +} diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index fb01339ead9..4d3dcf2bc9b 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -129,6 +129,7 @@ use tracing::instrument; use uuid::Uuid; pub mod builder; +pub mod distributed; pub mod io; mod partition_serde; pub mod v2; @@ -3325,60 +3326,23 @@ fn train_ivf_kmeans_step_arrow_array_no_loss( fn accumulate_refine_assignments( data: &FixedSizeListArray, centroids: &FixedSizeListArray, - cluster_sums: &mut [f32], - cluster_weights: &mut [f64], + accumulator: &mut Option, ) -> Result { - let dimension = data.value_length() as usize; - let kmeans = KMeans::with_centroids( - centroids.values().clone(), - dimension, - DistanceType::L2, - f64::MAX, - ); - let (membership, distances) = kmeans.compute_membership_and_distances(data)?; - let data_values = data.values().as_primitive::().values(); - let mut loss = 0.0; - - for row_idx in 0..data.len() { - let (Some(cluster_id), Some(distance)) = (membership[row_idx], distances[row_idx]) else { - continue; - }; - let cluster_id = cluster_id as usize; - cluster_weights[cluster_id] += 1.0; - loss += distance as f64; - let vector = &data_values[row_idx * dimension..(row_idx + 1) * dimension]; - let sum = &mut cluster_sums[cluster_id * dimension..(cluster_id + 1) * dimension]; - for (sum, value) in sum.iter_mut().zip(vector) { - *sum += *value; - } - } - - Ok(loss) + use lance_index::vector::kmeans::distributed::{compute_partial_stats, merge_partial_stats}; + let chunk = compute_partial_stats(centroids, data, DistanceType::L2)?; + let chunk_loss = chunk.total_loss(); + *accumulator = Some(match accumulator.take() { + Some(prev) => merge_partial_stats(prev, chunk)?, + None => chunk, + }); + Ok(chunk_loss) } fn update_refined_centroids( centroids: &FixedSizeListArray, - cluster_sums: &[f32], - cluster_weights: &[f64], + accumulator: lance_index::vector::kmeans::distributed::PartialStats, ) -> Result { - let dimension = centroids.value_length() as usize; - let mut next = centroids - .values() - .as_primitive::() - .values() - .to_vec(); - for cluster_id in 0..centroids.len() { - let weight = cluster_weights[cluster_id]; - if weight <= 0.0 { - continue; - } - let centroid = &mut next[cluster_id * dimension..(cluster_id + 1) * dimension]; - let sum = &cluster_sums[cluster_id * dimension..(cluster_id + 1) * dimension]; - for (value, sum) in centroid.iter_mut().zip(sum) { - *value = *sum / weight as f32; - } - } - f32_fsl_from_values(next, dimension) + lance_index::vector::kmeans::distributed::finalize_centroids(&accumulator, centroids) } async fn refine_streaming_f32_kmeans_with_sampler( @@ -3390,11 +3354,9 @@ async fn refine_streaming_f32_kmeans_with_sampler( passes: usize, on_progress: Arc, ) -> Result { - let dimension = initial_centroids.value_length() as usize; let mut centroids = initial_centroids.clone(); for pass in 1..=passes { - let mut cluster_sums = vec![0.0_f32; centroids.len() * dimension]; - let mut cluster_weights = vec![0.0_f64; centroids.len()]; + let mut accumulator: Option = None; let mut loss = 0.0; let mut row_offset = 0; while row_offset < sample_ranges.num_rows() { @@ -3412,21 +3374,19 @@ async fn refine_streaming_f32_kmeans_with_sampler( metric_type ))); } - loss += accumulate_refine_assignments( - &training_data, - ¢roids, - &mut cluster_sums, - &mut cluster_weights, - )?; + loss += accumulate_refine_assignments(&training_data, ¢roids, &mut accumulator)?; } - centroids = update_refined_centroids(¢roids, &cluster_sums, &cluster_weights)?; + let stats = accumulator.ok_or_else(|| { + Error::invalid_input( + "streaming IVF refinement: no training data assigned in this pass".to_string(), + ) + })?; + let assigned = stats.total_count() as usize; + centroids = update_refined_centroids(¢roids, stats)?; on_progress(pass as u32, passes as u32); info!( "Streaming IVF raw-vector refinement pass {} / {} assigned {} vectors; pre-update loss={}", - pass, - passes, - cluster_weights.iter().sum::() as usize, - loss + pass, passes, assigned, loss ); } Ok(centroids) @@ -3445,11 +3405,9 @@ async fn refine_streaming_f32_kmeans_with_resampling( passes: usize, on_progress: Arc, ) -> Result { - let dimension = initial_centroids.value_length() as usize; let mut centroids = initial_centroids.clone(); for pass in 1..=passes { - let mut cluster_sums = vec![0.0_f32; centroids.len() * dimension]; - let mut cluster_weights = vec![0.0_f64; centroids.len()]; + let mut accumulator: Option = None; let mut remaining_sample_rate = total_sample_rate; let mut loss = 0.0; while remaining_sample_rate > 0 { @@ -3474,22 +3432,20 @@ async fn refine_streaming_f32_kmeans_with_resampling( metric_type ))); } - loss += accumulate_refine_assignments( - &training_data, - ¢roids, - &mut cluster_sums, - &mut cluster_weights, - )?; + loss += accumulate_refine_assignments(&training_data, ¢roids, &mut accumulator)?; remaining_sample_rate -= step_sample_rate; } - centroids = update_refined_centroids(¢roids, &cluster_sums, &cluster_weights)?; + let stats = accumulator.ok_or_else(|| { + Error::invalid_input( + "streaming IVF refinement: no training data assigned in this pass".to_string(), + ) + })?; + let assigned = stats.total_count() as usize; + centroids = update_refined_centroids(¢roids, stats)?; on_progress(pass as u32, passes as u32); info!( "Streaming IVF resampled raw-vector refinement pass {} / {} assigned {} vectors; pre-update loss={}", - pass, - passes, - cluster_weights.iter().sum::() as usize, - loss + pass, passes, assigned, loss ); } Ok(centroids) @@ -4581,6 +4537,79 @@ mod tests { const DIM: usize = 32; + #[test] + fn test_streaming_refine_unchanged_partial_stats() { + // Regression contract for I6: the refactored streaming refine + // (compute_partial_stats + finalize_centroids) must produce centroids + // bit-identical to the legacy `cluster_sums / cluster_weights` path on + // the same input. Inline the legacy reference implementation so the + // test still works after the legacy helpers are deleted. + use arrow_array::cast::AsArray; + use arrow_array::types::Float32Type; + use arrow_array::{FixedSizeListArray, Float32Array}; + use lance_index::vector::kmeans::KMeans; + use lance_index::vector::kmeans::distributed::{compute_partial_stats, finalize_centroids}; + use lance_linalg::distance::DistanceType; + + let dim = 2; + let centroids = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 10.0, 0.0, 0.0, 10.0, 10.0, 10.0]), + dim as i32, + ) + .unwrap(); + let data = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![ + 0.1_f32, 0.0, 10.1, 0.1, 0.0, 9.9, 9.8, 10.2, 0.2, 0.1, 10.0, 0.0, + ]), + dim as i32, + ) + .unwrap(); + + let kmeans = + KMeans::with_centroids(centroids.values().clone(), dim, DistanceType::L2, f64::MAX); + let (membership, distances) = kmeans.compute_membership_and_distances(&data).unwrap(); + let data_values = data.values().as_primitive::().values(); + let k = centroids.len(); + let mut cluster_sums = vec![0.0_f32; k * dim]; + let mut cluster_weights = vec![0.0_f64; k]; + for i in 0..data.len() { + if let (Some(c), Some(_d)) = (membership[i], distances[i]) { + let ci = c as usize; + cluster_weights[ci] += 1.0; + let row = &data_values[i * dim..(i + 1) * dim]; + for d in 0..dim { + cluster_sums[ci * dim + d] += row[d]; + } + } + } + let mut baseline = centroids + .values() + .as_primitive::() + .values() + .to_vec(); + for ci in 0..k { + if cluster_weights[ci] > 0.0 { + for d in 0..dim { + baseline[ci * dim + d] = + cluster_sums[ci * dim + d] / cluster_weights[ci] as f32; + } + } + } + + let stats = compute_partial_stats(¢roids, &data, DistanceType::L2).unwrap(); + let new_new = finalize_centroids(&stats, ¢roids).unwrap(); + let v_new = new_new.values().as_primitive::().values(); + + for (a, b) in baseline.iter().zip(v_new.iter()) { + assert!( + (a - b).abs() < 1e-6, + "centroid mismatch: baseline={} new={}", + a, + b + ); + } + } + async fn compute_test_ivf_loss(dataset: &Dataset, column: &str, ivf: &IvfModel) -> f64 { let centroids = ivf .centroids_array() diff --git a/rust/lance/src/index/vector/ivf/distributed.rs b/rust/lance/src/index/vector/ivf/distributed.rs new file mode 100644 index 00000000000..c97f328536c --- /dev/null +++ b/rust/lance/src/index/vector/ivf/distributed.rs @@ -0,0 +1,465 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Layer 2 distributed IVF centroid training: dataset-aware async wrappers +//! around the pure primitives in `lance_index::vector::kmeans::distributed`. +//! +//! See `docs/superpowers/specs/2026-06-10-distributed-centroid-training-abstraction-design.md`. + +pub use lance_index::vector::kmeans::distributed::{ + PartialStats, bootstrap_centroids, finalize_centroids, merge_partial_stats, + reduce_partial_stats, select_initial_centroids, +}; + +use arrow_array::{Array, FixedSizeListArray, RecordBatch}; +use futures::StreamExt; +use lance_linalg::distance::DistanceType; +use lance_linalg::kernels::normalize_fsl_owned; + +use crate::Result; +use crate::dataset::Dataset; +use crate::index::vector::utils::{ + filter_finite_training_data, sample_training_data_stream, vector_column_to_fsl, +}; + +/// Round-0 worker entrypoint: reservoir-sample a Lance dataset slice. +/// +/// Streams the worker's projected training rows through a single +/// [`StreamingReservoir`](lance_index::vector::kmeans::distributed::StreamingReservoir) +/// so peak memory is bounded by the output sample +/// (target rows) rather than the worker's full fragment slice. The same +/// `rng_seed` is forwarded to both the upstream sampling stream and the +/// reservoir, making same-seed runs byte-deterministic for a given +/// `(dataset, column, fragments, target)` tuple. +/// +/// Internals mirror `build_ivf_model` (`rust/lance/src/index/vector/ivf.rs`): +/// 1. Stream raw rows via `sample_training_data_stream` (oversample to 2*target). +/// 2. Per batch: extract FSL, optionally L2-normalize for Cosine, drop non-finite rows. +/// 3. Feed each filtered chunk into the streaming reservoir. +pub async fn sample_round_0( + dataset: &Dataset, + column: &str, + fragment_ids: Option<&[u32]>, + target: usize, + distance_type: DistanceType, + rng_seed: u64, +) -> Result { + use lance_index::vector::kmeans::distributed::StreamingReservoir; + + // Round-0 oversamples to give the driver-side bootstrap enough material. + let mut stream = sample_training_data_stream( + dataset, + column, + target.saturating_mul(2), + fragment_ids, + Some(rng_seed), + ) + .await?; + + let mut reservoir = StreamingReservoir::new(target, rng_seed); + while let Some(batch) = stream.next().await { + let batch = batch?; + let fsl = vector_column_to_fsl(&batch, column)?; + let normalized = if distance_type == DistanceType::Cosine { + normalize_fsl_owned(fsl)? + } else { + fsl + }; + let filtered = filter_finite_training_data(normalized)?; + if filtered.is_empty() { + continue; + } + reservoir.feed(&filtered)?; + } + + reservoir.into_record_batch() +} + +/// Round-r worker entrypoint: scan the worker's fragment slice and produce a +/// `PartialStats` batch against the broadcast `centroids`. +/// +/// `centroids` are interpreted in the same dtype as the dataset's vector column +/// (caller is responsible for keeping them stable across rounds). +/// +/// Streams the projected training rows and accumulates per-batch +/// `PartialStats`, keeping peak memory at O(k·d) instead of O(N·d). +pub async fn compute_partial_stats( + dataset: &Dataset, + column: &str, + fragment_ids: Option<&[u32]>, + centroids: &FixedSizeListArray, + distance_type: DistanceType, +) -> Result { + use lance_index::vector::kmeans::distributed::{ + compute_centroids_fingerprint, compute_partial_stats as l1_compute_partial_stats, + merge_partial_stats, + }; + + // Worker is expected to have a small enough fragment slice that scanning + // it whole is cheap. `usize::MAX` skips sampling entirely; we just want + // every row streamed in. + let mut stream = + sample_training_data_stream(dataset, column, usize::MAX, fragment_ids, None).await?; + + let mut acc: Option = None; + while let Some(batch) = stream.next().await { + let batch = batch?; + let fsl = vector_column_to_fsl(&batch, column)?; + let normalized = if distance_type == DistanceType::Cosine { + normalize_fsl_owned(fsl)? + } else { + fsl + }; + let filtered = filter_finite_training_data(normalized)?; + if filtered.is_empty() { + continue; + } + let chunk = l1_compute_partial_stats(centroids, &filtered, distance_type)?; + acc = Some(match acc.take() { + None => chunk, + Some(prev) => merge_partial_stats(prev, chunk)?, + }); + } + + // Empty input (no rows / all filtered): synthesize an empty `PartialStats` + // with the right metadata so downstream merge/finalize works uniformly. + Ok(acc.unwrap_or_else(|| { + let dim = centroids.value_length() as usize; + let k = centroids.len(); + let fp = compute_centroids_fingerprint(centroids); + PartialStats::empty(k, dim, distance_type, fp) + })) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + use arrow_array::{RecordBatch, RecordBatchIterator}; + use arrow_schema::{DataType, Field, Schema}; + use lance_arrow::FixedSizeListArrayExt; + use lance_testing::datagen::generate_random_array_with_seed; + + use crate::dataset::Dataset; + + /// Build a small in-memory Lance dataset with `n` rows of FSL + /// across 4 fragments so distributed sampling has multiple slices to play with. + async fn make_vector_dataset(uri: &str, n: usize, dim: usize) -> Dataset { + let total = n * dim; + let values = + generate_random_array_with_seed::(total, [42; 32]); + let fsl = FixedSizeListArray::try_new_from_values(values, dim as i32).unwrap(); + + let schema = Arc::new(Schema::new(vec![Field::new( + "vec", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + dim as i32, + ), + true, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(fsl)]).unwrap(); + + let params = crate::dataset::WriteParams { + // Spread the rows across 4 fragments so worker partitioning has something to bite on. + max_rows_per_file: (n / 4).max(1), + max_rows_per_group: 256, + ..Default::default() + }; + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + Dataset::write(batches, uri, Some(params)).await.unwrap() + } + + #[tokio::test] + async fn test_sample_round_0_returns_target_rows() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let ds = make_vector_dataset(uri, 2_000, 8).await; + let sample = sample_round_0(&ds, "vec", None, 256, DistanceType::L2, 7) + .await + .unwrap(); + assert_eq!(sample.num_rows(), 256); + assert_eq!(sample.schema().field(0).name(), "vec"); + } + + /// Same `(dataset, fragments, target, seed)` must produce a byte-identical + /// sample, exercising the seed plumbed through `sample_training_data_stream` + /// and `StreamingReservoir`. + #[tokio::test] + async fn test_sample_round_0_is_deterministic_for_same_seed() { + use arrow_array::cast::AsArray; + use arrow_array::types::Float32Type; + + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let ds = make_vector_dataset(uri, 4_000, 8).await; + + let s1 = sample_round_0(&ds, "vec", None, 256, DistanceType::L2, 42) + .await + .unwrap(); + let s2 = sample_round_0(&ds, "vec", None, 256, DistanceType::L2, 42) + .await + .unwrap(); + + assert_eq!(s1.num_rows(), s2.num_rows()); + let v1 = s1 + .column(0) + .as_fixed_size_list() + .values() + .as_primitive::() + .values() + .to_vec(); + let v2 = s2 + .column(0) + .as_fixed_size_list() + .values() + .as_primitive::() + .values() + .to_vec(); + assert_eq!(v1, v2, "same seed must produce byte-identical samples"); + } + + /// Different seeds should not produce byte-identical samples on a dataset + /// large enough to make collisions astronomically unlikely. + #[tokio::test] + async fn test_sample_round_0_different_seed_changes_output() { + use arrow_array::cast::AsArray; + use arrow_array::types::Float32Type; + + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let ds = make_vector_dataset(uri, 4_000, 8).await; + + let mut samples = Vec::new(); + for seed in 0..5 { + let s = sample_round_0(&ds, "vec", None, 256, DistanceType::L2, seed) + .await + .unwrap(); + let v = s + .column(0) + .as_fixed_size_list() + .values() + .as_primitive::() + .values() + .to_vec(); + samples.push(v); + } + + let mut differs = false; + for i in 0..samples.len() { + for j in (i + 1)..samples.len() { + if samples[i] != samples[j] { + differs = true; + break; + } + } + if differs { + break; + } + } + assert!(differs, "different seeds should change the sample"); + } + + /// Streaming `compute_partial_stats` must produce identical statistics to + /// a materializing reference path that loads the whole dataset at once + /// (parity test for the Layer-2 streaming refactor). + #[tokio::test] + async fn test_compute_partial_stats_streaming_matches_materialized() { + use arrow_array::Float32Array; + use arrow_array::cast::AsArray; + use arrow_array::types::Float64Type; + use lance_index::vector::kmeans::distributed::compute_partial_stats as l1_compute; + + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let ds = make_vector_dataset(uri, 4_000, 8).await; + + // 4 fixed centroids inside the [0, 1)^8 unit cube generated by the + // dataset helper, so each one has at least one nearby vector. + let cs: Vec = vec![ + 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, // + 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, // + 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, 0.9, // + 0.2, 0.7, 0.3, 0.8, 0.4, 0.6, 0.5, 0.5, // + ]; + let centroids = FixedSizeListArray::try_new_from_values(Float32Array::from(cs), 8).unwrap(); + + let stream_stats = compute_partial_stats(&ds, "vec", None, ¢roids, DistanceType::L2) + .await + .unwrap(); + + // Reference: scan the whole dataset into one FSL and run Layer-1 + // `compute_partial_stats` directly. + let raw = + crate::index::vector::utils::maybe_sample_training_data(&ds, "vec", usize::MAX, None) + .await + .unwrap(); + let filtered = crate::index::vector::utils::filter_finite_training_data(raw).unwrap(); + let mat_stats = l1_compute(¢roids, &filtered, DistanceType::L2).unwrap(); + + assert_eq!(stream_stats.k(), mat_stats.k()); + assert_eq!(stream_stats.dim(), mat_stats.dim()); + assert_eq!(stream_stats.total_count(), mat_stats.total_count()); + // total loss can drift by a couple ULPs due to summation order across + // streamed chunks; allow a tiny relative tolerance. + let l_stream = stream_stats.total_loss(); + let l_mat = mat_stats.total_loss(); + assert!( + (l_stream - l_mat).abs() <= 1e-3 * l_mat.abs().max(1.0), + "loss drift: stream={} mat={}", + l_stream, + l_mat + ); + + let s_stream = stream_stats + .record_batch() + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_primitive::() + .values() + .to_vec(); + let s_mat = mat_stats + .record_batch() + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_primitive::() + .values() + .to_vec(); + for (a, b) in s_stream.iter().zip(s_mat.iter()) { + assert!( + (a - b).abs() <= 1e-3 * b.abs().max(1.0), + "sum drift: {} vs {}", + a, + b + ); + } + } + + #[tokio::test] + async fn test_layer2_compute_partial_stats_l2() { + use arrow_array::Float32Array; + + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let ds = make_vector_dataset(uri, 1_000, 4).await; + let centroids = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![0.0_f32, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0]), + 4, + ) + .unwrap(); + + let stats = compute_partial_stats(&ds, "vec", None, ¢roids, DistanceType::L2) + .await + .unwrap(); + assert_eq!(stats.k(), 2); + assert!(stats.total_count() > 0, "should have some assigned vectors"); + } + + #[tokio::test] + async fn test_end_to_end_four_workers_match_single_machine() { + use arrow_array::cast::AsArray; + use arrow_array::types::Float32Type; + use lance_index::vector::ivf::builder::IvfBuildParams; + + let dim = 8; + let k = 16; + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let ds = make_vector_dataset(uri, 4_000, dim).await; + let column = "vec"; + + // Pre-train a single-machine baseline using the existing path. + let baseline_ivf = crate::index::vector::ivf::build_ivf_model( + &ds, + column, + dim, + DistanceType::L2, + &IvfBuildParams { + num_partitions: Some(k), + sample_rate: 256, + max_iters: 5, + ..Default::default() + }, + None, + std::sync::Arc::new(lance_index::progress::NoopIndexBuildProgress), + ) + .await + .unwrap(); + let baseline_centroids = baseline_ivf + .centroids_array() + .expect("baseline IVF model should have centroids") + .clone(); + + // Simulate 4 workers via fragment id slicing. + let frags: Vec = ds.get_fragments().iter().map(|f| f.id() as u32).collect(); + assert!(frags.len() >= 4, "need >=4 fragments to simulate 4 workers"); + let groups: Vec> = frags + .chunks(frags.len().div_ceil(4)) + .map(|s| s.to_vec()) + .collect(); + + // Round 0: each "worker" reservoir-samples its slice; driver bootstraps. + let mut samples = Vec::new(); + for g in &groups { + let s = sample_round_0(&ds, column, Some(g), 256, DistanceType::L2, 42) + .await + .unwrap(); + samples.push(s); + } + let mut centroids = bootstrap_centroids(samples, k, DistanceType::L2, 42).unwrap(); + + // 5 Lloyd's rounds. + for _ in 0..5 { + let mut partials = Vec::new(); + for g in &groups { + let s = compute_partial_stats(&ds, column, Some(g), ¢roids, DistanceType::L2) + .await + .unwrap(); + partials.push(s); + } + let merged = reduce_partial_stats(partials).unwrap(); + centroids = finalize_centroids(&merged, ¢roids).unwrap(); + } + + // Centroid sets are close (allow re-permutation: assert each baseline centroid has a near + // neighbour in the distributed centroids). + let v_base = baseline_centroids + .values() + .as_primitive::() + .values() + .to_vec(); + let v_dist = centroids + .values() + .as_primitive::() + .values() + .to_vec(); + for ci in 0..k { + let base = &v_base[ci * dim..(ci + 1) * dim]; + let mut best = f32::INFINITY; + for cj in 0..k { + let dist_c = &v_dist[cj * dim..(cj + 1) * dim]; + let d: f32 = base + .iter() + .zip(dist_c.iter()) + .map(|(a, b)| (a - b).powi(2)) + .sum(); + best = best.min(d); + } + // Each baseline centroid should have a fairly close match in the distributed result; + // values are L2-squared in 8 dims so we accept a few units of slack. + assert!( + best < 25.0, + "no close match for centroid {}: best={}", + ci, + best + ); + } + } +} diff --git a/rust/lance/src/index/vector/utils.rs b/rust/lance/src/index/vector/utils.rs index 3046d0f3a83..6da68520d43 100644 --- a/rust/lance/src/index/vector/utils.rs +++ b/rust/lance/src/index/vector/utils.rs @@ -23,6 +23,19 @@ use tokio::sync::Mutex; use crate::dataset::{Dataset, ProjectionRequest, TakeBuilder, row_offsets_to_row_addresses}; use crate::{Error, Result}; +/// Construct a [`SmallRng`] from an optional seed. +/// +/// `Some(seed)` produces a deterministic RNG (used by reproducible-sampling +/// callers, e.g. distributed round-0). `None` falls back to the OS source — +/// matching the historical behavior of the callers that don't need +/// reproducibility. +fn seeded_rng(seed: Option) -> SmallRng { + match seed { + Some(s) => SmallRng::seed_from_u64(s), + None => SmallRng::from_os_rng(), + } +} + /// Helper function to extract a column from a RecordBatch, supporting nested field paths. /// /// This function handles: @@ -89,6 +102,7 @@ async fn estimate_multivector_vectors_per_row( column: &str, num_rows: usize, fragments: Option<&[u32]>, + seed: Option, ) -> Result { if num_rows == 0 { return Ok(1030); @@ -96,26 +110,34 @@ async fn estimate_multivector_vectors_per_row( let projection = dataset.schema().project(&[column])?; - // Try a few random samples first (fast path). - let sample_batch_size = std::cmp::min(64, num_rows); - for _ in 0..8 { - let batch = dataset - .sample(sample_batch_size, &projection, fragments) - .await?; - let array = get_column_from_batch(&batch, column)?; - let list_array = array.as_list::(); - for i in 0..list_array.len() { - if list_array.is_null(i) { - continue; - } - let len = list_array.value_length(i) as usize; - if len > 0 { - return Ok(len); + // Random-sample fast path: 8 small `Dataset::sample` rounds. This is fast + // when most rows are populated, but `Dataset::sample` itself uses + // `rand::rng()` and is non-deterministic. When the caller asks for + // determinism (e.g. distributed round-0), skip this path entirely and + // fall through to the scanner-based prefix below, which is already + // deterministic. + if seed.is_none() { + let sample_batch_size = std::cmp::min(64, num_rows); + for _ in 0..8 { + let batch = dataset + .sample(sample_batch_size, &projection, fragments) + .await?; + let array = get_column_from_batch(&batch, column)?; + let list_array = array.as_list::(); + for i in 0..list_array.len() { + if list_array.is_null(i) { + continue; + } + let len = list_array.value_length(i) as usize; + if len > 0 { + return Ok(len); + } } } } - // Fallback: scan a small prefix to find a non-null example. This avoids rare + // Fallback (and only path when `seed.is_some()`): scan a small prefix to + // find a non-null example. This is deterministic and avoids rare // flakiness when values are extremely sparse. let mut scanner = dataset.scan(); scanner.project(&[column])?; @@ -338,7 +360,7 @@ pub async fn maybe_sample_training_data( // Set a minimum sample size of 128 to avoid too small samples, // it's not a problem because 128 multivectors is just about 64 MiB let vectors_per_row = - estimate_multivector_vectors_per_row(dataset, column, num_rows, fragment_ids) + estimate_multivector_vectors_per_row(dataset, column, num_rows, fragment_ids, None) .await?; sample_size_hint.div_ceil(vectors_per_row).max(128) } @@ -412,7 +434,10 @@ impl PartitionLoadLock { /// /// Handles both regular vector columns (FixedSizeList) and multivector columns /// (List\), flattening the latter. -fn vector_column_to_fsl(batch: &RecordBatch, column: &str) -> Result { +pub(crate) fn vector_column_to_fsl( + batch: &RecordBatch, + column: &str, +) -> Result { let array = get_column_from_batch(batch, column)?; match array.data_type() { arrow::datatypes::DataType::FixedSizeList(_, _) => Ok(array.as_fixed_size_list().clone()), @@ -454,6 +479,118 @@ async fn scan_all_training_data( Ok(batch) } +/// Streaming counterpart of [`scan_all_training_data`]: yields the same +/// projected/optionally-null-filtered rows as a `DatasetRecordBatchStream` +/// without materializing them in memory. Used by Layer-2 distributed paths +/// that consume training rows incrementally. +pub(crate) async fn scan_all_training_data_stream( + dataset: &Dataset, + column: &str, + is_nullable: bool, + fragment_ids: Option<&[u32]>, +) -> Result { + let mut scanner = dataset.scan(); + scanner.project(&[column])?; + if let Some(fragment_ids) = fragment_ids { + scanner.with_fragments(resolve_scan_fragments(dataset, fragment_ids)?); + } + if is_nullable { + let column_expr = lance_datafusion::logical_expr::field_path_to_expr(column)?; + scanner.filter_expr(column_expr.is_not_null()); + } + let stream = scanner.try_into_stream().await?; + Ok(stream) +} + +/// Streaming counterpart of [`maybe_sample_training_data`]. +/// +/// Mirrors the same `should_sample` branching as the non-streaming variant +/// but yields a [`Pin>>`] of [`RecordBatch`] without ever +/// materializing the full sample in memory. The optional `seed` is forwarded +/// through to the random-range / fragment-shuffle samplers so callers (e.g. +/// distributed round-0) can reproduce the same byte-for-byte stream for a +/// given `(dataset, column, fragments, sample_size_hint, seed)` tuple. +/// +/// This intentionally does **not** use the `sample_fsl_uniform` chunked-take +/// fast path used by [`maybe_sample_training_data`]; that path returns an +/// already-materialized FSL and would require buffering the whole sample +/// before re-emitting it. The streaming callers either (a) consume every row +/// (Layer-2 partial stats) or (b) reservoir-sample (Layer-2 round 0), so the +/// loss of the dense-path optimisation is irrelevant. +pub(crate) async fn sample_training_data_stream( + dataset: &Dataset, + column: &str, + sample_size_hint: usize, + fragment_ids: Option<&[u32]>, + seed: Option, +) -> Result> + Send>>> { + let num_rows = count_rows(dataset, fragment_ids).await?; + + let vector_field = dataset.schema().field(column).ok_or(Error::index(format!( + "Sample training data (stream): column {} does not exist in schema", + column + )))?; + + if sample_size_hint == 0 { + info!("No sampling required (stream): yielding empty stream"); + return Ok(Box::pin(stream::empty())); + } + + let is_nullable = vector_field.nullable; + + let sample_size_hint = match vector_field.data_type() { + arrow::datatypes::DataType::List(_) => { + let vectors_per_row = + estimate_multivector_vectors_per_row(dataset, column, num_rows, fragment_ids, seed) + .await?; + sample_size_hint.div_ceil(vectors_per_row).max(128) + } + _ => sample_size_hint, + }; + + let byte_width = vector_field + .data_type() + .byte_width_opt() + .unwrap_or(4 * 1024); + + let should_sample = num_rows > sample_size_hint; + if !should_sample { + info!( + "Sample training data (stream): scanning all {} rows for column {}", + num_rows, column + ); + let s = scan_all_training_data_stream(dataset, column, is_nullable, fragment_ids).await?; + let mapped = s.map_err(Error::from); + return Ok(Box::pin(mapped)); + } + + info!( + "Sample training data (stream): sampling {} rows from {} rows for column {}", + sample_size_hint, num_rows, column + ); + + if let Some(fragment_ids) = fragment_ids { + return sample_training_data_scan_from_fragments( + dataset, + column, + sample_size_hint, + num_rows, + fragment_ids, + seed, + ); + } + + let scan = sample_training_data_scan( + dataset, + column, + sample_size_hint, + num_rows, + byte_width, + seed, + )?; + Ok(Box::pin(scan.map_err(Error::from))) +} + /// Sample training data from the dataset. /// /// Dispatches to the most efficient strategy based on column type and nullability: @@ -492,6 +629,7 @@ async fn sample_training_data( sample_size_hint, num_rows, fragment_ids, + None, )?; return match vector_field.data_type() { DataType::FixedSizeList(_, _) => { @@ -514,13 +652,25 @@ async fn sample_training_data( .await } DataType::FixedSizeList(_, _) => { - let scan = - sample_training_data_scan(dataset, column, sample_size_hint, num_rows, byte_width)?; + let scan = sample_training_data_scan( + dataset, + column, + sample_size_hint, + num_rows, + byte_width, + None, + )?; sample_nullable_fsl(column, sample_size_hint, byte_width, vector_field, scan).await } _ => { - let scan = - sample_training_data_scan(dataset, column, sample_size_hint, num_rows, byte_width)?; + let scan = sample_training_data_scan( + dataset, + column, + sample_size_hint, + num_rows, + byte_width, + None, + )?; sample_nullable_fallback(column, sample_size_hint, is_nullable, scan).await } } @@ -533,9 +683,10 @@ fn sample_training_data_scan( sample_size_hint: usize, num_rows: usize, byte_width: usize, + seed: Option, ) -> Result { let block_size = dataset.object_store.as_ref().block_size(); - let ranges = random_ranges(num_rows, sample_size_hint, block_size, byte_width); + let ranges = random_ranges(num_rows, sample_size_hint, block_size, byte_width, seed); Ok(dataset.take_scan( Box::pin(futures::stream::iter(ranges).map(Ok)), Arc::new(dataset.schema().project(&[column])?), @@ -556,6 +707,7 @@ fn sample_training_data_scan_from_fragments( sample_size_hint: usize, num_rows: usize, fragment_ids: &[u32], + seed: Option, ) -> Result> + Send>>> { if fragment_ids.is_empty() { return Err(Error::invalid_input( @@ -590,7 +742,7 @@ fn sample_training_data_scan_from_fragments( projection, selected_fragments, HashSet::::with_capacity(sample_size_hint.min(num_rows)), - SmallRng::from_os_rng(), + seeded_rng(seed), ), move |(dataset, projection, selected_fragments, mut seen_offsets, mut rng)| async move { if seen_offsets.len() >= num_rows { @@ -786,7 +938,7 @@ async fn sample_fsl_uniform( byte_width: usize, vector_field: &lance_core::datatypes::Field, ) -> Result { - let indices = generate_random_indices(num_rows, sample_size_hint); + let indices = generate_random_indices(num_rows, sample_size_hint, None); let projection = Arc::new(dataset.schema().project(&[column])?); let mut values_buf = MutableBuffer::with_capacity(sample_size_hint * byte_width); @@ -942,9 +1094,9 @@ fn filter_non_null_rows(array: ArrayRef, batch: RecordBatch) -> Result Vec { +fn generate_random_indices(num_rows: usize, k: usize, seed: Option) -> Vec { assert!(k <= num_rows); - let mut rng = SmallRng::from_os_rng(); + let mut rng = seeded_rng(seed); let mut indices = if k * 2 < num_rows { let mut set = std::collections::HashSet::with_capacity(k); while set.len() < k { @@ -984,9 +1136,10 @@ fn random_ranges( sample_size_hint: usize, block_size: usize, byte_width: usize, + seed: Option, ) -> impl Iterator> + Send { let rows_per_batch = 1.max(block_size / byte_width); - let mut rng = SmallRng::from_os_rng(); + let mut rng = seeded_rng(seed); let num_bins = num_rows.div_ceil(rows_per_batch); let bins_iter: Box + Send> = if sample_size_hint * 5 >= num_rows { @@ -1058,7 +1211,7 @@ mod tests { assert_eq!(bin_size, 10); let mut ranges = - random_ranges(num_rows, sample_size, block_size, byte_width).collect::>(); + random_ranges(num_rows, sample_size, block_size, byte_width, None).collect::>(); ranges.sort_by_key(|r| r.start); let expected = (0..num_rows as u64).step_by(bin_size).map(|start| { let end = std::cmp::min(start + bin_size as u64, num_rows as u64); @@ -1180,7 +1333,7 @@ mod tests { #[case::exact(100, 100)] #[test] fn test_generate_random_indices(#[case] num_rows: usize, #[case] k: usize) { - let indices = generate_random_indices(num_rows, k); + let indices = generate_random_indices(num_rows, k, None); assert_eq!(indices.len(), k); assert!(indices.windows(2).all(|w| w[0] < w[1])); assert!(indices.iter().all(|&i| (i as usize) < num_rows)); @@ -1251,12 +1404,50 @@ mod tests { .await .unwrap(); - let n = estimate_multivector_vectors_per_row(&dataset, "mv", nrows, None) + let n = estimate_multivector_vectors_per_row(&dataset, "mv", nrows, None, None) .await .unwrap(); assert_eq!(n, 1030); } + #[tokio::test] + async fn test_estimate_multivector_vectors_per_row_seeded_is_deterministic() { + let nrows: usize = 64; + let dims: u32 = 4; + + // Each multivector row holds between 2 and 3 inner vectors. No + // random-null thinning, so every row is populated and the scanner + // prefix will hit a non-empty row on its first batch. + let mv = array::cycle_vec_var( + array::rand_vec::(Dimension::from(dims)), + Dimension::from(2), + Dimension::from(3), + ); + + let data = gen_batch() + .col("mv", mv) + .into_batch_rows(RowCount::from(nrows as u64)) + .unwrap(); + + let dataset = InsertBuilder::new("memory://") + .execute(vec![data]) + .await + .unwrap(); + + let a = estimate_multivector_vectors_per_row(&dataset, "mv", nrows, None, Some(7)) + .await + .unwrap(); + let b = estimate_multivector_vectors_per_row(&dataset, "mv", nrows, None, Some(7)) + .await + .unwrap(); + assert_eq!(a, b, "seeded estimator must be deterministic"); + // Length came from the populated 2..3 range; sanity-check the value. + assert!( + (2..=3).contains(&a), + "expected inner length in [2,3], got {a}" + ); + } + // Creates a dataset with three fragments holding 100, 200, and 150 rows. async fn make_three_fragment_dataset() -> Dataset { use arrow_array::{RecordBatch, RecordBatchIterator};