diff --git a/python/python/lance/vector.py b/python/python/lance/vector.py index 37e331313dc..23a82253a23 100644 --- a/python/python/lance/vector.py +++ b/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( diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 1afbe6b86f9..b606df357d1 100644 --- a/python/python/tests/test_vector_index.py +++ b/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") diff --git a/rust/lance-index/src/lib.rs b/rust/lance-index/src/lib.rs index 52fa9e5a5de..3f23ff3a5f2 100644 --- a/rust/lance-index/src/lib.rs +++ b/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::*; diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index cb5f8818b59..b33f8be0788 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/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()), diff --git a/rust/lance-index/src/scalar/json.rs b/rust/lance-index/src/scalar/json.rs index 4ffbeac274d..1c2ee3ad4ff 100644 --- a/rust/lance-index/src/scalar/json.rs +++ b/rust/lance-index/src/scalar/json.rs @@ -1351,6 +1351,7 @@ mod tests { #[case] query: SargableQuery, #[case] expected_row_ids: Vec, ) { + 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, ) { - 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, diff --git a/rust/lance-index/src/scalar/rtree.rs b/rust/lance-index/src/scalar/rtree.rs index 13cc7265e34..af12c1438a8 100644 --- a/rust/lance-index/src/scalar/rtree.rs +++ b/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();