From 96ea8254dd593483f2159d532cf0e8eb4aa22086 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 31 Aug 2026 17:00:46 +0800 Subject: [PATCH] test(index): serialize spill-sort tests --- Cargo.lock | 1 + rust/lance-index/Cargo.toml | 1 + rust/lance-index/src/scalar/btree.rs | 4 ++++ rust/lance-index/src/scalar/json.rs | 15 +++++++-------- rust/lance-index/src/scalar/rtree.rs | 10 ++++++++++ 5 files changed, 23 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 56b142cf48a..9e8f82587e0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4868,6 +4868,7 @@ dependencies = [ "rstest", "serde", "serde_json", + "serial_test", "smallvec", "tempfile", "test-log", diff --git a/rust/lance-index/Cargo.toml b/rust/lance-index/Cargo.toml index 3ca392e4d31..3138c1f67cf 100644 --- a/rust/lance-index/Cargo.toml +++ b/rust/lance-index/Cargo.toml @@ -84,6 +84,7 @@ lance-datafusion = { workspace = true, features = ["datagen"] } lance-testing.workspace = true test-log.workspace = true rstest.workspace = true +serial_test.workspace = true chrono.workspace = true uuid.workspace = true diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index cb5f8818b59..f58548f00f9 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -5146,7 +5146,10 @@ mod tests { assert_eq!(original_data, remapped_data); } + // Spill-enabled index builds share the cached DataFusion memory pool within the + // test process, so keep them in one resource group. #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_update_ranged_index() { // Setup stores for both indexes let old_tmpdir = TempObjDir::default(); @@ -5297,6 +5300,7 @@ mod tests { } #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_update_with_exact_row_id_filter() { let old_tmpdir = TempObjDir::default(); let old_store = Arc::new(LanceIndexStore::new( diff --git a/rust/lance-index/src/scalar/json.rs b/rust/lance-index/src/scalar/json.rs index 4ffbeac274d..5ef2457b83f 100644 --- a/rust/lance-index/src/scalar/json.rs +++ b/rust/lance-index/src/scalar/json.rs @@ -1344,7 +1344,10 @@ mod tests { SargableQuery::Equals(ScalarValue::Float64(Some(2.0))), vec![1] )] + // Spill-enabled index builds share the cached DataFusion memory pool within the + // test process, so keep them in one resource group. #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_json_btree_update_uses_trained_target_type( #[case] initial_docs: &[&str], #[case] update_docs: &[&str], @@ -1443,6 +1446,7 @@ mod tests { } #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_json_btree_update_reports_type_drift() { let (source_store, _source_dir) = local_json_index_store(); let index = train_and_load_json_index( @@ -1470,6 +1474,7 @@ mod tests { } #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_json_derived_params_preserve_wrapper() { let (store, _tmpdir) = local_json_index_store(); let index = train_and_load_json_index( @@ -1537,12 +1542,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), @@ -1562,11 +1561,11 @@ mod tests { vec![0, 1, 2] )] #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_json_float_btree_index_unsorted_input( #[case] query: SargableQuery, #[case] expected: Vec, ) { - let _guard = FLOAT_INDEX_CASE_GUARD.lock().await; use crate::metrics::NoOpMetricsCollector; use lance_select::RowAddrTreeMap; @@ -1612,11 +1611,11 @@ mod tests { /// contains JSONB bytes, so conversion must use the accompanying type tag to turn it /// into an Arrow null before sorting and training the target index. #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_json_btree_index_null_at_path() { use crate::metrics::NoOpMetricsCollector; use lance_select::RowAddrTreeMap; - let _guard = FLOAT_INDEX_CASE_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..84cdd17f97c 100644 --- a/rust/lance-index/src/scalar/rtree.rs +++ b/rust/lance-index/src/scalar/rtree.rs @@ -1307,7 +1307,10 @@ mod tests { ) } + // Spill-enabled index builds share the cached DataFusion memory pool within the + // test process, so keep them in one resource group. #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_search_bbox() { let bbox_type = RectType::new(Dimension::XY, Default::default()); @@ -1353,6 +1356,7 @@ mod tests { } #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_search_null() { let point_type = PointType::new(Dimension::XY, Default::default()); @@ -1389,6 +1393,7 @@ mod tests { } #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_empty_geometries_are_not_indexed() { let line_string_type = LineStringType::new(Dimension::XY, Default::default()); let mut builder = LineStringBuilder::new(line_string_type); @@ -1449,6 +1454,7 @@ mod tests { } #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_non_finite_bounds_are_not_treated_as_empty() { let rect_type = RectType::new(Dimension::XY, Default::default()); let mut builder = RectBuilder::new(rect_type); @@ -1489,6 +1495,7 @@ mod tests { } #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_merge_rtree_indices_filters_rows_and_nulls() { let point_type = PointType::new(Dimension::XY, Default::default()); let mut first_builder = PointBuilder::new(point_type.clone()); @@ -1586,6 +1593,7 @@ mod tests { } #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_update_removes_pre_fix_empty_entries() { let rect_type = RectType::new(Dimension::XY, Default::default()); let mut builder = RectBuilder::new(rect_type); @@ -1651,6 +1659,7 @@ mod tests { } #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_update_and_search() { fn gen_data(num_items: u32, frag_id: u32, nulls_addrs: &mut RowAddrTreeMap) -> RectArray { let bbox_type = RectType::new(Dimension::XY, Default::default()); @@ -1748,6 +1757,7 @@ mod tests { } #[tokio::test] + #[serial_test::serial(LANCE_DF_SPILL_POOL)] async fn test_prewarm() { let point_type = PointType::new(Dimension::XY, Default::default());