bug: creating an IVF index on a float16 vector column hangs forever when values are large

Summary

Building any IVF index (IVF_FLAT, IVF_PQ, IVF_SQ, and the three IVF_HNSW_* variants) over a float16 vector column never returns — the process spins at 100% CPU with no error, no timeout and no progress. It is an infinite loop inside k-means, not slowness.

The trigger is numeric magnitude, not float16 itself and not a degenerate data shape: it fires when vectors per cluster × value magnitude exceeds the f16 range (65504). Normalized embeddings never hit it; unnormalized raw features with a smallish num_partitions do.

Reproduction

import numpy as np, pyarrow as pa, lance, tempfile

dim, n = 16, 1024
vals = np.random.default_rng(0).uniform(0, 1000, size=(n, dim)).astype(np.float16)
flat = pa.array(vals.reshape(-1), type=pa.float16())
tbl = pa.table({"vec": pa.FixedSizeListArray.from_arrays(flat, dim)})

ds = lance.write_dataset(tbl, tempfile.mkdtemp(), mode="overwrite")
ds.create_index("vec", index_type="IVF_FLAT", num_partitions=4, metric="l2")
# never returns

Any of these avoids it: num_partitions=64, uniform(0, 100), metric="cosine", or a float32 column.

The boundary lands exactly on 65504.

Root cause

Two independent defects in rust/lance-index/src/vector/kmeans.rs.

1. The centroid update accumulates in the storage type. KMeansAlgoFloat::to_kmeans allocates vec![T::Native::zero(); k * dimension] and sums vectors straight into it, so an f16 cluster sum saturates at 65504. Separately, T::Native::from_usize(cnt) is inf for a cluster larger than 65504, so 1 / cnt becomes 0 and the centroid is zeroed out — or turns into NaN when the sum had already reached inf.

2. split_clusters has an unbounded loop. Non-finite centroids make every L2 distance +inf, so argmin_value_float — which starts its running minimum at +inf and is documented to return None when everything is NaN/Inf — assigns nothing, and every cluster ends up empty. split_clusters then searches for a donor with

let p = (cnts[j] as f32 - 1.0) / (n - cnts.len()) as f32;
if rng.random::<f32>() < p { break; }

With every count at zero, p is negative, rng.random::<f32>() is in [0, 1), and the loop has no iteration cap. max_iters does not help because the spin happens inside a single iteration.

This second defect is not float16-specific. Any input where every distance is non-finite reaches it — an all-NaN float32 column hangs the same way.

Why cosine and dot escape

• cosine: vectors are normalized before training, so sums stay bounded.
• dot: the distance is -dot, so overflow yields -inf, and -inf < +inf holds, so argmin still returns a cluster.

L2 is the only metric whose overflow direction collides with argmin's initial +inf.

IVF_PQ + dot hangs anyway, one level down: use_residual is false for dot (vector/pq/builder.rs), so the sub-quantizers train on the raw vectors and hit the same overflow. Raising num_partitions does not help there, because that k-means uses k = 2^num_bits.

Environment

Reproduced end-to-end on pylance 7.0.0 / macOS arm64, and confirmed against main by reading the source; kmeans.rs is unchanged upstream since.

Start with rust/lance-index/src/vector/kmeans.rs, especially KMeansAlgoFloat::to_kmeans, split_clusters, and argmin_value_float; also inspect vector/pq/builder.rs for the IVF_PQ path. Reproduce the float16 L2 case, then verify that clustering remains finite and terminates for the reported inputs, including the all-non-finite case.