fix(index): reject unattainable partition counts in weighted k-means

[LuciferYang]

Closes #9345

train_weighted_hierarchical_f32_kmeans, the trainer behind the streaming coreset path, padded the cluster list with clones of the heaviest cluster whenever it formed fewer than target_k. A padded centroid is bit-identical to the one it came from, and assignment resolves a tie to the incumbent, so a duplicate never receives a vector; the refinement passes keep it that way, since a cluster with no members retains its previous centroid. The build therefore succeeded with partitions that could not hold a row, and a query near that centroid spent one of its nprobes on an empty partition because find_partitions takes the first nprobes of the sorted distances and the duplicate ties exactly with the original.

It now returns invalid_input naming how many clusters were formed, how many coreset rows they came from, and a num_partitions that would work. The flat hierarchical trainer already rejects the same request (lance-index/src/vector/kmeans.rs:1493), so this removes an inconsistency rather than inventing a policy — worth knowing that the two are not the same error kind, though: the flat path raises through ArrowError::InvalidArgumentError, which surfaces in Python as RuntimeError, while this one is Error::InvalidInput and surfaces as ValueError. Aligning them is a separate change.

The split loop now exhausts every splittable cluster before the shortfall check. It used to abandon the loop when the cluster at the top of the heap was a singleton, and priority there is loss rather than size, so a heavy singleton could outrank clusters that still had rows to give — the build was rejected while target_k was still reachable, and the suggested maximum in the message was lower than what the trainer could actually have formed. Only the singleton case changes: Ord already sorts finalized clusters last, so a finalized cluster at the top does mean nothing splittable is left, and breaking there stays correct.

The message says the remaining clusters "could not be split further" rather than blaming duplicate data, because the shortfall has two causes: coreset rows too alike to split, and a split loop that stops early. Naming only the first would send a reader with well-spread data down the wrong path.

clusters.truncate(target_k) right after the check became a no-op once the shortfall is an error — the loop stops at target_k and every step it takes is bounded by what is left — so it is now a debug_assert_eq! on that invariant instead of a silent guard.

How was this patch tested?

test_weighted_hierarchical_rejects_padding_with_duplicates covers both shapes that reach the shortfall, as an rstest pair. fewer_rows_than_partitions is the obvious one, 6 rows against target_k 16. enough_rows_all_identical is the one the only production caller can actually produce — 32 identical rows against 16 partitions — because the coreset always carries at least num_partitions rows, so a shortfall there comes from rows that cannot be split rather than from a row count. Both assert the coreset row count appears in the message, which is what separates the two cases in a bug report.

Restoring the padding loop fails both cases.

The singleton guard is reasoned from the heap ordering and the flat trainer's precedent rather than pinned by a test. I could not build a deterministic red case: heap priority comes from the initial clustering's cluster_losses, not from the losses argument a test passes in, so a fixture cannot force a heavy singleton to the top. A fixture that looked right turned out to pass with the old break too, so I removed it rather than ship a test with no discriminating power.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The duplicate-centroid failure is real, but the replacement must distinguish a truly exhausted splitter from one that stopped on a single unsplittable cluster. Exhaust all remaining non-finalized clusters before returning the error, as the flat trainer does, so supported streaming builds are not rejected and the suggested maximum partition count is truthful.

[lance-gatekeeper]

