feat(index): add strided KMeans data view to avoid PQ sub-vector copies

Problem

PQ codebook training currently materializes each sub-vector into a contiguous PrimitiveArray before running KMeans.

The current divide_to_subvectors implementation is costly, and there is a TODO comment.

This copy is expensive for large training sets because the original vectors are row-major:

row0: [sub0][sub1][sub2]...
row1: [sub0][sub1][sub2]...
row2: [sub0][sub1][sub2]...

But KMeans currently expects each training matrix to be contiguous:

[sub0_row0][sub0_row1][sub0_row2]...

As a result, PQ has to copy sub-vectors into a new contiguous layout before KMeans can consume them.

Proposal

Add an internal KMeans data view abstraction that can represent both contiguous training data and strided PQ sub-vector views.

For example:

trait KMeansData<T> {
    fn len(&self) -> usize;
    fn dimension(&self) -> usize;
    fn vector(&self, row: usize) -> &[T];
}

The existing contiguous case can be represented as:

struct ContiguousKMeansData<'a, T> {
    values: &'a [T],
    dimension: usize,
}

impl<'a, T> KMeansData<T> for ContiguousKMeansData<'a, T> {
    fn len(&self) -> usize {
        self.values.len() / self.dimension
    }

    fn dimension(&self) -> usize {
        self.dimension
    }

    fn vector(&self, row: usize) -> &[T] {
        let start = row * self.dimension;
        &self.values[start..start + self.dimension]
    }
}

PQ can then use a strided view over the original row-major values:

struct SubVectorKMeansData<'a, T> {
    values: &'a [T],
    full_dimension: usize,
    sub_dimension: usize,
    sub_vector_idx: usize,
    len: usize,
}

impl<'a, T> KMeansData<T> for SubVectorKMeansData<'a, T> {
    fn len(&self) -> usize {
        self.len
    }

    fn dimension(&self) -> usize {
        self.sub_dimension
    }

    fn vector(&self, row: usize) -> &[T] {
        let start = row * self.full_dimension
            + self.sub_vector_idx * self.sub_dimension;
        &self.values[start..start + self.sub_dimension]
    }
}

Then PQ training could avoid divide_to_subvectors entirely:

let values = data.values().as_primitive::<T>().values();
let full_dimension = data.value_length() as usize;
let sub_dimension = full_dimension / self.num_sub_vectors;

for sub_vector_idx in 0..self.num_sub_vectors {
    let training_data = SubVectorKMeansData {
        values,
        full_dimension,
        sub_dimension,
        sub_vector_idx,
        len: data.len().min(self.sample_rate * num_centroids),
    };

    train_kmeans_with_data::<T, _>(
        &training_data,
        params,
        num_centroids,
    )?;
}

Internally, KMeans hot loops would move from direct contiguous chunking:

data.par_chunks(dimension)
    .map(|vector| {
        l2_distance_batch(vector, centroids, dimension)
    })

to row-based access through the data view:

(0..data.len())
    .into_par_iter()
    .map(|row| {
        let vector = data.vector(row);
        l2_distance_batch(vector, centroids, data.dimension())
    })

Centroid recomputation would similarly use the view:

for row in 0..data.len() {
    if let Some(cluster_id) = membership[row] {
        let vector = data.vector(row);
        let centroid = &mut centroids[
            cluster_id as usize * dimension
                ..(cluster_id as usize + 1) * dimension
        ];

        centroid
            .iter_mut()
            .zip(vector)
            .for_each(|(c, v)| *c += *v);
    }
}

Any suggestions about this, or can I propose a PR directly?

Start in rust/lance-index/src/vector/pq/utils.rs at divide_to_subvectors and its TODO, then trace the KMeans training entry points that consume the copied data. Compare the contiguous and strided view examples in the issue, and consider how distance and centroid-recomputation loops would use row-based access. Done means PQ training no longer needs divide_to_subvectors while existing KMeans behavior is preserved.