diff --git a/rust/lance/src/dataset/tests/dataset_concurrency_store.rs b/rust/lance/src/dataset/tests/dataset_concurrency_store.rs index a9c2aa44c38..92bfc8f1b19 100644 --- a/rust/lance/src/dataset/tests/dataset_concurrency_store.rs +++ b/rust/lance/src/dataset/tests/dataset_concurrency_store.rs @@ -192,7 +192,7 @@ async fn test_add_bases() { use std::sync::Arc; // Create a test dataset - let test_uri = "memory://add_bases_test"; + let test_uri = "shared-memory://add_bases_test/primary"; let mut data_gen = BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); @@ -213,13 +213,13 @@ async fn test_add_bases() { let new_bases = vec![ BasePath::new( 0, - "memory://bucket1".to_string(), + "shared-memory://add_bases_test/bucket1".to_string(), Some("bucket1".to_string()), false, ), BasePath::new( 0, - "memory://bucket2".to_string(), + "shared-memory://add_bases_test/bucket2".to_string(), Some("bucket2".to_string()), true, ), @@ -243,9 +243,9 @@ async fn test_add_bases() { .find(|bp| bp.name == Some("bucket2".to_string())) .expect("bucket2 not found"); - assert_eq!(bucket1.path, "memory://bucket1"); + assert_eq!(bucket1.path, "shared-memory://add_bases_test/bucket1"); assert!(!bucket1.is_dataset_root); - assert_eq!(bucket2.path, "memory://bucket2"); + assert_eq!(bucket2.path, "shared-memory://add_bases_test/bucket2"); assert!(bucket2.is_dataset_root); let updated_dataset = Arc::new(updated_dataset); @@ -253,7 +253,7 @@ async fn test_add_bases() { // Test conflict detection - try to add a base with the same name let conflicting_bases = vec![BasePath::new( 0, - "memory://bucket3".to_string(), + "shared-memory://add_bases_test/bucket3".to_string(), Some("bucket1".to_string()), false, )]; @@ -270,7 +270,7 @@ async fn test_add_bases() { // Test conflict detection - try to add a base with the same path let conflicting_bases = vec![BasePath::new( 0, - "memory://bucket1".to_string(), + "shared-memory://add_bases_test/bucket1".to_string(), Some("bucket3".to_string()), false, )]; @@ -292,7 +292,7 @@ async fn test_concurrent_add_bases_conflict() { use std::sync::Arc; // Create a test dataset - let test_uri = "memory://concurrent_add_bases_test"; + let test_uri = "shared-memory://concurrent_add_bases_test/primary"; let mut data_gen = BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); @@ -314,7 +314,7 @@ async fn test_concurrent_add_bases_conflict() { // First transaction adds base1 let new_bases1 = vec![BasePath::new( 0, - "memory://bucket1".to_string(), + "shared-memory://concurrent_add_bases_test/bucket1".to_string(), Some("base1".to_string()), false, )]; @@ -325,7 +325,7 @@ async fn test_concurrent_add_bases_conflict() { // This should succeed as there's no conflict let new_bases2 = vec![BasePath::new( 0, - "memory://bucket2".to_string(), + "shared-memory://concurrent_add_bases_test/bucket2".to_string(), Some("base2".to_string()), false, )]; @@ -360,7 +360,7 @@ async fn test_concurrent_add_bases_name_conflict() { use std::sync::Arc; // Create a test dataset - let test_uri = "memory://concurrent_name_conflict_test"; + let test_uri = "shared-memory://concurrent_name_conflict_test/primary"; let mut data_gen = BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); @@ -383,7 +383,7 @@ async fn test_concurrent_add_bases_name_conflict() { // First transaction adds base with name "shared_base" let new_bases1 = vec![BasePath::new( 0, - "memory://bucket1".to_string(), + "shared-memory://concurrent_name_conflict_test/bucket1".to_string(), Some("shared_base".to_string()), false, )]; @@ -394,7 +394,7 @@ async fn test_concurrent_add_bases_name_conflict() { // This should fail due to name conflict let new_bases2 = vec![BasePath::new( 0, - "memory://bucket2".to_string(), + "shared-memory://concurrent_name_conflict_test/bucket2".to_string(), Some("shared_base".to_string()), false, )]; @@ -416,7 +416,7 @@ async fn test_concurrent_add_bases_path_conflict() { use std::sync::Arc; // Create a test dataset - let test_uri = "memory://concurrent_path_conflict_test"; + let test_uri = "shared-memory://concurrent_path_conflict_test/primary"; let mut data_gen = BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); @@ -436,10 +436,10 @@ async fn test_concurrent_add_bases_path_conflict() { let dataset = Arc::new(dataset); let dataset_clone = Arc::new(dataset_clone); - // First transaction adds base with path "memory://shared_path" + // First transaction adds a base at the shared path let new_bases1 = vec![BasePath::new( 0, - "memory://shared_path".to_string(), + "shared-memory://concurrent_path_conflict_test/shared_path".to_string(), Some("base1".to_string()), false, )]; @@ -450,7 +450,7 @@ async fn test_concurrent_add_bases_path_conflict() { // This should fail due to path conflict let new_bases2 = vec![BasePath::new( 0, - "memory://shared_path".to_string(), + "shared-memory://concurrent_path_conflict_test/shared_path".to_string(), Some("base2".to_string()), false, )]; @@ -472,7 +472,7 @@ async fn test_concurrent_add_bases_with_data_write() { use std::sync::Arc; // Create a test dataset - let test_uri = "memory://concurrent_write_test"; + let test_uri = "shared-memory://concurrent_write_test/primary"; let mut data_gen = BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); @@ -494,7 +494,7 @@ async fn test_concurrent_add_bases_with_data_write() { // First transaction adds a new base let new_bases = vec![BasePath::new( 0, - "memory://bucket1".to_string(), + "shared-memory://concurrent_write_test/bucket1".to_string(), Some("base1".to_string()), false, )]; diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index fffe71f930a..939b88cdcdc 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -1913,6 +1913,32 @@ mod tests { .await } + async fn scan_sorted_ids(dataset: &Dataset) -> Vec { + let batches = dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut ids: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + ids.sort_unstable(); + ids + } + #[test] fn test_auto_cleanup_disabled_by_default() { // Auto-cleanup must be off by default: the cleanup hook is expensive on @@ -2860,7 +2886,7 @@ mod tests { use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; // Create dataset with multi-base configuration - let test_uri = "memory://multi_base_test"; + let test_uri = "shared-memory://multi_base_test"; let primary_uri = format!("{}/primary", test_uri); let base1_uri = format!("{}/base1", test_uri); let base2_uri = format!("{}/base2", test_uri); @@ -2927,6 +2953,10 @@ mod tests { ); } + assert_eq!(scan_sorted_ids(&dataset).await, (0..5).collect::>()); + let reopened = Dataset::open(&primary_uri).await.unwrap(); + assert_eq!(scan_sorted_ids(&reopened).await, (0..5).collect::>()); + // Test validation: cannot specify both target_bases and target_base_names_or_paths let mut data_gen2 = BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); @@ -3037,7 +3067,7 @@ mod tests { use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; // Create initial dataset - let test_uri = "memory://overwrite_test"; + let test_uri = "shared-memory://overwrite_test"; let primary_uri = format!("{}/primary", test_uri); let base1_uri = format!("{}/base1", test_uri); let base2_uri = format!("{}/base2", test_uri); @@ -3120,6 +3150,9 @@ mod tests { .all(|f| f.metadata.files.iter().all(|file| file.base_id == Some(2))) ); + let reopened = Dataset::open(&primary_uri).await.unwrap(); + assert_eq!(scan_sorted_ids(&reopened).await, (0..2).collect::>()); + // Test validation: cannot specify initial_bases in OVERWRITE mode let mut data_gen3 = BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); @@ -3155,7 +3188,7 @@ mod tests { use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; // Create initial dataset with multi-base configuration - let test_uri = "memory://append_test"; + let test_uri = "shared-memory://append_test"; let primary_uri = format!("{}/primary", test_uri); let base1_uri = format!("{}/base1", test_uri); let base2_uri = format!("{}/base2", test_uri); @@ -3242,6 +3275,11 @@ mod tests { assert!(has_base1_data, "Should have data in base1"); assert!(has_base2_data, "Should have data in base2"); + let mut expected: Vec = (0..3).chain(0..2).chain(0..4).collect(); + expected.sort_unstable(); + let reopened = Dataset::open(&primary_uri).await.unwrap(); + assert_eq!(scan_sorted_ids(&reopened).await, expected); + // Test validation: cannot specify initial_bases in APPEND mode let mut data_gen4 = BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); @@ -4266,7 +4304,7 @@ mod tests { async fn test_multi_base_target_primary_and_bases() { use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; - let test_uri = "memory://primary_slot_test"; + let test_uri = "shared-memory://primary_slot_test"; let primary_uri = format!("{}/primary", test_uri); let base1_uri = format!("{}/base1", test_uri); let base2_uri = format!("{}/base2", test_uri); @@ -4358,6 +4396,11 @@ mod tests { assert_eq!(file_bases, vec![None, Some(2)]); assert_eq!(dataset.count_rows(None).await.unwrap(), 21); + + let mut expected: Vec = (0..6).chain(0..9).chain(0..6).collect(); + expected.sort_unstable(); + let reopened = Dataset::open(&primary_uri).await.unwrap(); + assert_eq!(scan_sorted_ids(&reopened).await, expected); } /// `target_all_bases` resolves to every registered base at execution @@ -4366,7 +4409,7 @@ mod tests { async fn test_multi_base_target_all_bases() { use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; - let test_uri = "memory://all_bases_test"; + let test_uri = "shared-memory://all_bases_test"; let primary_uri = format!("{}/primary", test_uri); let base1_uri = format!("{}/base1", test_uri); let base2_uri = format!("{}/base2", test_uri); @@ -4448,6 +4491,11 @@ mod tests { .collect(); assert_eq!(file_bases, vec![Some(1), Some(2)]); + let mut expected: Vec = (0..3).chain(0..9).chain(0..6).collect(); + expected.sort_unstable(); + let reopened = Dataset::open(&primary_uri).await.unwrap(); + assert_eq!(scan_sorted_ids(&reopened).await, expected); + // Cannot be combined with explicit target bases. let mut data_gen4 = BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); @@ -4473,7 +4521,7 @@ mod tests { // On a dataset with no registered bases: include_primary=true is a // no-op rotation over primary, false is rejected. - let plain_uri = "memory://all_bases_plain"; + let plain_uri = "shared-memory://all_bases_plain/primary"; let mut data_gen5 = BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); let plain = Dataset::write(data_gen5.batch(3), plain_uri, None) @@ -4524,12 +4572,13 @@ mod tests { // CREATE mode: initial_bases join the rotation before their ids are // committed to a manifest. - let create_uri = "memory://all_bases_create"; + let create_root = "shared-memory://all_bases_create"; + let create_uri = format!("{}/primary", create_root); let mut data_gen8 = BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); let dataset = Dataset::write( data_gen8.batch(9), - create_uri, + &create_uri, Some( WriteParams { mode: WriteMode::Create, @@ -4539,13 +4588,13 @@ mod tests { id: 0, name: Some("base1".to_string()), is_dataset_root: true, - path: format!("{}/base1", create_uri), + path: format!("{}/base1", create_root), }, BasePath { id: 0, name: Some("base2".to_string()), is_dataset_root: false, - path: format!("{}/base2", create_uri), + path: format!("{}/base2", create_root), }, ]), ..Default::default() @@ -4562,6 +4611,8 @@ mod tests { .flat_map(|f| f.metadata.files.iter().map(|file| file.base_id)) .collect(); assert_eq!(file_bases, vec![None, Some(1), Some(2)]); + let reopened = Dataset::open(&create_uri).await.unwrap(); + assert_eq!(scan_sorted_ids(&reopened).await, (0..9).collect::>()); } #[tokio::test]