This return can reject an attainable partition count. The loop stops as soon as the highest-loss heap item has one member, but inherited base_losses contribute to that ordering; a high-loss singleton can therefore stop training while other non-finalized clusters still contain distinct rows that can be split. At the production routing threshold this turns a valid 257-partition request into an error advising the user to reduce it to 16. Please continue past newly finalized singletons and return only once the heap has no splittable cluster left (the flat trainer's fill_missing_leaves follows that contract).

Reproducer

Add this test beside the new regression test:

#[test]
fn probe_weighted_hierarchical_production_threshold() {
    let dimension = 1;
    let num_points = 258;
    let target_k = 257;
    let mut values = Vec::with_capacity(num_points);
    values.push(1.0e9);
    values.extend((1..num_points).map(|i| i as f32 * 100.0));
    let data =
        FixedSizeListArray::try_new_from_values(Float32Array::from(values), dimension).unwrap();
    let params = WeightedHierarchicalKMeansParams {
        dimension: dimension as usize,
        target_k,
        metric_type: MetricType::L2,
        max_iters: 20,
        on_progress: Arc::new(|_, _| {}),
    };
    let mut losses = vec![0.0; num_points];
    losses[0] = 1.0e20;

    let result = train_weighted_hierarchical_f32_kmeans(
        &data,
        &vec![1.0; num_points],
        &losses,
        &params,
    );
    assert!(result.is_ok(), "unexpected shortfall: {result:?}");
}

Running cargo test -p lance probe_weighted_hierarchical_production_threshold -- --nocapture on this head fails with Cannot create 257 IVF partitions: ... could only form 16 non-empty clusters ... Reduce num_partitions to <= 16, even though all 258 coreset rows are distinct.

[lance-gatekeeper]

Fixed in 6488b60: a singleton is now finalized and requeued so the remaining non-finalized clusters can still split. The 257-of-258 production-threshold reproducer now passes.

[LuciferYang]

Reworked the loop rather than the wording — you were pointing at the same hazard my own review round flagged and I had left as a note.

The loop abandoned itself when the cluster at the top of the heap was a singleton. Priority there is loss, not size (Ord only stratifies on finalized, then compares loss), so a heavy singleton could outrank clusters that still had rows to give: the build was rejected while target_k was reachable, and the Reduce num_partitions to <= N in the message was lower than what the trainer could actually have formed. Only the singleton case changed. A finalized cluster at the top genuinely does mean nothing splittable is left, because Ord sorts finalized last, so breaking there is still right; the singleton is now finalized, pushed back so it sinks, and the loop keeps splitting the rest.

One thing I could not do, so it is not claimed anywhere: there is no test pinning that guard. Heap priority comes from the initial clustering's cluster_losses rather than from the losses argument a test supplies, so I could not force a heavy singleton to the top deterministically. A fixture that looked convincing passed with the old break as well, so I deleted it instead of shipping a test with no discriminating power. The two shortfall cases remain covered and both fail if the padding loop comes back.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The earlier correctness blocker is fixed. However, the latest author note leaves that repair without a regression guard. The repository requires every bugfix to have corresponding coverage, and the deterministic 257-of-258 case already distinguishes the old and current heads. Add that case to complete the path to acceptance.

[lance-gatekeeper]

This corrected branch has no regression test, so the subtle heap-ordering failure can return unnoticed despite the repository's explicit requirement that bugfixes carry coverage. The reason given in the author note does not hold for this case: assign_weighted_f32_points adds each supplied base_losses[row_idx] into initial.cluster_losses, so setting a large base loss deterministically raises the singleton's heap priority. The test below failed on 945c6f6 with “could only form 16” and passes on this head; please add it beside the existing shortfall test.

Regression test

#[test]
fn test_weighted_hierarchical_exhausts_splittable_clusters() {
    let dimension = 1;
    let num_points = 258;
    let target_k = 257;
    let mut values = Vec::with_capacity(num_points);
    values.push(1.0e9);
    values.extend((1..num_points).map(|i| i as f32 * 100.0));
    let data =
        FixedSizeListArray::try_new_from_values(Float32Array::from(values), dimension).unwrap();
    let params = WeightedHierarchicalKMeansParams {
        dimension: dimension as usize,
        target_k,
        metric_type: MetricType::L2,
        max_iters: 20,
        on_progress: Arc::new(|_, _| {}),
    };
    let mut losses = vec![0.0; num_points];
    losses[0] = 1.0e20;

    let centroids = train_weighted_hierarchical_f32_kmeans(
        &data,
        &vec![1.0; num_points],
        &losses,
        &params,
    )
    .unwrap();
    assert_eq!(centroids.len(), target_k);
}

Verified with cargo test -p lance weighted_hierarchical -- --nocapture; all three cases complete in 0.04s on the current head.