fix(index): serialize spill-sort tests and skip NaN centroid initialization

python/python/lance/vector.py

@@ -198,6 +198,49 @@ def train_pq_codebook_on_accelerator(
     return pq_codebook, kmeans_list
 
 
+def _sample_finite_vectors(
+    ds: Iterable[Union[dict[str, "torch.Tensor"], "torch.Tensor"]],
+    column: str,
+    k: int,
+) -> "torch.Tensor":
+    """Draw `k` vectors with no NaN/inf values from a TorchDataset stream.
+
+    A vector containing non-finite values is invalid kmeans input: chosen as an
+    initial centroid it makes every distance NaN, so training "converges" on
+    the first epoch without ever updating the centroids, and partition
+    assignment then drops every row (id -1) — downstream stages see an empty
+    dataset. Skip such rows instead of letting them become centroids.
+
+    Raises
+    ------
+    ValueError
+        If the stream yields fewer than `k` finite vectors, with the counts
+        needed to diagnose why (e.g. an all-NaN column).
+    """
+    import torch
+
+    valid: list["torch.Tensor"] = []
+    num_valid = 0
+    num_scanned = 0
+    for batch in ds:
+        vecs = batch[column] if isinstance(batch, dict) else batch
+        finite_rows = vecs.reshape(vecs.shape[0], -1).isfinite().all(dim=1)
+        kept = vecs[finite_rows] if not bool(finite_rows.all()) else vecs
+        valid.append(kept)
+        num_valid += kept.shape[0]
+        num_scanned += vecs.shape[0]
+        if num_valid >= k:
+            break
+
+    if num_valid < k:
+        raise ValueError(
+            f"only {num_valid} finite vectors available to initialize "
+            f"{k} centroids after scanning {num_scanned} rows; column "
+            f"'{column}' contains too many NaN/inf vectors"
+        )
+    return torch.cat(valid, dim=0)[:k]
+
+
 def train_ivf_centroids_on_accelerator(
     dataset: LanceDataset,
     column: str,
@@ -245,7 +288,10 @@ def train_ivf_centroids_on_accelerator(
         filter=filt,
     )
 
-    init_centroids = next(iter(ds))
+    if filter_nan:
+        init_centroids = _sample_finite_vectors(ds, column, k)
+    else:
+        init_centroids = next(iter(ds))
     LOGGER.info("Done sampling: centroids shape: %s", init_centroids.shape)
 
     ds = TorchDataset(

python/python/tests/test_vector_index.py

@@ -568,6 +568,76 @@ def test_torch_index_with_nans(tmp_path, index_file_version):
     validate_vector_index(dataset, "vector", sample_size=16)
 
 
+def test_torch_ivf_init_skips_nan_vectors(tmp_path):
+    torch = pytest.importorskip("torch")
+
+    from lance.vector import train_ivf_centroids_on_accelerator
+
+    # The first init batch in storage order is all-NaN, so an unfiltered init
+    # draw deterministically produces NaN centroids: every distance is NaN,
+    # kmeans "converges" on the first epoch without ever updating them, and
+    # partition assignment drops every row. The sampler must skip non-finite
+    # rows instead.
+    mat = np.concatenate(
+        [
+            np.full((4, 16), np.nan, dtype=np.float32),
+            np.random.randn(8, 16).astype(np.float32),
+        ]
+    )
+    dataset = lance.write_dataset(vec_to_table(data=mat), tmp_path)
+
+    centroids, _ = train_ivf_centroids_on_accelerator(
+        dataset,
+        "vector",
+        4,
+        "l2",
+        torch.device("cpu"),
+        max_iters=2,
+    )
+    assert centroids.shape == (4, 16)
+    assert np.isfinite(centroids).all()
+
+
+def test_torch_ivf_init_rejects_all_nan_column(tmp_path):
+    torch = pytest.importorskip("torch")
+
+    from lance.vector import train_ivf_centroids_on_accelerator
+
+    tbl = create_table(nvec=0, ndim=16, nans=50)
+    dataset = lance.write_dataset(tbl, tmp_path)
+
+    with pytest.raises(ValueError, match="finite vectors"):
+        train_ivf_centroids_on_accelerator(
+            dataset,
+            "vector",
+            4,
+            "l2",
+            torch.device("cpu"),
+            max_iters=2,
+        )
+
+
+def test_sample_finite_vectors_skips_nan_rows():
+    torch = pytest.importorskip("torch")
+
+    from lance.vector import _sample_finite_vectors
+
+    valid = torch.arange(8, dtype=torch.float32).reshape(4, 2)
+    nans = torch.full((3, 2), float("nan"))
+
+    # Tensor batches (single-column stream) and dict batches (multi-column
+    # stream) must both skip non-finite rows and preserve order.
+    out = _sample_finite_vectors(iter([nans, valid[:2], nans, valid[2:]]), "vector", 4)
+    assert torch.equal(out, valid)
+    out = _sample_finite_vectors(
+        iter([{"vector": nans}, {"vector": valid}]), "vector", 4
+    )
+    assert torch.equal(out, valid)
+
+    with pytest.raises(ValueError, match="finite vectors"):
+        _sample_finite_vectors(iter([nans]), "vector", 1)
+
+
 def test_index_with_no_centroid_movement(tmp_path):
     torch = pytest.importorskip("torch")
 

rust/lance-index/src/lib.rs

@@ -93,6 +93,18 @@ pub fn infer_system_index_type(
     }
 }
 
+/// Serializes tests that drive a spill-enabled DataFusion execution.
+///
+/// Each spill `SortExec` reserves a non-spillable merge buffer
+/// (`sort_spill_reservation_bytes`, 40MB with the default pool) from the
+/// process-wide cached memory pool (see `get_session_context`), and the default
+/// pool fits at most three concurrent reservations (150MB). Under nextest every
+/// test is its own process and this guard is a no-op, but under `cargo test`
+/// the whole test binary shares one process — and one pool — so unguarded
+/// tests must hold it to avoid spuriously exhausting the pool.
+#[cfg(test)]
+pub(crate) static SPILL_POOL_TEST_GUARD: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
+
 #[cfg(test)]
 mod tests {
     use super::*;

rust/lance-index/src/scalar/btree.rs

@@ -5148,6 +5148,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_update_ranged_index() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         // Setup stores for both indexes
         let old_tmpdir = TempObjDir::default();
         let old_store = Arc::new(LanceIndexStore::new(
@@ -5298,6 +5299,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_update_with_exact_row_id_filter() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         let old_tmpdir = TempObjDir::default();
         let old_store = Arc::new(LanceIndexStore::new(
             Arc::new(ObjectStore::local()),

rust/lance-index/src/scalar/json.rs

@@ -1351,6 +1351,7 @@ mod tests {
         #[case] query: SargableQuery,
         #[case] expected_row_ids: Vec<u64>,
     ) {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         use crate::metrics::NoOpMetricsCollector;
         use lance_select::RowAddrTreeMap;
 
@@ -1444,6 +1445,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_json_btree_update_reports_type_drift() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         let (source_store, _source_dir) = local_json_index_store();
         let index = train_and_load_json_index(
             source_store,
@@ -1471,6 +1473,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_json_derived_params_preserve_wrapper() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         let (store, _tmpdir) = local_json_index_store();
         let index = train_and_load_json_index(
             store,
@@ -1536,13 +1539,6 @@ mod tests {
     ///
     /// Rows are fed in raw storage order (not sorted by value) to simulate what an
     /// unordered scan would produce.
-    ///
-    /// Each case below runs a spilling `SortExec` that reserves a non-spillable merge
-    /// buffer from the process-wide cached DataFusion memory pool (see
-    /// `get_session_context`); running the cases concurrently contends for that shared
-    /// pool and can spuriously exhaust it, so this guard serializes them.
-    static FLOAT_INDEX_CASE_GUARD: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
-
     #[rstest]
     #[case::range_gt_zero(
         SargableQuery::Range(Bound::Excluded(ScalarValue::Float64(Some(0.0))), Bound::Unbounded),
@@ -1566,7 +1562,7 @@ mod tests {
         #[case] query: SargableQuery,
         #[case] expected: Vec<u64>,
     ) {
-        let _guard = FLOAT_INDEX_CASE_GUARD.lock().await;
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         use crate::metrics::NoOpMetricsCollector;
         use lance_select::RowAddrTreeMap;
 
@@ -1616,7 +1612,7 @@ mod tests {
         use crate::metrics::NoOpMetricsCollector;
         use lance_select::RowAddrTreeMap;
 
-        let _guard = FLOAT_INDEX_CASE_GUARD.lock().await;
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         let (store, _tmpdir) = local_json_index_store();
         let index = train_and_load_json_index(
             store,

rust/lance-index/src/scalar/rtree.rs

@@ -1309,6 +1309,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_search_bbox() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         let bbox_type = RectType::new(Dimension::XY, Default::default());
 
         let mut rng = rand::rng();
@@ -1354,6 +1355,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_search_null() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         let point_type = PointType::new(Dimension::XY, Default::default());
 
         let mut rng = rand::rng();
@@ -1390,6 +1392,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_empty_geometries_are_not_indexed() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         let line_string_type = LineStringType::new(Dimension::XY, Default::default());
         let mut builder = LineStringBuilder::new(line_string_type);
         builder
@@ -1450,6 +1453,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_non_finite_bounds_are_not_treated_as_empty() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         let rect_type = RectType::new(Dimension::XY, Default::default());
         let mut builder = RectBuilder::new(rect_type);
         builder.push_rect(Some(&Rect::new(
@@ -1490,6 +1494,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_merge_rtree_indices_filters_rows_and_nulls() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         let point_type = PointType::new(Dimension::XY, Default::default());
         let mut first_builder = PointBuilder::new(point_type.clone());
         first_builder.push_point(Some(&geo_types::point!(x: 10.0, y: 10.0)));
@@ -1587,6 +1592,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_update_removes_pre_fix_empty_entries() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         let rect_type = RectType::new(Dimension::XY, Default::default());
         let mut builder = RectBuilder::new(rect_type);
         builder.push_rect(Some(&BoundingBox::new()));
@@ -1652,6 +1658,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_update_and_search() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         fn gen_data(num_items: u32, frag_id: u32, nulls_addrs: &mut RowAddrTreeMap) -> RectArray {
             let bbox_type = RectType::new(Dimension::XY, Default::default());
 
@@ -1749,6 +1756,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_prewarm() {
+        let _guard = crate::SPILL_POOL_TEST_GUARD.lock().await;
         let point_type = PointType::new(Dimension::XY, Default::default());
 
         let mut rng = rand::rng();