From 269181c2c7170d0d690b295b430fe8b17156b169 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 3 Dec 2025 01:22:17 +0800 Subject: [PATCH 1/3] refactor: split dataset tests in a tests mod Signed-off-by: Xuanwo --- rust/lance/src/dataset.rs | 7359 +---------------- .../lance/src/dataset/tests/dataset_common.rs | 186 + .../tests/dataset_concurrency_store.rs | 508 ++ rust/lance/src/dataset/tests/dataset_geo.rs | 143 + rust/lance/src/dataset/tests/dataset_index.rs | 2419 ++++++ rust/lance/src/dataset/tests/dataset_io.rs | 1207 +++ .../src/dataset/tests/dataset_merge_update.rs | 1455 ++++ .../src/dataset/tests/dataset_migrations.rs | 358 + .../src/dataset/tests/dataset_transactions.rs | 344 + .../src/dataset/tests/dataset_versioning.rs | 738 ++ rust/lance/src/dataset/tests/mod.rs | 10 + 11 files changed, 7369 insertions(+), 7358 deletions(-) create mode 100644 rust/lance/src/dataset/tests/dataset_common.rs create mode 100644 rust/lance/src/dataset/tests/dataset_concurrency_store.rs create mode 100644 rust/lance/src/dataset/tests/dataset_geo.rs create mode 100644 rust/lance/src/dataset/tests/dataset_index.rs create mode 100644 rust/lance/src/dataset/tests/dataset_io.rs create mode 100644 rust/lance/src/dataset/tests/dataset_merge_update.rs create mode 100644 rust/lance/src/dataset/tests/dataset_migrations.rs create mode 100644 rust/lance/src/dataset/tests/dataset_transactions.rs create mode 100644 rust/lance/src/dataset/tests/dataset_versioning.rs create mode 100644 rust/lance/src/dataset/tests/mod.rs diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 4fecf55b25e..9de986ce022 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -2758,7361 +2758,4 @@ impl Projectable for Dataset { } #[cfg(test)] -mod tests { - use std::vec; - - use super::*; - use crate::dataset::optimize::{compact_files, CompactionOptions}; - use crate::dataset::transaction::DataReplacementGroup; - use crate::dataset::WriteMode::Overwrite; - use crate::index::vector::VectorIndexParams; - use crate::utils::test::copy_test_data_to_tmp; - use lance_arrow::FixedSizeListArrayExt; - use mock_instant::thread_local::MockClock; - - use crate::dataset::write::{CommitBuilder, InsertBuilder, WriteMode, WriteParams}; - use arrow::array::{as_struct_array, AsArray, GenericListBuilder, GenericStringBuilder}; - use arrow::compute::concat_batches; - use arrow::datatypes::UInt64Type; - use arrow_array::{ - builder::StringDictionaryBuilder, - cast::as_string_array, - types::{Float32Type, Int32Type}, - ArrayRef, DictionaryArray, Float32Array, Int32Array, Int64Array, Int8Array, - Int8DictionaryArray, ListArray, RecordBatchIterator, StringArray, UInt16Array, UInt32Array, - }; - use arrow_array::{ - Array, FixedSizeListArray, GenericStringArray, Int16Array, Int16DictionaryArray, - LargeBinaryArray, StructArray, UInt64Array, - }; - use arrow_ord::sort::sort_to_indices; - use arrow_schema::{ - DataType, Field as ArrowField, Field, Fields as ArrowFields, Schema as ArrowSchema, - }; - use lance_arrow::bfloat16::{self, BFLOAT16_EXT_NAME}; - use lance_arrow::{ARROW_EXT_META_KEY, ARROW_EXT_NAME_KEY, BLOB_META_KEY}; - use lance_core::utils::tempfile::{TempDir, TempStdDir, TempStrDir}; - use lance_datagen::{array, gen_batch, BatchCount, Dimension, RowCount}; - use lance_file::version::LanceFileVersion; - use lance_file::writer::FileWriter; - use lance_index::scalar::inverted::{ - query::{BooleanQuery, MatchQuery, Occur, Operator, PhraseQuery}, - tokenizer::InvertedIndexParams, - }; - use lance_index::scalar::FullTextSearchQuery; - use lance_index::{scalar::ScalarIndexParams, vector::DIST_COL, IndexType}; - use lance_io::assert_io_eq; - use lance_io::utils::CachedFileSize; - use lance_linalg::distance::MetricType; - use lance_table::feature_flags; - use lance_table::format::{DataFile, WriterVersion}; - - use crate::datafusion::LanceTableProvider; - use crate::dataset::refs::branch_contents_path; - use datafusion::common::{assert_contains, assert_not_contains}; - use datafusion::prelude::SessionContext; - use lance_arrow::json::ARROW_JSON_EXT_NAME; - use lance_datafusion::datagen::DatafusionDatagenExt; - use lance_datafusion::udf::register_functions; - use lance_index::scalar::inverted::query::{FtsQuery, MultiMatchQuery}; - use lance_testing::datagen::generate_random_array; - use pretty_assertions::assert_eq; - use rand::seq::SliceRandom; - use rand::Rng; - use rstest::rstest; - use std::cmp::Ordering; - - // Used to validate that futures returned are Send. - fn require_send(t: T) -> T { - t - } - - async fn create_file( - path: &std::path::Path, - mode: WriteMode, - data_storage_version: LanceFileVersion, - ) { - let fields = vec![ - ArrowField::new("i", DataType::Int32, false), - ArrowField::new( - "dict", - DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)), - false, - ), - ]; - let schema = Arc::new(ArrowSchema::new(fields)); - let dict_values = StringArray::from_iter_values(["a", "b", "c", "d", "e"]); - let batches: Vec = (0..20) - .map(|i| { - let mut arrays = - vec![Arc::new(Int32Array::from_iter_values(i * 20..(i + 1) * 20)) as ArrayRef]; - arrays.push(Arc::new( - DictionaryArray::try_new( - UInt16Array::from_iter_values((0_u16..20_u16).map(|v| v % 5)), - Arc::new(dict_values.clone()), - ) - .unwrap(), - )); - RecordBatch::try_new(schema.clone(), arrays).unwrap() - }) - .collect(); - let expected_batches = batches.clone(); - - let test_uri = path.to_str().unwrap(); - let write_params = WriteParams { - max_rows_per_file: 40, - max_rows_per_group: 10, - mode, - data_storage_version: Some(data_storage_version), - ..WriteParams::default() - }; - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - Dataset::write(reader, test_uri, Some(write_params)) - .await - .unwrap(); - - let actual_ds = Dataset::open(test_uri).await.unwrap(); - assert_eq!(actual_ds.version().version, 1); - assert_eq!( - actual_ds.manifest.writer_version, - Some(WriterVersion::default()) - ); - let actual_schema = ArrowSchema::from(actual_ds.schema()); - assert_eq!(&actual_schema, schema.as_ref()); - - let actual_batches = actual_ds - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - - // The batch size batches the group size. - // (the v2 writer has no concept of group size) - if data_storage_version == LanceFileVersion::Legacy { - for batch in &actual_batches { - assert_eq!(batch.num_rows(), 10); - } - } - - // sort - let actual_batch = concat_batches(&schema, &actual_batches).unwrap(); - let idx_arr = actual_batch.column_by_name("i").unwrap(); - let sorted_indices = sort_to_indices(idx_arr, None, None).unwrap(); - let struct_arr: StructArray = actual_batch.into(); - let sorted_arr = arrow_select::take::take(&struct_arr, &sorted_indices, None).unwrap(); - - let expected_struct_arr: StructArray = - concat_batches(&schema, &expected_batches).unwrap().into(); - assert_eq!(&expected_struct_arr, as_struct_array(sorted_arr.as_ref())); - - // Each fragments has different fragment ID - assert_eq!( - actual_ds - .fragments() - .iter() - .map(|f| f.id) - .collect::>(), - (0..10).collect::>() - ) - } - - #[rstest] - #[lance_test_macros::test(tokio::test)] - async fn test_create_dataset( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - // Appending / Overwriting a dataset that does not exist is treated as Create - for mode in [WriteMode::Create, WriteMode::Append, Overwrite] { - let test_dir = TempStdDir::default(); - create_file(&test_dir, mode, data_storage_version).await - } - } - - #[rstest] - #[lance_test_macros::test(tokio::test)] - async fn test_create_and_fill_empty_dataset( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - let test_uri = TempStrDir::default(); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let i32_array: ArrayRef = Arc::new(Int32Array::new(vec![].into(), None)); - let batch = RecordBatch::try_from_iter(vec![("i", i32_array)]).unwrap(); - let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); - // check schema of reader and original is same - assert_eq!(schema.as_ref(), reader.schema().as_ref()); - let result = Dataset::write( - reader, - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - ..Default::default() - }), - ) - .await - .unwrap(); - - // check dataset empty - assert_eq!(result.count_rows(None).await.unwrap(), 0); - // Since the dataset is empty, will return None. - assert_eq!(result.manifest.max_fragment_id(), None); - - // append rows to dataset - let mut write_params = WriteParams { - max_rows_per_file: 40, - max_rows_per_group: 10, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - // We should be able to append even if the metadata doesn't exactly match. - let schema_with_meta = Arc::new( - schema - .as_ref() - .clone() - .with_metadata([("key".to_string(), "value".to_string())].into()), - ); - let batches = vec![RecordBatch::try_new( - schema_with_meta, - vec![Arc::new(Int32Array::from_iter_values(0..10))], - ) - .unwrap()]; - write_params.mode = WriteMode::Append; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - Dataset::write(batches, &test_uri, Some(write_params)) - .await - .unwrap(); - - let expected_batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..10))], - ) - .unwrap(); - - // get actual dataset - let actual_ds = Dataset::open(&test_uri).await.unwrap(); - // confirm schema is same - let actual_schema = ArrowSchema::from(actual_ds.schema()); - assert_eq!(&actual_schema, schema.as_ref()); - // check num rows is 10 - assert_eq!(actual_ds.count_rows(None).await.unwrap(), 10); - // Max fragment id is still 0 since we only have 1 fragment. - assert_eq!(actual_ds.manifest.max_fragment_id(), Some(0)); - // check expected batch is correct - let actual_batches = actual_ds - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - // sort - let actual_batch = concat_batches(&schema, &actual_batches).unwrap(); - let idx_arr = actual_batch.column_by_name("i").unwrap(); - let sorted_indices = sort_to_indices(idx_arr, None, None).unwrap(); - let struct_arr: StructArray = actual_batch.into(); - let sorted_arr = arrow_select::take::take(&struct_arr, &sorted_indices, None).unwrap(); - let expected_struct_arr: StructArray = expected_batch.into(); - assert_eq!(&expected_struct_arr, as_struct_array(sorted_arr.as_ref())); - } - - #[rstest] - #[lance_test_macros::test(tokio::test)] - async fn test_create_with_empty_iter( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - let test_uri = TempStrDir::default(); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let reader = RecordBatchIterator::new(vec![].into_iter().map(Ok), schema.clone()); - // check schema of reader and original is same - assert_eq!(schema.as_ref(), reader.schema().as_ref()); - let write_params = Some(WriteParams { - data_storage_version: Some(data_storage_version), - ..Default::default() - }); - let result = Dataset::write(reader, &test_uri, write_params) - .await - .unwrap(); - - // check dataset empty - assert_eq!(result.count_rows(None).await.unwrap(), 0); - // Since the dataset is empty, will return None. - assert_eq!(result.manifest.max_fragment_id(), None); - } - - #[tokio::test] - async fn test_load_manifest_iops() { - // Use consistent session so memory store can be reused. - let session = Arc::new(Session::default()); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..10_i32))], - ) - .unwrap(); - let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); - let _original_ds = Dataset::write( - batches, - "memory://test", - Some(WriteParams { - session: Some(session.clone()), - ..Default::default() - }), - ) - .await - .unwrap(); - - let _ = _original_ds.object_store().io_stats_incremental(); //reset - - let _dataset = DatasetBuilder::from_uri("memory://test") - .with_session(session) - .load() - .await - .unwrap(); - - // There should be only two IOPS: - // 1. List _versions directory to get the latest manifest location - // 2. Read the manifest file. (The manifest is small enough to be read in one go. - // Larger manifests would result in more IOPS.) - let io_stats = _dataset.object_store().io_stats_incremental(); - assert_io_eq!(io_stats, read_iops, 2); - } - - #[rstest] - #[tokio::test] - async fn test_write_params( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - use fragment::FragReadConfig; - - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let num_rows: usize = 1_000; - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..num_rows as i32))], - ) - .unwrap()]; - - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - - let write_params = WriteParams { - max_rows_per_file: 100, - max_rows_per_group: 10, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - let dataset = Dataset::write(batches, &test_uri, Some(write_params)) - .await - .unwrap(); - - assert_eq!(dataset.count_rows(None).await.unwrap(), num_rows); - - let fragments = dataset.get_fragments(); - assert_eq!(fragments.len(), 10); - assert_eq!(dataset.count_fragments(), 10); - for fragment in &fragments { - assert_eq!(fragment.count_rows(None).await.unwrap(), 100); - let reader = fragment - .open(dataset.schema(), FragReadConfig::default()) - .await - .unwrap(); - // No group / batch concept in v2 - if data_storage_version == LanceFileVersion::Legacy { - assert_eq!(reader.legacy_num_batches(), 10); - for i in 0..reader.legacy_num_batches() as u32 { - assert_eq!(reader.legacy_num_rows_in_batch(i).unwrap(), 10); - } - } - } - } - - #[rstest] - #[tokio::test] - async fn test_write_manifest( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - use lance_table::feature_flags::FLAG_UNKNOWN; - - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..20))], - ) - .unwrap()]; - - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let write_fut = Dataset::write( - batches, - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - auto_cleanup: None, - ..Default::default() - }), - ); - let write_fut = require_send(write_fut); - let mut dataset = write_fut.await.unwrap(); - - // Check it has no flags - let manifest = read_manifest( - dataset.object_store(), - &dataset - .commit_handler - .resolve_latest_location(&dataset.base, dataset.object_store()) - .await - .unwrap() - .path, - None, - ) - .await - .unwrap(); - - assert_eq!( - manifest.data_storage_format, - DataStorageFormat::new(data_storage_version) - ); - assert_eq!(manifest.reader_feature_flags, 0); - - // Create one with deletions - dataset.delete("i < 10").await.unwrap(); - dataset.validate().await.unwrap(); - - // Check it set the flag - let mut manifest = read_manifest( - dataset.object_store(), - &dataset - .commit_handler - .resolve_latest_location(&dataset.base, dataset.object_store()) - .await - .unwrap() - .path, - None, - ) - .await - .unwrap(); - assert_eq!( - manifest.writer_feature_flags, - feature_flags::FLAG_DELETION_FILES - ); - assert_eq!( - manifest.reader_feature_flags, - feature_flags::FLAG_DELETION_FILES - ); - - // Write with custom manifest - manifest.writer_feature_flags |= FLAG_UNKNOWN; // Set another flag - manifest.reader_feature_flags |= FLAG_UNKNOWN; - manifest.version += 1; - write_manifest_file( - dataset.object_store(), - dataset.commit_handler.as_ref(), - &dataset.base, - &mut manifest, - None, - &ManifestWriteConfig { - auto_set_feature_flags: false, - timestamp: None, - use_stable_row_ids: false, - use_legacy_format: None, - storage_format: None, - disable_transaction_file: false, - }, - dataset.manifest_location.naming_scheme, - None, - ) - .await - .unwrap(); - - // Check it rejects reading it - let read_result = Dataset::open(&test_uri).await; - assert!(matches!(read_result, Err(Error::NotSupported { .. }))); - - // Check it rejects writing to it. - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..20))], - ) - .unwrap()]; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let write_result = Dataset::write( - batches, - &test_uri, - Some(WriteParams { - mode: WriteMode::Append, - data_storage_version: Some(data_storage_version), - ..Default::default() - }), - ) - .await; - - assert!(matches!(write_result, Err(Error::NotSupported { .. }))); - } - - #[rstest] - #[tokio::test] - async fn append_dataset( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..20))], - ) - .unwrap()]; - - let mut write_params = WriteParams { - max_rows_per_file: 40, - max_rows_per_group: 10, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - Dataset::write(batches, &test_uri, Some(write_params.clone())) - .await - .unwrap(); - - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(20..40))], - ) - .unwrap()]; - write_params.mode = WriteMode::Append; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - Dataset::write(batches, &test_uri, Some(write_params.clone())) - .await - .unwrap(); - - let expected_batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..40))], - ) - .unwrap(); - - let actual_ds = Dataset::open(&test_uri).await.unwrap(); - assert_eq!(actual_ds.version().version, 2); - let actual_schema = ArrowSchema::from(actual_ds.schema()); - assert_eq!(&actual_schema, schema.as_ref()); - - let actual_batches = actual_ds - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - // sort - let actual_batch = concat_batches(&schema, &actual_batches).unwrap(); - let idx_arr = actual_batch.column_by_name("i").unwrap(); - let sorted_indices = sort_to_indices(idx_arr, None, None).unwrap(); - let struct_arr: StructArray = actual_batch.into(); - let sorted_arr = arrow_select::take::take(&struct_arr, &sorted_indices, None).unwrap(); - - let expected_struct_arr: StructArray = expected_batch.into(); - assert_eq!(&expected_struct_arr, as_struct_array(sorted_arr.as_ref())); - - // Each fragments has different fragment ID - assert_eq!( - actual_ds - .fragments() - .iter() - .map(|f| f.id) - .collect::>(), - (0..2).collect::>() - ) - } - - #[rstest] - #[tokio::test] - async fn test_shallow_clone_with_hybrid_paths( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - let test_dir = TempStdDir::default(); - let base_dir = test_dir.join("base"); - let test_uri = base_dir.to_str().unwrap(); - let clone_dir = test_dir.join("clone"); - let cloned_uri = clone_dir.to_str().unwrap(); - - // Generate consistent test data batches - let generate_data = |prefix: &str, start_id: i32, row_count: u64| { - gen_batch() - .col("id", array::step_custom::(start_id, 1)) - .col("value", array::fill_utf8(format!("{prefix}_data"))) - .into_reader_rows(RowCount::from(row_count), BatchCount::from(1)) - }; - - // Reusable dataset writer with configurable mode - async fn write_dataset( - uri: &str, - data_reader: impl RecordBatchReader + Send + 'static, - mode: WriteMode, - version: LanceFileVersion, - ) -> Dataset { - let params = WriteParams { - max_rows_per_file: 100, - max_rows_per_group: 20, - data_storage_version: Some(version), - mode, - ..Default::default() - }; - Dataset::write(data_reader, uri, Some(params)) - .await - .unwrap() - } - - // Unified dataset scanning and row counting - async fn collect_rows(dataset: &Dataset) -> (usize, Vec) { - let batches = dataset - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - (batches.iter().map(|b| b.num_rows()).sum(), batches) - } - - // Create initial dataset - let mut dataset = write_dataset( - test_uri, - generate_data("initial", 0, 50), - WriteMode::Create, - data_storage_version, - ) - .await; - - // Store original state for comparison - let original_version = dataset.version().version; - let original_fragment_count = dataset.fragments().len(); - - // Create tag and shallow clone - dataset - .tags() - .create("test_tag", original_version) - .await - .unwrap(); - let cloned_dataset = dataset - .shallow_clone(cloned_uri, "test_tag", None) - .await - .unwrap(); - - // Verify cloned dataset state - let (cloned_rows, _) = collect_rows(&cloned_dataset).await; - assert_eq!(cloned_rows, 50); - assert_eq!(cloned_dataset.version().version, original_version); - - // Append data to cloned dataset - let updated_cloned = write_dataset( - cloned_uri, - generate_data("cloned_new", 50, 30), - WriteMode::Append, - data_storage_version, - ) - .await; - - // Verify updated cloned dataset - let (updated_cloned_rows, updated_batches) = collect_rows(&updated_cloned).await; - assert_eq!(updated_cloned_rows, 80); - assert_eq!(updated_cloned.version().version, original_version + 1); - - // Append data to original dataset - let updated_original = write_dataset( - test_uri, - generate_data("original_new", 50, 25), - WriteMode::Append, - data_storage_version, - ) - .await; - - // Verify updated original dataset - let (original_rows, _) = collect_rows(&updated_original).await; - assert_eq!(original_rows, 75); - assert_eq!(updated_original.version().version, original_version + 1); - - // Final validations - // Verify cloned dataset isolation - let final_cloned = Dataset::open(cloned_uri).await.unwrap(); - let (final_cloned_rows, _) = collect_rows(&final_cloned).await; - - // Data integrity check - let combined_batch = - concat_batches(&updated_batches[0].schema(), &updated_batches).unwrap(); - assert_eq!(combined_batch.column_by_name("id").unwrap().len(), 80); - assert_eq!(combined_batch.column_by_name("value").unwrap().len(), 80); - - // Fragment count validation - assert_eq!( - updated_original.fragments().len(), - original_fragment_count + 1 - ); - assert_eq!(final_cloned.fragments().len(), original_fragment_count + 1); - - // Final assertions - assert_eq!(final_cloned_rows, 80); - assert_eq!(final_cloned.version().version, original_version + 1); - } - - #[rstest] - #[tokio::test] - async fn test_shallow_clone_multiple_times( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - let test_uri = TempStrDir::default(); - let append_row_count = 36; - - // Async dataset writer function - async fn write_dataset( - dest: impl Into>, - row_count: u64, - mode: WriteMode, - version: LanceFileVersion, - ) -> Dataset { - let data = gen_batch() - .col("index", array::step::()) - .col("category", array::fill_utf8("base".to_string())) - .col("score", array::step_custom::(1.0, 0.5)); - Dataset::write( - data.into_reader_rows(RowCount::from(row_count), BatchCount::from(1)), - dest, - Some(WriteParams { - max_rows_per_file: 60, - max_rows_per_group: 12, - mode, - data_storage_version: Some(version), - ..Default::default() - }), - ) - .await - .unwrap() - } - - let mut current_dataset = write_dataset( - &test_uri, - append_row_count, - WriteMode::Create, - data_storage_version, - ) - .await; - - let test_round = 3; - // Generate clone paths - let clone_paths = (1..=test_round) - .map(|i| format!("{}/clone{}", test_uri, i)) - .collect::>(); - let mut cloned_datasets = Vec::with_capacity(test_round); - - // Unified cloning procedure, write a fragment to each cloned dataset. - for path in clone_paths.iter() { - current_dataset - .tags() - .create("v1", current_dataset.latest_version_id().await.unwrap()) - .await - .unwrap(); - - current_dataset = current_dataset - .shallow_clone(path, "v1", None) - .await - .unwrap(); - current_dataset = write_dataset( - Arc::new(current_dataset), - append_row_count, - WriteMode::Append, - data_storage_version, - ) - .await; - cloned_datasets.push(current_dataset.clone()); - } - - // Validation function - async fn validate_dataset( - dataset: &Dataset, - expected_rows: usize, - expected_fragments_count: usize, - expected_base_paths_count: usize, - ) { - let batches = dataset - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total_rows, expected_rows); - assert_eq!(dataset.fragments().len(), expected_fragments_count); - assert_eq!( - dataset.manifest().base_paths.len(), - expected_base_paths_count - ); - } - - // Verify cloned datasets row count, fragment count, base_path count - for (i, ds) in cloned_datasets.iter().enumerate() { - validate_dataset(ds, 36 * (i + 2), i + 2, i + 1).await; - } - - // Verify original dataset row count, fragment count, base_path count - let original = Dataset::open(&test_uri).await.unwrap(); - validate_dataset(&original, 36, 1, 0).await; - } - - #[rstest] - #[tokio::test] - async fn test_self_dataset_append( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..20))], - ) - .unwrap()]; - - let mut write_params = WriteParams { - max_rows_per_file: 40, - max_rows_per_group: 10, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let mut ds = Dataset::write(batches, &test_uri, Some(write_params.clone())) - .await - .unwrap(); - - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(20..40))], - ) - .unwrap()]; - write_params.mode = WriteMode::Append; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - - ds.append(batches, Some(write_params.clone())) - .await - .unwrap(); - - let expected_batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..40))], - ) - .unwrap(); - - let actual_ds = Dataset::open(&test_uri).await.unwrap(); - assert_eq!(actual_ds.version().version, 2); - // validate fragment ids - assert_eq!(actual_ds.fragments().len(), 2); - assert_eq!( - actual_ds - .fragments() - .iter() - .map(|f| f.id) - .collect::>(), - (0..2).collect::>() - ); - - let actual_schema = ArrowSchema::from(actual_ds.schema()); - assert_eq!(&actual_schema, schema.as_ref()); - - let actual_batches = actual_ds - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - // sort - let actual_batch = concat_batches(&schema, &actual_batches).unwrap(); - let idx_arr = actual_batch.column_by_name("i").unwrap(); - let sorted_indices = sort_to_indices(idx_arr, None, None).unwrap(); - let struct_arr: StructArray = actual_batch.into(); - let sorted_arr = arrow_select::take::take(&struct_arr, &sorted_indices, None).unwrap(); - - let expected_struct_arr: StructArray = expected_batch.into(); - assert_eq!(&expected_struct_arr, as_struct_array(sorted_arr.as_ref())); - - actual_ds.validate().await.unwrap(); - } - - #[rstest] - #[tokio::test] - async fn test_self_dataset_append_schema_different( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..20))], - ) - .unwrap()]; - - let other_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int64, - false, - )])); - let other_batches = vec![RecordBatch::try_new( - other_schema.clone(), - vec![Arc::new(Int64Array::from_iter_values(0..20))], - ) - .unwrap()]; - - let mut write_params = WriteParams { - max_rows_per_file: 40, - max_rows_per_group: 10, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let mut ds = Dataset::write(batches, &test_uri, Some(write_params.clone())) - .await - .unwrap(); - - write_params.mode = WriteMode::Append; - let other_batches = - RecordBatchIterator::new(other_batches.into_iter().map(Ok), other_schema.clone()); - - let result = ds.append(other_batches, Some(write_params.clone())).await; - // Error because schema is different - assert!(matches!(result, Err(Error::SchemaMismatch { .. }))) - } - - #[rstest] - #[tokio::test] - async fn append_dictionary( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - // We store the dictionary as part of the schema, so we check that the - // dictionary is consistent between appends. - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "x", - DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), - false, - )])); - let dictionary = Arc::new(StringArray::from(vec!["a", "b"])); - let indices = Int8Array::from(vec![0, 1, 0]); - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new( - Int8DictionaryArray::try_new(indices, dictionary.clone()).unwrap(), - )], - ) - .unwrap()]; - - let test_uri = TempStrDir::default(); - let mut write_params = WriteParams { - max_rows_per_file: 40, - max_rows_per_group: 10, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - Dataset::write(batches, &test_uri, Some(write_params.clone())) - .await - .unwrap(); - - // create a new one with same dictionary - let indices = Int8Array::from(vec![1, 0, 1]); - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new( - Int8DictionaryArray::try_new(indices, dictionary).unwrap(), - )], - ) - .unwrap()]; - - // Write to dataset (successful) - write_params.mode = WriteMode::Append; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - Dataset::write(batches, &test_uri, Some(write_params.clone())) - .await - .unwrap(); - - // Create a new one with *different* dictionary - let dictionary = Arc::new(StringArray::from(vec!["d", "c"])); - let indices = Int8Array::from(vec![1, 0, 1]); - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new( - Int8DictionaryArray::try_new(indices, dictionary).unwrap(), - )], - ) - .unwrap()]; - - // Try write to dataset (fails with legacy format) - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let result = Dataset::write(batches, &test_uri, Some(write_params)).await; - if data_storage_version == LanceFileVersion::Legacy { - assert!(result.is_err()); - } else { - assert!(result.is_ok()); - } - } - - #[rstest] - #[tokio::test] - async fn overwrite_dataset( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..20))], - ) - .unwrap()]; - - let mut write_params = WriteParams { - max_rows_per_file: 40, - max_rows_per_group: 10, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let dataset = Dataset::write(batches, &test_uri, Some(write_params.clone())) - .await - .unwrap(); - - let fragments = dataset.get_fragments(); - assert_eq!(fragments.len(), 1); - assert_eq!(dataset.manifest.max_fragment_id(), Some(0)); - - let new_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "s", - DataType::Utf8, - false, - )])); - let new_batches = vec![RecordBatch::try_new( - new_schema.clone(), - vec![Arc::new(StringArray::from_iter_values( - (20..40).map(|v| v.to_string()), - ))], - ) - .unwrap()]; - write_params.mode = Overwrite; - let new_batch_reader = - RecordBatchIterator::new(new_batches.into_iter().map(Ok), new_schema.clone()); - let dataset = Dataset::write(new_batch_reader, &test_uri, Some(write_params.clone())) - .await - .unwrap(); - - let fragments = dataset.get_fragments(); - assert_eq!(fragments.len(), 1); - // Fragment ids reset after overwrite. - assert_eq!(fragments[0].id(), 0); - assert_eq!(dataset.manifest.max_fragment_id(), Some(0)); - - let actual_ds = Dataset::open(&test_uri).await.unwrap(); - assert_eq!(actual_ds.version().version, 2); - let actual_schema = ArrowSchema::from(actual_ds.schema()); - assert_eq!(&actual_schema, new_schema.as_ref()); - - let actual_batches = actual_ds - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let actual_batch = concat_batches(&new_schema, &actual_batches).unwrap(); - - assert_eq!(new_schema.clone(), actual_batch.schema()); - let arr = actual_batch.column_by_name("s").unwrap(); - assert_eq!( - &StringArray::from_iter_values((20..40).map(|v| v.to_string())), - as_string_array(arr) - ); - assert_eq!(actual_ds.version().version, 2); - - // But we can still check out the first version - let first_ver = DatasetBuilder::from_uri(&test_uri) - .with_version(1) - .load() - .await - .unwrap(); - assert_eq!(first_ver.version().version, 1); - assert_eq!(&ArrowSchema::from(first_ver.schema()), schema.as_ref()); - } - - #[rstest] - #[tokio::test] - async fn test_fast_count_rows( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - - let batches: Vec = (0..20) - .map(|i| { - RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(i * 20..(i + 1) * 20))], - ) - .unwrap() - }) - .collect(); - - let write_params = WriteParams { - max_rows_per_file: 40, - max_rows_per_group: 10, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - Dataset::write(batches, &test_uri, Some(write_params)) - .await - .unwrap(); - - let dataset = Dataset::open(&test_uri).await.unwrap(); - dataset.validate().await.unwrap(); - assert_eq!(10, dataset.fragments().len()); - assert_eq!(400, dataset.count_rows(None).await.unwrap()); - assert_eq!( - 200, - dataset - .count_rows(Some("i < 200".to_string())) - .await - .unwrap() - ); - } - - #[rstest] - #[tokio::test] - async fn test_create_index( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - let test_uri = TempStrDir::default(); - - let dimension = 16; - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "embeddings", - DataType::FixedSizeList( - Arc::new(ArrowField::new("item", DataType::Float32, true)), - dimension, - ), - false, - )])); - - let float_arr = generate_random_array(512 * dimension as usize); - let vectors = Arc::new( - ::try_new_from_values( - float_arr, dimension, - ) - .unwrap(), - ); - let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors.clone()]).unwrap()]; - - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - - let mut dataset = Dataset::write( - reader, - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - ..Default::default() - }), - ) - .await - .unwrap(); - dataset.validate().await.unwrap(); - - // Make sure valid arguments should create index successfully - let params = VectorIndexParams::ivf_pq(10, 8, 2, MetricType::L2, 50); - dataset - .create_index(&["embeddings"], IndexType::Vector, None, ¶ms, true) - .await - .unwrap(); - dataset.validate().await.unwrap(); - - // The version should match the table version it was created from. - let indices = dataset.load_indices().await.unwrap(); - let actual = indices.first().unwrap().dataset_version; - let expected = dataset.manifest.version - 1; - assert_eq!(actual, expected); - let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); - assert_eq!(fragment_bitmap.len(), 1); - assert!(fragment_bitmap.contains(0)); - - // Append should inherit index - let write_params = WriteParams { - mode: WriteMode::Append, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors.clone()]).unwrap()]; - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let dataset = Dataset::write(reader, &test_uri, Some(write_params)) - .await - .unwrap(); - let indices = dataset.load_indices().await.unwrap(); - let actual = indices.first().unwrap().dataset_version; - let expected = dataset.manifest.version - 2; - assert_eq!(actual, expected); - dataset.validate().await.unwrap(); - // Fragment bitmap should show the original fragments, and not include - // the newly appended fragment. - let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); - assert_eq!(fragment_bitmap.len(), 1); - assert!(fragment_bitmap.contains(0)); - - let actual_statistics: serde_json::Value = - serde_json::from_str(&dataset.index_statistics("embeddings_idx").await.unwrap()) - .unwrap(); - let actual_statistics = actual_statistics.as_object().unwrap(); - assert_eq!(actual_statistics["index_type"].as_str().unwrap(), "IVF_PQ"); - - let deltas = actual_statistics["indices"].as_array().unwrap(); - assert_eq!(deltas.len(), 1); - assert_eq!(deltas[0]["metric_type"].as_str().unwrap(), "l2"); - assert_eq!(deltas[0]["num_partitions"].as_i64().unwrap(), 10); - - assert!(dataset.index_statistics("non-existent_idx").await.is_err()); - assert!(dataset.index_statistics("").await.is_err()); - - // Overwrite should invalidate index - let write_params = WriteParams { - mode: WriteMode::Overwrite, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors]).unwrap()]; - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let dataset = Dataset::write(reader, &test_uri, Some(write_params)) - .await - .unwrap(); - assert!(dataset.manifest.index_section.is_none()); - assert!(dataset.load_indices().await.unwrap().is_empty()); - dataset.validate().await.unwrap(); - - let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); - assert_eq!(fragment_bitmap.len(), 1); - assert!(fragment_bitmap.contains(0)); - } - - #[rstest] - #[tokio::test] - async fn test_create_scalar_index( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - #[values(false, true)] use_stable_row_id: bool, - ) { - let test_uri = TempStrDir::default(); - - let data = gen_batch().col("int", array::step::()); - // Write 64Ki rows. We should get 16 4Ki pages - let mut dataset = Dataset::write( - data.into_reader_rows(RowCount::from(16 * 1024), BatchCount::from(4)), - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - enable_stable_row_ids: use_stable_row_id, - ..Default::default() - }), - ) - .await - .unwrap(); - - let index_name = "my_index".to_string(); - - dataset - .create_index( - &["int"], - IndexType::Scalar, - Some(index_name.clone()), - &ScalarIndexParams::default(), - false, - ) - .await - .unwrap(); - - let indices = dataset.load_indices_by_name(&index_name).await.unwrap(); - - assert_eq!(indices.len(), 1); - assert_eq!(indices[0].dataset_version, 1); - assert_eq!(indices[0].fields, vec![0]); - assert_eq!(indices[0].name, index_name); - - dataset.index_statistics(&index_name).await.unwrap(); - } - - async fn create_bad_file(data_storage_version: LanceFileVersion) -> Result { - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "a.b.c", - DataType::Int32, - false, - )])); - - let batches: Vec = (0..20) - .map(|i| { - RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(i * 20..(i + 1) * 20))], - ) - .unwrap() - }) - .collect(); - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - Dataset::write( - reader, - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - ..Default::default() - }), - ) - .await - } - - #[tokio::test] - async fn test_create_fts_index_with_empty_table() { - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "text", - DataType::Utf8, - false, - )])); - - let batches: Vec = vec![]; - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let mut dataset = Dataset::write(reader, &test_uri, None) - .await - .expect("write dataset"); - - let params = InvertedIndexParams::default(); - dataset - .create_index(&["text"], IndexType::Inverted, None, ¶ms, true) - .await - .unwrap(); - - let batch = dataset - .scan() - .full_text_search(FullTextSearchQuery::new("lance".to_owned())) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(batch.num_rows(), 0); - } - - #[rstest] - #[tokio::test] - async fn test_create_int8_index( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - use lance_testing::datagen::generate_random_int8_array; - - let test_uri = TempStrDir::default(); - - let dimension = 16; - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "embeddings", - DataType::FixedSizeList( - Arc::new(ArrowField::new("item", DataType::Int8, true)), - dimension, - ), - false, - )])); - - let int8_arr = generate_random_int8_array(512 * dimension as usize); - let vectors = Arc::new( - ::try_new_from_values( - int8_arr, dimension, - ) - .unwrap(), - ); - let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors.clone()]).unwrap()]; - - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - - let mut dataset = Dataset::write( - reader, - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - ..Default::default() - }), - ) - .await - .unwrap(); - dataset.validate().await.unwrap(); - - // Make sure valid arguments should create index successfully - let params = VectorIndexParams::ivf_pq(10, 8, 2, MetricType::L2, 50); - dataset - .create_index(&["embeddings"], IndexType::Vector, None, ¶ms, true) - .await - .unwrap(); - dataset.validate().await.unwrap(); - - // The version should match the table version it was created from. - let indices = dataset.load_indices().await.unwrap(); - let actual = indices.first().unwrap().dataset_version; - let expected = dataset.manifest.version - 1; - assert_eq!(actual, expected); - let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); - assert_eq!(fragment_bitmap.len(), 1); - assert!(fragment_bitmap.contains(0)); - - // Append should inherit index - let write_params = WriteParams { - mode: WriteMode::Append, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors.clone()]).unwrap()]; - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let dataset = Dataset::write(reader, &test_uri, Some(write_params)) - .await - .unwrap(); - let indices = dataset.load_indices().await.unwrap(); - let actual = indices.first().unwrap().dataset_version; - let expected = dataset.manifest.version - 2; - assert_eq!(actual, expected); - dataset.validate().await.unwrap(); - // Fragment bitmap should show the original fragments, and not include - // the newly appended fragment. - let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); - assert_eq!(fragment_bitmap.len(), 1); - assert!(fragment_bitmap.contains(0)); - - let actual_statistics: serde_json::Value = - serde_json::from_str(&dataset.index_statistics("embeddings_idx").await.unwrap()) - .unwrap(); - let actual_statistics = actual_statistics.as_object().unwrap(); - assert_eq!(actual_statistics["index_type"].as_str().unwrap(), "IVF_PQ"); - - let deltas = actual_statistics["indices"].as_array().unwrap(); - assert_eq!(deltas.len(), 1); - assert_eq!(deltas[0]["metric_type"].as_str().unwrap(), "l2"); - assert_eq!(deltas[0]["num_partitions"].as_i64().unwrap(), 10); - - assert!(dataset.index_statistics("non-existent_idx").await.is_err()); - assert!(dataset.index_statistics("").await.is_err()); - - // Overwrite should invalidate index - let write_params = WriteParams { - mode: WriteMode::Overwrite, - data_storage_version: Some(data_storage_version), - ..Default::default() - }; - let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors]).unwrap()]; - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let dataset = Dataset::write(reader, &test_uri, Some(write_params)) - .await - .unwrap(); - assert!(dataset.manifest.index_section.is_none()); - assert!(dataset.load_indices().await.unwrap().is_empty()); - dataset.validate().await.unwrap(); - - let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); - assert_eq!(fragment_bitmap.len(), 1); - assert!(fragment_bitmap.contains(0)); - } - - #[tokio::test] - async fn test_create_fts_index_with_empty_strings() { - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "text", - DataType::Utf8, - false, - )])); - - let batches: Vec = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new(StringArray::from(vec!["", "", ""]))], - ) - .unwrap()]; - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - let mut dataset = Dataset::write(reader, &test_uri, None) - .await - .expect("write dataset"); - - let params = InvertedIndexParams::default(); - dataset - .create_index(&["text"], IndexType::Inverted, None, ¶ms, true) - .await - .unwrap(); - - let batch = dataset - .scan() - .full_text_search(FullTextSearchQuery::new("lance".to_owned())) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(batch.num_rows(), 0); - } - - #[rstest] - #[tokio::test] - async fn test_bad_field_name( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - // don't allow `.` in the field name - assert!(create_bad_file(data_storage_version).await.is_err()); - } - - #[tokio::test] - async fn test_open_dataset_not_found() { - let result = Dataset::open(".").await; - assert!(matches!(result.unwrap_err(), Error::DatasetNotFound { .. })); - } - - fn assert_all_manifests_use_scheme(test_dir: &TempStdDir, scheme: ManifestNamingScheme) { - let entries_names = test_dir - .join("_versions") - .read_dir() - .unwrap() - .map(|entry| entry.unwrap().file_name().into_string().unwrap()) - .collect::>(); - assert!( - entries_names - .iter() - .all(|name| ManifestNamingScheme::detect_scheme(name) == Some(scheme)), - "Entries: {:?}", - entries_names - ); - } - - #[tokio::test] - async fn test_v2_manifest_path_create() { - // Can create a dataset, using V2 paths - let data = lance_datagen::gen_batch() - .col("key", array::step::()) - .into_batch_rows(RowCount::from(10)) - .unwrap(); - let test_dir = TempStdDir::default(); - let test_uri = test_dir.to_str().unwrap(); - Dataset::write( - RecordBatchIterator::new([Ok(data.clone())], data.schema().clone()), - test_uri, - Some(WriteParams { - enable_v2_manifest_paths: true, - ..Default::default() - }), - ) - .await - .unwrap(); - - assert_all_manifests_use_scheme(&test_dir, ManifestNamingScheme::V2); - - // Appending to it will continue to use those paths - let dataset = Dataset::write( - RecordBatchIterator::new([Ok(data.clone())], data.schema().clone()), - test_uri, - Some(WriteParams { - mode: WriteMode::Append, - ..Default::default() - }), - ) - .await - .unwrap(); - - assert_all_manifests_use_scheme(&test_dir, ManifestNamingScheme::V2); - - UpdateBuilder::new(Arc::new(dataset)) - .update_where("key = 5") - .unwrap() - .set("key", "200") - .unwrap() - .build() - .unwrap() - .execute() - .await - .unwrap(); - - assert_all_manifests_use_scheme(&test_dir, ManifestNamingScheme::V2); - } - - #[tokio::test] - async fn test_v2_manifest_path_commit() { - let schema = Schema::try_from(&ArrowSchema::new(vec![ArrowField::new( - "x", - DataType::Int32, - false, - )])) - .unwrap(); - let operation = Operation::Overwrite { - fragments: vec![], - schema, - config_upsert_values: None, - initial_bases: None, - }; - let test_dir = TempStdDir::default(); - let test_uri = test_dir.to_str().unwrap(); - let dataset = Dataset::commit( - test_uri, - operation, - None, - None, - None, - Default::default(), - true, // enable_v2_manifest_paths - ) - .await - .unwrap(); - - assert!(dataset.manifest_location.naming_scheme == ManifestNamingScheme::V2); - - assert_all_manifests_use_scheme(&test_dir, ManifestNamingScheme::V2); - } - - #[tokio::test] - async fn test_strict_overwrite() { - let schema = Schema::try_from(&ArrowSchema::new(vec![ArrowField::new( - "x", - DataType::Int32, - false, - )])) - .unwrap(); - let operation = Operation::Overwrite { - fragments: vec![], - schema, - config_upsert_values: None, - initial_bases: None, - }; - let test_uri = TempStrDir::default(); - let read_version_0_transaction = Transaction::new(0, operation, None); - let strict_builder = CommitBuilder::new(&test_uri).with_max_retries(0); - let unstrict_builder = CommitBuilder::new(&test_uri).with_max_retries(1); - strict_builder - .clone() - .execute(read_version_0_transaction.clone()) - .await - .expect("Strict overwrite should succeed when writing a new dataset"); - strict_builder - .clone() - .execute(read_version_0_transaction.clone()) - .await - .expect_err("Strict overwrite should fail when committing to a stale version"); - unstrict_builder - .clone() - .execute(read_version_0_transaction.clone()) - .await - .expect("Unstrict overwrite should succeed when committing to a stale version"); - } - - #[rstest] - #[tokio::test] - async fn test_merge( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - #[values(false, true)] use_stable_row_id: bool, - ) { - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("i", DataType::Int32, false), - ArrowField::new("x", DataType::Float32, false), - ])); - let batch1 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![1, 2])), - Arc::new(Float32Array::from(vec![1.0, 2.0])), - ], - ) - .unwrap(); - let batch2 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![3, 2])), - Arc::new(Float32Array::from(vec![3.0, 4.0])), - ], - ) - .unwrap(); - - let test_uri = TempStrDir::default(); - - let write_params = WriteParams { - mode: WriteMode::Append, - data_storage_version: Some(data_storage_version), - enable_stable_row_ids: use_stable_row_id, - ..Default::default() - }; - - let batches = RecordBatchIterator::new(vec![batch1].into_iter().map(Ok), schema.clone()); - Dataset::write(batches, &test_uri, Some(write_params.clone())) - .await - .unwrap(); - - let batches = RecordBatchIterator::new(vec![batch2].into_iter().map(Ok), schema.clone()); - Dataset::write(batches, &test_uri, Some(write_params.clone())) - .await - .unwrap(); - - let dataset = Dataset::open(&test_uri).await.unwrap(); - assert_eq!(dataset.fragments().len(), 2); - assert_eq!(dataset.manifest.max_fragment_id(), Some(1)); - - let right_schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("i2", DataType::Int32, false), - ArrowField::new("y", DataType::Utf8, true), - ])); - let right_batch1 = RecordBatch::try_new( - right_schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![1, 2])), - Arc::new(StringArray::from(vec!["a", "b"])), - ], - ) - .unwrap(); - - let batches = - RecordBatchIterator::new(vec![right_batch1].into_iter().map(Ok), right_schema.clone()); - let mut dataset = Dataset::open(&test_uri).await.unwrap(); - dataset.merge(batches, "i", "i2").await.unwrap(); - dataset.validate().await.unwrap(); - - assert_eq!(dataset.version().version, 3); - assert_eq!(dataset.fragments().len(), 2); - assert_eq!(dataset.fragments()[0].files.len(), 2); - assert_eq!(dataset.fragments()[1].files.len(), 2); - assert_eq!(dataset.manifest.max_fragment_id(), Some(1)); - - let actual_batches = dataset - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let actual = concat_batches(&actual_batches[0].schema(), &actual_batches).unwrap(); - let expected = RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![ - ArrowField::new("i", DataType::Int32, false), - ArrowField::new("x", DataType::Float32, false), - ArrowField::new("y", DataType::Utf8, true), - ])), - vec![ - Arc::new(Int32Array::from(vec![1, 2, 3, 2])), - Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, 4.0])), - Arc::new(StringArray::from(vec![ - Some("a"), - Some("b"), - None, - Some("b"), - ])), - ], - ) - .unwrap(); - - assert_eq!(actual, expected); - - // Validate we can still read after re-instantiating dataset, which - // clears the cache. - let dataset = Dataset::open(&test_uri).await.unwrap(); - let actual_batches = dataset - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let actual = concat_batches(&actual_batches[0].schema(), &actual_batches).unwrap(); - assert_eq!(actual, expected); - } - - #[rstest] - #[tokio::test] - async fn test_large_merge( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - #[values(false, true)] use_stable_row_id: bool, - ) { - // Tests a merge that spans multiple batches within files - - // This test also tests "null filling" when merging (e.g. when keys do not match - // we need to insert nulls) - - let data = lance_datagen::gen_batch() - .col("key", array::step::()) - .col("value", array::fill_utf8("value".to_string())) - .into_reader_rows(RowCount::from(1_000), BatchCount::from(10)); - - let test_uri = TempStrDir::default(); - - let write_params = WriteParams { - mode: WriteMode::Append, - data_storage_version: Some(data_storage_version), - max_rows_per_file: 1024, - max_rows_per_group: 150, - enable_stable_row_ids: use_stable_row_id, - ..Default::default() - }; - Dataset::write(data, &test_uri, Some(write_params.clone())) - .await - .unwrap(); - - let mut dataset = Dataset::open(&test_uri).await.unwrap(); - assert_eq!(dataset.fragments().len(), 10); - assert_eq!(dataset.manifest.max_fragment_id(), Some(9)); - - let new_data = lance_datagen::gen_batch() - .col("key2", array::step_custom::(500, 1)) - .col("new_value", array::fill_utf8("new_value".to_string())) - .into_reader_rows(RowCount::from(1_000), BatchCount::from(10)); - - dataset.merge(new_data, "key", "key2").await.unwrap(); - dataset.validate().await.unwrap(); - } - - #[rstest] - #[tokio::test] - async fn test_merge_on_row_id( - #[values(LanceFileVersion::Stable)] data_storage_version: LanceFileVersion, - #[values(false, true)] use_stable_row_id: bool, - ) { - // Tests a merge on _rowid - - let data = lance_datagen::gen_batch() - .col("key", array::step::()) - .col("value", array::fill_utf8("value".to_string())) - .into_reader_rows(RowCount::from(1_000), BatchCount::from(10)); - - let write_params = WriteParams { - mode: WriteMode::Append, - data_storage_version: Some(data_storage_version), - max_rows_per_file: 1024, - max_rows_per_group: 150, - enable_stable_row_ids: use_stable_row_id, - ..Default::default() - }; - let mut dataset = Dataset::write(data, "memory://", Some(write_params.clone())) - .await - .unwrap(); - assert_eq!(dataset.fragments().len(), 10); - assert_eq!(dataset.manifest.max_fragment_id(), Some(9)); - - let data = dataset.scan().with_row_id().try_into_batch().await.unwrap(); - let row_ids: Arc = data[ROW_ID].clone(); - let key = data["key"].as_primitive::(); - let new_schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("rowid", DataType::UInt64, false), - ArrowField::new("new_value", DataType::Int32, false), - ])); - let new_value = Arc::new( - key.into_iter() - .map(|v| v.unwrap() + 1) - .collect::(), - ); - let len = new_value.len() as u32; - let new_batch = RecordBatch::try_new(new_schema.clone(), vec![row_ids, new_value]).unwrap(); - // shuffle new_batch - let mut rng = rand::rng(); - let mut indices: Vec = (0..len).collect(); - indices.shuffle(&mut rng); - let indices = arrow_array::UInt32Array::from_iter_values(indices); - let new_batch = arrow::compute::take_record_batch(&new_batch, &indices).unwrap(); - let new_data = RecordBatchIterator::new(vec![Ok(new_batch)], new_schema.clone()); - dataset.merge(new_data, ROW_ID, "rowid").await.unwrap(); - dataset.validate().await.unwrap(); - assert_eq!(dataset.schema().fields.len(), 3); - assert!(dataset.schema().field("key").is_some()); - assert!(dataset.schema().field("value").is_some()); - assert!(dataset.schema().field("new_value").is_some()); - let batch = dataset.scan().try_into_batch().await.unwrap(); - let key = batch["key"].as_primitive::(); - let new_value = batch["new_value"].as_primitive::(); - for i in 0..key.len() { - assert_eq!(key.value(i) + 1, new_value.value(i)); - } - } - - #[rstest] - #[tokio::test] - async fn test_merge_on_row_addr( - #[values(LanceFileVersion::Stable)] data_storage_version: LanceFileVersion, - #[values(false, true)] use_stable_row_id: bool, - ) { - // Tests a merge on _rowaddr - - let data = lance_datagen::gen_batch() - .col("key", array::step::()) - .col("value", array::fill_utf8("value".to_string())) - .into_reader_rows(RowCount::from(1_000), BatchCount::from(10)); - - let write_params = WriteParams { - mode: WriteMode::Append, - data_storage_version: Some(data_storage_version), - max_rows_per_file: 1024, - max_rows_per_group: 150, - enable_stable_row_ids: use_stable_row_id, - ..Default::default() - }; - let mut dataset = Dataset::write(data, "memory://", Some(write_params.clone())) - .await - .unwrap(); - - assert_eq!(dataset.fragments().len(), 10); - assert_eq!(dataset.manifest.max_fragment_id(), Some(9)); - - let data = dataset - .scan() - .with_row_address() - .try_into_batch() - .await - .unwrap(); - let row_addrs = data[ROW_ADDR].clone(); - let key = data["key"].as_primitive::(); - let new_schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("rowaddr", DataType::UInt64, false), - ArrowField::new("new_value", DataType::Int32, false), - ])); - let new_value = Arc::new( - key.into_iter() - .map(|v| v.unwrap() + 1) - .collect::(), - ); - let len = new_value.len() as u32; - let new_batch = - RecordBatch::try_new(new_schema.clone(), vec![row_addrs, new_value]).unwrap(); - // shuffle new_batch - let mut rng = rand::rng(); - let mut indices: Vec = (0..len).collect(); - indices.shuffle(&mut rng); - let indices = arrow_array::UInt32Array::from_iter_values(indices); - let new_batch = arrow::compute::take_record_batch(&new_batch, &indices).unwrap(); - let new_data = RecordBatchIterator::new(vec![Ok(new_batch)], new_schema.clone()); - dataset.merge(new_data, ROW_ADDR, "rowaddr").await.unwrap(); - dataset.validate().await.unwrap(); - assert_eq!(dataset.schema().fields.len(), 3); - assert!(dataset.schema().field("key").is_some()); - assert!(dataset.schema().field("value").is_some()); - assert!(dataset.schema().field("new_value").is_some()); - let batch = dataset.scan().try_into_batch().await.unwrap(); - let key = batch["key"].as_primitive::(); - let new_value = batch["new_value"].as_primitive::(); - for i in 0..key.len() { - assert_eq!(key.value(i) + 1, new_value.value(i)); - } - } - - #[rstest] - #[tokio::test] - async fn test_restore( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - // Create a table - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::UInt32, - false, - )])); - - let test_uri = TempStrDir::default(); - - let data = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(UInt32Array::from_iter_values(0..100))], - ); - let reader = RecordBatchIterator::new(vec![data.unwrap()].into_iter().map(Ok), schema); - let mut dataset = Dataset::write( - reader, - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - ..Default::default() - }), - ) - .await - .unwrap(); - assert_eq!(dataset.manifest.version, 1); - let original_manifest = dataset.manifest.clone(); - - // Delete some rows - dataset.delete("i > 50").await.unwrap(); - assert_eq!(dataset.manifest.version, 2); - - // Checkout a previous version - let mut dataset = dataset.checkout_version(1).await.unwrap(); - assert_eq!(dataset.manifest.version, 1); - let fragments = dataset.get_fragments(); - assert_eq!(fragments.len(), 1); - assert_eq!(dataset.count_fragments(), 1); - assert_eq!(fragments[0].metadata.deletion_file, None); - assert_eq!(dataset.manifest, original_manifest); - - // Checkout latest and then go back. - dataset.checkout_latest().await.unwrap(); - assert_eq!(dataset.manifest.version, 2); - let mut dataset = dataset.checkout_version(1).await.unwrap(); - - // Restore to a previous version - dataset.restore().await.unwrap(); - assert_eq!(dataset.manifest.version, 3); - assert_eq!(dataset.manifest.fragments, original_manifest.fragments); - assert_eq!(dataset.manifest.schema, original_manifest.schema); - - // Delete some rows again (make sure we can still write as usual) - dataset.delete("i > 30").await.unwrap(); - assert_eq!(dataset.manifest.version, 4); - let fragments = dataset.get_fragments(); - assert_eq!(fragments.len(), 1); - assert_eq!(dataset.count_fragments(), 1); - assert!(fragments[0].metadata.deletion_file.is_some()); - } - - #[rstest] - #[tokio::test] - async fn test_tag( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - // Create a table - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::UInt32, - false, - )])); - - let test_uri = TempStrDir::default(); - - let data = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(UInt32Array::from_iter_values(0..100))], - ); - let reader = RecordBatchIterator::new(vec![data.unwrap()].into_iter().map(Ok), schema); - let mut dataset = Dataset::write( - reader, - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - ..Default::default() - }), - ) - .await - .unwrap(); - assert_eq!(dataset.manifest.version, 1); - - // delete some rows - dataset.delete("i > 50").await.unwrap(); - assert_eq!(dataset.manifest.version, 2); - - assert_eq!(dataset.tags().list().await.unwrap().len(), 0); - - let bad_tag_creation = dataset.tags().create("tag1", 3).await; - assert_eq!( - bad_tag_creation.err().unwrap().to_string(), - "Version not found error: version Main::3 does not exist" - ); - - let bad_tag_deletion = dataset.tags().delete("tag1").await; - assert_eq!( - bad_tag_deletion.err().unwrap().to_string(), - "Ref not found error: tag tag1 does not exist" - ); - - dataset.tags().create("tag1", 1).await.unwrap(); - - assert_eq!(dataset.tags().list().await.unwrap().len(), 1); - - let another_bad_tag_creation = dataset.tags().create("tag1", 1).await; - assert_eq!( - another_bad_tag_creation.err().unwrap().to_string(), - "Ref conflict error: tag tag1 already exists" - ); - - dataset.tags().delete("tag1").await.unwrap(); - - assert_eq!(dataset.tags().list().await.unwrap().len(), 0); - - dataset.tags().create("tag1", 1).await.unwrap(); - dataset.tags().create("tag2", 1).await.unwrap(); - dataset.tags().create("v1.0.0-rc1", 2).await.unwrap(); - - let default_order = dataset.tags().list_tags_ordered(None).await.unwrap(); - let default_names: Vec<_> = default_order.iter().map(|t| &t.0).collect(); - assert_eq!( - default_names, - ["v1.0.0-rc1", "tag1", "tag2"], - "Default ordering mismatch" - ); - - let asc_order = dataset - .tags() - .list_tags_ordered(Some(Ordering::Less)) - .await - .unwrap(); - let asc_names: Vec<_> = asc_order.iter().map(|t| &t.0).collect(); - assert_eq!( - asc_names, - ["tag1", "tag2", "v1.0.0-rc1"], - "Ascending ordering mismatch" - ); - - let desc_order = dataset - .tags() - .list_tags_ordered(Some(Ordering::Greater)) - .await - .unwrap(); - let desc_names: Vec<_> = desc_order.iter().map(|t| &t.0).collect(); - assert_eq!( - desc_names, - ["v1.0.0-rc1", "tag1", "tag2"], - "Descending ordering mismatch" - ); - - assert_eq!(dataset.tags().list().await.unwrap().len(), 3); - - let bad_checkout = dataset.checkout_version("tag3").await; - assert_eq!( - bad_checkout.err().unwrap().to_string(), - "Ref not found error: tag tag3 does not exist" - ); - - dataset = dataset.checkout_version("tag1").await.unwrap(); - assert_eq!(dataset.manifest.version, 1); - - let first_ver = DatasetBuilder::from_uri(&test_uri) - .with_tag("tag1") - .load() - .await - .unwrap(); - assert_eq!(first_ver.version().version, 1); - - // test update tag - let bad_tag_update = dataset.tags().update("tag3", 1).await; - assert_eq!( - bad_tag_update.err().unwrap().to_string(), - "Ref not found error: tag tag3 does not exist" - ); - - let another_bad_tag_update = dataset.tags().update("tag1", 3).await; - assert_eq!( - another_bad_tag_update.err().unwrap().to_string(), - "Version not found error: version 3 does not exist" - ); - - dataset.tags().update("tag1", 2).await.unwrap(); - dataset = dataset.checkout_version("tag1").await.unwrap(); - assert_eq!(dataset.manifest.version, 2); - - dataset.tags().update("tag1", 1).await.unwrap(); - dataset = dataset.checkout_version("tag1").await.unwrap(); - assert_eq!(dataset.manifest.version, 1); - } - - #[rstest] - #[tokio::test] - async fn test_search_empty( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - // Create a table - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "vec", - DataType::FixedSizeList( - Arc::new(ArrowField::new("item", DataType::Float32, true)), - 128, - ), - false, - )])); - - let test_uri = TempStrDir::default(); - - let vectors = Arc::new( - ::try_new_from_values( - Float32Array::from_iter_values(vec![]), - 128, - ) - .unwrap(), - ); - - let data = RecordBatch::try_new(schema.clone(), vec![vectors]); - let reader = RecordBatchIterator::new(vec![data.unwrap()].into_iter().map(Ok), schema); - let dataset = Dataset::write( - reader, - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - ..Default::default() - }), - ) - .await - .unwrap(); - - let mut stream = dataset - .scan() - .nearest( - "vec", - &Float32Array::from_iter_values((0..128).map(|_| 0.1)), - 1, - ) - .unwrap() - .try_into_stream() - .await - .unwrap(); - - while let Some(batch) = stream.next().await { - let schema = batch.unwrap().schema(); - assert_eq!(schema.fields.len(), 2); - assert_eq!( - schema.field_with_name("vec").unwrap(), - &ArrowField::new( - "vec", - DataType::FixedSizeList( - Arc::new(ArrowField::new("item", DataType::Float32, true)), - 128 - ), - false, - ) - ); - assert_eq!( - schema.field_with_name(DIST_COL).unwrap(), - &ArrowField::new(DIST_COL, DataType::Float32, true) - ); - } - } - - #[rstest] - #[tokio::test] - async fn test_search_empty_after_delete( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - #[values(false, true)] use_stable_row_id: bool, - ) { - // Create a table - let test_uri = TempStrDir::default(); - - let data = gen_batch().col("vec", array::rand_vec::(Dimension::from(32))); - let reader = data.into_reader_rows(RowCount::from(500), BatchCount::from(1)); - let mut dataset = Dataset::write( - reader, - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - enable_stable_row_ids: use_stable_row_id, - ..Default::default() - }), - ) - .await - .unwrap(); - - let params = VectorIndexParams::ivf_pq(1, 8, 1, MetricType::L2, 50); - dataset - .create_index(&["vec"], IndexType::Vector, None, ¶ms, true) - .await - .unwrap(); - - dataset.delete("true").await.unwrap(); - - // This behavior will be re-introduced once we work on empty vector index handling. - // https://github.com/lance-format/lance/issues/4034 - // let indices = dataset.load_indices().await.unwrap(); - // // With the new retention behavior, indices are kept even when all fragments are deleted - // // This allows the index configuration to persist through data changes - // assert_eq!(indices.len(), 1); - - // // Verify the index has an empty effective fragment bitmap - // let index = &indices[0]; - // let effective_bitmap = index - // .effective_fragment_bitmap(&dataset.fragment_bitmap) - // .unwrap(); - // assert!(effective_bitmap.is_empty()); - - let mut stream = dataset - .scan() - .nearest( - "vec", - &Float32Array::from_iter_values((0..32).map(|_| 0.1)), - 1, - ) - .unwrap() - .try_into_stream() - .await - .unwrap(); - - while let Some(batch) = stream.next().await { - let schema = batch.unwrap().schema(); - assert_eq!(schema.fields.len(), 2); - assert_eq!( - schema.field_with_name("vec").unwrap(), - &ArrowField::new( - "vec", - DataType::FixedSizeList( - Arc::new(ArrowField::new("item", DataType::Float32, true)), - 32 - ), - false, - ) - ); - assert_eq!( - schema.field_with_name(DIST_COL).unwrap(), - &ArrowField::new(DIST_COL, DataType::Float32, true) - ); - } - - // predicate with redundant whitespace - dataset.delete(" True").await.unwrap(); - - let mut stream = dataset - .scan() - .nearest( - "vec", - &Float32Array::from_iter_values((0..32).map(|_| 0.1)), - 1, - ) - .unwrap() - .try_into_stream() - .await - .unwrap(); - - while let Some(batch) = stream.next().await { - let batch = batch.unwrap(); - let schema = batch.schema(); - assert_eq!(schema.fields.len(), 2); - assert_eq!( - schema.field_with_name("vec").unwrap(), - &ArrowField::new( - "vec", - DataType::FixedSizeList( - Arc::new(ArrowField::new("item", DataType::Float32, true)), - 32 - ), - false, - ) - ); - assert_eq!( - schema.field_with_name(DIST_COL).unwrap(), - &ArrowField::new(DIST_COL, DataType::Float32, true) - ); - assert_eq!(batch.num_rows(), 0, "Expected no results after delete"); - } - } - - #[rstest] - #[tokio::test] - async fn test_num_small_files( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - let test_uri = TempStrDir::default(); - let dimensions = 16; - let column_name = "vec"; - let field = ArrowField::new( - column_name, - DataType::FixedSizeList( - Arc::new(ArrowField::new("item", DataType::Float32, true)), - dimensions, - ), - false, - ); - - let schema = Arc::new(ArrowSchema::new(vec![field])); - - let float_arr = generate_random_array(512 * dimensions as usize); - let vectors = - arrow_array::FixedSizeListArray::try_new_from_values(float_arr, dimensions).unwrap(); - - let record_batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vectors)]).unwrap(); - - let reader = - RecordBatchIterator::new(vec![record_batch].into_iter().map(Ok), schema.clone()); - - let dataset = Dataset::write( - reader, - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - ..Default::default() - }), - ) - .await - .unwrap(); - dataset.validate().await.unwrap(); - - assert!(dataset.num_small_files(1024).await > 0); - assert!(dataset.num_small_files(512).await == 0); - } - - #[tokio::test] - async fn test_read_struct_of_dictionary_arrays() { - let test_uri = TempStrDir::default(); - - let arrow_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "s", - DataType::Struct(ArrowFields::from(vec![ArrowField::new( - "d", - DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), - true, - )])), - true, - )])); - - let mut batches: Vec = Vec::new(); - for _ in 1..2 { - let mut dict_builder = StringDictionaryBuilder::::new(); - dict_builder.append("a").unwrap(); - dict_builder.append("b").unwrap(); - dict_builder.append("c").unwrap(); - dict_builder.append("d").unwrap(); - - let struct_array = Arc::new(StructArray::from(vec![( - Arc::new(ArrowField::new( - "d", - DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), - true, - )), - Arc::new(dict_builder.finish()) as ArrayRef, - )])); - - let batch = - RecordBatch::try_new(arrow_schema.clone(), vec![struct_array.clone()]).unwrap(); - batches.push(batch); - } - - let batch_reader = - RecordBatchIterator::new(batches.clone().into_iter().map(Ok), arrow_schema.clone()); - Dataset::write(batch_reader, &test_uri, Some(WriteParams::default())) - .await - .unwrap(); - - let result = scan_dataset(&test_uri).await.unwrap(); - - assert_eq!(batches, result); - } - - async fn scan_dataset(uri: &str) -> Result> { - let results = Dataset::open(uri) - .await? - .scan() - .try_into_stream() - .await? - .try_collect::>() - .await?; - Ok(results) - } - - #[rstest] - #[tokio::test] - async fn test_v0_7_5_migration() { - // We migrate to add Fragment.physical_rows and DeletionFile.num_deletions - // after this version. - - // Copy over table - let test_dir = copy_test_data_to_tmp("v0.7.5/with_deletions").unwrap(); - let test_uri = test_dir.path_str(); - - // Assert num rows, deletions, and physical rows are all correct. - let dataset = Dataset::open(&test_uri).await.unwrap(); - assert_eq!(dataset.count_rows(None).await.unwrap(), 90); - assert_eq!(dataset.count_deleted_rows().await.unwrap(), 10); - let total_physical_rows = futures::stream::iter(dataset.get_fragments()) - .then(|f| async move { f.physical_rows().await }) - .try_fold(0, |acc, x| async move { Ok(acc + x) }) - .await - .unwrap(); - assert_eq!(total_physical_rows, 100); - - // Append 5 rows - let schema = Arc::new(ArrowSchema::from(dataset.schema())); - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int64Array::from_iter_values(100..105))], - ) - .unwrap(); - let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); - let write_params = WriteParams { - mode: WriteMode::Append, - ..Default::default() - }; - let dataset = Dataset::write(batches, &test_uri, Some(write_params)) - .await - .unwrap(); - - // Assert num rows, deletions, and physical rows are all correct. - assert_eq!(dataset.count_rows(None).await.unwrap(), 95); - assert_eq!(dataset.count_deleted_rows().await.unwrap(), 10); - let total_physical_rows = futures::stream::iter(dataset.get_fragments()) - .then(|f| async move { f.physical_rows().await }) - .try_fold(0, |acc, x| async move { Ok(acc + x) }) - .await - .unwrap(); - assert_eq!(total_physical_rows, 105); - - dataset.validate().await.unwrap(); - - // Scan data and assert it is as expected. - let expected = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int64Array::from_iter_values( - (0..10).chain(20..105), - ))], - ) - .unwrap(); - let actual_batches = dataset - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let actual = concat_batches(&actual_batches[0].schema(), &actual_batches).unwrap(); - assert_eq!(actual, expected); - } - - #[rstest] - #[tokio::test] - async fn test_fix_v0_8_0_broken_migration() { - // The migration from v0.7.5 was broken in 0.8.0. This validates we can - // automatically fix tables that have this problem. - - // Copy over table - let test_dir = copy_test_data_to_tmp("v0.8.0/migrated_from_v0.7.5").unwrap(); - let test_uri = test_dir.path_str(); - let test_uri = &test_uri; - - // Assert num rows, deletions, and physical rows are all correct, even - // though stats are bad. - let dataset = Dataset::open(test_uri).await.unwrap(); - assert_eq!(dataset.count_rows(None).await.unwrap(), 92); - assert_eq!(dataset.count_deleted_rows().await.unwrap(), 10); - let total_physical_rows = futures::stream::iter(dataset.get_fragments()) - .then(|f| async move { f.physical_rows().await }) - .try_fold(0, |acc, x| async move { Ok(acc + x) }) - .await - .unwrap(); - assert_eq!(total_physical_rows, 102); - - // Append 5 rows to table. - let schema = Arc::new(ArrowSchema::from(dataset.schema())); - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int64Array::from_iter_values(100..105))], - ) - .unwrap(); - let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); - let write_params = WriteParams { - mode: WriteMode::Append, - data_storage_version: Some(LanceFileVersion::Legacy), - ..Default::default() - }; - let dataset = Dataset::write(batches, test_uri, Some(write_params)) - .await - .unwrap(); - - // Assert statistics are all now correct. - let physical_rows: Vec<_> = dataset - .get_fragments() - .iter() - .map(|f| f.metadata.physical_rows) - .collect(); - assert_eq!(physical_rows, vec![Some(100), Some(2), Some(5)]); - let num_deletions: Vec<_> = dataset - .get_fragments() - .iter() - .map(|f| { - f.metadata - .deletion_file - .as_ref() - .and_then(|df| df.num_deleted_rows) - }) - .collect(); - assert_eq!(num_deletions, vec![Some(10), None, None]); - assert_eq!(dataset.count_rows(None).await.unwrap(), 97); - - // Scan data and assert it is as expected. - let expected = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int64Array::from_iter_values( - (0..10).chain(20..100).chain(0..2).chain(100..105), - ))], - ) - .unwrap(); - let actual_batches = dataset - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let actual = concat_batches(&actual_batches[0].schema(), &actual_batches).unwrap(); - assert_eq!(actual, expected); - } - - #[rstest] - #[tokio::test] - async fn test_v0_8_14_invalid_index_fragment_bitmap( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) { - // Old versions of lance could create an index whose fragment bitmap was - // invalid because it did not include fragments that were part of the index - // - // We need to make sure we do not rely on the fragment bitmap in these older - // versions and instead fall back to a slower legacy behavior - let test_dir = copy_test_data_to_tmp("v0.8.14/corrupt_index").unwrap(); - let test_uri = test_dir.path_str(); - let test_uri = &test_uri; - - let mut dataset = Dataset::open(test_uri).await.unwrap(); - - // Uncomment to reproduce the issue. The below query will panic - // let mut scan = dataset.scan(); - // let query_vec = Float32Array::from(vec![0_f32; 128]); - // let scan_fut = scan - // .nearest("vector", &query_vec, 2000) - // .unwrap() - // .nprobes(4) - // .prefilter(true) - // .try_into_stream() - // .await - // .unwrap() - // .try_collect::>() - // .await - // .unwrap(); - - // Add some data and recalculate the index, forcing a migration - let mut scan = dataset.scan(); - let data = scan - .limit(Some(10), None) - .unwrap() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let schema = data[0].schema(); - let data = RecordBatchIterator::new(data.into_iter().map(arrow::error::Result::Ok), schema); - - let broken_version = dataset.version().version; - - // Any transaction, no matter how simple, should trigger the fragment bitmap to be recalculated - dataset - .append( - data, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - ..Default::default() - }), - ) - .await - .unwrap(); - - for idx in dataset.load_indices().await.unwrap().iter() { - // The corrupt fragment_bitmap does not contain 0 but the - // restored one should - assert!(idx.fragment_bitmap.as_ref().unwrap().contains(0)); - } - - let mut dataset = dataset.checkout_version(broken_version).await.unwrap(); - dataset.restore().await.unwrap(); - - // Running compaction right away should work (this is verifying compaction - // is not broken by the potentially malformed fragment bitmaps) - compact_files(&mut dataset, CompactionOptions::default(), None) - .await - .unwrap(); - - for idx in dataset.load_indices().await.unwrap().iter() { - assert!(idx.fragment_bitmap.as_ref().unwrap().contains(0)); - } - - let mut scan = dataset.scan(); - let query_vec = Float32Array::from(vec![0_f32; 128]); - let batches = scan - .nearest("vector", &query_vec, 2000) - .unwrap() - .nprobes(4) - .prefilter(true) - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - - let row_count = batches.iter().map(|batch| batch.num_rows()).sum::(); - assert_eq!(row_count, 1900); - } - - #[tokio::test] - async fn test_fix_v0_10_5_corrupt_schema() { - // Schemas could be corrupted by successive calls to `add_columns` and - // `drop_columns`. We should be able to detect this by checking for - // duplicate field ids. We should be able to fix this in new commits - // by dropping unused data files and re-writing the schema. - - // Copy over table - let test_dir = copy_test_data_to_tmp("v0.10.5/corrupt_schema").unwrap(); - let test_uri = test_dir.path_str(); - let test_uri = &test_uri; - - let mut dataset = Dataset::open(test_uri).await.unwrap(); - - let validate_res = dataset.validate().await; - assert!(validate_res.is_err()); - - // Force a migration. - dataset.delete("false").await.unwrap(); - dataset.validate().await.unwrap(); - - let data = dataset.scan().try_into_batch().await.unwrap(); - assert_eq!( - data["b"] - .as_any() - .downcast_ref::() - .unwrap() - .values(), - &[0, 4, 8, 12] - ); - assert_eq!( - data["c"] - .as_any() - .downcast_ref::() - .unwrap() - .values(), - &[0, 5, 10, 15] - ); - } - - #[tokio::test] - async fn test_fix_v0_21_0_corrupt_fragment_bitmap() { - // In v0.21.0 and earlier, delta indices had a bug where the fragment bitmap - // could contain fragments that are part of other index deltas. - - // Copy over table - let test_dir = copy_test_data_to_tmp("v0.21.0/bad_index_fragment_bitmap").unwrap(); - let test_uri = test_dir.path_str(); - let test_uri = &test_uri; - - let mut dataset = Dataset::open(test_uri).await.unwrap(); - - let validate_res = dataset.validate().await; - assert!(validate_res.is_err()); - assert_eq!(dataset.load_indices().await.unwrap()[0].name, "vector_idx"); - - // Calling index statistics will force a migration - let stats = dataset.index_statistics("vector_idx").await.unwrap(); - let stats: serde_json::Value = serde_json::from_str(&stats).unwrap(); - assert_eq!(stats["num_indexed_fragments"], 2); - - dataset.checkout_latest().await.unwrap(); - dataset.validate().await.unwrap(); - - let indices = dataset.load_indices().await.unwrap(); - assert_eq!(indices.len(), 2); - fn get_bitmap(meta: &IndexMetadata) -> Vec { - meta.fragment_bitmap.as_ref().unwrap().iter().collect() - } - assert_eq!(get_bitmap(&indices[0]), vec![0]); - assert_eq!(get_bitmap(&indices[1]), vec![1]); - } - - #[tokio::test] - async fn test_max_fragment_id_migration() { - // v0.5.9 and earlier did not store the max fragment id in the manifest. - // This test ensures that we can read such datasets and migrate them to - // the latest version, which requires the max fragment id to be present. - { - let test_dir = copy_test_data_to_tmp("v0.5.9/no_fragments").unwrap(); - let test_uri = test_dir.path_str(); - let test_uri = &test_uri; - let dataset = Dataset::open(test_uri).await.unwrap(); - - assert_eq!(dataset.manifest.max_fragment_id, None); - assert_eq!(dataset.manifest.max_fragment_id(), None); - } - - { - let test_dir = copy_test_data_to_tmp("v0.5.9/dataset_with_fragments").unwrap(); - let test_uri = test_dir.path_str(); - let test_uri = &test_uri; - let dataset = Dataset::open(test_uri).await.unwrap(); - - assert_eq!(dataset.manifest.max_fragment_id, None); - assert_eq!(dataset.manifest.max_fragment_id(), Some(2)); - } - } - - #[rstest] - #[tokio::test] - async fn test_bfloat16_roundtrip( - #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] - data_storage_version: LanceFileVersion, - ) -> Result<()> { - let inner_field = Arc::new( - ArrowField::new("item", DataType::FixedSizeBinary(2), true).with_metadata( - [ - (ARROW_EXT_NAME_KEY.into(), BFLOAT16_EXT_NAME.into()), - (ARROW_EXT_META_KEY.into(), "".into()), - ] - .into(), - ), - ); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "fsl", - DataType::FixedSizeList(inner_field.clone(), 2), - false, - )])); - - let values = bfloat16::BFloat16Array::from_iter_values( - (0..6).map(|i| i as f32).map(half::bf16::from_f32), - ); - let vectors = FixedSizeListArray::new(inner_field, 2, Arc::new(values.into_inner()), None); - - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vectors)]).unwrap(); - - let test_uri = TempStrDir::default(); - - let dataset = Dataset::write( - RecordBatchIterator::new(vec![Ok(batch.clone())], schema.clone()), - &test_uri, - Some(WriteParams { - data_storage_version: Some(data_storage_version), - ..Default::default() - }), - ) - .await?; - - let data = dataset.scan().try_into_batch().await?; - assert_eq!(batch, data); - - Ok(()) - } - - #[tokio::test] - async fn test_overwrite_mixed_version() { - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "a", - DataType::Int32, - false, - )])); - let arr = Arc::new(Int32Array::from(vec![1, 2, 3])); - - let data = RecordBatch::try_new(schema.clone(), vec![arr]).unwrap(); - let reader = - RecordBatchIterator::new(vec![data.clone()].into_iter().map(Ok), schema.clone()); - - let dataset = Dataset::write( - reader, - &test_uri, - Some(WriteParams { - data_storage_version: Some(LanceFileVersion::Legacy), - ..Default::default() - }), - ) - .await - .unwrap(); - - assert_eq!( - dataset - .manifest - .data_storage_format - .lance_file_version() - .unwrap(), - LanceFileVersion::Legacy - ); - - let reader = RecordBatchIterator::new(vec![data].into_iter().map(Ok), schema); - let dataset = Dataset::write( - reader, - &test_uri, - Some(WriteParams { - mode: WriteMode::Overwrite, - ..Default::default() - }), - ) - .await - .unwrap(); - - assert_eq!( - dataset - .manifest - .data_storage_format - .lance_file_version() - .unwrap(), - LanceFileVersion::Legacy - ); - } - - // Bug: https://github.com/lancedb/lancedb/issues/1223 - #[tokio::test] - async fn test_open_nonexisting_dataset() { - let temp_dir = TempStdDir::default(); - let dataset_dir = temp_dir.join("non_existing"); - let dataset_uri = dataset_dir.to_str().unwrap(); - - let res = Dataset::open(dataset_uri).await; - assert!(res.is_err()); - - assert!(!dataset_dir.exists()); - } - - #[tokio::test] - async fn test_manifest_partially_fits() { - // This regresses a bug that occurred when the manifest file was over 4KiB but the manifest - // itself was less than 4KiB (due to a dictionary). 4KiB is important here because that's the - // block size we use when reading the "last block" - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "x", - DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), - false, - )])); - let dictionary = Arc::new(StringArray::from_iter_values( - (0..1000).map(|i| i.to_string()), - )); - let indices = Int16Array::from_iter_values(0..1000); - let batches = vec![RecordBatch::try_new( - schema.clone(), - vec![Arc::new( - Int16DictionaryArray::try_new(indices, dictionary.clone()).unwrap(), - )], - ) - .unwrap()]; - - let test_uri = TempStrDir::default(); - let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - Dataset::write(batches, &test_uri, None).await.unwrap(); - - let dataset = Dataset::open(&test_uri).await.unwrap(); - assert_eq!(1000, dataset.count_rows(None).await.unwrap()); - } - - #[tokio::test] - async fn test_dataset_uri_roundtrips() { - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "a", - DataType::Int32, - false, - )])); - - let test_uri = TempStrDir::default(); - let vectors = Arc::new(Int32Array::from_iter_values(vec![])); - - let data = RecordBatch::try_new(schema.clone(), vec![vectors]); - let reader = RecordBatchIterator::new(vec![data.unwrap()].into_iter().map(Ok), schema); - let dataset = Dataset::write( - reader, - &test_uri, - Some(WriteParams { - ..Default::default() - }), - ) - .await - .unwrap(); - - let uri = dataset.uri(); - assert_eq!(uri, test_uri.as_str()); - - let ds2 = Dataset::open(uri).await.unwrap(); - assert_eq!( - ds2.latest_version_id().await.unwrap(), - dataset.latest_version_id().await.unwrap() - ); - } - - #[tokio::test] - async fn test_fts_fuzzy_query() { - let params = InvertedIndexParams::default(); - let text_col = GenericStringArray::::from(vec![ - "fa", "fo", "fob", "focus", "foo", "food", "foul", // # spellchecker:disable-line - ]); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![arrow_schema::Field::new( - "text", - text_col.data_type().to_owned(), - false, - )]) - .into(), - vec![Arc::new(text_col) as ArrayRef], - ) - .unwrap(); - let schema = batch.schema(); - let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let test_uri = TempStrDir::default(); - let mut dataset = Dataset::write(batches, &test_uri, None).await.unwrap(); - dataset - .create_index(&["text"], IndexType::Inverted, None, ¶ms, true) - .await - .unwrap(); - let results = dataset - .scan() - .full_text_search(FullTextSearchQuery::new_fuzzy("foo".to_owned(), Some(1))) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 4); - let texts = results["text"] - .as_string::() - .iter() - .map(|s| s.unwrap().to_owned()) - .collect::>(); - assert_eq!( - texts, - vec![ - "foo".to_owned(), // 0 edits - "fo".to_owned(), // 1 deletion # spellchecker:disable-line - "fob".to_owned(), // 1 substitution # spellchecker:disable-line - "food".to_owned(), // 1 insertion # spellchecker:disable-line - ] - .into_iter() - .collect() - ); - } - - #[tokio::test] - async fn test_fts_on_multiple_columns() { - let params = InvertedIndexParams::default(); - let title_col = - GenericStringArray::::from(vec!["title common", "title hello", "title lance"]); - let content_col = GenericStringArray::::from(vec![ - "content world", - "content database", - "content common", - ]); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - arrow_schema::Field::new("title", title_col.data_type().to_owned(), false), - arrow_schema::Field::new("content", title_col.data_type().to_owned(), false), - ]) - .into(), - vec![ - Arc::new(title_col) as ArrayRef, - Arc::new(content_col) as ArrayRef, - ], - ) - .unwrap(); - let schema = batch.schema(); - let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let test_uri = TempStrDir::default(); - let mut dataset = Dataset::write(batches, &test_uri, None).await.unwrap(); - dataset - .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) - .await - .unwrap(); - dataset - .create_index(&["content"], IndexType::Inverted, None, ¶ms, true) - .await - .unwrap(); - - let results = dataset - .scan() - .full_text_search(FullTextSearchQuery::new("title".to_owned())) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 3); - - let results = dataset - .scan() - .full_text_search(FullTextSearchQuery::new("content".to_owned())) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 3); - - let results = dataset - .scan() - .full_text_search(FullTextSearchQuery::new("common".to_owned())) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 2); - - let results = dataset - .scan() - .full_text_search( - FullTextSearchQuery::new("common".to_owned()) - .with_column("title".to_owned()) - .unwrap(), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 1); - - let results = dataset - .scan() - .full_text_search( - FullTextSearchQuery::new("common".to_owned()) - .with_column("content".to_owned()) - .unwrap(), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 1); - } - - #[tokio::test] - async fn test_fts_unindexed_data() { - let params = InvertedIndexParams::default(); - let title_col = StringArray::from(vec!["title hello", "title lance", "title common"]); - let content_col = - StringArray::from(vec!["content world", "content database", "content common"]); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - Field::new("title", title_col.data_type().to_owned(), false), - Field::new("content", title_col.data_type().to_owned(), false), - ]) - .into(), - vec![ - Arc::new(title_col) as ArrayRef, - Arc::new(content_col) as ArrayRef, - ], - ) - .unwrap(); - let schema = batch.schema(); - let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let mut dataset = Dataset::write(batches, "memory://test.lance", None) - .await - .unwrap(); - dataset - .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) - .await - .unwrap(); - - let results = dataset - .scan() - .full_text_search(FullTextSearchQuery::new("title".to_owned())) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 3); - - // write new data - let title_col = StringArray::from(vec!["new title"]); - let content_col = StringArray::from(vec!["new content"]); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - Field::new("title", title_col.data_type().to_owned(), false), - Field::new("content", title_col.data_type().to_owned(), false), - ]) - .into(), - vec![ - Arc::new(title_col) as ArrayRef, - Arc::new(content_col) as ArrayRef, - ], - ) - .unwrap(); - let schema = batch.schema(); - let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - dataset.append(batches, None).await.unwrap(); - - let results = dataset - .scan() - .full_text_search(FullTextSearchQuery::new("title".to_owned())) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 4); - - let results = dataset - .scan() - .full_text_search(FullTextSearchQuery::new("new".to_owned())) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 1); - } - - #[tokio::test] - async fn test_fts_unindexed_data_on_empty_index() { - // Empty dataset with fts index - let params = InvertedIndexParams::default(); - let title_col = StringArray::from(Vec::<&str>::new()); - let content_col = StringArray::from(Vec::<&str>::new()); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - Field::new("title", title_col.data_type().to_owned(), false), - Field::new("content", title_col.data_type().to_owned(), false), - ]) - .into(), - vec![ - Arc::new(title_col) as ArrayRef, - Arc::new(content_col) as ArrayRef, - ], - ) - .unwrap(); - let schema = batch.schema(); - let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let mut dataset = Dataset::write(batches, "memory://test.lance", None) - .await - .unwrap(); - dataset - .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) - .await - .unwrap(); - - // Test fts search - let results = dataset - .scan() - .full_text_search(FullTextSearchQuery::new_query(FtsQuery::Match( - MatchQuery::new("title".to_owned()).with_column(Some("title".to_owned())), - ))) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 0); - - // write new data - let title_col = StringArray::from(vec!["title hello", "title lance", "title common"]); - let content_col = - StringArray::from(vec!["content world", "content database", "content common"]); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - Field::new("title", title_col.data_type().to_owned(), false), - Field::new("content", title_col.data_type().to_owned(), false), - ]) - .into(), - vec![ - Arc::new(title_col) as ArrayRef, - Arc::new(content_col) as ArrayRef, - ], - ) - .unwrap(); - let schema = batch.schema(); - let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - dataset.append(batches, None).await.unwrap(); - - let results = dataset - .scan() - .full_text_search(FullTextSearchQuery::new_query(FtsQuery::Match( - MatchQuery::new("title".to_owned()).with_column(Some("title".to_owned())), - ))) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 3); - } - - #[tokio::test] - async fn test_fts_without_index() { - // create table without index - let title_col = StringArray::from(vec!["title hello", "title lance", "title common"]); - let content_col = - StringArray::from(vec!["content world", "content database", "content common"]); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - Field::new("title", title_col.data_type().to_owned(), false), - Field::new("content", title_col.data_type().to_owned(), false), - ]) - .into(), - vec![ - Arc::new(title_col) as ArrayRef, - Arc::new(content_col) as ArrayRef, - ], - ) - .unwrap(); - let schema = batch.schema(); - let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let mut dataset = Dataset::write(batches, "memory://test.lance", None) - .await - .unwrap(); - - // match query on title and content - let results = dataset - .scan() - .full_text_search( - FullTextSearchQuery::new("title".to_owned()) - .with_columns(&["title".to_string(), "content".to_string()]) - .unwrap(), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 3); - - // write new data - let title_col = StringArray::from(vec!["new title"]); - let content_col = StringArray::from(vec!["new content"]); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - Field::new("title", title_col.data_type().to_owned(), false), - Field::new("content", title_col.data_type().to_owned(), false), - ]) - .into(), - vec![ - Arc::new(title_col) as ArrayRef, - Arc::new(content_col) as ArrayRef, - ], - ) - .unwrap(); - let schema = batch.schema(); - let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - dataset.append(batches, None).await.unwrap(); - - // match query on title and content - let results = dataset - .scan() - .full_text_search( - FullTextSearchQuery::new("title".to_owned()) - .with_columns(&["title".to_string(), "content".to_string()]) - .unwrap(), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 4); - - let results = dataset - .scan() - .full_text_search( - FullTextSearchQuery::new("new".to_owned()) - .with_columns(&["title".to_string(), "content".to_string()]) - .unwrap(), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 1); - } - - #[tokio::test] - async fn test_fts_rank() { - let params = InvertedIndexParams::default(); - let text_col = - GenericStringArray::::from(vec!["score", "find score", "try to find score"]); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![arrow_schema::Field::new( - "text", - text_col.data_type().to_owned(), - false, - )]) - .into(), - vec![Arc::new(text_col) as ArrayRef], - ) - .unwrap(); - let schema = batch.schema(); - let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let test_uri = TempStrDir::default(); - let mut dataset = Dataset::write(batches, &test_uri, None).await.unwrap(); - dataset - .create_index(&["text"], IndexType::Inverted, None, ¶ms, true) - .await - .unwrap(); - - let results = dataset - .scan() - .with_row_id() - .full_text_search(FullTextSearchQuery::new("score".to_owned())) - .unwrap() - .limit(Some(3), None) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 3); - let row_ids = results[ROW_ID].as_primitive::().values(); - assert_eq!(row_ids, &[0, 1, 2]); - - let results = dataset - .scan() - .with_row_id() - .full_text_search(FullTextSearchQuery::new("score".to_owned())) - .unwrap() - .limit(Some(2), None) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 2); - let row_ids = results[ROW_ID].as_primitive::().values(); - assert_eq!(row_ids, &[0, 1]); - - let results = dataset - .scan() - .with_row_id() - .full_text_search(FullTextSearchQuery::new("score".to_owned())) - .unwrap() - .limit(Some(1), None) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(results.num_rows(), 1); - let row_ids = results[ROW_ID].as_primitive::().values(); - assert_eq!(row_ids, &[0]); - } - - async fn create_fts_dataset< - Offset: arrow::array::OffsetSizeTrait, - ListOffset: arrow::array::OffsetSizeTrait, - >( - is_list: bool, - with_position: bool, - params: InvertedIndexParams, - ) -> Dataset { - let tempdir = TempStrDir::default(); - let uri = tempdir.to_owned(); - drop(tempdir); - - let params = params.with_position(with_position); - let doc_col: Arc = if is_list { - let string_builder = GenericStringBuilder::::new(); - let mut list_col = GenericListBuilder::::new(string_builder); - // Create a list of strings - list_col.values().append_value("lance database the search"); // for testing phrase query - list_col.append(true); - list_col.values().append_value("lance database"); // for testing phrase query - list_col.append(true); - list_col.values().append_value("lance search"); - list_col.append(true); - list_col.values().append_value("database"); - list_col.values().append_value("search"); - list_col.append(true); - list_col.values().append_value("unrelated doc"); - list_col.append(true); - list_col.values().append_value("unrelated"); - list_col.append(true); - list_col.values().append_value("mots"); - list_col.values().append_value("accentués"); - list_col.append(true); - list_col - .values() - .append_value("lance database full text search"); - list_col.append(true); - - // for testing null - list_col.append(false); - - Arc::new(list_col.finish()) - } else { - Arc::new(GenericStringArray::::from(vec![ - "lance database the search", - "lance database", - "lance search", - "database search", - "unrelated doc", - "unrelated", - "mots accentués", - "lance database full text search", - ])) - }; - let ids = UInt64Array::from_iter_values(0..doc_col.len() as u64); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - arrow_schema::Field::new("doc", doc_col.data_type().to_owned(), true), - arrow_schema::Field::new("id", DataType::UInt64, false), - ]) - .into(), - vec![Arc::new(doc_col) as ArrayRef, Arc::new(ids) as ArrayRef], - ) - .unwrap(); - let schema = batch.schema(); - let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let mut dataset = Dataset::write(batches, &uri, None).await.unwrap(); - - dataset - .create_index(&["doc"], IndexType::Inverted, None, ¶ms, true) - .await - .unwrap(); - - dataset - } - - async fn test_fts_index< - Offset: arrow::array::OffsetSizeTrait, - ListOffset: arrow::array::OffsetSizeTrait, - >( - is_list: bool, - ) { - let ds = create_fts_dataset::( - is_list, - false, - InvertedIndexParams::default(), - ) - .await; - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new("lance".to_owned()).limit(Some(3))) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 3, "{:?}", result); - let ids = result["id"].as_primitive::().values(); - assert!(ids.contains(&0), "{:?}", result); - assert!(ids.contains(&1), "{:?}", result); - assert!(ids.contains(&2), "{:?}", result); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new("database".to_owned()).limit(Some(3))) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 3); - let ids = result["id"].as_primitive::().values(); - assert!(ids.contains(&0), "{:?}", result); - assert!(ids.contains(&1), "{:?}", result); - assert!(ids.contains(&3), "{:?}", result); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search( - FullTextSearchQuery::new_query( - MatchQuery::new("lance database".to_owned()) - .with_operator(Operator::And) - .into(), - ) - .limit(Some(5)), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 3, "{:?}", result); - let ids = result["id"].as_primitive::().values(); - assert!(ids.contains(&0), "{:?}", result); - assert!(ids.contains(&1), "{:?}", result); - assert!(ids.contains(&7), "{:?}", result); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new("unknown null".to_owned()).limit(Some(3))) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 0); - - // test phrase query - // for non-phrasal query, the order of the tokens doesn't matter - // so there should be 4 documents that contain "database" or "lance" - - // we built the index without position, so the phrase query will not work - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search( - FullTextSearchQuery::new_query( - PhraseQuery::new("lance database".to_owned()).into(), - ) - .limit(Some(10)), - ) - .unwrap() - .try_into_batch() - .await; - let err = result.unwrap_err().to_string(); - assert!(err.contains("position is not found but required for phrase queries, try recreating the index with position"),"{}",err); - - // recreate the index with position - let ds = - create_fts_dataset::(is_list, true, InvertedIndexParams::default()) - .await; - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new("lance database".to_owned()).limit(Some(10))) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 5, "{:?}", result); - let ids = result["id"].as_primitive::().values(); - assert!(ids.contains(&0)); - assert!(ids.contains(&1)); - assert!(ids.contains(&2)); - assert!(ids.contains(&3)); - assert!(ids.contains(&7)); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search( - FullTextSearchQuery::new_query( - PhraseQuery::new("lance database".to_owned()).into(), - ) - .limit(Some(10)), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - let ids = result["id"].as_primitive::().values(); - assert_eq!(result.num_rows(), 3, "{:?}", ids); - assert!(ids.contains(&0)); - assert!(ids.contains(&1)); - assert!(ids.contains(&7)); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search( - FullTextSearchQuery::new_query( - PhraseQuery::new("database lance".to_owned()).into(), - ) - .limit(Some(10)), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 0); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search( - FullTextSearchQuery::new_query(PhraseQuery::new("lance unknown".to_owned()).into()) - .limit(Some(10)), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 0); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search( - FullTextSearchQuery::new_query(PhraseQuery::new("unknown null".to_owned()).into()) - .limit(Some(3)), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 0); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search( - FullTextSearchQuery::new_query(PhraseQuery::new("lance search".to_owned()).into()) - .limit(Some(3)), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 1); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search( - FullTextSearchQuery::new_query( - PhraseQuery::new("lance search".to_owned()) - .with_slop(2) - .into(), - ) - .limit(Some(3)), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 2); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search( - FullTextSearchQuery::new_query( - PhraseQuery::new("search lance".to_owned()) - .with_slop(2) - .into(), - ) - .limit(Some(3)), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 0); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search( - // must contain "lance" and "database", and may contain "search" - FullTextSearchQuery::new_query( - BooleanQuery::new([ - ( - Occur::Should, - MatchQuery::new("search".to_owned()) - .with_operator(Operator::And) - .into(), - ), - ( - Occur::Must, - MatchQuery::new("lance database".to_owned()) - .with_operator(Operator::And) - .into(), - ), - ]) - .into(), - ) - .limit(Some(3)), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 3, "{:?}", result); - let ids = result["id"].as_primitive::().values(); - assert!(ids.contains(&0), "{:?}", result); - assert!(ids.contains(&1), "{:?}", result); - assert!(ids.contains(&7), "{:?}", result); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search( - // must contain "lance" and "database", and may contain "search" - FullTextSearchQuery::new_query( - BooleanQuery::new([ - ( - Occur::Should, - MatchQuery::new("search".to_owned()) - .with_operator(Operator::And) - .into(), - ), - ( - Occur::Must, - MatchQuery::new("lance database".to_owned()) - .with_operator(Operator::And) - .into(), - ), - ( - Occur::MustNot, - MatchQuery::new("full text".to_owned()).into(), - ), - ]) - .into(), - ) - .limit(Some(3)), - ) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 2, "{:?}", result); - let ids = result["id"].as_primitive::().values(); - assert!(ids.contains(&0), "{:?}", result); - assert!(ids.contains(&1), "{:?}", result); - } - - #[tokio::test] - async fn test_fts_index_with_string() { - test_fts_index::(false).await; - test_fts_index::(true).await; - test_fts_index::(true).await; - } - - #[tokio::test] - async fn test_fts_index_with_large_string() { - test_fts_index::(false).await; - test_fts_index::(true).await; - test_fts_index::(true).await; - } - - #[tokio::test] - async fn test_fts_accented_chars() { - let ds = create_fts_dataset::(false, false, InvertedIndexParams::default()).await; - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new("accentués".to_owned()).limit(Some(3))) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 1); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new("accentues".to_owned()).limit(Some(3))) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 0); - - // with ascii folding enabled, the search should be accent-insensitive - let ds = create_fts_dataset::( - false, - false, - InvertedIndexParams::default() - .stem(false) - .ascii_folding(true), - ) - .await; - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new("accentués".to_owned()).limit(Some(3))) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 1); - - let result = ds - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new("accentues".to_owned()).limit(Some(3))) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 1); - } - - #[tokio::test] - async fn test_fts_phrase_query() { - let tmpdir = TempStrDir::default(); - let uri = tmpdir.to_owned(); - drop(tmpdir); - - let words = ["lance", "full", "text", "search"]; - let mut lance_search_count = 0; - let mut full_text_count = 0; - let mut doc_array = (0..4096) - .map(|_| { - let mut rng = rand::rng(); - let mut text = String::with_capacity(512); - let len = rng.random_range(127..512); - for i in 0..len { - if i > 0 { - text.push(' '); - } - text.push_str(words[rng.random_range(0..words.len())]); - } - if text.contains("lance search") { - lance_search_count += 1; - } - if text.contains("full text") { - full_text_count += 1; - } - text - }) - .collect_vec(); - // Ensure at least one doc matches each phrase deterministically - doc_array.push("lance search".to_owned()); - lance_search_count += 1; - doc_array.push("full text".to_owned()); - full_text_count += 1; - doc_array.push("position for phrase query".to_owned()); - - // 1) Build index without positions and assert phrase query errors - let params_no_pos = InvertedIndexParams::default().with_position(false); - let doc_col: Arc = Arc::new(GenericStringArray::::from(doc_array.clone())); - let ids = UInt64Array::from_iter_values(0..doc_col.len() as u64); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - arrow_schema::Field::new("doc", doc_col.data_type().to_owned(), true), - arrow_schema::Field::new("id", DataType::UInt64, false), - ]) - .into(), - vec![Arc::new(doc_col) as ArrayRef, Arc::new(ids) as ArrayRef], - ) - .unwrap(); - let schema = batch.schema(); - let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let mut dataset = Dataset::write(batches, &uri, None).await.unwrap(); - dataset - .create_index(&["doc"], IndexType::Inverted, None, ¶ms_no_pos, true) - .await - .unwrap(); - - let err = dataset - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new_query( - PhraseQuery::new("lance search".to_owned()).into(), - )) - .unwrap() - .try_into_batch() - .await - .unwrap_err() - .to_string(); - assert!(err.contains("position is not found but required for phrase queries, try recreating the index with position"), "{}", err); - assert!(err.starts_with("Invalid user input: "), "{}", err); - - // 2) Recreate index with positions and assert phrase query works - let params_with_pos = InvertedIndexParams::default().with_position(true); - dataset - .create_index(&["doc"], IndexType::Inverted, None, ¶ms_with_pos, true) - .await - .unwrap(); - - let result = dataset - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new_query( - PhraseQuery::new("lance search".to_owned()).into(), - )) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), lance_search_count); - - let result = dataset - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new_query( - PhraseQuery::new("full text".to_owned()).into(), - )) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), full_text_count); - - let result = dataset - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new_query( - PhraseQuery::new("phrase query".to_owned()).into(), - )) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 1); - - let result = dataset - .scan() - .project(&["id"]) - .unwrap() - .full_text_search(FullTextSearchQuery::new_query( - PhraseQuery::new("".to_owned()).into(), - )) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(result.num_rows(), 0); - } - - #[tokio::test] - async fn concurrent_create() { - async fn write(uri: &str) -> Result<()> { - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "a", - DataType::Int32, - false, - )])); - let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); - Dataset::write(empty_reader, uri, None).await?; - Ok(()) - } - - for _ in 0..5 { - let test_uri = TempStrDir::default(); - - let (res1, res2) = tokio::join!(write(&test_uri), write(&test_uri)); - - assert!(res1.is_ok() || res2.is_ok()); - if res1.is_err() { - assert!( - matches!(res1, Err(Error::DatasetAlreadyExists { .. })), - "{:?}", - res1 - ); - } else if res2.is_err() { - assert!( - matches!(res2, Err(Error::DatasetAlreadyExists { .. })), - "{:?}", - res2 - ); - } else { - assert!(res1.is_ok() && res2.is_ok()); - } - } - } - - #[tokio::test] - async fn test_read_transaction_properties() { - const LANCE_COMMIT_MESSAGE_KEY: &str = "__lance_commit_message"; - // Create a test dataset - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int32, false), - ArrowField::new("value", DataType::Utf8, false), - ])); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![1, 2, 3])), - Arc::new(StringArray::from(vec!["a", "b", "c"])), - ], - ) - .unwrap(); - - let test_uri = TempStrDir::default(); - - // Create WriteParams with properties - let mut properties1 = HashMap::new(); - properties1.insert( - LANCE_COMMIT_MESSAGE_KEY.to_string(), - "First commit".to_string(), - ); - properties1.insert("custom_prop".to_string(), "custom_value".to_string()); - - let write_params = WriteParams { - transaction_properties: Some(Arc::new(properties1)), - ..Default::default() - }; - - let dataset = Dataset::write( - RecordBatchIterator::new([Ok(batch.clone())], schema.clone()), - &test_uri, - Some(write_params), - ) - .await - .unwrap(); - - let transaction = dataset.read_transaction_by_version(1).await.unwrap(); - assert!(transaction.is_some()); - let props = transaction.unwrap().transaction_properties.unwrap(); - assert_eq!(props.len(), 2); - assert_eq!( - props.get(LANCE_COMMIT_MESSAGE_KEY), - Some(&"First commit".to_string()) - ); - assert_eq!(props.get("custom_prop"), Some(&"custom_value".to_string())); - - let mut properties2 = HashMap::new(); - properties2.insert( - LANCE_COMMIT_MESSAGE_KEY.to_string(), - "Second commit".to_string(), - ); - properties2.insert("another_prop".to_string(), "another_value".to_string()); - - let write_params = WriteParams { - transaction_properties: Some(Arc::new(properties2)), - mode: WriteMode::Append, - ..Default::default() - }; - - let batch2 = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![4, 5])), - Arc::new(StringArray::from(vec!["d", "e"])), - ], - ) - .unwrap(); - - let mut dataset = dataset; - dataset - .append( - RecordBatchIterator::new([Ok(batch2)], schema.clone()), - Some(write_params), - ) - .await - .unwrap(); - - let transaction = dataset.read_transaction_by_version(2).await.unwrap(); - assert!(transaction.is_some()); - let props = transaction.unwrap().transaction_properties.unwrap(); - assert_eq!(props.len(), 2); - assert_eq!( - props.get(LANCE_COMMIT_MESSAGE_KEY), - Some(&"Second commit".to_string()) - ); - assert_eq!( - props.get("another_prop"), - Some(&"another_value".to_string()) - ); - - let transaction = dataset.read_transaction_by_version(1).await.unwrap(); - assert!(transaction.is_some()); - let props = transaction.unwrap().transaction_properties.unwrap(); - assert_eq!(props.len(), 2); - assert_eq!( - props.get(LANCE_COMMIT_MESSAGE_KEY), - Some(&"First commit".to_string()) - ); - assert_eq!(props.get("custom_prop"), Some(&"custom_value".to_string())); - - let result = dataset.read_transaction_by_version(999).await; - assert!(result.is_err()); - } - - #[tokio::test] - async fn test_insert_subschema() { - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("a", DataType::Int32, false), - ArrowField::new("b", DataType::Int32, true), - ])); - let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); - let mut dataset = Dataset::write(empty_reader, "memory://", None) - .await - .unwrap(); - dataset.validate().await.unwrap(); - - // If missing columns that aren't nullable, will return an error - // TODO: provide alternative default than null. - let just_b = Arc::new(schema.project(&[1]).unwrap()); - let batch = RecordBatch::try_new(just_b.clone(), vec![Arc::new(Int32Array::from(vec![1]))]) - .unwrap(); - let reader = RecordBatchIterator::new(vec![Ok(batch)], just_b.clone()); - let res = dataset.append(reader, None).await; - assert!( - matches!(res, Err(Error::SchemaMismatch { .. })), - "Expected Error::SchemaMismatch, got {:?}", - res - ); - - // If missing columns that are nullable, the write succeeds. - let just_a = Arc::new(schema.project(&[0]).unwrap()); - let batch = RecordBatch::try_new(just_a.clone(), vec![Arc::new(Int32Array::from(vec![1]))]) - .unwrap(); - let reader = RecordBatchIterator::new(vec![Ok(batch)], just_a.clone()); - dataset.append(reader, None).await.unwrap(); - dataset.validate().await.unwrap(); - assert_eq!(dataset.count_rows(None).await.unwrap(), 1); - - // Looking at the fragments, there is no data file with the missing field - let fragments = dataset.get_fragments(); - assert_eq!(fragments.len(), 1); - assert_eq!(fragments[0].metadata.files.len(), 1); - assert_eq!(&fragments[0].metadata.files[0].fields, &[0]); - - // When reading back, columns that are missing are null - let data = dataset.scan().try_into_batch().await.unwrap(); - let expected = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![1])), - Arc::new(Int32Array::from(vec![None])), - ], - ) - .unwrap(); - assert_eq!(data, expected); - - // Can still insert all columns - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![2])), - Arc::new(Int32Array::from(vec![3])), - ], - ) - .unwrap(); - let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], schema.clone()); - dataset.append(reader, None).await.unwrap(); - dataset.validate().await.unwrap(); - assert_eq!(dataset.count_rows(None).await.unwrap(), 2); - - // When reading back, only missing data is null, otherwise is filled in - let data = dataset.scan().try_into_batch().await.unwrap(); - let expected = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![1, 2])), - Arc::new(Int32Array::from(vec![None, Some(3)])), - ], - ) - .unwrap(); - assert_eq!(data, expected); - - // Can run compaction. All files should now have all fields. - compact_files(&mut dataset, CompactionOptions::default(), None) - .await - .unwrap(); - dataset.validate().await.unwrap(); - let fragments = dataset.get_fragments(); - assert_eq!(fragments.len(), 1); - assert_eq!(fragments[0].metadata.files.len(), 1); - assert_eq!(&fragments[0].metadata.files[0].fields, &[0, 1]); - - // Can scan and get expected data. - let data = dataset.scan().try_into_batch().await.unwrap(); - assert_eq!(data, expected); - } - - #[tokio::test] - async fn test_insert_nested_subschemas() { - // Test subschemas at struct level - // Test different orders - // Test the Dataset::write() path - // Test Take across fragments with different field id sets - let test_uri = TempStrDir::default(); - - let field_a = Arc::new(ArrowField::new("a", DataType::Int32, true)); - let field_b = Arc::new(ArrowField::new("b", DataType::Int32, false)); - let field_c = Arc::new(ArrowField::new("c", DataType::Int32, true)); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "s", - DataType::Struct(vec![field_a.clone(), field_b.clone(), field_c.clone()].into()), - true, - )])); - let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); - let dataset = Dataset::write(empty_reader, &test_uri, None).await.unwrap(); - dataset.validate().await.unwrap(); - - let append_options = WriteParams { - mode: WriteMode::Append, - ..Default::default() - }; - // Can insert b, a - let just_b_a = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "s", - DataType::Struct(vec![field_b.clone(), field_a.clone()].into()), - true, - )])); - let batch = RecordBatch::try_new( - just_b_a.clone(), - vec![Arc::new(StructArray::from(vec![ - ( - field_b.clone(), - Arc::new(Int32Array::from(vec![1])) as ArrayRef, - ), - (field_a.clone(), Arc::new(Int32Array::from(vec![2]))), - ]))], - ) - .unwrap(); - let reader = RecordBatchIterator::new(vec![Ok(batch)], just_b_a.clone()); - let dataset = Dataset::write(reader, &test_uri, Some(append_options.clone())) - .await - .unwrap(); - dataset.validate().await.unwrap(); - let fragments = dataset.get_fragments(); - assert_eq!(fragments.len(), 1); - assert_eq!(fragments[0].metadata.files.len(), 1); - assert_eq!(&fragments[0].metadata.files[0].fields, &[0, 2, 1]); - assert_eq!(&fragments[0].metadata.files[0].column_indices, &[0, 1, 2]); - - // Can insert c, b - let just_c_b = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "s", - DataType::Struct(vec![field_c.clone(), field_b.clone()].into()), - true, - )])); - let batch = RecordBatch::try_new( - just_c_b.clone(), - vec![Arc::new(StructArray::from(vec![ - ( - field_c.clone(), - Arc::new(Int32Array::from(vec![4])) as ArrayRef, - ), - (field_b.clone(), Arc::new(Int32Array::from(vec![3]))), - ]))], - ) - .unwrap(); - let reader = RecordBatchIterator::new(vec![Ok(batch)], just_c_b.clone()); - let dataset = Dataset::write(reader, &test_uri, Some(append_options.clone())) - .await - .unwrap(); - dataset.validate().await.unwrap(); - let fragments = dataset.get_fragments(); - assert_eq!(fragments.len(), 2); - assert_eq!(fragments[1].metadata.files.len(), 1); - assert_eq!(&fragments[1].metadata.files[0].fields, &[0, 3, 2]); - assert_eq!(&fragments[1].metadata.files[0].column_indices, &[0, 1, 2]); - - // Can't insert a, c (b is non-nullable) - let just_a_c = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "s", - DataType::Struct(vec![field_a.clone(), field_c.clone()].into()), - true, - )])); - let batch = RecordBatch::try_new( - just_a_c.clone(), - vec![Arc::new(StructArray::from(vec![ - ( - field_a.clone(), - Arc::new(Int32Array::from(vec![5])) as ArrayRef, - ), - (field_c.clone(), Arc::new(Int32Array::from(vec![6]))), - ]))], - ) - .unwrap(); - let reader = RecordBatchIterator::new(vec![Ok(batch)], just_a_c.clone()); - let res = Dataset::write(reader, &test_uri, Some(append_options)).await; - assert!( - matches!(res, Err(Error::SchemaMismatch { .. })), - "Expected Error::SchemaMismatch, got {:?}", - res - ); - - // Can scan and get all data - let data = dataset.scan().try_into_batch().await.unwrap(); - let expected = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(StructArray::from(vec![ - ( - field_a.clone(), - Arc::new(Int32Array::from(vec![Some(2), None])) as ArrayRef, - ), - (field_b.clone(), Arc::new(Int32Array::from(vec![1, 3]))), - ( - field_c.clone(), - Arc::new(Int32Array::from(vec![None, Some(4)])), - ), - ]))], - ) - .unwrap(); - assert_eq!(data, expected); - - // Can call take and get rows from all three back in one batch - let result = dataset - .take(&[1, 0], Arc::new(dataset.schema().clone())) - .await - .unwrap(); - let expected = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(StructArray::from(vec![ - ( - field_a.clone(), - Arc::new(Int32Array::from(vec![None, Some(2)])) as ArrayRef, - ), - (field_b.clone(), Arc::new(Int32Array::from(vec![3, 1]))), - ( - field_c.clone(), - Arc::new(Int32Array::from(vec![Some(4), None])), - ), - ]))], - ) - .unwrap(); - assert_eq!(result, expected); - } - - #[tokio::test] - async fn test_insert_balanced_subschemas() { - let test_uri = TempStrDir::default(); - - let field_a = ArrowField::new("a", DataType::Int32, true); - let field_b = ArrowField::new("b", DataType::LargeBinary, true); - let schema = Arc::new(ArrowSchema::new(vec![ - field_a.clone(), - field_b - .clone() - .with_metadata([(BLOB_META_KEY.to_string(), "true".to_string())].into()), - ])); - let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); - let options = WriteParams { - enable_stable_row_ids: true, - enable_v2_manifest_paths: true, - ..Default::default() - }; - let mut dataset = Dataset::write(empty_reader, &test_uri, Some(options)) - .await - .unwrap(); - dataset.validate().await.unwrap(); - - // Insert left side - let just_a = Arc::new(ArrowSchema::new(vec![field_a.clone()])); - let batch = RecordBatch::try_new(just_a.clone(), vec![Arc::new(Int32Array::from(vec![1]))]) - .unwrap(); - let reader = RecordBatchIterator::new(vec![Ok(batch)], just_a.clone()); - dataset.append(reader, None).await.unwrap(); - dataset.validate().await.unwrap(); - - let fragments = dataset.get_fragments(); - assert_eq!(fragments.len(), 1); - assert_eq!(fragments[0].metadata.files.len(), 1); - assert_eq!(&fragments[0].metadata.files[0].fields, &[0]); - - // Insert right side - let just_b = Arc::new(ArrowSchema::new(vec![field_b.clone()])); - let batch = RecordBatch::try_new( - just_b.clone(), - vec![Arc::new(LargeBinaryArray::from_iter(vec![Some(vec![2u8])]))], - ) - .unwrap(); - let reader = RecordBatchIterator::new(vec![Ok(batch)], just_b.clone()); - dataset.append(reader, None).await.unwrap(); - dataset.validate().await.unwrap(); - - let fragments = dataset.get_fragments(); - assert_eq!(fragments.len(), 2); - assert_eq!(fragments[1].metadata.files.len(), 1); - assert_eq!(&fragments[1].metadata.files[0].fields, &[1]); - - let data = dataset - .take( - &[0, 1], - ProjectionRequest::from_columns(["a"], dataset.schema()), - ) - .await - .unwrap(); - assert_eq!(data.num_rows(), 2); - let a_column = data.column(0).as_primitive::(); - assert_eq!(a_column.value(0), 1); - assert!(a_column.is_null(1)); - - let blob_batch = dataset - .take( - &[0, 1], - ProjectionRequest::from_columns(["b"], dataset.schema()), - ) - .await - .unwrap(); - let blob_descriptions = blob_batch.column(0).as_struct(); - assert!(blob_descriptions.is_null(0)); - assert!(blob_descriptions.is_valid(1)); - } - - #[tokio::test] - async fn test_datafile_replacement() { - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "a", - DataType::Int32, - true, - )])); - let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); - let dataset = Arc::new( - Dataset::write(empty_reader, "memory://", None) - .await - .unwrap(), - ); - dataset.validate().await.unwrap(); - - // Test empty replacement should commit a new manifest and do nothing - let mut dataset = Dataset::commit( - WriteDestination::Dataset(dataset.clone()), - Operation::DataReplacement { - replacements: vec![], - }, - Some(1), - None, - None, - Arc::new(Default::default()), - false, - ) - .await - .unwrap(); - dataset.validate().await.unwrap(); - - assert_eq!(dataset.version().version, 2); - assert_eq!(dataset.get_fragments().len(), 0); - - // try the same thing on a non-empty dataset - let vals: Int32Array = vec![1, 2, 3].into(); - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vals)]).unwrap(); - dataset - .append( - RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), - None, - ) - .await - .unwrap(); - - let dataset = Dataset::commit( - WriteDestination::Dataset(Arc::new(dataset)), - Operation::DataReplacement { - replacements: vec![], - }, - Some(3), - None, - None, - Arc::new(Default::default()), - false, - ) - .await - .unwrap(); - dataset.validate().await.unwrap(); - - assert_eq!(dataset.version().version, 4); - assert_eq!(dataset.get_fragments().len(), 1); - - let batch = dataset.scan().try_into_batch().await.unwrap(); - assert_eq!(batch.num_rows(), 3); - assert_eq!( - batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values(), - &[1, 2, 3] - ); - - // write a new datafile - let object_writer = dataset - .object_store - .create(&Path::from("data/test.lance")) - .await - .unwrap(); - let mut writer = FileWriter::try_new( - object_writer, - schema.as_ref().try_into().unwrap(), - Default::default(), - ) - .unwrap(); - - let vals: Int32Array = vec![4, 5, 6].into(); - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vals)]).unwrap(); - writer.write_batch(&batch).await.unwrap(); - writer.finish().await.unwrap(); - - // find the datafile we want to replace - let frag = dataset.get_fragment(0).unwrap(); - let data_file = frag.data_file_for_field(0).unwrap(); - let mut new_data_file = data_file.clone(); - new_data_file.path = "test.lance".to_string(); - - let dataset = Dataset::commit( - WriteDestination::Dataset(Arc::new(dataset)), - Operation::DataReplacement { - replacements: vec![DataReplacementGroup(0, new_data_file)], - }, - Some(4), - None, - None, - Arc::new(Default::default()), - false, - ) - .await - .unwrap(); - - assert_eq!(dataset.version().version, 5); - assert_eq!(dataset.get_fragments().len(), 1); - assert_eq!(dataset.get_fragments()[0].metadata.files.len(), 1); - - let batch = dataset.scan().try_into_batch().await.unwrap(); - assert_eq!(batch.num_rows(), 3); - assert_eq!( - batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values(), - &[4, 5, 6] - ); - } - - #[tokio::test] - async fn test_datafile_partial_replacement() { - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "a", - DataType::Int32, - true, - )])); - let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); - let mut dataset = Dataset::write(empty_reader, "memory://", None) - .await - .unwrap(); - dataset.validate().await.unwrap(); - - let vals: Int32Array = vec![1, 2, 3].into(); - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vals)]).unwrap(); - dataset - .append( - RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), - None, - ) - .await - .unwrap(); - - let fragment = dataset.get_fragments().pop().unwrap().metadata; - - let extended_schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("a", DataType::Int32, true), - ArrowField::new("b", DataType::Int32, true), - ])); - - // add all null column - let dataset = Dataset::commit( - WriteDestination::Dataset(Arc::new(dataset)), - Operation::Merge { - fragments: vec![fragment], - schema: extended_schema.as_ref().try_into().unwrap(), - }, - Some(2), - None, - None, - Arc::new(Default::default()), - false, - ) - .await - .unwrap(); - - let partial_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "b", - DataType::Int32, - true, - )])); - - // write a new datafile - let object_writer = dataset - .object_store - .create(&Path::from("data/test.lance")) - .await - .unwrap(); - let mut writer = FileWriter::try_new( - object_writer, - partial_schema.as_ref().try_into().unwrap(), - Default::default(), - ) - .unwrap(); - - let vals: Int32Array = vec![4, 5, 6].into(); - let batch = RecordBatch::try_new(partial_schema.clone(), vec![Arc::new(vals)]).unwrap(); - writer.write_batch(&batch).await.unwrap(); - writer.finish().await.unwrap(); - - let (major, minor) = lance_file::version::LanceFileVersion::Stable.to_numbers(); - - // find the datafile we want to replace - let new_data_file = DataFile { - path: "test.lance".to_string(), - // the second column in the dataset - fields: vec![1], - // is located in the first column of this datafile - column_indices: vec![0], - file_major_version: major, - file_minor_version: minor, - file_size_bytes: CachedFileSize::unknown(), - base_id: None, - }; - - let dataset = Dataset::commit( - WriteDestination::Dataset(Arc::new(dataset)), - Operation::DataReplacement { - replacements: vec![DataReplacementGroup(0, new_data_file)], - }, - Some(3), - None, - None, - Arc::new(Default::default()), - false, - ) - .await - .unwrap(); - - assert_eq!(dataset.version().version, 4); - assert_eq!(dataset.get_fragments().len(), 1); - assert_eq!(dataset.get_fragments()[0].metadata.files.len(), 2); - assert_eq!(dataset.get_fragments()[0].metadata.files[0].fields, vec![0]); - assert_eq!(dataset.get_fragments()[0].metadata.files[1].fields, vec![1]); - - let batch = dataset.scan().try_into_batch().await.unwrap(); - assert_eq!(batch.num_rows(), 3); - assert_eq!( - batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values(), - &[1, 2, 3] - ); - assert_eq!( - batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap() - .values(), - &[4, 5, 6] - ); - - // do it again but on the first column - // find the datafile we want to replace - let new_data_file = DataFile { - path: "test.lance".to_string(), - // the first column in the dataset - fields: vec![0], - // is located in the first column of this datafile - column_indices: vec![0], - file_major_version: major, - file_minor_version: minor, - file_size_bytes: CachedFileSize::unknown(), - base_id: None, - }; - - let dataset = Dataset::commit( - WriteDestination::Dataset(Arc::new(dataset)), - Operation::DataReplacement { - replacements: vec![DataReplacementGroup(0, new_data_file)], - }, - Some(4), - None, - None, - Arc::new(Default::default()), - false, - ) - .await - .unwrap(); - - assert_eq!(dataset.version().version, 5); - assert_eq!(dataset.get_fragments().len(), 1); - assert_eq!(dataset.get_fragments()[0].metadata.files.len(), 2); - - let batch = dataset.scan().try_into_batch().await.unwrap(); - assert_eq!(batch.num_rows(), 3); - assert_eq!( - batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values(), - &[4, 5, 6] - ); - assert_eq!( - batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap() - .values(), - &[4, 5, 6] - ); - } - - #[tokio::test] - async fn test_datafile_replacement_error() { - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "a", - DataType::Int32, - true, - )])); - let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); - let mut dataset = Dataset::write(empty_reader, "memory://", None) - .await - .unwrap(); - dataset.validate().await.unwrap(); - - let vals: Int32Array = vec![1, 2, 3].into(); - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vals)]).unwrap(); - dataset - .append( - RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), - None, - ) - .await - .unwrap(); - - let fragment = dataset.get_fragments().pop().unwrap().metadata; - - let extended_schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("a", DataType::Int32, true), - ArrowField::new("b", DataType::Int32, true), - ])); - - // add all null column - let dataset = Dataset::commit( - WriteDestination::Dataset(Arc::new(dataset)), - Operation::Merge { - fragments: vec![fragment], - schema: extended_schema.as_ref().try_into().unwrap(), - }, - Some(2), - None, - None, - Arc::new(Default::default()), - false, - ) - .await - .unwrap(); - - // find the datafile we want to replace - let new_data_file = DataFile { - path: "test.lance".to_string(), - // the second column in the dataset - fields: vec![1], - // is located in the first column of this datafile - column_indices: vec![0], - file_major_version: 2, - file_minor_version: 0, - file_size_bytes: CachedFileSize::unknown(), - base_id: None, - }; - - let new_data_file = DataFile { - fields: vec![0, 1], - ..new_data_file - }; - - let err = Dataset::commit( - WriteDestination::Dataset(Arc::new(dataset.clone())), - Operation::DataReplacement { - replacements: vec![DataReplacementGroup(0, new_data_file)], - }, - Some(2), - None, - None, - Arc::new(Default::default()), - false, - ) - .await - .unwrap_err(); - assert!( - err.to_string() - .contains("Expected to modify the fragment but no changes were made"), - "Expected Error::DataFileReplacementError, got {:?}", - err - ); - } - - #[tokio::test] - async fn test_replace_dataset() { - let test_dir = TempDir::default(); - let test_uri = test_dir.path_str(); - let test_path = test_dir.obj_path(); - - let data = gen_batch() - .col("int", array::step::()) - .into_batch_rows(RowCount::from(20)) - .unwrap(); - let data1 = data.slice(0, 10); - let data2 = data.slice(10, 10); - let mut ds = InsertBuilder::new(&test_uri) - .execute(vec![data1]) - .await - .unwrap(); - - ds.object_store().remove_dir_all(test_path).await.unwrap(); - - let ds2 = InsertBuilder::new(&test_uri) - .execute(vec![data2.clone()]) - .await - .unwrap(); - - ds.checkout_latest().await.unwrap(); - let roundtripped = ds.scan().try_into_batch().await.unwrap(); - assert_eq!(roundtripped, data2); - - ds.validate().await.unwrap(); - ds2.validate().await.unwrap(); - assert_eq!(ds.manifest.version, 1); - assert_eq!(ds2.manifest.version, 1); - } - - #[tokio::test] - async fn test_session_store_registry() { - // Create a session - let session = Arc::new(Session::default()); - let registry = session.store_registry(); - assert!(registry.active_stores().is_empty()); - - // Create a dataset with memory store - let write_params = WriteParams { - session: Some(session.clone()), - ..Default::default() - }; - let batch = RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![ArrowField::new( - "a", - DataType::Int32, - false, - )])), - vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], - ) - .unwrap(); - let dataset = InsertBuilder::new("memory://test") - .with_params(&write_params) - .execute(vec![batch.clone()]) - .await - .unwrap(); - - // Assert there is one active store. - assert_eq!(registry.active_stores().len(), 1); - - // If we create another dataset also in memory, it should re-use the - // existing store. - let dataset2 = InsertBuilder::new("memory://test2") - .with_params(&write_params) - .execute(vec![batch.clone()]) - .await - .unwrap(); - assert_eq!(registry.active_stores().len(), 1); - assert_eq!( - Arc::as_ptr(&dataset.object_store().inner), - Arc::as_ptr(&dataset2.object_store().inner) - ); - - // If we create another with **different parameters**, it should create a new store. - let write_params2 = WriteParams { - session: Some(session.clone()), - store_params: Some(ObjectStoreParams { - block_size: Some(10_000), - ..Default::default() - }), - ..Default::default() - }; - let dataset3 = InsertBuilder::new("memory://test3") - .with_params(&write_params2) - .execute(vec![batch.clone()]) - .await - .unwrap(); - assert_eq!(registry.active_stores().len(), 2); - assert_ne!( - Arc::as_ptr(&dataset.object_store().inner), - Arc::as_ptr(&dataset3.object_store().inner) - ); - - // Remove both datasets - drop(dataset3); - assert_eq!(registry.active_stores().len(), 1); - drop(dataset2); - drop(dataset); - assert_eq!(registry.active_stores().len(), 0); - } - - #[tokio::test] - async fn test_migrate_v2_manifest_paths() { - let test_uri = TempStrDir::default(); - - let data = lance_datagen::gen_batch() - .col("key", array::step::()) - .into_reader_rows(RowCount::from(10), BatchCount::from(1)); - let mut dataset = Dataset::write(data, &test_uri, None).await.unwrap(); - assert_eq!( - dataset.manifest_location().naming_scheme, - ManifestNamingScheme::V1 - ); - - dataset.migrate_manifest_paths_v2().await.unwrap(); - assert_eq!( - dataset.manifest_location().naming_scheme, - ManifestNamingScheme::V2 - ); - } - - #[rstest] - #[tokio::test] - async fn test_fragment_id_zero_not_reused() { - // Test case 1: Fragment id zero isn't re-used - // 1. Create a dataset with 1 fragment - // 2. Delete all rows - // 3. Append another fragment - // 4. Assert new fragment has id 1 not 0 - - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::UInt32, - false, - )])); - - // Create dataset with 1 fragment - let data = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(UInt32Array::from_iter_values(0..10))], - ) - .unwrap(); - let batches = RecordBatchIterator::new(vec![data].into_iter().map(Ok), schema.clone()); - let mut dataset = Dataset::write(batches, &test_uri, None).await.unwrap(); - - // Verify we have 1 fragment with id 0 - assert_eq!(dataset.get_fragments().len(), 1); - assert_eq!(dataset.get_fragments()[0].id(), 0); - assert_eq!(dataset.manifest.max_fragment_id(), Some(0)); - - // Delete all rows - dataset.delete("true").await.unwrap(); - - // After deletion, dataset should be empty but max_fragment_id preserved - assert_eq!(dataset.get_fragments().len(), 0); - assert_eq!(dataset.count_rows(None).await.unwrap(), 0); - assert_eq!(dataset.manifest.max_fragment_id(), Some(0)); - - // Append another fragment - let data = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(UInt32Array::from_iter_values(20..30))], - ) - .unwrap(); - let batches = RecordBatchIterator::new(vec![data].into_iter().map(Ok), schema.clone()); - let write_params = WriteParams { - mode: WriteMode::Append, - ..Default::default() - }; - let dataset = Dataset::write(batches, &test_uri, Some(write_params)) - .await - .unwrap(); - - // Assert new fragment has id 1, not 0 - assert_eq!(dataset.get_fragments().len(), 1); - assert_eq!(dataset.get_fragments()[0].id(), 1); - assert_eq!(dataset.manifest.max_fragment_id(), Some(1)); - } - - #[rstest] - #[tokio::test] - async fn test_fragment_id_never_reset() { - // Test case 2: Fragment id is never reset, even if all rows are deleted - // 1. Create dataset with N fragments - // 2. Delete all rows - // 3. Append more fragments - // 4. Assert new fragments have ids >= N - - let test_uri = TempStrDir::default(); - - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::UInt32, - false, - )])); - - // Create dataset with 3 fragments (N=3) - let data = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(UInt32Array::from_iter_values(0..30))], - ) - .unwrap(); - let batches = RecordBatchIterator::new(vec![Ok(data)], schema.clone()); - let write_params = WriteParams { - max_rows_per_file: 10, // Force multiple fragments - ..Default::default() - }; - let mut dataset = Dataset::write(batches, &test_uri, Some(write_params)) - .await - .unwrap(); - - // Verify we have 3 fragments with ids 0, 1, 2 - assert_eq!(dataset.get_fragments().len(), 3); - assert_eq!(dataset.get_fragments()[0].id(), 0); - assert_eq!(dataset.get_fragments()[1].id(), 1); - assert_eq!(dataset.get_fragments()[2].id(), 2); - assert_eq!(dataset.manifest.max_fragment_id(), Some(2)); - - // Delete all rows - dataset.delete("true").await.unwrap(); - - // After deletion, dataset should be empty but max_fragment_id preserved - assert_eq!(dataset.get_fragments().len(), 0); - assert_eq!(dataset.count_rows(None).await.unwrap(), 0); - assert_eq!(dataset.manifest.max_fragment_id(), Some(2)); - - // Append more fragments (2 new fragments) - let data = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(UInt32Array::from_iter_values(100..120))], - ) - .unwrap(); - let batches = RecordBatchIterator::new(vec![Ok(data)], schema.clone()); - let write_params = WriteParams { - mode: WriteMode::Append, - max_rows_per_file: 10, // Force multiple fragments - ..Default::default() - }; - let dataset = Dataset::write(batches, &test_uri, Some(write_params)) - .await - .unwrap(); - - // Assert new fragments have ids >= N (3, 4) - assert_eq!(dataset.get_fragments().len(), 2); - assert_eq!(dataset.get_fragments()[0].id(), 3); - assert_eq!(dataset.get_fragments()[1].id(), 4); - assert_eq!(dataset.manifest.max_fragment_id(), Some(4)); - } - - #[tokio::test] - async fn test_insert_skip_auto_cleanup() { - let test_uri = TempStrDir::default(); - - // Create initial dataset with aggressive auto cleanup (interval=1, older_than=1ms) - let data = gen_batch() - .col("id", array::step::()) - .into_reader_rows(RowCount::from(100), BatchCount::from(1)); - - let write_params = WriteParams { - mode: WriteMode::Create, - auto_cleanup: Some(AutoCleanupParams { - interval: 1, - older_than: chrono::TimeDelta::try_milliseconds(0).unwrap(), // Cleanup versions older than 0ms - }), - ..Default::default() - }; - - // Start at 1 second after epoch - MockClock::set_system_time(std::time::Duration::from_secs(1)); - - let dataset = Dataset::write(data, &test_uri, Some(write_params)) - .await - .unwrap(); - assert_eq!(dataset.version().version, 1); - - // Advance time by 1 second - MockClock::set_system_time(std::time::Duration::from_secs(2)); - - // First append WITHOUT skip_auto_cleanup - should trigger cleanup - let data1 = gen_batch() - .col("id", array::step::()) - .into_df_stream(RowCount::from(50), BatchCount::from(1)); - - let write_params1 = WriteParams { - mode: WriteMode::Append, - skip_auto_cleanup: false, - ..Default::default() - }; - - let dataset2 = InsertBuilder::new(WriteDestination::Dataset(Arc::new(dataset))) - .with_params(&write_params1) - .execute_stream(data1) - .await - .unwrap(); - - assert_eq!(dataset2.version().version, 2); - - // Advance time - MockClock::set_system_time(std::time::Duration::from_secs(3)); - - // Need to do another commit for cleanup to take effect since cleanup runs on the old dataset - let data1_extra = gen_batch() - .col("id", array::step::()) - .into_df_stream(RowCount::from(10), BatchCount::from(1)); - - let dataset2_extra = InsertBuilder::new(WriteDestination::Dataset(Arc::new(dataset2))) - .with_params(&write_params1) - .execute_stream(data1_extra) - .await - .unwrap(); - - assert_eq!(dataset2_extra.version().version, 3); - - // Version 1 should be cleaned up due to auto cleanup (cleanup runs every version) - assert!( - dataset2_extra.checkout_version(1).await.is_err(), - "Version 1 should have been cleaned up" - ); - // Version 2 should still exist - assert!( - dataset2_extra.checkout_version(2).await.is_ok(), - "Version 2 should still exist" - ); - - // Advance time - MockClock::set_system_time(std::time::Duration::from_secs(4)); - - // Second append WITH skip_auto_cleanup - should NOT trigger cleanup - let data2 = gen_batch() - .col("id", array::step::()) - .into_df_stream(RowCount::from(30), BatchCount::from(1)); - - let write_params2 = WriteParams { - mode: WriteMode::Append, - skip_auto_cleanup: true, // Skip auto cleanup - ..Default::default() - }; - - let dataset3 = InsertBuilder::new(WriteDestination::Dataset(Arc::new(dataset2_extra))) - .with_params(&write_params2) - .execute_stream(data2) - .await - .unwrap(); - - assert_eq!(dataset3.version().version, 4); - - // Version 2 should still exist because skip_auto_cleanup was enabled - assert!( - dataset3.checkout_version(2).await.is_ok(), - "Version 2 should still exist because skip_auto_cleanup was enabled" - ); - // Version 3 should also still exist - assert!( - dataset3.checkout_version(3).await.is_ok(), - "Version 3 should still exist" - ); - } - - #[tokio::test] - async fn test_nullable_struct_v2_1_issue_4385() { - // Test for issue #4385: nullable struct should preserve null values in v2.1 format - use arrow_array::cast::AsArray; - use arrow_schema::Fields; - - // Create a struct field with nullable float field - let struct_fields = Fields::from(vec![ArrowField::new("x", DataType::Float32, true)]); - - // Create outer struct with the nullable struct as a field (not root) - let outer_fields = Fields::from(vec![ - ArrowField::new("id", DataType::Int32, false), - ArrowField::new("data", DataType::Struct(struct_fields.clone()), true), - ]); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "record", - DataType::Struct(outer_fields.clone()), - false, - )])); - - // Create data with null struct - let id_values = Int32Array::from(vec![1, 2, 3]); - let x_values = Float32Array::from(vec![Some(1.0), Some(2.0), Some(3.0)]); - let inner_struct_array = StructArray::new( - struct_fields, - vec![Arc::new(x_values) as ArrayRef], - Some(vec![true, false, true].into()), // Second struct is null - ); - - let outer_struct_array = StructArray::new( - outer_fields, - vec![ - Arc::new(id_values) as ArrayRef, - Arc::new(inner_struct_array.clone()) as ArrayRef, - ], - None, // Outer struct is not nullable - ); - - let batch = - RecordBatch::try_new(schema.clone(), vec![Arc::new(outer_struct_array)]).unwrap(); - - // Write dataset with v2.1 format - let test_uri = TempStrDir::default(); - - let write_params = WriteParams { - mode: WriteMode::Create, - data_storage_version: Some(LanceFileVersion::V2_1), - ..Default::default() - }; - - let batches = vec![batch.clone()]; - let batch_reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - - Dataset::write(batch_reader, &test_uri, Some(write_params)) - .await - .unwrap(); - - // Read back the dataset - let dataset = Dataset::open(&test_uri).await.unwrap(); - let scanner = dataset.scan(); - let result_batches = scanner - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - - assert_eq!(result_batches.len(), 1); - let result_batch = &result_batches[0]; - let read_outer_struct = result_batch.column(0).as_struct(); - let read_inner_struct = read_outer_struct.column(1).as_struct(); // "data" field - - // The bug: null struct is not preserved - assert!( - read_inner_struct.is_null(1), - "Second struct should be null but it's not. Read value: {:?}", - read_inner_struct - ); - - // Verify the null count is preserved - assert_eq!( - inner_struct_array.null_count(), - read_inner_struct.null_count(), - "Null count should be preserved" - ); - } - - #[tokio::test] - async fn test_issue_4902_packed_struct_v2_1_read_error() { - use std::collections::HashMap; - - use arrow_array::{ArrayRef, Int32Array, RecordBatchIterator, StructArray, UInt32Array}; - use arrow_schema::{Field as ArrowField, Fields, Schema as ArrowSchema}; - - let struct_fields = Fields::from(vec![ - ArrowField::new("x", DataType::UInt32, false), - ArrowField::new("y", DataType::UInt32, false), - ]); - let mut packed_metadata = HashMap::new(); - packed_metadata.insert("packed".to_string(), "true".to_string()); - - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("int_col", DataType::Int32, false), - ArrowField::new("struct_col", DataType::Struct(struct_fields.clone()), false) - .with_metadata(packed_metadata), - ])); - - let int_values = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8])); - let x_values = Arc::new(UInt32Array::from(vec![1, 4, 7, 10, 13, 16, 19, 22])); - let y_values = Arc::new(UInt32Array::from(vec![2, 5, 8, 11, 14, 17, 20, 23])); - let struct_array = Arc::new(StructArray::new( - struct_fields, - vec![x_values.clone() as ArrayRef, y_values.clone() as ArrayRef], - None, - )); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - int_values.clone() as ArrayRef, - struct_array.clone() as ArrayRef, - ], - ) - .unwrap(); - - let test_uri = TempStrDir::default(); - let write_params = WriteParams { - mode: WriteMode::Create, - data_storage_version: Some(LanceFileVersion::V2_1), - ..Default::default() - }; - let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], schema.clone()); - Dataset::write(reader, &test_uri, Some(write_params)) - .await - .unwrap(); - - let dataset = Dataset::open(&test_uri).await.unwrap(); - - let result_batches = dataset - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - assert_eq!(result_batches, vec![batch.clone()]); - - let struct_batches = dataset - .scan() - .project(&["struct_col"]) - .unwrap() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - assert_eq!(struct_batches.len(), 1); - let read_struct = struct_batches[0].column(0).as_struct(); - assert_eq!(read_struct, struct_array.as_ref()); - } - - #[tokio::test] - async fn test_issue_4429_nested_struct_encoding_v2_1_with_over_65k_structs() { - // Regression test for miniblock 16KB limit with nested struct patterns - // Tests encoding behavior when a nested struct> contains - // large amounts of data that exceeds miniblock encoding limits - - // Create a struct with multiple fields that will trigger miniblock encoding - // Each field is 4 bytes, making the struct narrow enough for miniblock - let measurement_fields = vec![ - ArrowField::new("val_a", DataType::Float32, true), - ArrowField::new("val_b", DataType::Float32, true), - ArrowField::new("val_c", DataType::Float32, true), - ArrowField::new("val_d", DataType::Float32, true), - ArrowField::new("seq_high", DataType::Int32, true), - ArrowField::new("seq_low", DataType::Int32, true), - ]; - let measurement_type = DataType::Struct(measurement_fields.clone().into()); - - // Create nested schema: struct> - // This pattern can trigger encoding issues with large data volumes - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "data", - DataType::Struct( - vec![ArrowField::new( - "measurements", - DataType::List(Arc::new(ArrowField::new( - "item", - measurement_type.clone(), - true, - ))), - true, - )] - .into(), - ), - true, - )])); - - // Create large number of measurements that will exceed encoding limits - // Using 70,520 to match the exact problematic size - const NUM_MEASUREMENTS: usize = 70_520; - - // Generate data for two full sets (rows 0 and 2 will have data, row 1 empty) - const TOTAL_MEASUREMENTS: usize = NUM_MEASUREMENTS * 2; - - // Create arrays with realistic values - let val_a_array = Float32Array::from_iter( - (0..TOTAL_MEASUREMENTS).map(|i| Some(16.66 + (i as f32 * 0.0001))), - ); - let val_b_array = Float32Array::from_iter( - (0..TOTAL_MEASUREMENTS).map(|i| Some(-3.54 + (i as f32 * 0.0002))), - ); - let val_c_array = Float32Array::from_iter( - (0..TOTAL_MEASUREMENTS).map(|i| Some(2.94 + (i as f32 * 0.0001))), - ); - let val_d_array = - Float32Array::from_iter((0..TOTAL_MEASUREMENTS).map(|i| Some(((i % 50) + 10) as f32))); - let seq_high_array = - Int32Array::from_iter((0..TOTAL_MEASUREMENTS).map(|_| Some(1736962329))); - let seq_low_array = Int32Array::from_iter( - (0..TOTAL_MEASUREMENTS).map(|i| Some(304403000 + (i * 1000) as i32)), - ); - - // Create the struct array with all measurements - let struct_array = StructArray::from(vec![ - ( - Arc::new(ArrowField::new("val_a", DataType::Float32, true)), - Arc::new(val_a_array) as ArrayRef, - ), - ( - Arc::new(ArrowField::new("val_b", DataType::Float32, true)), - Arc::new(val_b_array) as ArrayRef, - ), - ( - Arc::new(ArrowField::new("val_c", DataType::Float32, true)), - Arc::new(val_c_array) as ArrayRef, - ), - ( - Arc::new(ArrowField::new("val_d", DataType::Float32, true)), - Arc::new(val_d_array) as ArrayRef, - ), - ( - Arc::new(ArrowField::new("seq_high", DataType::Int32, true)), - Arc::new(seq_high_array) as ArrayRef, - ), - ( - Arc::new(ArrowField::new("seq_low", DataType::Int32, true)), - Arc::new(seq_low_array) as ArrayRef, - ), - ]); - - // Create list array with pattern: [70520 items, 0 items, 70520 items] - // This pattern triggers the issue with V2.1 encoding - let offsets = vec![ - 0i32, - NUM_MEASUREMENTS as i32, // End of row 0 - NUM_MEASUREMENTS as i32, // End of row 1 (empty) - (NUM_MEASUREMENTS * 2) as i32, // End of row 2 - ]; - let list_array = ListArray::try_new( - Arc::new(ArrowField::new("item", measurement_type, true)), - arrow_buffer::OffsetBuffer::new(arrow_buffer::ScalarBuffer::from(offsets)), - Arc::new(struct_array) as ArrayRef, - None, - ) - .unwrap(); - - // Create the outer struct wrapping the list - let data_struct = StructArray::from(vec![( - Arc::new(ArrowField::new( - "measurements", - DataType::List(Arc::new(ArrowField::new( - "item", - DataType::Struct(measurement_fields.into()), - true, - ))), - true, - )), - Arc::new(list_array) as ArrayRef, - )]); - - // Create the final record batch with 3 rows - let batch = - RecordBatch::try_new(schema.clone(), vec![Arc::new(data_struct) as ArrayRef]).unwrap(); - - assert_eq!(batch.num_rows(), 3, "Should have exactly 3 rows"); - - let test_uri = TempStrDir::default(); - - // Test with V2.1 format which has different encoding behavior - let batches = vec![batch]; - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); - - // V2.1 format triggers miniblock encoding for narrow structs - let write_params = WriteParams { - data_storage_version: Some(lance_file::version::LanceFileVersion::V2_1), - ..Default::default() - }; - - // Write dataset - this will panic with miniblock 16KB assertion - let dataset = Dataset::write(reader, &test_uri, Some(write_params)) - .await - .unwrap(); - - dataset.validate().await.unwrap(); - assert_eq!(dataset.count_rows(None).await.unwrap(), 3); - } - - async fn prepare_json_dataset() -> (Dataset, String) { - let text_col = Arc::new(StringArray::from(vec![ - r#"{ - "Title": "HarryPotter Chapter One", - "Content": "Mr. and Mrs. Dursley, of number four, Privet Drive, were proud to say...", - "Author": "J.K. Rowling", - "Price": 128, - "Language": ["english", "chinese"] - }"#, - r#"{ - "Title": "Fairy Talest", - "Content": "Once upon a time, on a bitterly cold New Year's Eve, a little girl...", - "Author": "ANDERSEN", - "Price": 50, - "Language": ["english", "chinese"] - }"#, - ])); - let json_col = "json_field".to_string(); - - // Prepare dataset - let mut metadata = HashMap::new(); - metadata.insert( - ARROW_EXT_NAME_KEY.to_string(), - ARROW_JSON_EXT_NAME.to_string(), - ); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - Field::new(&json_col, DataType::Utf8, false).with_metadata(metadata) - ]) - .into(), - vec![text_col.clone()], - ) - .unwrap(); - let schema = batch.schema(); - let stream = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let dataset = Dataset::write(stream, "memory://test/table", None) - .await - .unwrap(); - - (dataset, json_col) - } - - #[tokio::test] - async fn test_json_inverted_fuzziness_query() { - let (mut dataset, json_col) = prepare_json_dataset().await; - - // Create inverted index for json col - dataset - .create_index( - &[&json_col], - IndexType::Inverted, - None, - &InvertedIndexParams::default().lance_tokenizer("json".to_string()), - true, - ) - .await - .unwrap(); - - // Match query with fuzziness - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Content,str,Dursley".to_string()) - .with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(1, batch.num_rows()); - - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Content,str,Bursley".to_string()) - .with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(0, batch.num_rows()); - - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Content,str,Bursley".to_string()) - .with_column(Some(json_col.clone())) - .with_fuzziness(Some(1)), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(1, batch.num_rows()); - - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Content,str,ABursley".to_string()) - .with_column(Some(json_col.clone())) - .with_fuzziness(Some(1)), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(0, batch.num_rows()); - - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Content,str,ABursley".to_string()) - .with_column(Some(json_col.clone())) - .with_fuzziness(Some(2)), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(1, batch.num_rows()); - - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Dontent,str,Bursley".to_string()) - .with_column(Some(json_col.clone())) - .with_fuzziness(Some(2)), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(0, batch.num_rows()); - } - - #[tokio::test] - async fn test_json_inverted_match_query() { - let (mut dataset, json_col) = prepare_json_dataset().await; - - // Create inverted index for json col, with max token len 10 and enable stemming, - // lower case, and remove stop words - dataset - .create_index( - &[&json_col], - IndexType::Inverted, - None, - &InvertedIndexParams::default() - .lance_tokenizer("json".to_string()) - .max_token_length(Some(10)) - .stem(true) - .lower_case(true) - .remove_stop_words(true), - true, - ) - .await - .unwrap(); - - // Match query with token length exceed max token length - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Title,str,harrypotter".to_string()) - .with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(0, batch.num_rows()); - - // Match query with stemming - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Content,str,onc".to_string()).with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(1, batch.num_rows()); - - // Match query with lower case - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Content,str,DURSLEY".to_string()) - .with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(1, batch.num_rows()); - - // Match query with stop word - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Content,str,and".to_string()).with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(0, batch.num_rows()); - } - - #[tokio::test] - async fn test_json_inverted_flat_match_query() { - let (mut dataset, json_col) = prepare_json_dataset().await; - - // Create inverted index for json col - dataset - .create_index( - &[&json_col], - IndexType::Inverted, - None, - &InvertedIndexParams::default() - .lance_tokenizer("json".to_string()) - .stem(false), - true, - ) - .await - .unwrap(); - - // Append data - let text_col = Arc::new(StringArray::from(vec![ - r#"{ - "Title": "HarryPotter Chapter Two", - "Content": "Nearly ten years had passed since the Dursleys had woken up...", - "Author": "J.K. Rowling", - "Price": 128, - "Language": ["english", "chinese"] - }"#, - ])); - - let mut metadata = HashMap::new(); - metadata.insert( - ARROW_EXT_NAME_KEY.to_string(), - ARROW_JSON_EXT_NAME.to_string(), - ); - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![ - Field::new(&json_col, DataType::Utf8, false).with_metadata(metadata) - ]) - .into(), - vec![text_col.clone()], - ) - .unwrap(); - let schema = batch.schema(); - let stream = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - dataset.append(stream, None).await.unwrap(); - - // Test match query - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Title,str,harrypotter".to_string()) - .with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(2, batch.num_rows()); - } - - #[tokio::test] - async fn test_json_inverted_phrase_query() { - // Prepare json dataset - let (mut dataset, json_col) = prepare_json_dataset().await; - - // Create inverted index for json col - dataset - .create_index( - &[&json_col], - IndexType::Inverted, - None, - &InvertedIndexParams::default() - .lance_tokenizer("json".to_string()) - .stem(false) - .with_position(true), - true, - ) - .await - .unwrap(); - - // Test phrase query - let query = FullTextSearchQuery { - query: FtsQuery::Phrase( - PhraseQuery::new("Title,str,harrypotter one chapter".to_string()) - .with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(0, batch.num_rows()); - - let query = FullTextSearchQuery { - query: FtsQuery::Phrase( - PhraseQuery::new("Title,str,harrypotter chapter one".to_string()) - .with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(1, batch.num_rows()); - } - - #[tokio::test] - async fn test_json_inverted_multimatch_query() { - // Prepare json dataset - let (mut dataset, json_col) = prepare_json_dataset().await; - - // Create inverted index for json col - dataset - .create_index( - &[&json_col], - IndexType::Inverted, - None, - &InvertedIndexParams::default() - .lance_tokenizer("json".to_string()) - .stem(false), - true, - ) - .await - .unwrap(); - - // Test multi match query - let query = FullTextSearchQuery { - query: FtsQuery::MultiMatch(MultiMatchQuery { - match_queries: vec![ - MatchQuery::new("Title,str,harrypotter".to_string()) - .with_column(Some(json_col.clone())), - MatchQuery::new("Language,str,english".to_string()) - .with_column(Some(json_col.clone())), - ], - }), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(2, batch.num_rows()); - } - - #[tokio::test] - async fn test_json_inverted_boolean_query() { - // Prepare json dataset - let (mut dataset, json_col) = prepare_json_dataset().await; - - // Create inverted index for json col - dataset - .create_index( - &[&json_col], - IndexType::Inverted, - None, - &InvertedIndexParams::default() - .lance_tokenizer("json".to_string()) - .stem(false), - true, - ) - .await - .unwrap(); - - // Test boolean query - let query = FullTextSearchQuery { - query: FtsQuery::Boolean(BooleanQuery { - should: vec![], - must: vec![ - FtsQuery::Match( - MatchQuery::new("Language,str,english".to_string()) - .with_column(Some(json_col.clone())), - ), - FtsQuery::Match( - MatchQuery::new("Title,str,harrypotter".to_string()) - .with_column(Some(json_col.clone())), - ), - ], - must_not: vec![], - }), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(1, batch.num_rows()); - } - - #[tokio::test] - async fn test_sql_contains_tokens() { - let text_col = Arc::new(StringArray::from(vec![ - "a cat catch a fish", - "a fish catch a cat", - "a white cat catch a big fish", - "cat catchup fish", - "cat fish catch", - ])); - - // Prepare dataset - let batch = RecordBatch::try_new( - arrow_schema::Schema::new(vec![Field::new("text", DataType::Utf8, false)]).into(), - vec![text_col.clone()], - ) - .unwrap(); - let schema = batch.schema(); - let stream = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); - let mut dataset = Dataset::write(stream, "memory://test/table", None) - .await - .unwrap(); - - // Test without fts index - let results = execute_sql( - "select * from foo where contains_tokens(text, 'cat catch fish')", - "foo".to_string(), - Arc::new(dataset.clone()), - ) - .await - .unwrap(); - - assert_results( - results, - &StringArray::from(vec![ - "a cat catch a fish", - "a fish catch a cat", - "a white cat catch a big fish", - "cat fish catch", - ]), - ); - - // Verify plan, should not contain ScalarIndexQuery. - let results = execute_sql( - "explain select * from foo where contains_tokens(text, 'cat catch fish')", - "foo".to_string(), - Arc::new(dataset.clone()), - ) - .await - .unwrap(); - let plan = format!("{:?}", results); - assert_not_contains!(&plan, "ScalarIndexQuery"); - - // Test with unsuitable fts index - dataset - .create_index( - &["text"], - IndexType::Inverted, - None, - &InvertedIndexParams::default().base_tokenizer("raw".to_string()), - true, - ) - .await - .unwrap(); - - let results = execute_sql( - "select * from foo where contains_tokens(text, 'cat catch fish')", - "foo".to_string(), - Arc::new(dataset.clone()), - ) - .await - .unwrap(); - - assert_results( - results, - &StringArray::from(vec![ - "a cat catch a fish", - "a fish catch a cat", - "a white cat catch a big fish", - "cat fish catch", - ]), - ); - - // Verify plan, should not contain ScalarIndexQuery because fts index is not unsuitable. - let results = execute_sql( - "explain select * from foo where contains_tokens(text, 'cat catch fish')", - "foo".to_string(), - Arc::new(dataset.clone()), - ) - .await - .unwrap(); - let plan = format!("{:?}", results); - assert_not_contains!(&plan, "ScalarIndexQuery"); - - // Test with suitable fts index - dataset - .create_index( - &["text"], - IndexType::Inverted, - None, - &InvertedIndexParams::default() - .max_token_length(None) - .stem(false), - true, - ) - .await - .unwrap(); - - let results = execute_sql( - "select * from foo where contains_tokens(text, 'cat catch fish')", - "foo".to_string(), - Arc::new(dataset.clone()), - ) - .await - .unwrap(); - - assert_results( - results, - &StringArray::from(vec![ - "a cat catch a fish", - "a fish catch a cat", - "a white cat catch a big fish", - "cat fish catch", - ]), - ); - - // Verify plan, should contain ScalarIndexQuery. - let results = execute_sql( - "explain select * from foo where contains_tokens(text, 'cat catch fish')", - "foo".to_string(), - Arc::new(dataset.clone()), - ) - .await - .unwrap(); - let plan = format!("{:?}", results); - assert_contains!(&plan, "ScalarIndexQuery"); - } - - async fn execute_sql( - sql: &str, - table: String, - dataset: Arc, - ) -> Result> { - let ctx = SessionContext::new(); - ctx.register_table( - table, - Arc::new(LanceTableProvider::new(dataset, false, false)), - )?; - register_functions(&ctx); - - let df = ctx.sql(sql).await?; - Ok(df - .execute_stream() - .await - .unwrap() - .try_collect::>() - .await?) - } - - fn assert_results(results: Vec, values: &T) { - assert_eq!(results.len(), 1); - let results = results.into_iter().next().unwrap(); - assert_eq!(results.num_columns(), 1); - - assert_eq!( - results.column(0).as_any().downcast_ref::().unwrap(), - values - ) - } - - // Test coverage: - // Case 1: delete external transaction file → read_transaction should prioritize inline and succeed. - // Case 2: reading small manifest caches transaction data, eliminating transaction reading IO. - // Case 3: manifest does not contain inline → read_transaction should fall back to external transaction file and succeed. - #[tokio::test] - async fn test_inline_transaction() { - use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator}; - use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; - use std::sync::Arc; - - async fn create_dataset(rows: i32) -> Arc { - let dir = TempDir::default(); - let uri = dir.path_str(); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "i", - DataType::Int32, - false, - )])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(Int32Array::from_iter_values(0..rows))], - ) - .unwrap(); - let ds = Dataset::write( - RecordBatchIterator::new(vec![Ok(batch)], schema), - uri.as_str(), - None, - ) - .await - .unwrap(); - Arc::new(ds) - } - - fn make_tx(read_version: u64) -> Transaction { - Transaction::new(read_version, Operation::Append { fragments: vec![] }, None) - } - - async fn delete_external_tx_file(ds: &Dataset) { - if let Some(tx_file) = ds.manifest.transaction_file.as_ref() { - let tx_path = ds.base.child("_transactions").child(tx_file.as_str()); - let _ = ds.object_store.inner.delete(&tx_path).await; // ignore errors - } - } - - let session = Arc::new(Session::default()); - - // Case 1: Default write_flag=true, delete external transaction file, read should use inline transaction - let ds = create_dataset(5).await; - let read_version = ds.manifest().version; - let tx = make_tx(read_version); - let ds2 = CommitBuilder::new(ds.clone()) - .execute(tx.clone()) - .await - .unwrap(); - delete_external_tx_file(&ds2).await; - let read_tx = ds2.read_transaction().await.unwrap().unwrap(); - assert_eq!(read_tx, tx.clone()); - - // Case 2: reading small manifest caches transaction data, eliminating transaction reading IO. - let read_ds2 = DatasetBuilder::from_uri(ds2.uri.clone()) - .with_session(session.clone()) - .load() - .await - .unwrap(); - let stats = read_ds2.object_store().io_stats_incremental(); // Reset - assert!(stats.read_bytes < 64 * 1024); - // Because the manifest is so small, we should have opportunistically - // cached the transaction in memory already. - let inline_tx = read_ds2.read_transaction().await.unwrap().unwrap(); - let stats = read_ds2.object_store().io_stats_incremental(); - assert_eq!(stats.read_iops, 0); - assert_eq!(stats.read_bytes, 0); - assert_eq!(inline_tx, tx); - - // Case 3: manifest does not contain inline transaction, read should fall back to external transaction file - let ds = create_dataset(2).await; - let tx = make_tx(ds.manifest().version); - let tx_file = crate::io::commit::write_transaction_file(ds.object_store(), &ds.base, &tx) - .await - .unwrap(); - let (mut manifest, indices) = tx - .build_manifest( - Some(ds.manifest.as_ref()), - ds.load_indices().await.unwrap().as_ref().clone(), - &tx_file, - &ManifestWriteConfig::default(), - ) - .unwrap(); - let location = write_manifest_file( - ds.object_store(), - ds.commit_handler.as_ref(), - &ds.base, - &mut manifest, - if indices.is_empty() { - None - } else { - Some(indices.clone()) - }, - &ManifestWriteConfig::default(), - ds.manifest_location.naming_scheme, - None, - ) - .await - .unwrap(); - let ds_new = ds.checkout_version(location.version).await.unwrap(); - assert!(ds_new.manifest.transaction_section.is_none()); - assert!(ds_new.manifest.transaction_file.is_some()); - let read_tx = ds_new.read_transaction().await.unwrap().unwrap(); - assert_eq!(read_tx, tx); - } - - #[tokio::test] - async fn test_limit_pushdown_in_physical_plan() -> Result<()> { - use tempfile::tempdir; - let temp_dir = tempdir()?; - - let dataset_path = temp_dir.path().join("limit_pushdown_dataset"); - let values: Vec = (0..1000).collect(); - let array = Int32Array::from(values); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "value", - DataType::Int32, - false, - )])); - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)])?; - - let write_params = WriteParams { - mode: WriteMode::Create, - max_rows_per_file: 100, - ..Default::default() - }; - - let batch_reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); - Dataset::write( - batch_reader, - dataset_path.to_str().unwrap(), - Some(write_params), - ) - .await?; - - let mut dataset = Dataset::open(dataset_path.to_str().unwrap()).await?; - - dataset - .create_index( - &["value"], - IndexType::Scalar, - None, - &ScalarIndexParams::default(), - false, - ) - .await?; - - // Test 1: No filter with limit - { - let mut scanner = dataset.scan(); - scanner.limit(Some(100), None)?; - let plan = scanner.explain_plan(true).await?; - - assert!(plan.contains("range_before=Some(0..100)")); - assert!(plan.contains("range_after=None")); - - let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(100, total_rows); - } - - // Test 2: Indexed filter with limit - { - let mut scanner = dataset.scan(); - scanner.filter("value >= 500")?.limit(Some(50), None)?; - let plan = scanner.explain_plan(true).await?; - - assert!(plan.contains("range_after=Some(0..50)")); - assert!(plan.contains("range_before=None")); - - let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(50, total_rows); - } - - // Test 3: Offset + Limit - { - let mut scanner = dataset.scan(); - scanner.filter("value < 500")?.limit(Some(30), Some(20))?; - let plan = scanner.explain_plan(true).await?; - - assert!(plan.contains("GlobalLimitExec: skip=20, fetch=30")); - assert!(plan.contains("range_after=Some(0..50)")); - - let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(30, total_rows); - - // Verify exact values (should be 20..50) - let all_values: Vec = batches - .iter() - .flat_map(|batch| { - batch - .column_by_name("value") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .values() - .iter() - .copied() - .collect::>() - }) - .collect(); - assert_eq!(all_values, (20..50).collect::>()); - } - - // Test 4: Large limit exceeding data - { - let mut scanner = dataset.scan(); - scanner.limit(Some(5000), None)?; - let plan = scanner.explain_plan(true).await?; - - assert!(plan.contains("range_before=Some(0..1000)")); - - let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(1000, total_rows); - } - - // Test 5: Cross-fragment filter with limit - { - let mut scanner = dataset.scan(); - scanner - .filter("value >= 95 AND value <= 205")? - .limit(Some(50), None)?; - let plan = scanner.explain_plan(true).await?; - - assert!(plan.contains("range_after=Some(0..50)")); - - let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(50, total_rows); - } - - Ok(()) - } - - #[tokio::test] - async fn test_index_take_batch_size() -> Result<()> { - use tempfile::tempdir; - let temp_dir = tempdir()?; - - let dataset_path = temp_dir.path().join("ints_dataset"); - let values: Vec = (0..1024).collect(); - let array = Int32Array::from(values); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "ints", - DataType::Int32, - false, - )])); - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)])?; - let write_params = WriteParams { - mode: WriteMode::Create, - max_rows_per_file: 100, - ..Default::default() - }; - let batch_reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); - Dataset::write( - batch_reader, - dataset_path.to_str().unwrap(), - Some(write_params), - ) - .await?; - let mut dataset = Dataset::open(dataset_path.to_str().unwrap()).await?; - dataset - .create_index( - &["ints"], - IndexType::Scalar, - None, - &ScalarIndexParams::default(), - false, - ) - .await?; - - let mut scanner = dataset.scan(); - scanner.batch_size(50).filter("ints > 0")?.with_row_id(); - let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(1023, total_rows); - assert_eq!(21, batches.len()); - - let mut scanner = dataset.scan(); - scanner - .batch_size(50) - .filter("ints > 0")? - .limit(Some(1024), None)? - .with_row_id(); - let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(1023, total_rows); - assert_eq!(21, batches.len()); - - let dataset_path2 = temp_dir.path().join("strings_dataset"); - let strings: Vec = (0..1024).map(|i| format!("string-{}", i)).collect(); - let string_array = StringArray::from(strings); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "strings", - DataType::Utf8, - false, - )])); - let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(string_array)])?; - let write_params = WriteParams { - mode: WriteMode::Create, - max_rows_per_file: 100, - ..Default::default() - }; - let batch_reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); - Dataset::write( - batch_reader, - dataset_path2.to_str().unwrap(), - Some(write_params), - ) - .await?; - let mut dataset2 = Dataset::open(dataset_path2.to_str().unwrap()).await?; - dataset2 - .create_index( - &["strings"], - IndexType::Scalar, - None, - &ScalarIndexParams::default(), - false, - ) - .await?; - - let mut scanner = dataset2.scan(); - scanner - .batch_size(50) - .filter("contains(strings, 'ing')")? - .limit(Some(1024), None)? - .with_row_id(); - let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(1024, total_rows); - assert_eq!(21, batches.len()); - - Ok(()) - } - - // This test covers - // 1. Create branch from main, a branch and a global tag - // 2. Write to each created branch and verify data - // 2. Load branch from nested uris - // 3. Checkout branch from main, a branch and a global tag - // 4. List branches and verify branch metadata - // 5. Delete branches - // 6. Delete zombie branches - #[tokio::test] - async fn test_branch() { - let tempdir = TempDir::default(); - let test_uri = tempdir.path_str(); - let data_storage_version = LanceFileVersion::Stable; - - // Generate consistent test data batches - let generate_data = |prefix: &str, start_id: i32, row_count: u64| { - gen_batch() - .col("id", array::step_custom::(start_id, 1)) - .col("value", array::fill_utf8(format!("{prefix}_data"))) - .into_reader_rows(RowCount::from(row_count), BatchCount::from(1)) - }; - - // Reusable dataset writer with configurable mode - async fn write_dataset( - uri: &str, - data_reader: impl RecordBatchReader + Send + 'static, - mode: WriteMode, - version: LanceFileVersion, - ) -> Dataset { - let params = WriteParams { - max_rows_per_file: 100, - max_rows_per_group: 20, - data_storage_version: Some(version), - mode, - ..Default::default() - }; - Dataset::write(data_reader, uri, Some(params)) - .await - .unwrap() - } - - // Unified dataset scanning and row counting - async fn collect_rows(dataset: &Dataset) -> (usize, Vec) { - let batches = dataset - .scan() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - (batches.iter().map(|b| b.num_rows()).sum(), batches) - } - - // Phase 1: Create empty dataset, write data batch 1, create branch1 based on version_number, write data batch 2 - let mut dataset = write_dataset( - &test_uri, - generate_data("batch1", 0, 50), - WriteMode::Create, - data_storage_version, - ) - .await; - - let original_version = dataset.version().version; - assert_eq!(original_version, 1); - - // Create branch1 on the latest version and write data batch 2 - let mut branch1_dataset = dataset - .create_branch("branch1", original_version, None) - .await - .unwrap(); - assert_eq!(branch1_dataset.uri, format!("{}/tree/branch1", test_uri)); - - branch1_dataset = write_dataset( - branch1_dataset.uri(), - generate_data("batch2", 50, 30), - WriteMode::Append, - data_storage_version, - ) - .await; - - // Phase 2: Create branch2 based on branch1's latest version_number, write data batch 3 - let mut branch2_dataset = branch1_dataset - .create_branch( - "dev/branch2", - ("branch1", branch1_dataset.version().version), - None, - ) - .await - .unwrap(); - assert_eq!( - branch2_dataset.uri, - format!("{}/tree/dev/branch2", test_uri) - ); - - branch2_dataset = write_dataset( - branch2_dataset.uri(), - generate_data("batch3", 80, 20), - WriteMode::Append, - data_storage_version, - ) - .await; - - // Phase 3: Create a tag on branch2, the actual tag content is under root dataset - // create branch3 based on that tag, write data batch 4 - branch2_dataset - .tags() - .create_on_branch( - "tag1", - branch2_dataset.version().version, - Some("dev/branch2"), - ) - .await - .unwrap(); - - let mut branch3_dataset = branch2_dataset - .create_branch("feature/nathan/branch3", "tag1", None) - .await - .unwrap(); - assert_eq!( - branch3_dataset.uri, - format!("{}/tree/feature/nathan/branch3", test_uri) - ); - - branch3_dataset = write_dataset( - branch3_dataset.uri(), - generate_data("batch4", 100, 25), - WriteMode::Append, - data_storage_version, - ) - .await; - - // Verify data correctness and independence of each branch - // Main branch only has data 1 (50 rows) - let main_dataset = Dataset::open(&test_uri).await.unwrap(); - let (main_rows, _) = collect_rows(&main_dataset).await; - assert_eq!(main_rows, 50); // only batch1 - assert_eq!(main_dataset.version().version, 1); - - // branch1 has data 1 + 2 (80 rows) - let updated_branch1 = Dataset::open(branch1_dataset.uri()).await.unwrap(); - let (branch1_rows, _) = collect_rows(&updated_branch1).await; - assert_eq!(branch1_rows, 80); // batch1+batch2 - assert_eq!(updated_branch1.version().version, 2); - - // branch2 has data 1 + 2 + 3 (100 rows) - let updated_branch2 = Dataset::open(branch2_dataset.uri()).await.unwrap(); - let (branch2_rows, _) = collect_rows(&updated_branch2).await; - assert_eq!(branch2_rows, 100); // batch1+batch2+batch3 - assert_eq!(updated_branch2.version().version, 3); - - // branch3 has data 1 + 2 + 3 + 4 (125 rows) - let updated_branch3 = Dataset::open(branch3_dataset.uri()).await.unwrap(); - let (branch3_rows, _) = collect_rows(&updated_branch3).await; - assert_eq!(branch3_rows, 125); // batch1+batch2+batch3+batch4 - assert_eq!(updated_branch3.version().version, 4); - - // Use list_branches to get branch list and verify each field of branch_content - let branches = dataset.list_branches().await.unwrap(); - assert_eq!(branches.len(), 3); - assert!(branches.contains_key("branch1")); - assert!(branches.contains_key("dev/branch2")); - assert!(branches.contains_key("feature/nathan/branch3")); - - // Verify branch1 content - let branch1_content = branches.get("branch1").unwrap(); - assert_eq!(branch1_content.parent_branch, None); // Created based on main branch - assert_eq!(branch1_content.parent_version, 1); - assert!(branch1_content.create_at > 0); - assert!(branch1_content.manifest_size > 0); - - // Verify branch2 content - let branch2_content = branches.get("dev/branch2").unwrap(); - assert_eq!(branch2_content.parent_branch.as_deref().unwrap(), "branch1"); - assert_eq!(branch2_content.parent_version, 2); - assert!(branch2_content.create_at > 0); - assert!(branch2_content.manifest_size > 0); - assert!(branch2_content.create_at >= branch1_content.create_at); - - // Verify branch3 content - let branch3_content = branches.get("feature/nathan/branch3").unwrap(); - // Created based on tag pointed to branch2 - assert_eq!( - branch3_content.parent_branch.as_deref().unwrap(), - "dev/branch2" - ); - assert_eq!(branch3_content.parent_version, 3); - assert!(branch3_content.create_at > 0); - assert!(branch3_content.manifest_size > 0); - assert!(branch3_content.create_at >= branch2_content.create_at); - - // Verify checkout_branch - let checkout_branch1 = main_dataset.checkout_branch("branch1").await.unwrap(); - let checkout_branch2 = checkout_branch1 - .checkout_branch("dev/branch2") - .await - .unwrap(); - let checkout_branch2_tag = checkout_branch1.checkout_version("tag1").await.unwrap(); - let checkout_branch3 = checkout_branch2_tag - .checkout_branch("feature/nathan/branch3") - .await - .unwrap(); - let checkout_branch3_at_version3 = checkout_branch2 - .checkout_version(("feature/nathan/branch3", 3)) - .await - .unwrap(); - assert_eq!(checkout_branch3.version().version, 4); - assert_eq!(checkout_branch3_at_version3.version().version, 3); - assert_eq!(checkout_branch2.version().version, 3); - assert_eq!(checkout_branch2_tag.version().version, 3); - assert_eq!(checkout_branch1.version().version, 2); - assert_eq!(checkout_branch3.count_rows(None).await.unwrap(), 125); - assert_eq!( - checkout_branch3_at_version3.count_rows(None).await.unwrap(), - 100 - ); - assert_eq!(checkout_branch2.count_rows(None).await.unwrap(), 100); - assert_eq!(checkout_branch2_tag.count_rows(None).await.unwrap(), 100); - assert_eq!(checkout_branch1.count_rows(None).await.unwrap(), 80); - assert_eq!( - checkout_branch3.manifest.branch.as_deref().unwrap(), - "feature/nathan/branch3" - ); - assert_eq!( - checkout_branch3_at_version3 - .manifest - .branch - .as_deref() - .unwrap(), - "feature/nathan/branch3" - ); - assert_eq!( - checkout_branch2.manifest.branch.as_deref().unwrap(), - "dev/branch2" - ); - assert_eq!( - checkout_branch2_tag.manifest.branch.as_deref().unwrap(), - "dev/branch2" - ); - assert_eq!( - checkout_branch1.manifest.branch.as_deref().unwrap(), - "branch1" - ); - - let mut dataset = main_dataset; - // Finally delete all branches - dataset.delete_branch("branch1").await.unwrap(); - dataset.delete_branch("dev/branch2").await.unwrap(); - // Test deleting zombie branch - let root_location = dataset.refs.root().unwrap(); - let branch_file = branch_contents_path(&root_location.path, "feature/nathan/branch3"); - dataset.object_store.delete(&branch_file).await.unwrap(); - // Now "feature/nathan/branch3" is a zombie branch - // Use delete_branch to verify if the directory is cleaned up - dataset - .force_delete_branch("feature/nathan/branch3") - .await - .unwrap(); - let cleaned_path = Path::parse(format!("{}/tree/feature", test_uri)).unwrap(); - assert!(!dataset.object_store.exists(&cleaned_path).await.unwrap()); - - // Verify list_branches is empty - let branches_after_delete = dataset.list_branches().await.unwrap(); - assert!(branches_after_delete.is_empty()); - - // Verify branch directories are all deleted cleanly - let test_path = tempdir.obj_path(); - let branches = dataset - .object_store - .read_dir(test_path.child("tree")) - .await - .unwrap(); - assert!(branches.is_empty()); - } - - #[tokio::test] - async fn test_add_bases() { - use lance_table::format::BasePath; - use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; - use std::sync::Arc; - - // Create a test dataset - let test_uri = "memory://add_bases_test"; - let mut data_gen = - BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); - - let dataset = Dataset::write( - data_gen.batch(5), - test_uri, - Some(WriteParams { - mode: WriteMode::Create, - ..Default::default() - }), - ) - .await - .unwrap(); - - let dataset = Arc::new(dataset); - - // Test adding new base paths - let new_bases = vec![ - BasePath::new( - 0, - "memory://bucket1".to_string(), - Some("bucket1".to_string()), - false, - ), - BasePath::new( - 0, - "memory://bucket2".to_string(), - Some("bucket2".to_string()), - true, - ), - ]; - - let updated_dataset = dataset.add_bases(new_bases, None).await.unwrap(); - - // Verify the base paths were added - assert_eq!(updated_dataset.manifest.base_paths.len(), 2); - - let bucket1 = updated_dataset - .manifest - .base_paths - .values() - .find(|bp| bp.name == Some("bucket1".to_string())) - .expect("bucket1 not found"); - let bucket2 = updated_dataset - .manifest - .base_paths - .values() - .find(|bp| bp.name == Some("bucket2".to_string())) - .expect("bucket2 not found"); - - assert_eq!(bucket1.path, "memory://bucket1"); - assert!(!bucket1.is_dataset_root); - assert_eq!(bucket2.path, "memory://bucket2"); - assert!(bucket2.is_dataset_root); - - let updated_dataset = Arc::new(updated_dataset); - - // Test conflict detection - try to add a base with the same name - let conflicting_bases = vec![BasePath::new( - 0, - "memory://bucket3".to_string(), - Some("bucket1".to_string()), - false, - )]; - - let result = updated_dataset.add_bases(conflicting_bases, None).await; - assert!(result.is_err()); - assert!(result - .unwrap_err() - .to_string() - .contains("Conflict detected")); - - // Test conflict detection - try to add a base with the same path - let conflicting_bases = vec![BasePath::new( - 0, - "memory://bucket1".to_string(), - Some("bucket3".to_string()), - false, - )]; - - let result = updated_dataset.add_bases(conflicting_bases, None).await; - assert!(result.is_err()); - assert!(result - .unwrap_err() - .to_string() - .contains("Conflict detected")); - } - - #[tokio::test] - async fn test_concurrent_add_bases_conflict() { - use lance_table::format::BasePath; - use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; - use std::sync::Arc; - - // Create a test dataset - let test_uri = "memory://concurrent_add_bases_test"; - let mut data_gen = - BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); - - let dataset = Dataset::write( - data_gen.batch(5), - test_uri, - Some(WriteParams { - mode: WriteMode::Create, - ..Default::default() - }), - ) - .await - .unwrap(); - - // Clone the dataset to simulate concurrent access - let dataset = Arc::new(dataset); - let dataset_clone = Arc::new(dataset.clone()); - - // First transaction adds base1 - let new_bases1 = vec![BasePath::new( - 0, - "memory://bucket1".to_string(), - Some("base1".to_string()), - false, - )]; - - let updated_dataset = dataset.add_bases(new_bases1, None).await.unwrap(); - - // Second transaction tries to add a different base (base2) - // This should succeed as there's no conflict - let new_bases2 = vec![BasePath::new( - 0, - "memory://bucket2".to_string(), - Some("base2".to_string()), - false, - )]; - - let result = dataset_clone.add_bases(new_bases2, None).await; - assert!(result.is_ok()); - - // Verify both bases are present after conflict resolution - let mut final_dataset = updated_dataset; - final_dataset.checkout_latest().await.unwrap(); - assert_eq!(final_dataset.manifest.base_paths.len(), 2); - - let base1 = final_dataset - .manifest - .base_paths - .values() - .find(|bp| bp.name == Some("base1".to_string())); - let base2 = final_dataset - .manifest - .base_paths - .values() - .find(|bp| bp.name == Some("base2".to_string())); - - assert!(base1.is_some()); - assert!(base2.is_some()); - } - - #[tokio::test] - async fn test_concurrent_add_bases_name_conflict() { - use lance_table::format::BasePath; - use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; - use std::sync::Arc; - - // Create a test dataset - let test_uri = "memory://concurrent_name_conflict_test"; - let mut data_gen = - BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); - - let dataset = Dataset::write( - data_gen.batch(5), - test_uri, - Some(WriteParams { - mode: WriteMode::Create, - ..Default::default() - }), - ) - .await - .unwrap(); - - // Clone the dataset to simulate concurrent access - let dataset_clone = dataset.clone(); - let dataset = Arc::new(dataset); - let dataset_clone = Arc::new(dataset_clone); - - // First transaction adds base with name "shared_base" - let new_bases1 = vec![BasePath::new( - 0, - "memory://bucket1".to_string(), - Some("shared_base".to_string()), - false, - )]; - - let _updated_dataset = dataset.add_bases(new_bases1, None).await.unwrap(); - - // Second transaction tries to add a different base with same name - // This should fail due to name conflict - let new_bases2 = vec![BasePath::new( - 0, - "memory://bucket2".to_string(), - Some("shared_base".to_string()), - false, - )]; - - let result = dataset_clone.add_bases(new_bases2, None).await; - assert!(result.is_err()); - assert!(result - .unwrap_err() - .to_string() - .contains("incompatible with concurrent transaction")); - } - - #[tokio::test] - async fn test_concurrent_add_bases_path_conflict() { - use lance_table::format::BasePath; - use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; - use std::sync::Arc; - - // Create a test dataset - let test_uri = "memory://concurrent_path_conflict_test"; - let mut data_gen = - BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); - - let dataset = Dataset::write( - data_gen.batch(5), - test_uri, - Some(WriteParams { - mode: WriteMode::Create, - ..Default::default() - }), - ) - .await - .unwrap(); - - // Clone the dataset to simulate concurrent access - let dataset_clone = dataset.clone(); - let dataset = Arc::new(dataset); - let dataset_clone = Arc::new(dataset_clone); - - // First transaction adds base with path "memory://shared_path" - let new_bases1 = vec![BasePath::new( - 0, - "memory://shared_path".to_string(), - Some("base1".to_string()), - false, - )]; - - let _updated_dataset = dataset.add_bases(new_bases1, None).await.unwrap(); - - // Second transaction tries to add a different base with same path - // This should fail due to path conflict - let new_bases2 = vec![BasePath::new( - 0, - "memory://shared_path".to_string(), - Some("base2".to_string()), - false, - )]; - - let result = dataset_clone.add_bases(new_bases2, None).await; - assert!(result.is_err()); - assert!(result - .unwrap_err() - .to_string() - .contains("incompatible with concurrent transaction")); - } - - #[tokio::test] - async fn test_concurrent_add_bases_with_data_write() { - use lance_table::format::BasePath; - use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; - use std::sync::Arc; - - // Create a test dataset - let test_uri = "memory://concurrent_write_test"; - let mut data_gen = - BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); - - let dataset = Dataset::write( - data_gen.batch(5), - test_uri, - Some(WriteParams { - mode: WriteMode::Create, - ..Default::default() - }), - ) - .await - .unwrap(); - - // Clone the dataset to simulate concurrent access - let dataset_clone = dataset.clone(); - let dataset = Arc::new(dataset); - - // First transaction adds a new base - let new_bases = vec![BasePath::new( - 0, - "memory://bucket1".to_string(), - Some("base1".to_string()), - false, - )]; - - let updated_dataset = dataset.add_bases(new_bases, None).await.unwrap(); - - // Concurrent transaction appends data - // This should succeed as add_bases doesn't conflict with data writes - let result = Dataset::write( - data_gen.batch(5), - WriteDestination::Dataset(Arc::new(dataset_clone)), - Some(WriteParams { - mode: WriteMode::Append, - ..Default::default() - }), - ) - .await; - - assert!(result.is_ok()); - - // Verify both operations are reflected - let mut final_dataset = updated_dataset; - final_dataset.checkout_latest().await.unwrap(); - - // Should have the new base - assert_eq!(final_dataset.manifest.base_paths.len(), 1); - assert!(final_dataset - .manifest - .base_paths - .values() - .any(|bp| bp.name == Some("base1".to_string()))); - - // Should have both data writes (10 rows total) - assert_eq!(final_dataset.count_rows(None).await.unwrap(), 10); - } - - #[tokio::test] - async fn test_auto_infer_lance_tokenizer() { - let (mut dataset, json_col) = prepare_json_dataset().await; - - // Create inverted index for json col. Expect auto-infer 'json' for lance tokenizer. - dataset - .create_index( - &[&json_col], - IndexType::Inverted, - None, - &InvertedIndexParams::default(), - true, - ) - .await - .unwrap(); - - // Match query succeed only when lance tokenizer is 'json' - let query = FullTextSearchQuery { - query: FtsQuery::Match( - MatchQuery::new("Content,str,once".to_string()).with_column(Some(json_col.clone())), - ), - limit: None, - wand_factor: None, - }; - let batch = dataset - .scan() - .full_text_search(query) - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(1, batch.num_rows()); - } - - #[tokio::test] - async fn test_geo_types() { - use geo_types::{coord, line_string, Rect}; - use geoarrow_array::{ - builder::{LineStringBuilder, PointBuilder, PolygonBuilder}, - GeoArrowArray, - }; - use geoarrow_schema::{Dimension, LineStringType, PointType, PolygonType}; - - // 1. Creates arrow table with spatial data. - let point_type = PointType::new(Dimension::XY, Default::default()); - let line_string_type = LineStringType::new(Dimension::XY, Default::default()); - let polygon_type = PolygonType::new(Dimension::XY, Default::default()); - - let schema = arrow_schema::Schema::new(vec![ - point_type.clone().to_field("point", true), - line_string_type.clone().to_field("linestring", true), - polygon_type.clone().to_field("polygon", true), - ]); - let schema = Arc::new(schema) as arrow_schema::SchemaRef; - - let mut point_builder = PointBuilder::new(point_type.clone()); - point_builder.push_point(Some(&geo_types::point!(x: -72.1235, y: 42.3521))); - let point_arr = point_builder.finish(); - - let mut line_string_builder = LineStringBuilder::new(line_string_type.clone()); - line_string_builder - .push_line_string(Some(&line_string![ - (x: -72.1260, y: 42.45), - (x: -72.123, y: 42.1546), - (x: -73.123, y: 43.1546), - ])) - .unwrap(); - let line_arr = line_string_builder.finish(); - - let mut polygon_builder = PolygonBuilder::new(polygon_type.clone()); - let rect = Rect::new( - coord! { x: -72.123, y: 42.146 }, - coord! { x: -72.126, y: 42.45 }, - ); - polygon_builder.push_rect(Some(&rect)).unwrap(); - let polygon_arr = polygon_builder.finish(); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - point_arr.to_array_ref(), - line_arr.to_array_ref(), - polygon_arr.to_array_ref(), - ], - ) - .unwrap(); - - // 2. Write to lance - let lance_path = TempStrDir::default(); - let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); - let dataset = Dataset::write(reader, &lance_path, Some(Default::default())) - .await - .unwrap(); - - // 3. Verifies that the schema fields and extension metadata are preserved - assert_eq!(dataset.schema().fields.len(), 3); - let fields = &dataset.schema().fields; - assert_eq!( - fields.first().unwrap().metadata.get("ARROW:extension:name"), - Some(&"geoarrow.point".to_owned()) - ); - assert_eq!( - fields.get(1).unwrap().metadata.get("ARROW:extension:name"), - Some(&"geoarrow.linestring".to_owned()) - ); - assert_eq!( - fields.get(2).unwrap().metadata.get("ARROW:extension:name"), - Some(&"geoarrow.polygon".to_owned()) - ); - } - - #[tokio::test] - async fn test_geo_sql() { - use arrow_array::types::Float64Type; - use geo_types::line_string; - use geoarrow_array::{ - builder::{LineStringBuilder, PointBuilder}, - GeoArrowArray, - }; - use geoarrow_schema::{Dimension, LineStringType, PointType}; - - // 1. Creates arrow table with point and linestring spatial data - let point_type = PointType::new(Dimension::XY, Default::default()); - let line_string_type = LineStringType::new(Dimension::XY, Default::default()); - - let schema = arrow_schema::Schema::new(vec![ - point_type.clone().to_field("point", true), - line_string_type.clone().to_field("linestring", true), - ]); - let schema = Arc::new(schema) as arrow_schema::SchemaRef; - - let mut point_builder = PointBuilder::new(point_type.clone()); - point_builder.push_point(Some(&geo_types::point!(x: -72.1235, y: 42.3521))); - let point_arr = point_builder.finish(); - - let mut line_string_builder = LineStringBuilder::new(line_string_type.clone()); - line_string_builder - .push_line_string(Some(&line_string![ - (x: -72.1260, y: 42.45), - (x: -72.123, y: 42.1546), - (x: -73.123, y: 43.1546), - ])) - .unwrap(); - let line_arr = line_string_builder.finish(); - - let batch = RecordBatch::try_new( - schema.clone(), - vec![point_arr.to_array_ref(), line_arr.to_array_ref()], - ) - .unwrap(); - - // 2. Write to lance - let lance_path = TempStrDir::default(); - let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); - let dataset = Dataset::write(reader, &lance_path, Some(Default::default())) - .await - .unwrap(); - - // 3. Executes a SQL query with St_Distance function - let batches = execute_sql( - "SELECT ST_Distance(point, linestring) AS dist FROM dataset", - "dataset".to_owned(), - Arc::new(dataset.clone()), - ) - .await - .unwrap(); - assert_eq!(batches.len(), 1); - let batch = batches.first().unwrap(); - assert_eq!(batch.num_columns(), 1); - assert_eq!(batch.num_rows(), 1); - approx::assert_relative_eq!( - batch.column(0).as_primitive::().value(0), - 0.0015056772638228177 - ); - } -} +mod tests; diff --git a/rust/lance/src/dataset/tests/dataset_common.rs b/rust/lance/src/dataset/tests/dataset_common.rs new file mode 100644 index 00000000000..20adaf5dff8 --- /dev/null +++ b/rust/lance/src/dataset/tests/dataset_common.rs @@ -0,0 +1,186 @@ +#![allow(clippy::redundant_pub_crate)] +pub(crate) use std::collections::{HashMap, HashSet}; +pub(crate) use std::sync::Arc; +pub(crate) use std::vec; + +pub(crate) use crate::dataset::builder::DatasetBuilder; +pub(crate) use crate::dataset::tests::dataset_migrations::scan_dataset; +pub(crate) use crate::dataset::tests::dataset_transactions::{assert_results, execute_sql}; +pub(crate) use crate::dataset::{ + AutoCleanupParams, + ManifestWriteConfig, + ProjectionRequest, + write_manifest_file, +}; +pub(crate) use crate::{Dataset, Error, Result}; +pub(crate) use crate::session::Session; +pub(crate) use crate::dataset::optimize::{compact_files, CompactionOptions}; +pub(crate) use crate::dataset::transaction::{DataReplacementGroup, Operation, Transaction}; +pub(crate) use crate::dataset::WriteMode::Overwrite; +pub(crate) use crate::dataset::ROW_ID; +pub(crate) use crate::datatypes::Schema; +pub(crate) use crate::io::ObjectStoreParams; +pub(crate) use lance_core::ROW_ADDR; +pub(crate) use lance_table::format::{DataStorageFormat, IndexMetadata}; +pub(crate) use lance_table::io::commit::ManifestNamingScheme; +pub(crate) use crate::dataset::WriteDestination; +pub(crate) use crate::dataset::UpdateBuilder; +pub(crate) use crate::index::vector::VectorIndexParams; +pub(crate) use crate::utils::test::copy_test_data_to_tmp; +pub(crate) use lance_arrow::FixedSizeListArrayExt; +pub(crate) use mock_instant::thread_local::MockClock; + +pub(crate) use crate::dataset::write::{CommitBuilder, InsertBuilder, WriteMode, WriteParams}; +pub(crate) use arrow::array::{as_struct_array, AsArray, GenericListBuilder, GenericStringBuilder}; +pub(crate) use arrow::compute::concat_batches; +pub(crate) use arrow::datatypes::UInt64Type; +pub(crate) use arrow_array::{ + builder::StringDictionaryBuilder, + cast::as_string_array, + types::{Float32Type, Int32Type}, + ArrayRef, DictionaryArray, Float32Array, Int32Array, Int64Array, Int8Array, + Int8DictionaryArray, ListArray, RecordBatchIterator, StringArray, UInt16Array, UInt32Array, +}; +pub(crate) use arrow_array::{ + Array, FixedSizeListArray, GenericStringArray, Int16Array, Int16DictionaryArray, + LargeBinaryArray, StructArray, UInt64Array, +}; +pub(crate) use arrow_array::RecordBatch; +pub(crate) use arrow_array::RecordBatchReader; +pub(crate) use arrow_ord::sort::sort_to_indices; +pub(crate) use arrow_schema::{ + DataType, Field as ArrowField, Field, Fields as ArrowFields, Schema as ArrowSchema, +}; +pub(crate) use lance_arrow::bfloat16::{self, BFLOAT16_EXT_NAME}; +pub(crate) use lance_arrow::{ARROW_EXT_META_KEY, ARROW_EXT_NAME_KEY, BLOB_META_KEY}; +pub(crate) use lance_core::utils::tempfile::{TempDir, TempStdDir, TempStrDir}; +pub(crate) use lance_datagen::{array, gen_batch, BatchCount, Dimension, RowCount}; +pub(crate) use lance_file::version::LanceFileVersion; +pub(crate) use lance_file::writer::FileWriter; +pub(crate) use lance_index::scalar::inverted::{ + query::{BooleanQuery, MatchQuery, Occur, Operator, PhraseQuery}, + tokenizer::InvertedIndexParams, +}; +pub(crate) use lance_index::DatasetIndexExt; +pub(crate) use lance_index::scalar::FullTextSearchQuery; +pub(crate) use lance_index::{scalar::ScalarIndexParams, vector::DIST_COL, IndexType}; +pub(crate) use lance_io::assert_io_eq; +pub(crate) use lance_io::utils::CachedFileSize; +pub(crate) use lance_linalg::distance::MetricType; +pub(crate) use lance_table::feature_flags; +pub(crate) use lance_table::format::{DataFile, WriterVersion}; + +pub(crate) use crate::datafusion::LanceTableProvider; +pub(crate) use crate::dataset::refs::branch_contents_path; +pub(crate) use datafusion::common::{assert_contains, assert_not_contains}; +pub(crate) use datafusion::prelude::SessionContext; +pub(crate) use lance_arrow::json::ARROW_JSON_EXT_NAME; +pub(crate) use lance_datafusion::datagen::DatafusionDatagenExt; +pub(crate) use lance_datafusion::udf::register_functions; +pub(crate) use lance_index::scalar::inverted::query::{FtsQuery, MultiMatchQuery}; +pub(crate) use lance_testing::datagen::generate_random_array; +pub(crate) use itertools::Itertools; +pub(crate) use rand::seq::SliceRandom; +pub(crate) use rand::Rng; +pub(crate) use rstest::rstest; +pub(crate) use futures::{StreamExt, TryStreamExt}; +pub(crate) use std::cmp::Ordering; +pub(crate) use object_store::path::Path; +pub(crate) use lance_table::io::manifest::read_manifest; + +// Used to validate that futures returned are Send. +pub(crate) fn require_send(t: T) -> T { + t +} + +pub(crate) async fn create_file( + path: &std::path::Path, + mode: WriteMode, + data_storage_version: LanceFileVersion, +) { + let fields = vec![ + ArrowField::new("i", DataType::Int32, false), + ArrowField::new( + "dict", + DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)), + false, + ), + ]; + let schema = Arc::new(ArrowSchema::new(fields)); + let dict_values = StringArray::from_iter_values(["a", "b", "c", "d", "e"]); + let batches: Vec = (0..20) + .map(|i| { + let mut arrays = + vec![Arc::new(Int32Array::from_iter_values(i * 20..(i + 1) * 20)) as ArrayRef]; + arrays.push(Arc::new( + DictionaryArray::try_new( + UInt16Array::from_iter_values((0_u16..20_u16).map(|v| v % 5)), + Arc::new(dict_values.clone()), + ) + .unwrap(), + )); + RecordBatch::try_new(schema.clone(), arrays).unwrap() + }) + .collect(); + let expected_batches = batches.clone(); + + let test_uri = path.to_str().unwrap(); + let write_params = WriteParams { + max_rows_per_file: 40, + max_rows_per_group: 10, + mode, + data_storage_version: Some(data_storage_version), + ..WriteParams::default() + }; + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + Dataset::write(reader, test_uri, Some(write_params)) + .await + .unwrap(); + + let actual_ds = Dataset::open(test_uri).await.unwrap(); + assert_eq!(actual_ds.version().version, 1); + assert_eq!( + actual_ds.manifest.writer_version, + Some(WriterVersion::default()) + ); + let actual_schema = ArrowSchema::from(actual_ds.schema()); + assert_eq!(&actual_schema, schema.as_ref()); + + let actual_batches = actual_ds + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + + // The batch size batches the group size. + // (the v2 writer has no concept of group size) + if data_storage_version == LanceFileVersion::Legacy { + for batch in &actual_batches { + assert_eq!(batch.num_rows(), 10); + } + } + + // sort + let actual_batch = concat_batches(&schema, &actual_batches).unwrap(); + let idx_arr = actual_batch.column_by_name("i").unwrap(); + let sorted_indices = sort_to_indices(idx_arr, None, None).unwrap(); + let struct_arr: StructArray = actual_batch.into(); + let sorted_arr = arrow_select::take::take(&struct_arr, &sorted_indices, None).unwrap(); + + let expected_struct_arr: StructArray = + concat_batches(&schema, &expected_batches).unwrap().into(); + assert_eq!(&expected_struct_arr, as_struct_array(sorted_arr.as_ref())); + + // Each fragments has different fragment ID + assert_eq!( + actual_ds + .fragments() + .iter() + .map(|f| f.id) + .collect::>(), + (0..10).collect::>() + ) +} diff --git a/rust/lance/src/dataset/tests/dataset_concurrency_store.rs b/rust/lance/src/dataset/tests/dataset_concurrency_store.rs new file mode 100644 index 00000000000..182f9d0e4e3 --- /dev/null +++ b/rust/lance/src/dataset/tests/dataset_concurrency_store.rs @@ -0,0 +1,508 @@ +use super::dataset_common::*; + +#[tokio::test] +async fn concurrent_create() { + async fn write(uri: &str) -> Result<()> { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "a", + DataType::Int32, + false, + )])); + let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); + Dataset::write(empty_reader, uri, None).await?; + Ok(()) + } + + for _ in 0..5 { + let test_uri = TempStrDir::default(); + + let (res1, res2) = tokio::join!(write(&test_uri), write(&test_uri)); + + assert!(res1.is_ok() || res2.is_ok()); + if res1.is_err() { + assert!( + matches!(res1, Err(Error::DatasetAlreadyExists { .. })), + "{:?}", + res1 + ); + } else if res2.is_err() { + assert!( + matches!(res2, Err(Error::DatasetAlreadyExists { .. })), + "{:?}", + res2 + ); + } else { + assert!(res1.is_ok() && res2.is_ok()); + } + } +} + +#[tokio::test] +async fn test_limit_pushdown_in_physical_plan() -> Result<()> { + use tempfile::tempdir; + let temp_dir = tempdir()?; + + let dataset_path = temp_dir.path().join("limit_pushdown_dataset"); + let values: Vec = (0..1000).collect(); + let array = Int32Array::from(values); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "value", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)])?; + + let write_params = WriteParams { + mode: WriteMode::Create, + max_rows_per_file: 100, + ..Default::default() + }; + + let batch_reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + Dataset::write( + batch_reader, + dataset_path.to_str().unwrap(), + Some(write_params), + ) + .await?; + + let mut dataset = Dataset::open(dataset_path.to_str().unwrap()).await?; + + dataset + .create_index( + &["value"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + false, + ) + .await?; + + // Test 1: No filter with limit + { + let mut scanner = dataset.scan(); + scanner.limit(Some(100), None)?; + let plan = scanner.explain_plan(true).await?; + + assert!(plan.contains("range_before=Some(0..100)")); + assert!(plan.contains("range_after=None")); + + let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(100, total_rows); + } + + // Test 2: Indexed filter with limit + { + let mut scanner = dataset.scan(); + scanner.filter("value >= 500")?.limit(Some(50), None)?; + let plan = scanner.explain_plan(true).await?; + + assert!(plan.contains("range_after=Some(0..50)")); + assert!(plan.contains("range_before=None")); + + let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(50, total_rows); + } + + // Test 3: Offset + Limit + { + let mut scanner = dataset.scan(); + scanner.filter("value < 500")?.limit(Some(30), Some(20))?; + let plan = scanner.explain_plan(true).await?; + + assert!(plan.contains("GlobalLimitExec: skip=20, fetch=30")); + assert!(plan.contains("range_after=Some(0..50)")); + + let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(30, total_rows); + + // Verify exact values (should be 20..50) + let all_values: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + .collect::>() + }) + .collect(); + assert_eq!(all_values, (20..50).collect::>()); + } + + // Test 4: Large limit exceeding data + { + let mut scanner = dataset.scan(); + scanner.limit(Some(5000), None)?; + let plan = scanner.explain_plan(true).await?; + + assert!(plan.contains("range_before=Some(0..1000)")); + + let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(1000, total_rows); + } + + // Test 5: Cross-fragment filter with limit + { + let mut scanner = dataset.scan(); + scanner + .filter("value >= 95 AND value <= 205")? + .limit(Some(50), None)?; + let plan = scanner.explain_plan(true).await?; + + assert!(plan.contains("range_after=Some(0..50)")); + + let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(50, total_rows); + } + + Ok(()) +} + +#[tokio::test] +async fn test_add_bases() { + use lance_table::format::BasePath; + use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; + use std::sync::Arc; + + // Create a test dataset + let test_uri = "memory://add_bases_test"; + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + + let dataset = Dataset::write( + data_gen.batch(5), + test_uri, + Some(WriteParams { + mode: WriteMode::Create, + ..Default::default() + }), + ) + .await + .unwrap(); + + let dataset = Arc::new(dataset); + + // Test adding new base paths + let new_bases = vec![ + BasePath::new( + 0, + "memory://bucket1".to_string(), + Some("bucket1".to_string()), + false, + ), + BasePath::new( + 0, + "memory://bucket2".to_string(), + Some("bucket2".to_string()), + true, + ), + ]; + + let updated_dataset = dataset.add_bases(new_bases, None).await.unwrap(); + + // Verify the base paths were added + assert_eq!(updated_dataset.manifest.base_paths.len(), 2); + + let bucket1 = updated_dataset + .manifest + .base_paths + .values() + .find(|bp| bp.name == Some("bucket1".to_string())) + .expect("bucket1 not found"); + let bucket2 = updated_dataset + .manifest + .base_paths + .values() + .find(|bp| bp.name == Some("bucket2".to_string())) + .expect("bucket2 not found"); + + assert_eq!(bucket1.path, "memory://bucket1"); + assert!(!bucket1.is_dataset_root); + assert_eq!(bucket2.path, "memory://bucket2"); + assert!(bucket2.is_dataset_root); + + let updated_dataset = Arc::new(updated_dataset); + + // Test conflict detection - try to add a base with the same name + let conflicting_bases = vec![BasePath::new( + 0, + "memory://bucket3".to_string(), + Some("bucket1".to_string()), + false, + )]; + + let result = updated_dataset.add_bases(conflicting_bases, None).await; + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("Conflict detected")); + + // Test conflict detection - try to add a base with the same path + let conflicting_bases = vec![BasePath::new( + 0, + "memory://bucket1".to_string(), + Some("bucket3".to_string()), + false, + )]; + + let result = updated_dataset.add_bases(conflicting_bases, None).await; + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("Conflict detected")); +} + +#[tokio::test] +async fn test_concurrent_add_bases_conflict() { + use lance_table::format::BasePath; + use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; + use std::sync::Arc; + + // Create a test dataset + let test_uri = "memory://concurrent_add_bases_test"; + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + + let dataset = Dataset::write( + data_gen.batch(5), + test_uri, + Some(WriteParams { + mode: WriteMode::Create, + ..Default::default() + }), + ) + .await + .unwrap(); + + // Clone the dataset to simulate concurrent access + let dataset = Arc::new(dataset); + let dataset_clone = Arc::new(dataset.clone()); + + // First transaction adds base1 + let new_bases1 = vec![BasePath::new( + 0, + "memory://bucket1".to_string(), + Some("base1".to_string()), + false, + )]; + + let updated_dataset = dataset.add_bases(new_bases1, None).await.unwrap(); + + // Second transaction tries to add a different base (base2) + // This should succeed as there's no conflict + let new_bases2 = vec![BasePath::new( + 0, + "memory://bucket2".to_string(), + Some("base2".to_string()), + false, + )]; + + let result = dataset_clone.add_bases(new_bases2, None).await; + assert!(result.is_ok()); + + // Verify both bases are present after conflict resolution + let mut final_dataset = updated_dataset; + final_dataset.checkout_latest().await.unwrap(); + assert_eq!(final_dataset.manifest.base_paths.len(), 2); + + let base1 = final_dataset + .manifest + .base_paths + .values() + .find(|bp| bp.name == Some("base1".to_string())); + let base2 = final_dataset + .manifest + .base_paths + .values() + .find(|bp| bp.name == Some("base2".to_string())); + + assert!(base1.is_some()); + assert!(base2.is_some()); +} + +#[tokio::test] +async fn test_concurrent_add_bases_name_conflict() { + use lance_table::format::BasePath; + use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; + use std::sync::Arc; + + // Create a test dataset + let test_uri = "memory://concurrent_name_conflict_test"; + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + + let dataset = Dataset::write( + data_gen.batch(5), + test_uri, + Some(WriteParams { + mode: WriteMode::Create, + ..Default::default() + }), + ) + .await + .unwrap(); + + // Clone the dataset to simulate concurrent access + let dataset_clone = dataset.clone(); + let dataset = Arc::new(dataset); + let dataset_clone = Arc::new(dataset_clone); + + // First transaction adds base with name "shared_base" + let new_bases1 = vec![BasePath::new( + 0, + "memory://bucket1".to_string(), + Some("shared_base".to_string()), + false, + )]; + + let _updated_dataset = dataset.add_bases(new_bases1, None).await.unwrap(); + + // Second transaction tries to add a different base with same name + // This should fail due to name conflict + let new_bases2 = vec![BasePath::new( + 0, + "memory://bucket2".to_string(), + Some("shared_base".to_string()), + false, + )]; + + let result = dataset_clone.add_bases(new_bases2, None).await; + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("incompatible with concurrent transaction")); +} + +#[tokio::test] +async fn test_concurrent_add_bases_path_conflict() { + use lance_table::format::BasePath; + use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; + use std::sync::Arc; + + // Create a test dataset + let test_uri = "memory://concurrent_path_conflict_test"; + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + + let dataset = Dataset::write( + data_gen.batch(5), + test_uri, + Some(WriteParams { + mode: WriteMode::Create, + ..Default::default() + }), + ) + .await + .unwrap(); + + // Clone the dataset to simulate concurrent access + let dataset_clone = dataset.clone(); + let dataset = Arc::new(dataset); + let dataset_clone = Arc::new(dataset_clone); + + // First transaction adds base with path "memory://shared_path" + let new_bases1 = vec![BasePath::new( + 0, + "memory://shared_path".to_string(), + Some("base1".to_string()), + false, + )]; + + let _updated_dataset = dataset.add_bases(new_bases1, None).await.unwrap(); + + // Second transaction tries to add a different base with same path + // This should fail due to path conflict + let new_bases2 = vec![BasePath::new( + 0, + "memory://shared_path".to_string(), + Some("base2".to_string()), + false, + )]; + + let result = dataset_clone.add_bases(new_bases2, None).await; + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("incompatible with concurrent transaction")); +} + +#[tokio::test] +async fn test_concurrent_add_bases_with_data_write() { + use lance_table::format::BasePath; + use lance_testing::datagen::{BatchGenerator, IncrementingInt32}; + use std::sync::Arc; + + // Create a test dataset + let test_uri = "memory://concurrent_write_test"; + let mut data_gen = + BatchGenerator::new().col(Box::new(IncrementingInt32::new().named("id".to_owned()))); + + let dataset = Dataset::write( + data_gen.batch(5), + test_uri, + Some(WriteParams { + mode: WriteMode::Create, + ..Default::default() + }), + ) + .await + .unwrap(); + + // Clone the dataset to simulate concurrent access + let dataset_clone = dataset.clone(); + let dataset = Arc::new(dataset); + + // First transaction adds a new base + let new_bases = vec![BasePath::new( + 0, + "memory://bucket1".to_string(), + Some("base1".to_string()), + false, + )]; + + let updated_dataset = dataset.add_bases(new_bases, None).await.unwrap(); + + // Concurrent transaction appends data + // This should succeed as add_bases doesn't conflict with data writes + let result = Dataset::write( + data_gen.batch(5), + WriteDestination::Dataset(Arc::new(dataset_clone)), + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await; + + assert!(result.is_ok()); + + // Verify both operations are reflected + let mut final_dataset = updated_dataset; + final_dataset.checkout_latest().await.unwrap(); + + // Should have the new base + assert_eq!(final_dataset.manifest.base_paths.len(), 1); + assert!(final_dataset + .manifest + .base_paths + .values() + .any(|bp| bp.name == Some("base1".to_string()))); + + // Should have both data writes (10 rows total) + assert_eq!(final_dataset.count_rows(None).await.unwrap(), 10); +} diff --git a/rust/lance/src/dataset/tests/dataset_geo.rs b/rust/lance/src/dataset/tests/dataset_geo.rs new file mode 100644 index 00000000000..379fb2ae036 --- /dev/null +++ b/rust/lance/src/dataset/tests/dataset_geo.rs @@ -0,0 +1,143 @@ +use super::dataset_common::*; + +#[tokio::test] +async fn test_geo_types() { + use geo_types::{coord, line_string, Rect}; + use geoarrow_array::{ + builder::{LineStringBuilder, PointBuilder, PolygonBuilder}, + GeoArrowArray, + }; + use geoarrow_schema::{Dimension, LineStringType, PointType, PolygonType}; + + // 1. Creates arrow table with spatial data. + let point_type = PointType::new(Dimension::XY, Default::default()); + let line_string_type = LineStringType::new(Dimension::XY, Default::default()); + let polygon_type = PolygonType::new(Dimension::XY, Default::default()); + + let schema = arrow_schema::Schema::new(vec![ + point_type.clone().to_field("point", true), + line_string_type.clone().to_field("linestring", true), + polygon_type.clone().to_field("polygon", true), + ]); + let schema = Arc::new(schema) as arrow_schema::SchemaRef; + + let mut point_builder = PointBuilder::new(point_type.clone()); + point_builder.push_point(Some(&geo_types::point!(x: -72.1235, y: 42.3521))); + let point_arr = point_builder.finish(); + + let mut line_string_builder = LineStringBuilder::new(line_string_type.clone()); + line_string_builder + .push_line_string(Some(&line_string![ + (x: -72.1260, y: 42.45), + (x: -72.123, y: 42.1546), + (x: -73.123, y: 43.1546), + ])) + .unwrap(); + let line_arr = line_string_builder.finish(); + + let mut polygon_builder = PolygonBuilder::new(polygon_type.clone()); + let rect = Rect::new( + coord! { x: -72.123, y: 42.146 }, + coord! { x: -72.126, y: 42.45 }, + ); + polygon_builder.push_rect(Some(&rect)).unwrap(); + let polygon_arr = polygon_builder.finish(); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + point_arr.to_array_ref(), + line_arr.to_array_ref(), + polygon_arr.to_array_ref(), + ], + ) + .unwrap(); + + // 2. Write to lance + let lance_path = TempStrDir::default(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + let dataset = Dataset::write(reader, &lance_path, Some(Default::default())) + .await + .unwrap(); + + // 3. Verifies that the schema fields and extension metadata are preserved + assert_eq!(dataset.schema().fields.len(), 3); + let fields = &dataset.schema().fields; + assert_eq!( + fields.first().unwrap().metadata.get("ARROW:extension:name"), + Some(&"geoarrow.point".to_owned()) + ); + assert_eq!( + fields.get(1).unwrap().metadata.get("ARROW:extension:name"), + Some(&"geoarrow.linestring".to_owned()) + ); + assert_eq!( + fields.get(2).unwrap().metadata.get("ARROW:extension:name"), + Some(&"geoarrow.polygon".to_owned()) + ); +} + +#[tokio::test] +async fn test_geo_sql() { + use arrow_array::types::Float64Type; + use geo_types::line_string; + use geoarrow_array::{ + builder::{LineStringBuilder, PointBuilder}, + GeoArrowArray, + }; + use geoarrow_schema::{Dimension, LineStringType, PointType}; + + // 1. Creates arrow table with point and linestring spatial data + let point_type = PointType::new(Dimension::XY, Default::default()); + let line_string_type = LineStringType::new(Dimension::XY, Default::default()); + + let schema = arrow_schema::Schema::new(vec![ + point_type.clone().to_field("point", true), + line_string_type.clone().to_field("linestring", true), + ]); + let schema = Arc::new(schema) as arrow_schema::SchemaRef; + + let mut point_builder = PointBuilder::new(point_type.clone()); + point_builder.push_point(Some(&geo_types::point!(x: -72.1235, y: 42.3521))); + let point_arr = point_builder.finish(); + + let mut line_string_builder = LineStringBuilder::new(line_string_type.clone()); + line_string_builder + .push_line_string(Some(&line_string![ + (x: -72.1260, y: 42.45), + (x: -72.123, y: 42.1546), + (x: -73.123, y: 43.1546), + ])) + .unwrap(); + let line_arr = line_string_builder.finish(); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![point_arr.to_array_ref(), line_arr.to_array_ref()], + ) + .unwrap(); + + // 2. Write to lance + let lance_path = TempStrDir::default(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + let dataset = Dataset::write(reader, &lance_path, Some(Default::default())) + .await + .unwrap(); + + // 3. Executes a SQL query with St_Distance function + let batches = execute_sql( + "SELECT ST_Distance(point, linestring) AS dist FROM dataset", + "dataset".to_owned(), + Arc::new(dataset.clone()), + ) + .await + .unwrap(); + assert_eq!(batches.len(), 1); + let batch = batches.first().unwrap(); + assert_eq!(batch.num_columns(), 1); + assert_eq!(batch.num_rows(), 1); + approx::assert_relative_eq!( + batch.column(0).as_primitive::().value(0), + 0.0015056772638228177 + ); +} diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs new file mode 100644 index 00000000000..921cd8360aa --- /dev/null +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -0,0 +1,2419 @@ +use super::dataset_common::*; + +#[rstest] +#[tokio::test] +async fn test_create_index( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + let test_uri = TempStrDir::default(); + + let dimension = 16; + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "embeddings", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + dimension, + ), + false, + )])); + + let float_arr = generate_random_array(512 * dimension as usize); + let vectors = Arc::new( + ::try_new_from_values( + float_arr, dimension, + ) + .unwrap(), + ); + let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors.clone()]).unwrap()]; + + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + + let mut dataset = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await + .unwrap(); + dataset.validate().await.unwrap(); + + // Make sure valid arguments should create index successfully + let params = VectorIndexParams::ivf_pq(10, 8, 2, MetricType::L2, 50); + dataset + .create_index(&["embeddings"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + dataset.validate().await.unwrap(); + + // The version should match the table version it was created from. + let indices = dataset.load_indices().await.unwrap(); + let actual = indices.first().unwrap().dataset_version; + let expected = dataset.manifest.version - 1; + assert_eq!(actual, expected); + let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); + assert_eq!(fragment_bitmap.len(), 1); + assert!(fragment_bitmap.contains(0)); + + // Append should inherit index + let write_params = WriteParams { + mode: WriteMode::Append, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors.clone()]).unwrap()]; + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let dataset = Dataset::write(reader, &test_uri, Some(write_params)) + .await + .unwrap(); + let indices = dataset.load_indices().await.unwrap(); + let actual = indices.first().unwrap().dataset_version; + let expected = dataset.manifest.version - 2; + assert_eq!(actual, expected); + dataset.validate().await.unwrap(); + // Fragment bitmap should show the original fragments, and not include + // the newly appended fragment. + let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); + assert_eq!(fragment_bitmap.len(), 1); + assert!(fragment_bitmap.contains(0)); + + let actual_statistics: serde_json::Value = + serde_json::from_str(&dataset.index_statistics("embeddings_idx").await.unwrap()).unwrap(); + let actual_statistics = actual_statistics.as_object().unwrap(); + assert_eq!(actual_statistics["index_type"].as_str().unwrap(), "IVF_PQ"); + + let deltas = actual_statistics["indices"].as_array().unwrap(); + assert_eq!(deltas.len(), 1); + assert_eq!(deltas[0]["metric_type"].as_str().unwrap(), "l2"); + assert_eq!(deltas[0]["num_partitions"].as_i64().unwrap(), 10); + + assert!(dataset.index_statistics("non-existent_idx").await.is_err()); + assert!(dataset.index_statistics("").await.is_err()); + + // Overwrite should invalidate index + let write_params = WriteParams { + mode: WriteMode::Overwrite, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors]).unwrap()]; + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let dataset = Dataset::write(reader, &test_uri, Some(write_params)) + .await + .unwrap(); + assert!(dataset.manifest.index_section.is_none()); + assert!(dataset.load_indices().await.unwrap().is_empty()); + dataset.validate().await.unwrap(); + + let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); + assert_eq!(fragment_bitmap.len(), 1); + assert!(fragment_bitmap.contains(0)); +} + +#[rstest] +#[tokio::test] +async fn test_create_scalar_index( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, + #[values(false, true)] use_stable_row_id: bool, +) { + let test_uri = TempStrDir::default(); + + let data = gen_batch().col("int", array::step::()); + // Write 64Ki rows. We should get 16 4Ki pages + let mut dataset = Dataset::write( + data.into_reader_rows(RowCount::from(16 * 1024), BatchCount::from(4)), + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + enable_stable_row_ids: use_stable_row_id, + ..Default::default() + }), + ) + .await + .unwrap(); + + let index_name = "my_index".to_string(); + + dataset + .create_index( + &["int"], + IndexType::Scalar, + Some(index_name.clone()), + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + + let indices = dataset.load_indices_by_name(&index_name).await.unwrap(); + + assert_eq!(indices.len(), 1); + assert_eq!(indices[0].dataset_version, 1); + assert_eq!(indices[0].fields, vec![0]); + assert_eq!(indices[0].name, index_name); + + dataset.index_statistics(&index_name).await.unwrap(); +} + +async fn create_bad_file(data_storage_version: LanceFileVersion) -> Result { + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "a.b.c", + DataType::Int32, + false, + )])); + + let batches: Vec = (0..20) + .map(|i| { + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(i * 20..(i + 1) * 20))], + ) + .unwrap() + }) + .collect(); + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + Dataset::write( + reader, + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await +} + +#[tokio::test] +async fn test_create_fts_index_with_empty_table() { + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "text", + DataType::Utf8, + false, + )])); + + let batches: Vec = vec![]; + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let mut dataset = Dataset::write(reader, &test_uri, None) + .await + .expect("write dataset"); + + let params = InvertedIndexParams::default(); + dataset + .create_index(&["text"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let batch = dataset + .scan() + .full_text_search(FullTextSearchQuery::new("lance".to_owned())) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(batch.num_rows(), 0); +} + +#[rstest] +#[tokio::test] +async fn test_create_int8_index( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + use lance_testing::datagen::generate_random_int8_array; + + let test_uri = TempStrDir::default(); + + let dimension = 16; + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "embeddings", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Int8, true)), + dimension, + ), + false, + )])); + + let int8_arr = generate_random_int8_array(512 * dimension as usize); + let vectors = Arc::new( + ::try_new_from_values( + int8_arr, dimension, + ) + .unwrap(), + ); + let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors.clone()]).unwrap()]; + + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + + let mut dataset = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await + .unwrap(); + dataset.validate().await.unwrap(); + + // Make sure valid arguments should create index successfully + let params = VectorIndexParams::ivf_pq(10, 8, 2, MetricType::L2, 50); + dataset + .create_index(&["embeddings"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + dataset.validate().await.unwrap(); + + // The version should match the table version it was created from. + let indices = dataset.load_indices().await.unwrap(); + let actual = indices.first().unwrap().dataset_version; + let expected = dataset.manifest.version - 1; + assert_eq!(actual, expected); + let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); + assert_eq!(fragment_bitmap.len(), 1); + assert!(fragment_bitmap.contains(0)); + + // Append should inherit index + let write_params = WriteParams { + mode: WriteMode::Append, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors.clone()]).unwrap()]; + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let dataset = Dataset::write(reader, &test_uri, Some(write_params)) + .await + .unwrap(); + let indices = dataset.load_indices().await.unwrap(); + let actual = indices.first().unwrap().dataset_version; + let expected = dataset.manifest.version - 2; + assert_eq!(actual, expected); + dataset.validate().await.unwrap(); + // Fragment bitmap should show the original fragments, and not include + // the newly appended fragment. + let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); + assert_eq!(fragment_bitmap.len(), 1); + assert!(fragment_bitmap.contains(0)); + + let actual_statistics: serde_json::Value = + serde_json::from_str(&dataset.index_statistics("embeddings_idx").await.unwrap()).unwrap(); + let actual_statistics = actual_statistics.as_object().unwrap(); + assert_eq!(actual_statistics["index_type"].as_str().unwrap(), "IVF_PQ"); + + let deltas = actual_statistics["indices"].as_array().unwrap(); + assert_eq!(deltas.len(), 1); + assert_eq!(deltas[0]["metric_type"].as_str().unwrap(), "l2"); + assert_eq!(deltas[0]["num_partitions"].as_i64().unwrap(), 10); + + assert!(dataset.index_statistics("non-existent_idx").await.is_err()); + assert!(dataset.index_statistics("").await.is_err()); + + // Overwrite should invalidate index + let write_params = WriteParams { + mode: WriteMode::Overwrite, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let batches = vec![RecordBatch::try_new(schema.clone(), vec![vectors]).unwrap()]; + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let dataset = Dataset::write(reader, &test_uri, Some(write_params)) + .await + .unwrap(); + assert!(dataset.manifest.index_section.is_none()); + assert!(dataset.load_indices().await.unwrap().is_empty()); + dataset.validate().await.unwrap(); + + let fragment_bitmap = indices.first().unwrap().fragment_bitmap.as_ref().unwrap(); + assert_eq!(fragment_bitmap.len(), 1); + assert!(fragment_bitmap.contains(0)); +} + +#[tokio::test] +async fn test_create_fts_index_with_empty_strings() { + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "text", + DataType::Utf8, + false, + )])); + + let batches: Vec = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StringArray::from(vec!["", "", ""]))], + ) + .unwrap()]; + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let mut dataset = Dataset::write(reader, &test_uri, None) + .await + .expect("write dataset"); + + let params = InvertedIndexParams::default(); + dataset + .create_index(&["text"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let batch = dataset + .scan() + .full_text_search(FullTextSearchQuery::new("lance".to_owned())) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(batch.num_rows(), 0); +} + +#[rstest] +#[tokio::test] +async fn test_bad_field_name( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + // don't allow `.` in the field name + assert!(create_bad_file(data_storage_version).await.is_err()); +} + +#[tokio::test] +async fn test_open_dataset_not_found() { + let result = Dataset::open(".").await; + assert!(matches!(result.unwrap_err(), Error::DatasetNotFound { .. })); +} + +#[rstest] +#[tokio::test] +async fn test_search_empty( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + // Create a table + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 128, + ), + false, + )])); + + let test_uri = TempStrDir::default(); + + let vectors = Arc::new( + ::try_new_from_values( + Float32Array::from_iter_values(vec![]), + 128, + ) + .unwrap(), + ); + + let data = RecordBatch::try_new(schema.clone(), vec![vectors]); + let reader = RecordBatchIterator::new(vec![data.unwrap()].into_iter().map(Ok), schema); + let dataset = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await + .unwrap(); + + let mut stream = dataset + .scan() + .nearest( + "vec", + &Float32Array::from_iter_values((0..128).map(|_| 0.1)), + 1, + ) + .unwrap() + .try_into_stream() + .await + .unwrap(); + + while let Some(batch) = stream.next().await { + let schema = batch.unwrap().schema(); + assert_eq!(schema.fields.len(), 2); + assert_eq!( + schema.field_with_name("vec").unwrap(), + &ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 128 + ), + false, + ) + ); + assert_eq!( + schema.field_with_name(DIST_COL).unwrap(), + &ArrowField::new(DIST_COL, DataType::Float32, true) + ); + } +} + +#[rstest] +#[tokio::test] +async fn test_search_empty_after_delete( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, + #[values(false, true)] use_stable_row_id: bool, +) { + // Create a table + let test_uri = TempStrDir::default(); + + let data = gen_batch().col("vec", array::rand_vec::(Dimension::from(32))); + let reader = data.into_reader_rows(RowCount::from(500), BatchCount::from(1)); + let mut dataset = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + enable_stable_row_ids: use_stable_row_id, + ..Default::default() + }), + ) + .await + .unwrap(); + + let params = VectorIndexParams::ivf_pq(1, 8, 1, MetricType::L2, 50); + dataset + .create_index(&["vec"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + dataset.delete("true").await.unwrap(); + + // This behavior will be re-introduced once we work on empty vector index handling. + // https://github.com/lance-format/lance/issues/4034 + // let indices = dataset.load_indices().await.unwrap(); + // // With the new retention behavior, indices are kept even when all fragments are deleted + // // This allows the index configuration to persist through data changes + // assert_eq!(indices.len(), 1); + + // // Verify the index has an empty effective fragment bitmap + // let index = &indices[0]; + // let effective_bitmap = index + // .effective_fragment_bitmap(&dataset.fragment_bitmap) + // .unwrap(); + // assert!(effective_bitmap.is_empty()); + + let mut stream = dataset + .scan() + .nearest( + "vec", + &Float32Array::from_iter_values((0..32).map(|_| 0.1)), + 1, + ) + .unwrap() + .try_into_stream() + .await + .unwrap(); + + while let Some(batch) = stream.next().await { + let schema = batch.unwrap().schema(); + assert_eq!(schema.fields.len(), 2); + assert_eq!( + schema.field_with_name("vec").unwrap(), + &ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 32 + ), + false, + ) + ); + assert_eq!( + schema.field_with_name(DIST_COL).unwrap(), + &ArrowField::new(DIST_COL, DataType::Float32, true) + ); + } + + // predicate with redundant whitespace + dataset.delete(" True").await.unwrap(); + + let mut stream = dataset + .scan() + .nearest( + "vec", + &Float32Array::from_iter_values((0..32).map(|_| 0.1)), + 1, + ) + .unwrap() + .try_into_stream() + .await + .unwrap(); + + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + let schema = batch.schema(); + assert_eq!(schema.fields.len(), 2); + assert_eq!( + schema.field_with_name("vec").unwrap(), + &ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 32 + ), + false, + ) + ); + assert_eq!( + schema.field_with_name(DIST_COL).unwrap(), + &ArrowField::new(DIST_COL, DataType::Float32, true) + ); + assert_eq!(batch.num_rows(), 0, "Expected no results after delete"); + } +} + +#[rstest] +#[tokio::test] +async fn test_num_small_files( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + let test_uri = TempStrDir::default(); + let dimensions = 16; + let column_name = "vec"; + let field = ArrowField::new( + column_name, + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + dimensions, + ), + false, + ); + + let schema = Arc::new(ArrowSchema::new(vec![field])); + + let float_arr = generate_random_array(512 * dimensions as usize); + let vectors = + arrow_array::FixedSizeListArray::try_new_from_values(float_arr, dimensions).unwrap(); + + let record_batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vectors)]).unwrap(); + + let reader = RecordBatchIterator::new(vec![record_batch].into_iter().map(Ok), schema.clone()); + + let dataset = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await + .unwrap(); + dataset.validate().await.unwrap(); + + assert!(dataset.num_small_files(1024).await > 0); + assert!(dataset.num_small_files(512).await == 0); +} + +#[tokio::test] +async fn test_read_struct_of_dictionary_arrays() { + let test_uri = TempStrDir::default(); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "s", + DataType::Struct(ArrowFields::from(vec![ArrowField::new( + "d", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + )])), + true, + )])); + + let mut batches: Vec = Vec::new(); + for _ in 1..2 { + let mut dict_builder = StringDictionaryBuilder::::new(); + dict_builder.append("a").unwrap(); + dict_builder.append("b").unwrap(); + dict_builder.append("c").unwrap(); + dict_builder.append("d").unwrap(); + + let struct_array = Arc::new(StructArray::from(vec![( + Arc::new(ArrowField::new( + "d", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + )), + Arc::new(dict_builder.finish()) as ArrayRef, + )])); + + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![struct_array.clone()]).unwrap(); + batches.push(batch); + } + + let batch_reader = + RecordBatchIterator::new(batches.clone().into_iter().map(Ok), arrow_schema.clone()); + Dataset::write(batch_reader, &test_uri, Some(WriteParams::default())) + .await + .unwrap(); + + let result = scan_dataset(&test_uri).await.unwrap(); + + assert_eq!(batches, result); +} + +#[tokio::test] +async fn test_fts_fuzzy_query() { + let params = InvertedIndexParams::default(); + let text_col = GenericStringArray::::from(vec![ + "fa", "fo", "fob", "focus", "foo", "food", "foul", // # spellchecker:disable-line + ]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "text", + text_col.data_type().to_owned(), + false, + )]) + .into(), + vec![Arc::new(text_col) as ArrayRef], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write(batches, &test_uri, None).await.unwrap(); + dataset + .create_index(&["text"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + let results = dataset + .scan() + .full_text_search(FullTextSearchQuery::new_fuzzy("foo".to_owned(), Some(1))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 4); + let texts = results["text"] + .as_string::() + .iter() + .map(|s| s.unwrap().to_owned()) + .collect::>(); + assert_eq!( + texts, + vec![ + "foo".to_owned(), // 0 edits + "fo".to_owned(), // 1 deletion # spellchecker:disable-line + "fob".to_owned(), // 1 substitution # spellchecker:disable-line + "food".to_owned(), // 1 insertion # spellchecker:disable-line + ] + .into_iter() + .collect() + ); +} + +#[tokio::test] +async fn test_fts_on_multiple_columns() { + let params = InvertedIndexParams::default(); + let title_col = + GenericStringArray::::from(vec!["title common", "title hello", "title lance"]); + let content_col = GenericStringArray::::from(vec![ + "content world", + "content database", + "content common", + ]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("title", title_col.data_type().to_owned(), false), + arrow_schema::Field::new("content", title_col.data_type().to_owned(), false), + ]) + .into(), + vec![ + Arc::new(title_col) as ArrayRef, + Arc::new(content_col) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write(batches, &test_uri, None).await.unwrap(); + dataset + .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + dataset + .create_index(&["content"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let results = dataset + .scan() + .full_text_search(FullTextSearchQuery::new("title".to_owned())) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 3); + + let results = dataset + .scan() + .full_text_search(FullTextSearchQuery::new("content".to_owned())) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 3); + + let results = dataset + .scan() + .full_text_search(FullTextSearchQuery::new("common".to_owned())) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 2); + + let results = dataset + .scan() + .full_text_search( + FullTextSearchQuery::new("common".to_owned()) + .with_column("title".to_owned()) + .unwrap(), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 1); + + let results = dataset + .scan() + .full_text_search( + FullTextSearchQuery::new("common".to_owned()) + .with_column("content".to_owned()) + .unwrap(), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 1); +} + +#[tokio::test] +async fn test_fts_unindexed_data() { + let params = InvertedIndexParams::default(); + let title_col = StringArray::from(vec!["title hello", "title lance", "title common"]); + let content_col = + StringArray::from(vec!["content world", "content database", "content common"]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + Field::new("title", title_col.data_type().to_owned(), false), + Field::new("content", title_col.data_type().to_owned(), false), + ]) + .into(), + vec![ + Arc::new(title_col) as ArrayRef, + Arc::new(content_col) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(batches, "memory://test.lance", None) + .await + .unwrap(); + dataset + .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let results = dataset + .scan() + .full_text_search(FullTextSearchQuery::new("title".to_owned())) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 3); + + // write new data + let title_col = StringArray::from(vec!["new title"]); + let content_col = StringArray::from(vec!["new content"]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + Field::new("title", title_col.data_type().to_owned(), false), + Field::new("content", title_col.data_type().to_owned(), false), + ]) + .into(), + vec![ + Arc::new(title_col) as ArrayRef, + Arc::new(content_col) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + dataset.append(batches, None).await.unwrap(); + + let results = dataset + .scan() + .full_text_search(FullTextSearchQuery::new("title".to_owned())) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 4); + + let results = dataset + .scan() + .full_text_search(FullTextSearchQuery::new("new".to_owned())) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 1); +} + +#[tokio::test] +async fn test_fts_unindexed_data_on_empty_index() { + // Empty dataset with fts index + let params = InvertedIndexParams::default(); + let title_col = StringArray::from(Vec::<&str>::new()); + let content_col = StringArray::from(Vec::<&str>::new()); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + Field::new("title", title_col.data_type().to_owned(), false), + Field::new("content", title_col.data_type().to_owned(), false), + ]) + .into(), + vec![ + Arc::new(title_col) as ArrayRef, + Arc::new(content_col) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(batches, "memory://test.lance", None) + .await + .unwrap(); + dataset + .create_index(&["title"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + // Test fts search + let results = dataset + .scan() + .full_text_search(FullTextSearchQuery::new_query(FtsQuery::Match( + MatchQuery::new("title".to_owned()).with_column(Some("title".to_owned())), + ))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 0); + + // write new data + let title_col = StringArray::from(vec!["title hello", "title lance", "title common"]); + let content_col = + StringArray::from(vec!["content world", "content database", "content common"]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + Field::new("title", title_col.data_type().to_owned(), false), + Field::new("content", title_col.data_type().to_owned(), false), + ]) + .into(), + vec![ + Arc::new(title_col) as ArrayRef, + Arc::new(content_col) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + dataset.append(batches, None).await.unwrap(); + + let results = dataset + .scan() + .full_text_search(FullTextSearchQuery::new_query(FtsQuery::Match( + MatchQuery::new("title".to_owned()).with_column(Some("title".to_owned())), + ))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 3); +} + +#[tokio::test] +async fn test_fts_without_index() { + // create table without index + let title_col = StringArray::from(vec!["title hello", "title lance", "title common"]); + let content_col = + StringArray::from(vec!["content world", "content database", "content common"]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + Field::new("title", title_col.data_type().to_owned(), false), + Field::new("content", title_col.data_type().to_owned(), false), + ]) + .into(), + vec![ + Arc::new(title_col) as ArrayRef, + Arc::new(content_col) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(batches, "memory://test.lance", None) + .await + .unwrap(); + + // match query on title and content + let results = dataset + .scan() + .full_text_search( + FullTextSearchQuery::new("title".to_owned()) + .with_columns(&["title".to_string(), "content".to_string()]) + .unwrap(), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 3); + + // write new data + let title_col = StringArray::from(vec!["new title"]); + let content_col = StringArray::from(vec!["new content"]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + Field::new("title", title_col.data_type().to_owned(), false), + Field::new("content", title_col.data_type().to_owned(), false), + ]) + .into(), + vec![ + Arc::new(title_col) as ArrayRef, + Arc::new(content_col) as ArrayRef, + ], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + dataset.append(batches, None).await.unwrap(); + + // match query on title and content + let results = dataset + .scan() + .full_text_search( + FullTextSearchQuery::new("title".to_owned()) + .with_columns(&["title".to_string(), "content".to_string()]) + .unwrap(), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 4); + + let results = dataset + .scan() + .full_text_search( + FullTextSearchQuery::new("new".to_owned()) + .with_columns(&["title".to_string(), "content".to_string()]) + .unwrap(), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 1); +} + +#[tokio::test] +async fn test_fts_rank() { + let params = InvertedIndexParams::default(); + let text_col = + GenericStringArray::::from(vec!["score", "find score", "try to find score"]); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "text", + text_col.data_type().to_owned(), + false, + )]) + .into(), + vec![Arc::new(text_col) as ArrayRef], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write(batches, &test_uri, None).await.unwrap(); + dataset + .create_index(&["text"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + let results = dataset + .scan() + .with_row_id() + .full_text_search(FullTextSearchQuery::new("score".to_owned())) + .unwrap() + .limit(Some(3), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 3); + let row_ids = results[ROW_ID].as_primitive::().values(); + assert_eq!(row_ids, &[0, 1, 2]); + + let results = dataset + .scan() + .with_row_id() + .full_text_search(FullTextSearchQuery::new("score".to_owned())) + .unwrap() + .limit(Some(2), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 2); + let row_ids = results[ROW_ID].as_primitive::().values(); + assert_eq!(row_ids, &[0, 1]); + + let results = dataset + .scan() + .with_row_id() + .full_text_search(FullTextSearchQuery::new("score".to_owned())) + .unwrap() + .limit(Some(1), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(results.num_rows(), 1); + let row_ids = results[ROW_ID].as_primitive::().values(); + assert_eq!(row_ids, &[0]); +} + +async fn create_fts_dataset< + Offset: arrow::array::OffsetSizeTrait, + ListOffset: arrow::array::OffsetSizeTrait, +>( + is_list: bool, + with_position: bool, + params: InvertedIndexParams, +) -> Dataset { + let tempdir = TempStrDir::default(); + let uri = tempdir.to_owned(); + drop(tempdir); + + let params = params.with_position(with_position); + let doc_col: Arc = if is_list { + let string_builder = GenericStringBuilder::::new(); + let mut list_col = GenericListBuilder::::new(string_builder); + // Create a list of strings + list_col.values().append_value("lance database the search"); // for testing phrase query + list_col.append(true); + list_col.values().append_value("lance database"); // for testing phrase query + list_col.append(true); + list_col.values().append_value("lance search"); + list_col.append(true); + list_col.values().append_value("database"); + list_col.values().append_value("search"); + list_col.append(true); + list_col.values().append_value("unrelated doc"); + list_col.append(true); + list_col.values().append_value("unrelated"); + list_col.append(true); + list_col.values().append_value("mots"); + list_col.values().append_value("accentués"); + list_col.append(true); + list_col + .values() + .append_value("lance database full text search"); + list_col.append(true); + + // for testing null + list_col.append(false); + + Arc::new(list_col.finish()) + } else { + Arc::new(GenericStringArray::::from(vec![ + "lance database the search", + "lance database", + "lance search", + "database search", + "unrelated doc", + "unrelated", + "mots accentués", + "lance database full text search", + ])) + }; + let ids = UInt64Array::from_iter_values(0..doc_col.len() as u64); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("doc", doc_col.data_type().to_owned(), true), + arrow_schema::Field::new("id", DataType::UInt64, false), + ]) + .into(), + vec![Arc::new(doc_col) as ArrayRef, Arc::new(ids) as ArrayRef], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(batches, &uri, None).await.unwrap(); + + dataset + .create_index(&["doc"], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + + dataset +} + +async fn test_fts_index< + Offset: arrow::array::OffsetSizeTrait, + ListOffset: arrow::array::OffsetSizeTrait, +>( + is_list: bool, +) { + let ds = + create_fts_dataset::(is_list, false, InvertedIndexParams::default()) + .await; + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("lance".to_owned()).limit(Some(3))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 3, "{:?}", result); + let ids = result["id"].as_primitive::().values(); + assert!(ids.contains(&0), "{:?}", result); + assert!(ids.contains(&1), "{:?}", result); + assert!(ids.contains(&2), "{:?}", result); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("database".to_owned()).limit(Some(3))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 3); + let ids = result["id"].as_primitive::().values(); + assert!(ids.contains(&0), "{:?}", result); + assert!(ids.contains(&1), "{:?}", result); + assert!(ids.contains(&3), "{:?}", result); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query( + MatchQuery::new("lance database".to_owned()) + .with_operator(Operator::And) + .into(), + ) + .limit(Some(5)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 3, "{:?}", result); + let ids = result["id"].as_primitive::().values(); + assert!(ids.contains(&0), "{:?}", result); + assert!(ids.contains(&1), "{:?}", result); + assert!(ids.contains(&7), "{:?}", result); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("unknown null".to_owned()).limit(Some(3))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 0); + + // test phrase query + // for non-phrasal query, the order of the tokens doesn't matter + // so there should be 4 documents that contain "database" or "lance" + + // we built the index without position, so the phrase query will not work + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("lance database".to_owned()).into()) + .limit(Some(10)), + ) + .unwrap() + .try_into_batch() + .await; + let err = result.unwrap_err().to_string(); + assert!(err.contains("position is not found but required for phrase queries, try recreating the index with position"),"{}",err); + + // recreate the index with position + let ds = + create_fts_dataset::(is_list, true, InvertedIndexParams::default()) + .await; + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("lance database".to_owned()).limit(Some(10))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 5, "{:?}", result); + let ids = result["id"].as_primitive::().values(); + assert!(ids.contains(&0)); + assert!(ids.contains(&1)); + assert!(ids.contains(&2)); + assert!(ids.contains(&3)); + assert!(ids.contains(&7)); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("lance database".to_owned()).into()) + .limit(Some(10)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let ids = result["id"].as_primitive::().values(); + assert_eq!(result.num_rows(), 3, "{:?}", ids); + assert!(ids.contains(&0)); + assert!(ids.contains(&1)); + assert!(ids.contains(&7)); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("database lance".to_owned()).into()) + .limit(Some(10)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 0); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("lance unknown".to_owned()).into()) + .limit(Some(10)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 0); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("unknown null".to_owned()).into()) + .limit(Some(3)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 0); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query(PhraseQuery::new("lance search".to_owned()).into()) + .limit(Some(3)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 1); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query( + PhraseQuery::new("lance search".to_owned()) + .with_slop(2) + .into(), + ) + .limit(Some(3)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 2); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + FullTextSearchQuery::new_query( + PhraseQuery::new("search lance".to_owned()) + .with_slop(2) + .into(), + ) + .limit(Some(3)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 0); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + // must contain "lance" and "database", and may contain "search" + FullTextSearchQuery::new_query( + BooleanQuery::new([ + ( + Occur::Should, + MatchQuery::new("search".to_owned()) + .with_operator(Operator::And) + .into(), + ), + ( + Occur::Must, + MatchQuery::new("lance database".to_owned()) + .with_operator(Operator::And) + .into(), + ), + ]) + .into(), + ) + .limit(Some(3)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 3, "{:?}", result); + let ids = result["id"].as_primitive::().values(); + assert!(ids.contains(&0), "{:?}", result); + assert!(ids.contains(&1), "{:?}", result); + assert!(ids.contains(&7), "{:?}", result); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search( + // must contain "lance" and "database", and may contain "search" + FullTextSearchQuery::new_query( + BooleanQuery::new([ + ( + Occur::Should, + MatchQuery::new("search".to_owned()) + .with_operator(Operator::And) + .into(), + ), + ( + Occur::Must, + MatchQuery::new("lance database".to_owned()) + .with_operator(Operator::And) + .into(), + ), + ( + Occur::MustNot, + MatchQuery::new("full text".to_owned()).into(), + ), + ]) + .into(), + ) + .limit(Some(3)), + ) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 2, "{:?}", result); + let ids = result["id"].as_primitive::().values(); + assert!(ids.contains(&0), "{:?}", result); + assert!(ids.contains(&1), "{:?}", result); +} + +#[tokio::test] +async fn test_fts_index_with_string() { + test_fts_index::(false).await; + test_fts_index::(true).await; + test_fts_index::(true).await; +} + +#[tokio::test] +async fn test_fts_index_with_large_string() { + test_fts_index::(false).await; + test_fts_index::(true).await; + test_fts_index::(true).await; +} + +#[tokio::test] +async fn test_fts_accented_chars() { + let ds = create_fts_dataset::(false, false, InvertedIndexParams::default()).await; + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("accentués".to_owned()).limit(Some(3))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 1); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("accentues".to_owned()).limit(Some(3))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 0); + + // with ascii folding enabled, the search should be accent-insensitive + let ds = create_fts_dataset::( + false, + false, + InvertedIndexParams::default() + .stem(false) + .ascii_folding(true), + ) + .await; + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("accentués".to_owned()).limit(Some(3))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 1); + + let result = ds + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new("accentues".to_owned()).limit(Some(3))) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 1); +} + +#[tokio::test] +async fn test_fts_phrase_query() { + let tmpdir = TempStrDir::default(); + let uri = tmpdir.to_owned(); + drop(tmpdir); + + let words = ["lance", "full", "text", "search"]; + let mut lance_search_count = 0; + let mut full_text_count = 0; + let mut doc_array = (0..4096) + .map(|_| { + let mut rng = rand::rng(); + let mut text = String::with_capacity(512); + let len = rng.random_range(127..512); + for i in 0..len { + if i > 0 { + text.push(' '); + } + text.push_str(words[rng.random_range(0..words.len())]); + } + if text.contains("lance search") { + lance_search_count += 1; + } + if text.contains("full text") { + full_text_count += 1; + } + text + }) + .collect_vec(); + // Ensure at least one doc matches each phrase deterministically + doc_array.push("lance search".to_owned()); + lance_search_count += 1; + doc_array.push("full text".to_owned()); + full_text_count += 1; + doc_array.push("position for phrase query".to_owned()); + + // 1) Build index without positions and assert phrase query errors + let params_no_pos = InvertedIndexParams::default().with_position(false); + let doc_col: Arc = Arc::new(GenericStringArray::::from(doc_array.clone())); + let ids = UInt64Array::from_iter_values(0..doc_col.len() as u64); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("doc", doc_col.data_type().to_owned(), true), + arrow_schema::Field::new("id", DataType::UInt64, false), + ]) + .into(), + vec![Arc::new(doc_col) as ArrayRef, Arc::new(ids) as ArrayRef], + ) + .unwrap(); + let schema = batch.schema(); + let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(batches, &uri, None).await.unwrap(); + dataset + .create_index(&["doc"], IndexType::Inverted, None, ¶ms_no_pos, true) + .await + .unwrap(); + + let err = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query( + PhraseQuery::new("lance search".to_owned()).into(), + )) + .unwrap() + .try_into_batch() + .await + .unwrap_err() + .to_string(); + assert!(err.contains("position is not found but required for phrase queries, try recreating the index with position"), "{}", err); + assert!(err.starts_with("Invalid user input: "), "{}", err); + + // 2) Recreate index with positions and assert phrase query works + let params_with_pos = InvertedIndexParams::default().with_position(true); + dataset + .create_index(&["doc"], IndexType::Inverted, None, ¶ms_with_pos, true) + .await + .unwrap(); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query( + PhraseQuery::new("lance search".to_owned()).into(), + )) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), lance_search_count); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query( + PhraseQuery::new("full text".to_owned()).into(), + )) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), full_text_count); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query( + PhraseQuery::new("phrase query".to_owned()).into(), + )) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 1); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query( + PhraseQuery::new("".to_owned()).into(), + )) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(result.num_rows(), 0); +} + +async fn prepare_json_dataset() -> (Dataset, String) { + let text_col = Arc::new(StringArray::from(vec![ + r#"{ + "Title": "HarryPotter Chapter One", + "Content": "Mr. and Mrs. Dursley, of number four, Privet Drive, were proud to say...", + "Author": "J.K. Rowling", + "Price": 128, + "Language": ["english", "chinese"] + }"#, + r#"{ + "Title": "Fairy Talest", + "Content": "Once upon a time, on a bitterly cold New Year's Eve, a little girl...", + "Author": "ANDERSEN", + "Price": 50, + "Language": ["english", "chinese"] + }"#, + ])); + let json_col = "json_field".to_string(); + + // Prepare dataset + let mut metadata = HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + Field::new(&json_col, DataType::Utf8, false).with_metadata(metadata) + ]) + .into(), + vec![text_col.clone()], + ) + .unwrap(); + let schema = batch.schema(); + let stream = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let dataset = Dataset::write(stream, "memory://test/table", None) + .await + .unwrap(); + + (dataset, json_col) +} + +#[tokio::test] +async fn test_json_inverted_fuzziness_query() { + let (mut dataset, json_col) = prepare_json_dataset().await; + + // Create inverted index for json col + dataset + .create_index( + &[&json_col], + IndexType::Inverted, + None, + &InvertedIndexParams::default().lance_tokenizer("json".to_string()), + true, + ) + .await + .unwrap(); + + // Match query with fuzziness + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Content,str,Dursley".to_string()).with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(1, batch.num_rows()); + + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Content,str,Bursley".to_string()).with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(0, batch.num_rows()); + + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Content,str,Bursley".to_string()) + .with_column(Some(json_col.clone())) + .with_fuzziness(Some(1)), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(1, batch.num_rows()); + + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Content,str,ABursley".to_string()) + .with_column(Some(json_col.clone())) + .with_fuzziness(Some(1)), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(0, batch.num_rows()); + + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Content,str,ABursley".to_string()) + .with_column(Some(json_col.clone())) + .with_fuzziness(Some(2)), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(1, batch.num_rows()); + + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Dontent,str,Bursley".to_string()) + .with_column(Some(json_col.clone())) + .with_fuzziness(Some(2)), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(0, batch.num_rows()); +} + +#[tokio::test] +async fn test_json_inverted_match_query() { + let (mut dataset, json_col) = prepare_json_dataset().await; + + // Create inverted index for json col, with max token len 10 and enable stemming, + // lower case, and remove stop words + dataset + .create_index( + &[&json_col], + IndexType::Inverted, + None, + &InvertedIndexParams::default() + .lance_tokenizer("json".to_string()) + .max_token_length(Some(10)) + .stem(true) + .lower_case(true) + .remove_stop_words(true), + true, + ) + .await + .unwrap(); + + // Match query with token length exceed max token length + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Title,str,harrypotter".to_string()) + .with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(0, batch.num_rows()); + + // Match query with stemming + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Content,str,onc".to_string()).with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(1, batch.num_rows()); + + // Match query with lower case + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Content,str,DURSLEY".to_string()).with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(1, batch.num_rows()); + + // Match query with stop word + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Content,str,and".to_string()).with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(0, batch.num_rows()); +} + +#[tokio::test] +async fn test_json_inverted_flat_match_query() { + let (mut dataset, json_col) = prepare_json_dataset().await; + + // Create inverted index for json col + dataset + .create_index( + &[&json_col], + IndexType::Inverted, + None, + &InvertedIndexParams::default() + .lance_tokenizer("json".to_string()) + .stem(false), + true, + ) + .await + .unwrap(); + + // Append data + let text_col = Arc::new(StringArray::from(vec![ + r#"{ + "Title": "HarryPotter Chapter Two", + "Content": "Nearly ten years had passed since the Dursleys had woken up...", + "Author": "J.K. Rowling", + "Price": 128, + "Language": ["english", "chinese"] + }"#, + ])); + + let mut metadata = HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![ + Field::new(&json_col, DataType::Utf8, false).with_metadata(metadata) + ]) + .into(), + vec![text_col.clone()], + ) + .unwrap(); + let schema = batch.schema(); + let stream = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + dataset.append(stream, None).await.unwrap(); + + // Test match query + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Title,str,harrypotter".to_string()) + .with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(2, batch.num_rows()); +} + +#[tokio::test] +async fn test_json_inverted_phrase_query() { + // Prepare json dataset + let (mut dataset, json_col) = prepare_json_dataset().await; + + // Create inverted index for json col + dataset + .create_index( + &[&json_col], + IndexType::Inverted, + None, + &InvertedIndexParams::default() + .lance_tokenizer("json".to_string()) + .stem(false) + .with_position(true), + true, + ) + .await + .unwrap(); + + // Test phrase query + let query = FullTextSearchQuery { + query: FtsQuery::Phrase( + PhraseQuery::new("Title,str,harrypotter one chapter".to_string()) + .with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(0, batch.num_rows()); + + let query = FullTextSearchQuery { + query: FtsQuery::Phrase( + PhraseQuery::new("Title,str,harrypotter chapter one".to_string()) + .with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(1, batch.num_rows()); +} + +#[tokio::test] +async fn test_json_inverted_multimatch_query() { + // Prepare json dataset + let (mut dataset, json_col) = prepare_json_dataset().await; + + // Create inverted index for json col + dataset + .create_index( + &[&json_col], + IndexType::Inverted, + None, + &InvertedIndexParams::default() + .lance_tokenizer("json".to_string()) + .stem(false), + true, + ) + .await + .unwrap(); + + // Test multi match query + let query = FullTextSearchQuery { + query: FtsQuery::MultiMatch(MultiMatchQuery { + match_queries: vec![ + MatchQuery::new("Title,str,harrypotter".to_string()) + .with_column(Some(json_col.clone())), + MatchQuery::new("Language,str,english".to_string()) + .with_column(Some(json_col.clone())), + ], + }), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(2, batch.num_rows()); +} + +#[tokio::test] +async fn test_json_inverted_boolean_query() { + // Prepare json dataset + let (mut dataset, json_col) = prepare_json_dataset().await; + + // Create inverted index for json col + dataset + .create_index( + &[&json_col], + IndexType::Inverted, + None, + &InvertedIndexParams::default() + .lance_tokenizer("json".to_string()) + .stem(false), + true, + ) + .await + .unwrap(); + + // Test boolean query + let query = FullTextSearchQuery { + query: FtsQuery::Boolean(BooleanQuery { + should: vec![], + must: vec![ + FtsQuery::Match( + MatchQuery::new("Language,str,english".to_string()) + .with_column(Some(json_col.clone())), + ), + FtsQuery::Match( + MatchQuery::new("Title,str,harrypotter".to_string()) + .with_column(Some(json_col.clone())), + ), + ], + must_not: vec![], + }), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(1, batch.num_rows()); +} + +#[tokio::test] +async fn test_sql_contains_tokens() { + let text_col = Arc::new(StringArray::from(vec![ + "a cat catch a fish", + "a fish catch a cat", + "a white cat catch a big fish", + "cat catchup fish", + "cat fish catch", + ])); + + // Prepare dataset + let batch = RecordBatch::try_new( + arrow_schema::Schema::new(vec![Field::new("text", DataType::Utf8, false)]).into(), + vec![text_col.clone()], + ) + .unwrap(); + let schema = batch.schema(); + let stream = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); + let mut dataset = Dataset::write(stream, "memory://test/table", None) + .await + .unwrap(); + + // Test without fts index + let results = execute_sql( + "select * from foo where contains_tokens(text, 'cat catch fish')", + "foo".to_string(), + Arc::new(dataset.clone()), + ) + .await + .unwrap(); + + assert_results( + results, + &StringArray::from(vec![ + "a cat catch a fish", + "a fish catch a cat", + "a white cat catch a big fish", + "cat fish catch", + ]), + ); + + // Verify plan, should not contain ScalarIndexQuery. + let results = execute_sql( + "explain select * from foo where contains_tokens(text, 'cat catch fish')", + "foo".to_string(), + Arc::new(dataset.clone()), + ) + .await + .unwrap(); + let plan = format!("{:?}", results); + assert_not_contains!(&plan, "ScalarIndexQuery"); + + // Test with unsuitable fts index + dataset + .create_index( + &["text"], + IndexType::Inverted, + None, + &InvertedIndexParams::default().base_tokenizer("raw".to_string()), + true, + ) + .await + .unwrap(); + + let results = execute_sql( + "select * from foo where contains_tokens(text, 'cat catch fish')", + "foo".to_string(), + Arc::new(dataset.clone()), + ) + .await + .unwrap(); + + assert_results( + results, + &StringArray::from(vec![ + "a cat catch a fish", + "a fish catch a cat", + "a white cat catch a big fish", + "cat fish catch", + ]), + ); + + // Verify plan, should not contain ScalarIndexQuery because fts index is not unsuitable. + let results = execute_sql( + "explain select * from foo where contains_tokens(text, 'cat catch fish')", + "foo".to_string(), + Arc::new(dataset.clone()), + ) + .await + .unwrap(); + let plan = format!("{:?}", results); + assert_not_contains!(&plan, "ScalarIndexQuery"); + + // Test with suitable fts index + dataset + .create_index( + &["text"], + IndexType::Inverted, + None, + &InvertedIndexParams::default() + .max_token_length(None) + .stem(false), + true, + ) + .await + .unwrap(); + + let results = execute_sql( + "select * from foo where contains_tokens(text, 'cat catch fish')", + "foo".to_string(), + Arc::new(dataset.clone()), + ) + .await + .unwrap(); + + assert_results( + results, + &StringArray::from(vec![ + "a cat catch a fish", + "a fish catch a cat", + "a white cat catch a big fish", + "cat fish catch", + ]), + ); + + // Verify plan, should contain ScalarIndexQuery. + let results = execute_sql( + "explain select * from foo where contains_tokens(text, 'cat catch fish')", + "foo".to_string(), + Arc::new(dataset.clone()), + ) + .await + .unwrap(); + let plan = format!("{:?}", results); + assert_contains!(&plan, "ScalarIndexQuery"); +} + +#[tokio::test] +async fn test_index_take_batch_size() -> Result<()> { + use tempfile::tempdir; + let temp_dir = tempdir()?; + + let dataset_path = temp_dir.path().join("ints_dataset"); + let values: Vec = (0..1024).collect(); + let array = Int32Array::from(values); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "ints", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)])?; + let write_params = WriteParams { + mode: WriteMode::Create, + max_rows_per_file: 100, + ..Default::default() + }; + let batch_reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + Dataset::write( + batch_reader, + dataset_path.to_str().unwrap(), + Some(write_params), + ) + .await?; + let mut dataset = Dataset::open(dataset_path.to_str().unwrap()).await?; + dataset + .create_index( + &["ints"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + false, + ) + .await?; + + let mut scanner = dataset.scan(); + scanner.batch_size(50).filter("ints > 0")?.with_row_id(); + let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(1023, total_rows); + assert_eq!(21, batches.len()); + + let mut scanner = dataset.scan(); + scanner + .batch_size(50) + .filter("ints > 0")? + .limit(Some(1024), None)? + .with_row_id(); + let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(1023, total_rows); + assert_eq!(21, batches.len()); + + let dataset_path2 = temp_dir.path().join("strings_dataset"); + let strings: Vec = (0..1024).map(|i| format!("string-{}", i)).collect(); + let string_array = StringArray::from(strings); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "strings", + DataType::Utf8, + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(string_array)])?; + let write_params = WriteParams { + mode: WriteMode::Create, + max_rows_per_file: 100, + ..Default::default() + }; + let batch_reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + Dataset::write( + batch_reader, + dataset_path2.to_str().unwrap(), + Some(write_params), + ) + .await?; + let mut dataset2 = Dataset::open(dataset_path2.to_str().unwrap()).await?; + dataset2 + .create_index( + &["strings"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + false, + ) + .await?; + + let mut scanner = dataset2.scan(); + scanner + .batch_size(50) + .filter("contains(strings, 'ing')")? + .limit(Some(1024), None)? + .with_row_id(); + let batches: Vec = scanner.try_into_stream().await?.try_collect().await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(1024, total_rows); + assert_eq!(21, batches.len()); + + Ok(()) +} + +#[tokio::test] +async fn test_auto_infer_lance_tokenizer() { + let (mut dataset, json_col) = prepare_json_dataset().await; + + // Create inverted index for json col. Expect auto-infer 'json' for lance tokenizer. + dataset + .create_index( + &[&json_col], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + true, + ) + .await + .unwrap(); + + // Match query succeed only when lance tokenizer is 'json' + let query = FullTextSearchQuery { + query: FtsQuery::Match( + MatchQuery::new("Content,str,once".to_string()).with_column(Some(json_col.clone())), + ), + limit: None, + wand_factor: None, + }; + let batch = dataset + .scan() + .full_text_search(query) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(1, batch.num_rows()); +} diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs new file mode 100644 index 00000000000..6377a9cb016 --- /dev/null +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -0,0 +1,1207 @@ +use super::dataset_common::*; + +#[rstest] +#[lance_test_macros::test(tokio::test)] +async fn test_create_dataset( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + // Appending / Overwriting a dataset that does not exist is treated as Create + for mode in [WriteMode::Create, WriteMode::Append, Overwrite] { + let test_dir = TempStdDir::default(); + create_file(&test_dir, mode, data_storage_version).await + } +} + +#[rstest] +#[lance_test_macros::test(tokio::test)] +async fn test_create_and_fill_empty_dataset( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + let test_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let i32_array: ArrayRef = Arc::new(Int32Array::new(vec![].into(), None)); + let batch = RecordBatch::try_from_iter(vec![("i", i32_array)]).unwrap(); + let reader = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema.clone()); + // check schema of reader and original is same + assert_eq!(schema.as_ref(), reader.schema().as_ref()); + let result = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await + .unwrap(); + + // check dataset empty + assert_eq!(result.count_rows(None).await.unwrap(), 0); + // Since the dataset is empty, will return None. + assert_eq!(result.manifest.max_fragment_id(), None); + + // append rows to dataset + let mut write_params = WriteParams { + max_rows_per_file: 40, + max_rows_per_group: 10, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + // We should be able to append even if the metadata doesn't exactly match. + let schema_with_meta = Arc::new( + schema + .as_ref() + .clone() + .with_metadata([("key".to_string(), "value".to_string())].into()), + ); + let batches = vec![RecordBatch::try_new( + schema_with_meta, + vec![Arc::new(Int32Array::from_iter_values(0..10))], + ) + .unwrap()]; + write_params.mode = WriteMode::Append; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + Dataset::write(batches, &test_uri, Some(write_params)) + .await + .unwrap(); + + let expected_batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..10))], + ) + .unwrap(); + + // get actual dataset + let actual_ds = Dataset::open(&test_uri).await.unwrap(); + // confirm schema is same + let actual_schema = ArrowSchema::from(actual_ds.schema()); + assert_eq!(&actual_schema, schema.as_ref()); + // check num rows is 10 + assert_eq!(actual_ds.count_rows(None).await.unwrap(), 10); + // Max fragment id is still 0 since we only have 1 fragment. + assert_eq!(actual_ds.manifest.max_fragment_id(), Some(0)); + // check expected batch is correct + let actual_batches = actual_ds + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + // sort + let actual_batch = concat_batches(&schema, &actual_batches).unwrap(); + let idx_arr = actual_batch.column_by_name("i").unwrap(); + let sorted_indices = sort_to_indices(idx_arr, None, None).unwrap(); + let struct_arr: StructArray = actual_batch.into(); + let sorted_arr = arrow_select::take::take(&struct_arr, &sorted_indices, None).unwrap(); + let expected_struct_arr: StructArray = expected_batch.into(); + assert_eq!(&expected_struct_arr, as_struct_array(sorted_arr.as_ref())); +} + +#[rstest] +#[lance_test_macros::test(tokio::test)] +async fn test_create_with_empty_iter( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + let test_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let reader = RecordBatchIterator::new(vec![].into_iter().map(Ok), schema.clone()); + // check schema of reader and original is same + assert_eq!(schema.as_ref(), reader.schema().as_ref()); + let write_params = Some(WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }); + let result = Dataset::write(reader, &test_uri, write_params) + .await + .unwrap(); + + // check dataset empty + assert_eq!(result.count_rows(None).await.unwrap(), 0); + // Since the dataset is empty, will return None. + assert_eq!(result.manifest.max_fragment_id(), None); +} + +#[tokio::test] +async fn test_load_manifest_iops() { + // Use consistent session so memory store can be reused. + let session = Arc::new(Session::default()); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..10_i32))], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let _original_ds = Dataset::write( + batches, + "memory://test", + Some(WriteParams { + session: Some(session.clone()), + ..Default::default() + }), + ) + .await + .unwrap(); + + let _ = _original_ds.object_store().io_stats_incremental(); //reset + + let _dataset = DatasetBuilder::from_uri("memory://test") + .with_session(session) + .load() + .await + .unwrap(); + + // There should be only two IOPS: + // 1. List _versions directory to get the latest manifest location + // 2. Read the manifest file. (The manifest is small enough to be read in one go. + // Larger manifests would result in more IOPS.) + let io_stats = _dataset.object_store().io_stats_incremental(); + assert_io_eq!(io_stats, read_iops, 2); +} + +#[rstest] +#[tokio::test] +async fn test_write_params( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + use crate::dataset::fragment::FragReadConfig; + + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let num_rows: usize = 1_000; + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..num_rows as i32))], + ) + .unwrap()]; + + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + + let write_params = WriteParams { + max_rows_per_file: 100, + max_rows_per_group: 10, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let dataset = Dataset::write(batches, &test_uri, Some(write_params)) + .await + .unwrap(); + + assert_eq!(dataset.count_rows(None).await.unwrap(), num_rows); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 10); + assert_eq!(dataset.count_fragments(), 10); + for fragment in &fragments { + assert_eq!(fragment.count_rows(None).await.unwrap(), 100); + let reader = fragment + .open(dataset.schema(), FragReadConfig::default()) + .await + .unwrap(); + // No group / batch concept in v2 + if data_storage_version == LanceFileVersion::Legacy { + assert_eq!(reader.legacy_num_batches(), 10); + for i in 0..reader.legacy_num_batches() as u32 { + assert_eq!(reader.legacy_num_rows_in_batch(i).unwrap(), 10); + } + } + } +} + +#[rstest] +#[tokio::test] +async fn test_write_manifest( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + use lance_table::feature_flags::FLAG_UNKNOWN; + + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..20))], + ) + .unwrap()]; + + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let write_fut = Dataset::write( + batches, + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + auto_cleanup: None, + ..Default::default() + }), + ); + let write_fut = require_send(write_fut); + let mut dataset = write_fut.await.unwrap(); + + // Check it has no flags + let manifest = read_manifest( + dataset.object_store(), + &dataset + .commit_handler + .resolve_latest_location(&dataset.base, dataset.object_store()) + .await + .unwrap() + .path, + None, + ) + .await + .unwrap(); + + assert_eq!( + manifest.data_storage_format, + DataStorageFormat::new(data_storage_version) + ); + assert_eq!(manifest.reader_feature_flags, 0); + + // Create one with deletions + dataset.delete("i < 10").await.unwrap(); + dataset.validate().await.unwrap(); + + // Check it set the flag + let mut manifest = read_manifest( + dataset.object_store(), + &dataset + .commit_handler + .resolve_latest_location(&dataset.base, dataset.object_store()) + .await + .unwrap() + .path, + None, + ) + .await + .unwrap(); + assert_eq!( + manifest.writer_feature_flags, + feature_flags::FLAG_DELETION_FILES + ); + assert_eq!( + manifest.reader_feature_flags, + feature_flags::FLAG_DELETION_FILES + ); + + // Write with custom manifest + manifest.writer_feature_flags |= FLAG_UNKNOWN; // Set another flag + manifest.reader_feature_flags |= FLAG_UNKNOWN; + manifest.version += 1; + write_manifest_file( + dataset.object_store(), + dataset.commit_handler.as_ref(), + &dataset.base, + &mut manifest, + None, + &ManifestWriteConfig { + auto_set_feature_flags: false, + timestamp: None, + use_stable_row_ids: false, + use_legacy_format: None, + storage_format: None, + disable_transaction_file: false, + }, + dataset.manifest_location.naming_scheme, + None, + ) + .await + .unwrap(); + + // Check it rejects reading it + let read_result = Dataset::open(&test_uri).await; + assert!(matches!(read_result, Err(Error::NotSupported { .. }))); + + // Check it rejects writing to it. + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..20))], + ) + .unwrap()]; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let write_result = Dataset::write( + batches, + &test_uri, + Some(WriteParams { + mode: WriteMode::Append, + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await; + + assert!(matches!(write_result, Err(Error::NotSupported { .. }))); +} + +#[rstest] +#[tokio::test] +async fn append_dataset( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..20))], + ) + .unwrap()]; + + let mut write_params = WriteParams { + max_rows_per_file: 40, + max_rows_per_group: 10, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + Dataset::write(batches, &test_uri, Some(write_params.clone())) + .await + .unwrap(); + + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(20..40))], + ) + .unwrap()]; + write_params.mode = WriteMode::Append; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + Dataset::write(batches, &test_uri, Some(write_params.clone())) + .await + .unwrap(); + + let expected_batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..40))], + ) + .unwrap(); + + let actual_ds = Dataset::open(&test_uri).await.unwrap(); + assert_eq!(actual_ds.version().version, 2); + let actual_schema = ArrowSchema::from(actual_ds.schema()); + assert_eq!(&actual_schema, schema.as_ref()); + + let actual_batches = actual_ds + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + // sort + let actual_batch = concat_batches(&schema, &actual_batches).unwrap(); + let idx_arr = actual_batch.column_by_name("i").unwrap(); + let sorted_indices = sort_to_indices(idx_arr, None, None).unwrap(); + let struct_arr: StructArray = actual_batch.into(); + let sorted_arr = arrow_select::take::take(&struct_arr, &sorted_indices, None).unwrap(); + + let expected_struct_arr: StructArray = expected_batch.into(); + assert_eq!(&expected_struct_arr, as_struct_array(sorted_arr.as_ref())); + + // Each fragments has different fragment ID + assert_eq!( + actual_ds + .fragments() + .iter() + .map(|f| f.id) + .collect::>(), + (0..2).collect::>() + ) +} + +#[rstest] +#[tokio::test] +async fn test_shallow_clone_with_hybrid_paths( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + let test_dir = TempStdDir::default(); + let base_dir = test_dir.join("base"); + let test_uri = base_dir.to_str().unwrap(); + let clone_dir = test_dir.join("clone"); + let cloned_uri = clone_dir.to_str().unwrap(); + + // Generate consistent test data batches + let generate_data = |prefix: &str, start_id: i32, row_count: u64| { + gen_batch() + .col("id", array::step_custom::(start_id, 1)) + .col("value", array::fill_utf8(format!("{prefix}_data"))) + .into_reader_rows(RowCount::from(row_count), BatchCount::from(1)) + }; + + // Reusable dataset writer with configurable mode + async fn write_dataset( + uri: &str, + data_reader: impl RecordBatchReader + Send + 'static, + mode: WriteMode, + version: LanceFileVersion, + ) -> Dataset { + let params = WriteParams { + max_rows_per_file: 100, + max_rows_per_group: 20, + data_storage_version: Some(version), + mode, + ..Default::default() + }; + Dataset::write(data_reader, uri, Some(params)) + .await + .unwrap() + } + + // Unified dataset scanning and row counting + async fn collect_rows(dataset: &Dataset) -> (usize, Vec) { + let batches = dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + (batches.iter().map(|b| b.num_rows()).sum(), batches) + } + + // Create initial dataset + let mut dataset = write_dataset( + test_uri, + generate_data("initial", 0, 50), + WriteMode::Create, + data_storage_version, + ) + .await; + + // Store original state for comparison + let original_version = dataset.version().version; + let original_fragment_count = dataset.fragments().len(); + + // Create tag and shallow clone + dataset + .tags() + .create("test_tag", original_version) + .await + .unwrap(); + let cloned_dataset = dataset + .shallow_clone(cloned_uri, "test_tag", None) + .await + .unwrap(); + + // Verify cloned dataset state + let (cloned_rows, _) = collect_rows(&cloned_dataset).await; + assert_eq!(cloned_rows, 50); + assert_eq!(cloned_dataset.version().version, original_version); + + // Append data to cloned dataset + let updated_cloned = write_dataset( + cloned_uri, + generate_data("cloned_new", 50, 30), + WriteMode::Append, + data_storage_version, + ) + .await; + + // Verify updated cloned dataset + let (updated_cloned_rows, updated_batches) = collect_rows(&updated_cloned).await; + assert_eq!(updated_cloned_rows, 80); + assert_eq!(updated_cloned.version().version, original_version + 1); + + // Append data to original dataset + let updated_original = write_dataset( + test_uri, + generate_data("original_new", 50, 25), + WriteMode::Append, + data_storage_version, + ) + .await; + + // Verify updated original dataset + let (original_rows, _) = collect_rows(&updated_original).await; + assert_eq!(original_rows, 75); + assert_eq!(updated_original.version().version, original_version + 1); + + // Final validations + // Verify cloned dataset isolation + let final_cloned = Dataset::open(cloned_uri).await.unwrap(); + let (final_cloned_rows, _) = collect_rows(&final_cloned).await; + + // Data integrity check + let combined_batch = concat_batches(&updated_batches[0].schema(), &updated_batches).unwrap(); + assert_eq!(combined_batch.column_by_name("id").unwrap().len(), 80); + assert_eq!(combined_batch.column_by_name("value").unwrap().len(), 80); + + // Fragment count validation + assert_eq!( + updated_original.fragments().len(), + original_fragment_count + 1 + ); + assert_eq!(final_cloned.fragments().len(), original_fragment_count + 1); + + // Final assertions + assert_eq!(final_cloned_rows, 80); + assert_eq!(final_cloned.version().version, original_version + 1); +} + +#[rstest] +#[tokio::test] +async fn test_shallow_clone_multiple_times( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + let test_uri = TempStrDir::default(); + let append_row_count = 36; + + // Async dataset writer function + async fn write_dataset( + dest: impl Into>, + row_count: u64, + mode: WriteMode, + version: LanceFileVersion, + ) -> Dataset { + let data = gen_batch() + .col("index", array::step::()) + .col("category", array::fill_utf8("base".to_string())) + .col("score", array::step_custom::(1.0, 0.5)); + Dataset::write( + data.into_reader_rows(RowCount::from(row_count), BatchCount::from(1)), + dest, + Some(WriteParams { + max_rows_per_file: 60, + max_rows_per_group: 12, + mode, + data_storage_version: Some(version), + ..Default::default() + }), + ) + .await + .unwrap() + } + + let mut current_dataset = write_dataset( + &test_uri, + append_row_count, + WriteMode::Create, + data_storage_version, + ) + .await; + + let test_round = 3; + // Generate clone paths + let clone_paths = (1..=test_round) + .map(|i| format!("{}/clone{}", test_uri, i)) + .collect::>(); + let mut cloned_datasets = Vec::with_capacity(test_round); + + // Unified cloning procedure, write a fragment to each cloned dataset. + for path in clone_paths.iter() { + current_dataset + .tags() + .create("v1", current_dataset.latest_version_id().await.unwrap()) + .await + .unwrap(); + + current_dataset = current_dataset + .shallow_clone(path, "v1", None) + .await + .unwrap(); + current_dataset = write_dataset( + Arc::new(current_dataset), + append_row_count, + WriteMode::Append, + data_storage_version, + ) + .await; + cloned_datasets.push(current_dataset.clone()); + } + + // Validation function + async fn validate_dataset( + dataset: &Dataset, + expected_rows: usize, + expected_fragments_count: usize, + expected_base_paths_count: usize, + ) { + let batches = dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, expected_rows); + assert_eq!(dataset.fragments().len(), expected_fragments_count); + assert_eq!( + dataset.manifest().base_paths.len(), + expected_base_paths_count + ); + } + + // Verify cloned datasets row count, fragment count, base_path count + for (i, ds) in cloned_datasets.iter().enumerate() { + validate_dataset(ds, 36 * (i + 2), i + 2, i + 1).await; + } + + // Verify original dataset row count, fragment count, base_path count + let original = Dataset::open(&test_uri).await.unwrap(); + validate_dataset(&original, 36, 1, 0).await; +} + +#[rstest] +#[tokio::test] +async fn test_self_dataset_append( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..20))], + ) + .unwrap()]; + + let mut write_params = WriteParams { + max_rows_per_file: 40, + max_rows_per_group: 10, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let mut ds = Dataset::write(batches, &test_uri, Some(write_params.clone())) + .await + .unwrap(); + + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(20..40))], + ) + .unwrap()]; + write_params.mode = WriteMode::Append; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + + ds.append(batches, Some(write_params.clone())) + .await + .unwrap(); + + let expected_batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..40))], + ) + .unwrap(); + + let actual_ds = Dataset::open(&test_uri).await.unwrap(); + assert_eq!(actual_ds.version().version, 2); + // validate fragment ids + assert_eq!(actual_ds.fragments().len(), 2); + assert_eq!( + actual_ds + .fragments() + .iter() + .map(|f| f.id) + .collect::>(), + (0..2).collect::>() + ); + + let actual_schema = ArrowSchema::from(actual_ds.schema()); + assert_eq!(&actual_schema, schema.as_ref()); + + let actual_batches = actual_ds + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + // sort + let actual_batch = concat_batches(&schema, &actual_batches).unwrap(); + let idx_arr = actual_batch.column_by_name("i").unwrap(); + let sorted_indices = sort_to_indices(idx_arr, None, None).unwrap(); + let struct_arr: StructArray = actual_batch.into(); + let sorted_arr = arrow_select::take::take(&struct_arr, &sorted_indices, None).unwrap(); + + let expected_struct_arr: StructArray = expected_batch.into(); + assert_eq!(&expected_struct_arr, as_struct_array(sorted_arr.as_ref())); + + actual_ds.validate().await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn test_self_dataset_append_schema_different( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..20))], + ) + .unwrap()]; + + let other_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int64, + false, + )])); + let other_batches = vec![RecordBatch::try_new( + other_schema.clone(), + vec![Arc::new(Int64Array::from_iter_values(0..20))], + ) + .unwrap()]; + + let mut write_params = WriteParams { + max_rows_per_file: 40, + max_rows_per_group: 10, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let mut ds = Dataset::write(batches, &test_uri, Some(write_params.clone())) + .await + .unwrap(); + + write_params.mode = WriteMode::Append; + let other_batches = + RecordBatchIterator::new(other_batches.into_iter().map(Ok), other_schema.clone()); + + let result = ds.append(other_batches, Some(write_params.clone())).await; + // Error because schema is different + assert!(matches!(result, Err(Error::SchemaMismatch { .. }))) +} + +#[rstest] +#[tokio::test] +async fn append_dictionary( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + // We store the dictionary as part of the schema, so we check that the + // dictionary is consistent between appends. + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "x", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + false, + )])); + let dictionary = Arc::new(StringArray::from(vec!["a", "b"])); + let indices = Int8Array::from(vec![0, 1, 0]); + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new( + Int8DictionaryArray::try_new(indices, dictionary.clone()).unwrap(), + )], + ) + .unwrap()]; + + let test_uri = TempStrDir::default(); + let mut write_params = WriteParams { + max_rows_per_file: 40, + max_rows_per_group: 10, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + Dataset::write(batches, &test_uri, Some(write_params.clone())) + .await + .unwrap(); + + // create a new one with same dictionary + let indices = Int8Array::from(vec![1, 0, 1]); + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new( + Int8DictionaryArray::try_new(indices, dictionary).unwrap(), + )], + ) + .unwrap()]; + + // Write to dataset (successful) + write_params.mode = WriteMode::Append; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + Dataset::write(batches, &test_uri, Some(write_params.clone())) + .await + .unwrap(); + + // Create a new one with *different* dictionary + let dictionary = Arc::new(StringArray::from(vec!["d", "c"])); + let indices = Int8Array::from(vec![1, 0, 1]); + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new( + Int8DictionaryArray::try_new(indices, dictionary).unwrap(), + )], + ) + .unwrap()]; + + // Try write to dataset (fails with legacy format) + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let result = Dataset::write(batches, &test_uri, Some(write_params)).await; + if data_storage_version == LanceFileVersion::Legacy { + assert!(result.is_err()); + } else { + assert!(result.is_ok()); + } +} + +#[rstest] +#[tokio::test] +async fn overwrite_dataset( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..20))], + ) + .unwrap()]; + + let mut write_params = WriteParams { + max_rows_per_file: 40, + max_rows_per_group: 10, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + let dataset = Dataset::write(batches, &test_uri, Some(write_params.clone())) + .await + .unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1); + assert_eq!(dataset.manifest.max_fragment_id(), Some(0)); + + let new_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "s", + DataType::Utf8, + false, + )])); + let new_batches = vec![RecordBatch::try_new( + new_schema.clone(), + vec![Arc::new(StringArray::from_iter_values( + (20..40).map(|v| v.to_string()), + ))], + ) + .unwrap()]; + write_params.mode = Overwrite; + let new_batch_reader = + RecordBatchIterator::new(new_batches.into_iter().map(Ok), new_schema.clone()); + let dataset = Dataset::write(new_batch_reader, &test_uri, Some(write_params.clone())) + .await + .unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1); + // Fragment ids reset after overwrite. + assert_eq!(fragments[0].id(), 0); + assert_eq!(dataset.manifest.max_fragment_id(), Some(0)); + + let actual_ds = Dataset::open(&test_uri).await.unwrap(); + assert_eq!(actual_ds.version().version, 2); + let actual_schema = ArrowSchema::from(actual_ds.schema()); + assert_eq!(&actual_schema, new_schema.as_ref()); + + let actual_batches = actual_ds + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let actual_batch = concat_batches(&new_schema, &actual_batches).unwrap(); + + assert_eq!(new_schema.clone(), actual_batch.schema()); + let arr = actual_batch.column_by_name("s").unwrap(); + assert_eq!( + &StringArray::from_iter_values((20..40).map(|v| v.to_string())), + as_string_array(arr) + ); + assert_eq!(actual_ds.version().version, 2); + + // But we can still check out the first version + let first_ver = DatasetBuilder::from_uri(&test_uri) + .with_version(1) + .load() + .await + .unwrap(); + assert_eq!(first_ver.version().version, 1); + assert_eq!(&ArrowSchema::from(first_ver.schema()), schema.as_ref()); +} + +#[rstest] +#[tokio::test] +async fn test_fast_count_rows( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + + let batches: Vec = (0..20) + .map(|i| { + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(i * 20..(i + 1) * 20))], + ) + .unwrap() + }) + .collect(); + + let write_params = WriteParams { + max_rows_per_file: 40, + max_rows_per_group: 10, + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + Dataset::write(batches, &test_uri, Some(write_params)) + .await + .unwrap(); + + let dataset = Dataset::open(&test_uri).await.unwrap(); + dataset.validate().await.unwrap(); + assert_eq!(10, dataset.fragments().len()); + assert_eq!(400, dataset.count_rows(None).await.unwrap()); + assert_eq!( + 200, + dataset + .count_rows(Some("i < 200".to_string())) + .await + .unwrap() + ); +} + +#[rstest] +#[tokio::test] +async fn test_bfloat16_roundtrip( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) -> Result<()> { + let inner_field = Arc::new( + ArrowField::new("item", DataType::FixedSizeBinary(2), true).with_metadata( + [ + (ARROW_EXT_NAME_KEY.into(), BFLOAT16_EXT_NAME.into()), + (ARROW_EXT_META_KEY.into(), "".into()), + ] + .into(), + ), + ); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "fsl", + DataType::FixedSizeList(inner_field.clone(), 2), + false, + )])); + + let values = bfloat16::BFloat16Array::from_iter_values( + (0..6).map(|i| i as f32).map(half::bf16::from_f32), + ); + let vectors = FixedSizeListArray::new(inner_field, 2, Arc::new(values.into_inner()), None); + + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vectors)]).unwrap(); + + let test_uri = TempStrDir::default(); + + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch.clone())], schema.clone()), + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await?; + + let data = dataset.scan().try_into_batch().await?; + assert_eq!(batch, data); + + Ok(()) +} + +#[tokio::test] +async fn test_overwrite_mixed_version() { + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "a", + DataType::Int32, + false, + )])); + let arr = Arc::new(Int32Array::from(vec![1, 2, 3])); + + let data = RecordBatch::try_new(schema.clone(), vec![arr]).unwrap(); + let reader = RecordBatchIterator::new(vec![data.clone()].into_iter().map(Ok), schema.clone()); + + let dataset = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::Legacy), + ..Default::default() + }), + ) + .await + .unwrap(); + + assert_eq!( + dataset + .manifest + .data_storage_format + .lance_file_version() + .unwrap(), + LanceFileVersion::Legacy + ); + + let reader = RecordBatchIterator::new(vec![data].into_iter().map(Ok), schema); + let dataset = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + mode: WriteMode::Overwrite, + ..Default::default() + }), + ) + .await + .unwrap(); + + assert_eq!( + dataset + .manifest + .data_storage_format + .lance_file_version() + .unwrap(), + LanceFileVersion::Legacy + ); +} + +#[tokio::test] +async fn test_open_nonexisting_dataset() { + let temp_dir = TempStdDir::default(); + let dataset_dir = temp_dir.join("non_existing"); + let dataset_uri = dataset_dir.to_str().unwrap(); + + let res = Dataset::open(dataset_uri).await; + assert!(res.is_err()); + + assert!(!dataset_dir.exists()); +} + +#[tokio::test] +async fn test_manifest_partially_fits() { + // This regresses a bug that occurred when the manifest file was over 4KiB but the manifest + // itself was less than 4KiB (due to a dictionary). 4KiB is important here because that's the + // block size we use when reading the "last block" + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "x", + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), + false, + )])); + let dictionary = Arc::new(StringArray::from_iter_values( + (0..1000).map(|i| i.to_string()), + )); + let indices = Int16Array::from_iter_values(0..1000); + let batches = vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new( + Int16DictionaryArray::try_new(indices, dictionary.clone()).unwrap(), + )], + ) + .unwrap()]; + + let test_uri = TempStrDir::default(); + let batches = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + Dataset::write(batches, &test_uri, None).await.unwrap(); + + let dataset = Dataset::open(&test_uri).await.unwrap(); + assert_eq!(1000, dataset.count_rows(None).await.unwrap()); +} + +#[tokio::test] +async fn test_dataset_uri_roundtrips() { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "a", + DataType::Int32, + false, + )])); + + let test_uri = TempStrDir::default(); + let vectors = Arc::new(Int32Array::from_iter_values(vec![])); + + let data = RecordBatch::try_new(schema.clone(), vec![vectors]); + let reader = RecordBatchIterator::new(vec![data.unwrap()].into_iter().map(Ok), schema); + let dataset = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + ..Default::default() + }), + ) + .await + .unwrap(); + + let uri = dataset.uri(); + assert_eq!(uri, test_uri.as_str()); + + let ds2 = Dataset::open(uri).await.unwrap(); + assert_eq!( + ds2.latest_version_id().await.unwrap(), + dataset.latest_version_id().await.unwrap() + ); +} diff --git a/rust/lance/src/dataset/tests/dataset_merge_update.rs b/rust/lance/src/dataset/tests/dataset_merge_update.rs new file mode 100644 index 00000000000..b4dab03341f --- /dev/null +++ b/rust/lance/src/dataset/tests/dataset_merge_update.rs @@ -0,0 +1,1455 @@ +use super::dataset_common::*; + +#[rstest] +#[tokio::test] +async fn test_merge( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, + #[values(false, true)] use_stable_row_id: bool, +) { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("i", DataType::Int32, false), + ArrowField::new("x", DataType::Float32, false), + ])); + let batch1 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Float32Array::from(vec![1.0, 2.0])), + ], + ) + .unwrap(); + let batch2 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![3, 2])), + Arc::new(Float32Array::from(vec![3.0, 4.0])), + ], + ) + .unwrap(); + + let test_uri = TempStrDir::default(); + + let write_params = WriteParams { + mode: WriteMode::Append, + data_storage_version: Some(data_storage_version), + enable_stable_row_ids: use_stable_row_id, + ..Default::default() + }; + + let batches = RecordBatchIterator::new(vec![batch1].into_iter().map(Ok), schema.clone()); + Dataset::write(batches, &test_uri, Some(write_params.clone())) + .await + .unwrap(); + + let batches = RecordBatchIterator::new(vec![batch2].into_iter().map(Ok), schema.clone()); + Dataset::write(batches, &test_uri, Some(write_params.clone())) + .await + .unwrap(); + + let dataset = Dataset::open(&test_uri).await.unwrap(); + assert_eq!(dataset.fragments().len(), 2); + assert_eq!(dataset.manifest.max_fragment_id(), Some(1)); + + let right_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("i2", DataType::Int32, false), + ArrowField::new("y", DataType::Utf8, true), + ])); + let right_batch1 = RecordBatch::try_new( + right_schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec!["a", "b"])), + ], + ) + .unwrap(); + + let batches = + RecordBatchIterator::new(vec![right_batch1].into_iter().map(Ok), right_schema.clone()); + let mut dataset = Dataset::open(&test_uri).await.unwrap(); + dataset.merge(batches, "i", "i2").await.unwrap(); + dataset.validate().await.unwrap(); + + assert_eq!(dataset.version().version, 3); + assert_eq!(dataset.fragments().len(), 2); + assert_eq!(dataset.fragments()[0].files.len(), 2); + assert_eq!(dataset.fragments()[1].files.len(), 2); + assert_eq!(dataset.manifest.max_fragment_id(), Some(1)); + + let actual_batches = dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let actual = concat_batches(&actual_batches[0].schema(), &actual_batches).unwrap(); + let expected = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("i", DataType::Int32, false), + ArrowField::new("x", DataType::Float32, false), + ArrowField::new("y", DataType::Utf8, true), + ])), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 2])), + Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, 4.0])), + Arc::new(StringArray::from(vec![ + Some("a"), + Some("b"), + None, + Some("b"), + ])), + ], + ) + .unwrap(); + + assert_eq!(actual, expected); + + // Validate we can still read after re-instantiating dataset, which + // clears the cache. + let dataset = Dataset::open(&test_uri).await.unwrap(); + let actual_batches = dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let actual = concat_batches(&actual_batches[0].schema(), &actual_batches).unwrap(); + assert_eq!(actual, expected); +} + +#[rstest] +#[tokio::test] +async fn test_large_merge( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, + #[values(false, true)] use_stable_row_id: bool, +) { + // Tests a merge that spans multiple batches within files + + // This test also tests "null filling" when merging (e.g. when keys do not match + // we need to insert nulls) + + let data = lance_datagen::gen_batch() + .col("key", array::step::()) + .col("value", array::fill_utf8("value".to_string())) + .into_reader_rows(RowCount::from(1_000), BatchCount::from(10)); + + let test_uri = TempStrDir::default(); + + let write_params = WriteParams { + mode: WriteMode::Append, + data_storage_version: Some(data_storage_version), + max_rows_per_file: 1024, + max_rows_per_group: 150, + enable_stable_row_ids: use_stable_row_id, + ..Default::default() + }; + Dataset::write(data, &test_uri, Some(write_params.clone())) + .await + .unwrap(); + + let mut dataset = Dataset::open(&test_uri).await.unwrap(); + assert_eq!(dataset.fragments().len(), 10); + assert_eq!(dataset.manifest.max_fragment_id(), Some(9)); + + let new_data = lance_datagen::gen_batch() + .col("key2", array::step_custom::(500, 1)) + .col("new_value", array::fill_utf8("new_value".to_string())) + .into_reader_rows(RowCount::from(1_000), BatchCount::from(10)); + + dataset.merge(new_data, "key", "key2").await.unwrap(); + dataset.validate().await.unwrap(); +} + +#[rstest] +#[tokio::test] +async fn test_merge_on_row_id( + #[values(LanceFileVersion::Stable)] data_storage_version: LanceFileVersion, + #[values(false, true)] use_stable_row_id: bool, +) { + // Tests a merge on _rowid + + let data = lance_datagen::gen_batch() + .col("key", array::step::()) + .col("value", array::fill_utf8("value".to_string())) + .into_reader_rows(RowCount::from(1_000), BatchCount::from(10)); + + let write_params = WriteParams { + mode: WriteMode::Append, + data_storage_version: Some(data_storage_version), + max_rows_per_file: 1024, + max_rows_per_group: 150, + enable_stable_row_ids: use_stable_row_id, + ..Default::default() + }; + let mut dataset = Dataset::write(data, "memory://", Some(write_params.clone())) + .await + .unwrap(); + assert_eq!(dataset.fragments().len(), 10); + assert_eq!(dataset.manifest.max_fragment_id(), Some(9)); + + let data = dataset.scan().with_row_id().try_into_batch().await.unwrap(); + let row_ids: Arc = data[ROW_ID].clone(); + let key = data["key"].as_primitive::(); + let new_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("rowid", DataType::UInt64, false), + ArrowField::new("new_value", DataType::Int32, false), + ])); + let new_value = Arc::new( + key.into_iter() + .map(|v| v.unwrap() + 1) + .collect::(), + ); + let len = new_value.len() as u32; + let new_batch = RecordBatch::try_new(new_schema.clone(), vec![row_ids, new_value]).unwrap(); + // shuffle new_batch + let mut rng = rand::rng(); + let mut indices: Vec = (0..len).collect(); + indices.shuffle(&mut rng); + let indices = arrow_array::UInt32Array::from_iter_values(indices); + let new_batch = arrow::compute::take_record_batch(&new_batch, &indices).unwrap(); + let new_data = RecordBatchIterator::new(vec![Ok(new_batch)], new_schema.clone()); + dataset.merge(new_data, ROW_ID, "rowid").await.unwrap(); + dataset.validate().await.unwrap(); + assert_eq!(dataset.schema().fields.len(), 3); + assert!(dataset.schema().field("key").is_some()); + assert!(dataset.schema().field("value").is_some()); + assert!(dataset.schema().field("new_value").is_some()); + let batch = dataset.scan().try_into_batch().await.unwrap(); + let key = batch["key"].as_primitive::(); + let new_value = batch["new_value"].as_primitive::(); + for i in 0..key.len() { + assert_eq!(key.value(i) + 1, new_value.value(i)); + } +} + +#[rstest] +#[tokio::test] +async fn test_merge_on_row_addr( + #[values(LanceFileVersion::Stable)] data_storage_version: LanceFileVersion, + #[values(false, true)] use_stable_row_id: bool, +) { + // Tests a merge on _rowaddr + + let data = lance_datagen::gen_batch() + .col("key", array::step::()) + .col("value", array::fill_utf8("value".to_string())) + .into_reader_rows(RowCount::from(1_000), BatchCount::from(10)); + + let write_params = WriteParams { + mode: WriteMode::Append, + data_storage_version: Some(data_storage_version), + max_rows_per_file: 1024, + max_rows_per_group: 150, + enable_stable_row_ids: use_stable_row_id, + ..Default::default() + }; + let mut dataset = Dataset::write(data, "memory://", Some(write_params.clone())) + .await + .unwrap(); + + assert_eq!(dataset.fragments().len(), 10); + assert_eq!(dataset.manifest.max_fragment_id(), Some(9)); + + let data = dataset + .scan() + .with_row_address() + .try_into_batch() + .await + .unwrap(); + let row_addrs = data[ROW_ADDR].clone(); + let key = data["key"].as_primitive::(); + let new_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("rowaddr", DataType::UInt64, false), + ArrowField::new("new_value", DataType::Int32, false), + ])); + let new_value = Arc::new( + key.into_iter() + .map(|v| v.unwrap() + 1) + .collect::(), + ); + let len = new_value.len() as u32; + let new_batch = RecordBatch::try_new(new_schema.clone(), vec![row_addrs, new_value]).unwrap(); + // shuffle new_batch + let mut rng = rand::rng(); + let mut indices: Vec = (0..len).collect(); + indices.shuffle(&mut rng); + let indices = arrow_array::UInt32Array::from_iter_values(indices); + let new_batch = arrow::compute::take_record_batch(&new_batch, &indices).unwrap(); + let new_data = RecordBatchIterator::new(vec![Ok(new_batch)], new_schema.clone()); + dataset.merge(new_data, ROW_ADDR, "rowaddr").await.unwrap(); + dataset.validate().await.unwrap(); + assert_eq!(dataset.schema().fields.len(), 3); + assert!(dataset.schema().field("key").is_some()); + assert!(dataset.schema().field("value").is_some()); + assert!(dataset.schema().field("new_value").is_some()); + let batch = dataset.scan().try_into_batch().await.unwrap(); + let key = batch["key"].as_primitive::(); + let new_value = batch["new_value"].as_primitive::(); + for i in 0..key.len() { + assert_eq!(key.value(i) + 1, new_value.value(i)); + } +} + +#[tokio::test] +async fn test_insert_subschema() { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, false), + ArrowField::new("b", DataType::Int32, true), + ])); + let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); + let mut dataset = Dataset::write(empty_reader, "memory://", None) + .await + .unwrap(); + dataset.validate().await.unwrap(); + + // If missing columns that aren't nullable, will return an error + // TODO: provide alternative default than null. + let just_b = Arc::new(schema.project(&[1]).unwrap()); + let batch = + RecordBatch::try_new(just_b.clone(), vec![Arc::new(Int32Array::from(vec![1]))]).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], just_b.clone()); + let res = dataset.append(reader, None).await; + assert!( + matches!(res, Err(Error::SchemaMismatch { .. })), + "Expected Error::SchemaMismatch, got {:?}", + res + ); + + // If missing columns that are nullable, the write succeeds. + let just_a = Arc::new(schema.project(&[0]).unwrap()); + let batch = + RecordBatch::try_new(just_a.clone(), vec![Arc::new(Int32Array::from(vec![1]))]).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], just_a.clone()); + dataset.append(reader, None).await.unwrap(); + dataset.validate().await.unwrap(); + assert_eq!(dataset.count_rows(None).await.unwrap(), 1); + + // Looking at the fragments, there is no data file with the missing field + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1); + assert_eq!(fragments[0].metadata.files.len(), 1); + assert_eq!(&fragments[0].metadata.files[0].fields, &[0]); + + // When reading back, columns that are missing are null + let data = dataset.scan().try_into_batch().await.unwrap(); + let expected = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![None])), + ], + ) + .unwrap(); + assert_eq!(data, expected); + + // Can still insert all columns + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![2])), + Arc::new(Int32Array::from(vec![3])), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], schema.clone()); + dataset.append(reader, None).await.unwrap(); + dataset.validate().await.unwrap(); + assert_eq!(dataset.count_rows(None).await.unwrap(), 2); + + // When reading back, only missing data is null, otherwise is filled in + let data = dataset.scan().try_into_batch().await.unwrap(); + let expected = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![None, Some(3)])), + ], + ) + .unwrap(); + assert_eq!(data, expected); + + // Can run compaction. All files should now have all fields. + compact_files(&mut dataset, CompactionOptions::default(), None) + .await + .unwrap(); + dataset.validate().await.unwrap(); + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1); + assert_eq!(fragments[0].metadata.files.len(), 1); + assert_eq!(&fragments[0].metadata.files[0].fields, &[0, 1]); + + // Can scan and get expected data. + let data = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!(data, expected); +} + +#[tokio::test] +async fn test_insert_nested_subschemas() { + // Test subschemas at struct level + // Test different orders + // Test the Dataset::write() path + // Test Take across fragments with different field id sets + let test_uri = TempStrDir::default(); + + let field_a = Arc::new(ArrowField::new("a", DataType::Int32, true)); + let field_b = Arc::new(ArrowField::new("b", DataType::Int32, false)); + let field_c = Arc::new(ArrowField::new("c", DataType::Int32, true)); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "s", + DataType::Struct(vec![field_a.clone(), field_b.clone(), field_c.clone()].into()), + true, + )])); + let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); + let dataset = Dataset::write(empty_reader, &test_uri, None).await.unwrap(); + dataset.validate().await.unwrap(); + + let append_options = WriteParams { + mode: WriteMode::Append, + ..Default::default() + }; + // Can insert b, a + let just_b_a = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "s", + DataType::Struct(vec![field_b.clone(), field_a.clone()].into()), + true, + )])); + let batch = RecordBatch::try_new( + just_b_a.clone(), + vec![Arc::new(StructArray::from(vec![ + ( + field_b.clone(), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + ), + (field_a.clone(), Arc::new(Int32Array::from(vec![2]))), + ]))], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], just_b_a.clone()); + let dataset = Dataset::write(reader, &test_uri, Some(append_options.clone())) + .await + .unwrap(); + dataset.validate().await.unwrap(); + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1); + assert_eq!(fragments[0].metadata.files.len(), 1); + assert_eq!(&fragments[0].metadata.files[0].fields, &[0, 2, 1]); + assert_eq!(&fragments[0].metadata.files[0].column_indices, &[0, 1, 2]); + + // Can insert c, b + let just_c_b = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "s", + DataType::Struct(vec![field_c.clone(), field_b.clone()].into()), + true, + )])); + let batch = RecordBatch::try_new( + just_c_b.clone(), + vec![Arc::new(StructArray::from(vec![ + ( + field_c.clone(), + Arc::new(Int32Array::from(vec![4])) as ArrayRef, + ), + (field_b.clone(), Arc::new(Int32Array::from(vec![3]))), + ]))], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], just_c_b.clone()); + let dataset = Dataset::write(reader, &test_uri, Some(append_options.clone())) + .await + .unwrap(); + dataset.validate().await.unwrap(); + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 2); + assert_eq!(fragments[1].metadata.files.len(), 1); + assert_eq!(&fragments[1].metadata.files[0].fields, &[0, 3, 2]); + assert_eq!(&fragments[1].metadata.files[0].column_indices, &[0, 1, 2]); + + // Can't insert a, c (b is non-nullable) + let just_a_c = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "s", + DataType::Struct(vec![field_a.clone(), field_c.clone()].into()), + true, + )])); + let batch = RecordBatch::try_new( + just_a_c.clone(), + vec![Arc::new(StructArray::from(vec![ + ( + field_a.clone(), + Arc::new(Int32Array::from(vec![5])) as ArrayRef, + ), + (field_c.clone(), Arc::new(Int32Array::from(vec![6]))), + ]))], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], just_a_c.clone()); + let res = Dataset::write(reader, &test_uri, Some(append_options)).await; + assert!( + matches!(res, Err(Error::SchemaMismatch { .. })), + "Expected Error::SchemaMismatch, got {:?}", + res + ); + + // Can scan and get all data + let data = dataset.scan().try_into_batch().await.unwrap(); + let expected = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StructArray::from(vec![ + ( + field_a.clone(), + Arc::new(Int32Array::from(vec![Some(2), None])) as ArrayRef, + ), + (field_b.clone(), Arc::new(Int32Array::from(vec![1, 3]))), + ( + field_c.clone(), + Arc::new(Int32Array::from(vec![None, Some(4)])), + ), + ]))], + ) + .unwrap(); + assert_eq!(data, expected); + + // Can call take and get rows from all three back in one batch + let result = dataset + .take(&[1, 0], Arc::new(dataset.schema().clone())) + .await + .unwrap(); + let expected = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StructArray::from(vec![ + ( + field_a.clone(), + Arc::new(Int32Array::from(vec![None, Some(2)])) as ArrayRef, + ), + (field_b.clone(), Arc::new(Int32Array::from(vec![3, 1]))), + ( + field_c.clone(), + Arc::new(Int32Array::from(vec![Some(4), None])), + ), + ]))], + ) + .unwrap(); + assert_eq!(result, expected); +} + +#[tokio::test] +async fn test_insert_balanced_subschemas() { + let test_uri = TempStrDir::default(); + + let field_a = ArrowField::new("a", DataType::Int32, true); + let field_b = ArrowField::new("b", DataType::LargeBinary, true); + let schema = Arc::new(ArrowSchema::new(vec![ + field_a.clone(), + field_b + .clone() + .with_metadata([(BLOB_META_KEY.to_string(), "true".to_string())].into()), + ])); + let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); + let options = WriteParams { + enable_stable_row_ids: true, + enable_v2_manifest_paths: true, + ..Default::default() + }; + let mut dataset = Dataset::write(empty_reader, &test_uri, Some(options)) + .await + .unwrap(); + dataset.validate().await.unwrap(); + + // Insert left side + let just_a = Arc::new(ArrowSchema::new(vec![field_a.clone()])); + let batch = + RecordBatch::try_new(just_a.clone(), vec![Arc::new(Int32Array::from(vec![1]))]).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], just_a.clone()); + dataset.append(reader, None).await.unwrap(); + dataset.validate().await.unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1); + assert_eq!(fragments[0].metadata.files.len(), 1); + assert_eq!(&fragments[0].metadata.files[0].fields, &[0]); + + // Insert right side + let just_b = Arc::new(ArrowSchema::new(vec![field_b.clone()])); + let batch = RecordBatch::try_new( + just_b.clone(), + vec![Arc::new(LargeBinaryArray::from_iter(vec![Some(vec![2u8])]))], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], just_b.clone()); + dataset.append(reader, None).await.unwrap(); + dataset.validate().await.unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 2); + assert_eq!(fragments[1].metadata.files.len(), 1); + assert_eq!(&fragments[1].metadata.files[0].fields, &[1]); + + let data = dataset + .take( + &[0, 1], + ProjectionRequest::from_columns(["a"], dataset.schema()), + ) + .await + .unwrap(); + assert_eq!(data.num_rows(), 2); + let a_column = data.column(0).as_primitive::(); + assert_eq!(a_column.value(0), 1); + assert!(a_column.is_null(1)); + + let blob_batch = dataset + .take( + &[0, 1], + ProjectionRequest::from_columns(["b"], dataset.schema()), + ) + .await + .unwrap(); + let blob_descriptions = blob_batch.column(0).as_struct(); + assert!(blob_descriptions.is_null(0)); + assert!(blob_descriptions.is_valid(1)); +} + +#[tokio::test] +async fn test_datafile_replacement() { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "a", + DataType::Int32, + true, + )])); + let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); + let dataset = Arc::new( + Dataset::write(empty_reader, "memory://", None) + .await + .unwrap(), + ); + dataset.validate().await.unwrap(); + + // Test empty replacement should commit a new manifest and do nothing + let mut dataset = Dataset::commit( + WriteDestination::Dataset(dataset.clone()), + Operation::DataReplacement { + replacements: vec![], + }, + Some(1), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + dataset.validate().await.unwrap(); + + assert_eq!(dataset.version().version, 2); + assert_eq!(dataset.get_fragments().len(), 0); + + // try the same thing on a non-empty dataset + let vals: Int32Array = vec![1, 2, 3].into(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vals)]).unwrap(); + dataset + .append( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + None, + ) + .await + .unwrap(); + + let dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataReplacement { + replacements: vec![], + }, + Some(3), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + dataset.validate().await.unwrap(); + + assert_eq!(dataset.version().version, 4); + assert_eq!(dataset.get_fragments().len(), 1); + + let batch = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 3); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[1, 2, 3] + ); + + // write a new datafile + let object_writer = dataset + .object_store + .create(&Path::from("data/test.lance")) + .await + .unwrap(); + let mut writer = FileWriter::try_new( + object_writer, + schema.as_ref().try_into().unwrap(), + Default::default(), + ) + .unwrap(); + + let vals: Int32Array = vec![4, 5, 6].into(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vals)]).unwrap(); + writer.write_batch(&batch).await.unwrap(); + writer.finish().await.unwrap(); + + // find the datafile we want to replace + let frag = dataset.get_fragment(0).unwrap(); + let data_file = frag.data_file_for_field(0).unwrap(); + let mut new_data_file = data_file.clone(); + new_data_file.path = "test.lance".to_string(); + + let dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, new_data_file)], + }, + Some(4), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + + assert_eq!(dataset.version().version, 5); + assert_eq!(dataset.get_fragments().len(), 1); + assert_eq!(dataset.get_fragments()[0].metadata.files.len(), 1); + + let batch = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 3); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[4, 5, 6] + ); +} + +#[tokio::test] +async fn test_datafile_partial_replacement() { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "a", + DataType::Int32, + true, + )])); + let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); + let mut dataset = Dataset::write(empty_reader, "memory://", None) + .await + .unwrap(); + dataset.validate().await.unwrap(); + + let vals: Int32Array = vec![1, 2, 3].into(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vals)]).unwrap(); + dataset + .append( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + None, + ) + .await + .unwrap(); + + let fragment = dataset.get_fragments().pop().unwrap().metadata; + + let extended_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])); + + // add all null column + let dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::Merge { + fragments: vec![fragment], + schema: extended_schema.as_ref().try_into().unwrap(), + }, + Some(2), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + + let partial_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "b", + DataType::Int32, + true, + )])); + + // write a new datafile + let object_writer = dataset + .object_store + .create(&Path::from("data/test.lance")) + .await + .unwrap(); + let mut writer = FileWriter::try_new( + object_writer, + partial_schema.as_ref().try_into().unwrap(), + Default::default(), + ) + .unwrap(); + + let vals: Int32Array = vec![4, 5, 6].into(); + let batch = RecordBatch::try_new(partial_schema.clone(), vec![Arc::new(vals)]).unwrap(); + writer.write_batch(&batch).await.unwrap(); + writer.finish().await.unwrap(); + + let (major, minor) = lance_file::version::LanceFileVersion::Stable.to_numbers(); + + // find the datafile we want to replace + let new_data_file = DataFile { + path: "test.lance".to_string(), + // the second column in the dataset + fields: vec![1], + // is located in the first column of this datafile + column_indices: vec![0], + file_major_version: major, + file_minor_version: minor, + file_size_bytes: CachedFileSize::unknown(), + base_id: None, + }; + + let dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, new_data_file)], + }, + Some(3), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + + assert_eq!(dataset.version().version, 4); + assert_eq!(dataset.get_fragments().len(), 1); + assert_eq!(dataset.get_fragments()[0].metadata.files.len(), 2); + assert_eq!(dataset.get_fragments()[0].metadata.files[0].fields, vec![0]); + assert_eq!(dataset.get_fragments()[0].metadata.files[1].fields, vec![1]); + + let batch = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 3); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[1, 2, 3] + ); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[4, 5, 6] + ); + + // do it again but on the first column + // find the datafile we want to replace + let new_data_file = DataFile { + path: "test.lance".to_string(), + // the first column in the dataset + fields: vec![0], + // is located in the first column of this datafile + column_indices: vec![0], + file_major_version: major, + file_minor_version: minor, + file_size_bytes: CachedFileSize::unknown(), + base_id: None, + }; + + let dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, new_data_file)], + }, + Some(4), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + + assert_eq!(dataset.version().version, 5); + assert_eq!(dataset.get_fragments().len(), 1); + assert_eq!(dataset.get_fragments()[0].metadata.files.len(), 2); + + let batch = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 3); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[4, 5, 6] + ); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[4, 5, 6] + ); +} + +#[tokio::test] +async fn test_datafile_replacement_error() { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "a", + DataType::Int32, + true, + )])); + let empty_reader = RecordBatchIterator::new(vec![], schema.clone()); + let mut dataset = Dataset::write(empty_reader, "memory://", None) + .await + .unwrap(); + dataset.validate().await.unwrap(); + + let vals: Int32Array = vec![1, 2, 3].into(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vals)]).unwrap(); + dataset + .append( + RecordBatchIterator::new(vec![Ok(batch)], schema.clone()), + None, + ) + .await + .unwrap(); + + let fragment = dataset.get_fragments().pop().unwrap().metadata; + + let extended_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])); + + // add all null column + let dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::Merge { + fragments: vec![fragment], + schema: extended_schema.as_ref().try_into().unwrap(), + }, + Some(2), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + + // find the datafile we want to replace + let new_data_file = DataFile { + path: "test.lance".to_string(), + // the second column in the dataset + fields: vec![1], + // is located in the first column of this datafile + column_indices: vec![0], + file_major_version: 2, + file_minor_version: 0, + file_size_bytes: CachedFileSize::unknown(), + base_id: None, + }; + + let new_data_file = DataFile { + fields: vec![0, 1], + ..new_data_file + }; + + let err = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset.clone())), + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, new_data_file)], + }, + Some(2), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("Expected to modify the fragment but no changes were made"), + "Expected Error::DataFileReplacementError, got {:?}", + err + ); +} + +#[tokio::test] +async fn test_replace_dataset() { + let test_dir = TempDir::default(); + let test_uri = test_dir.path_str(); + let test_path = test_dir.obj_path(); + + let data = gen_batch() + .col("int", array::step::()) + .into_batch_rows(RowCount::from(20)) + .unwrap(); + let data1 = data.slice(0, 10); + let data2 = data.slice(10, 10); + let mut ds = InsertBuilder::new(&test_uri) + .execute(vec![data1]) + .await + .unwrap(); + + ds.object_store().remove_dir_all(test_path).await.unwrap(); + + let ds2 = InsertBuilder::new(&test_uri) + .execute(vec![data2.clone()]) + .await + .unwrap(); + + ds.checkout_latest().await.unwrap(); + let roundtripped = ds.scan().try_into_batch().await.unwrap(); + assert_eq!(roundtripped, data2); + + ds.validate().await.unwrap(); + ds2.validate().await.unwrap(); + assert_eq!(ds.manifest.version, 1); + assert_eq!(ds2.manifest.version, 1); +} + +#[tokio::test] +async fn test_insert_skip_auto_cleanup() { + let test_uri = TempStrDir::default(); + + // Create initial dataset with aggressive auto cleanup (interval=1, older_than=1ms) + let data = gen_batch() + .col("id", array::step::()) + .into_reader_rows(RowCount::from(100), BatchCount::from(1)); + + let write_params = WriteParams { + mode: WriteMode::Create, + auto_cleanup: Some(AutoCleanupParams { + interval: 1, + older_than: chrono::TimeDelta::try_milliseconds(0).unwrap(), // Cleanup versions older than 0ms + }), + ..Default::default() + }; + + // Start at 1 second after epoch + MockClock::set_system_time(std::time::Duration::from_secs(1)); + + let dataset = Dataset::write(data, &test_uri, Some(write_params)) + .await + .unwrap(); + assert_eq!(dataset.version().version, 1); + + // Advance time by 1 second + MockClock::set_system_time(std::time::Duration::from_secs(2)); + + // First append WITHOUT skip_auto_cleanup - should trigger cleanup + let data1 = gen_batch() + .col("id", array::step::()) + .into_df_stream(RowCount::from(50), BatchCount::from(1)); + + let write_params1 = WriteParams { + mode: WriteMode::Append, + skip_auto_cleanup: false, + ..Default::default() + }; + + let dataset2 = InsertBuilder::new(WriteDestination::Dataset(Arc::new(dataset))) + .with_params(&write_params1) + .execute_stream(data1) + .await + .unwrap(); + + assert_eq!(dataset2.version().version, 2); + + // Advance time + MockClock::set_system_time(std::time::Duration::from_secs(3)); + + // Need to do another commit for cleanup to take effect since cleanup runs on the old dataset + let data1_extra = gen_batch() + .col("id", array::step::()) + .into_df_stream(RowCount::from(10), BatchCount::from(1)); + + let dataset2_extra = InsertBuilder::new(WriteDestination::Dataset(Arc::new(dataset2))) + .with_params(&write_params1) + .execute_stream(data1_extra) + .await + .unwrap(); + + assert_eq!(dataset2_extra.version().version, 3); + + // Version 1 should be cleaned up due to auto cleanup (cleanup runs every version) + assert!( + dataset2_extra.checkout_version(1).await.is_err(), + "Version 1 should have been cleaned up" + ); + // Version 2 should still exist + assert!( + dataset2_extra.checkout_version(2).await.is_ok(), + "Version 2 should still exist" + ); + + // Advance time + MockClock::set_system_time(std::time::Duration::from_secs(4)); + + // Second append WITH skip_auto_cleanup - should NOT trigger cleanup + let data2 = gen_batch() + .col("id", array::step::()) + .into_df_stream(RowCount::from(30), BatchCount::from(1)); + + let write_params2 = WriteParams { + mode: WriteMode::Append, + skip_auto_cleanup: true, // Skip auto cleanup + ..Default::default() + }; + + let dataset3 = InsertBuilder::new(WriteDestination::Dataset(Arc::new(dataset2_extra))) + .with_params(&write_params2) + .execute_stream(data2) + .await + .unwrap(); + + assert_eq!(dataset3.version().version, 4); + + // Version 2 should still exist because skip_auto_cleanup was enabled + assert!( + dataset3.checkout_version(2).await.is_ok(), + "Version 2 should still exist because skip_auto_cleanup was enabled" + ); + // Version 3 should also still exist + assert!( + dataset3.checkout_version(3).await.is_ok(), + "Version 3 should still exist" + ); +} + +#[tokio::test] +async fn test_nullable_struct_v2_1_issue_4385() { + // Test for issue #4385: nullable struct should preserve null values in v2.1 format + use arrow_array::cast::AsArray; + use arrow_schema::Fields; + + // Create a struct field with nullable float field + let struct_fields = Fields::from(vec![ArrowField::new("x", DataType::Float32, true)]); + + // Create outer struct with the nullable struct as a field (not root) + let outer_fields = Fields::from(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("data", DataType::Struct(struct_fields.clone()), true), + ]); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "record", + DataType::Struct(outer_fields.clone()), + false, + )])); + + // Create data with null struct + let id_values = Int32Array::from(vec![1, 2, 3]); + let x_values = Float32Array::from(vec![Some(1.0), Some(2.0), Some(3.0)]); + let inner_struct_array = StructArray::new( + struct_fields, + vec![Arc::new(x_values) as ArrayRef], + Some(vec![true, false, true].into()), // Second struct is null + ); + + let outer_struct_array = StructArray::new( + outer_fields, + vec![ + Arc::new(id_values) as ArrayRef, + Arc::new(inner_struct_array.clone()) as ArrayRef, + ], + None, // Outer struct is not nullable + ); + + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(outer_struct_array)]).unwrap(); + + // Write dataset with v2.1 format + let test_uri = TempStrDir::default(); + + let write_params = WriteParams { + mode: WriteMode::Create, + data_storage_version: Some(LanceFileVersion::V2_1), + ..Default::default() + }; + + let batches = vec![batch.clone()]; + let batch_reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + + Dataset::write(batch_reader, &test_uri, Some(write_params)) + .await + .unwrap(); + + // Read back the dataset + let dataset = Dataset::open(&test_uri).await.unwrap(); + let scanner = dataset.scan(); + let result_batches = scanner + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(result_batches.len(), 1); + let result_batch = &result_batches[0]; + let read_outer_struct = result_batch.column(0).as_struct(); + let read_inner_struct = read_outer_struct.column(1).as_struct(); // "data" field + + // The bug: null struct is not preserved + assert!( + read_inner_struct.is_null(1), + "Second struct should be null but it's not. Read value: {:?}", + read_inner_struct + ); + + // Verify the null count is preserved + assert_eq!( + inner_struct_array.null_count(), + read_inner_struct.null_count(), + "Null count should be preserved" + ); +} + +#[tokio::test] +async fn test_issue_4902_packed_struct_v2_1_read_error() { + use std::collections::HashMap; + + use arrow_array::{ArrayRef, Int32Array, RecordBatchIterator, StructArray, UInt32Array}; + use arrow_schema::{Field as ArrowField, Fields, Schema as ArrowSchema}; + + let struct_fields = Fields::from(vec![ + ArrowField::new("x", DataType::UInt32, false), + ArrowField::new("y", DataType::UInt32, false), + ]); + let mut packed_metadata = HashMap::new(); + packed_metadata.insert("packed".to_string(), "true".to_string()); + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("int_col", DataType::Int32, false), + ArrowField::new("struct_col", DataType::Struct(struct_fields.clone()), false) + .with_metadata(packed_metadata), + ])); + + let int_values = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8])); + let x_values = Arc::new(UInt32Array::from(vec![1, 4, 7, 10, 13, 16, 19, 22])); + let y_values = Arc::new(UInt32Array::from(vec![2, 5, 8, 11, 14, 17, 20, 23])); + let struct_array = Arc::new(StructArray::new( + struct_fields, + vec![x_values.clone() as ArrayRef, y_values.clone() as ArrayRef], + None, + )); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + int_values.clone() as ArrayRef, + struct_array.clone() as ArrayRef, + ], + ) + .unwrap(); + + let test_uri = TempStrDir::default(); + let write_params = WriteParams { + mode: WriteMode::Create, + data_storage_version: Some(LanceFileVersion::V2_1), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], schema.clone()); + Dataset::write(reader, &test_uri, Some(write_params)) + .await + .unwrap(); + + let dataset = Dataset::open(&test_uri).await.unwrap(); + + let result_batches = dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(result_batches, vec![batch.clone()]); + + let struct_batches = dataset + .scan() + .project(&["struct_col"]) + .unwrap() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(struct_batches.len(), 1); + let read_struct = struct_batches[0].column(0).as_struct(); + assert_eq!(read_struct, struct_array.as_ref()); +} + +#[tokio::test] +async fn test_issue_4429_nested_struct_encoding_v2_1_with_over_65k_structs() { + // Regression test for miniblock 16KB limit with nested struct patterns + // Tests encoding behavior when a nested struct> contains + // large amounts of data that exceeds miniblock encoding limits + + // Create a struct with multiple fields that will trigger miniblock encoding + // Each field is 4 bytes, making the struct narrow enough for miniblock + let measurement_fields = vec![ + ArrowField::new("val_a", DataType::Float32, true), + ArrowField::new("val_b", DataType::Float32, true), + ArrowField::new("val_c", DataType::Float32, true), + ArrowField::new("val_d", DataType::Float32, true), + ArrowField::new("seq_high", DataType::Int32, true), + ArrowField::new("seq_low", DataType::Int32, true), + ]; + let measurement_type = DataType::Struct(measurement_fields.clone().into()); + + // Create nested schema: struct> + // This pattern can trigger encoding issues with large data volumes + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "data", + DataType::Struct( + vec![ArrowField::new( + "measurements", + DataType::List(Arc::new(ArrowField::new( + "item", + measurement_type.clone(), + true, + ))), + true, + )] + .into(), + ), + true, + )])); + + // Create large number of measurements that will exceed encoding limits + // Using 70,520 to match the exact problematic size + const NUM_MEASUREMENTS: usize = 70_520; + + // Generate data for two full sets (rows 0 and 2 will have data, row 1 empty) + const TOTAL_MEASUREMENTS: usize = NUM_MEASUREMENTS * 2; + + // Create arrays with realistic values + let val_a_array = + Float32Array::from_iter((0..TOTAL_MEASUREMENTS).map(|i| Some(16.66 + (i as f32 * 0.0001)))); + let val_b_array = + Float32Array::from_iter((0..TOTAL_MEASUREMENTS).map(|i| Some(-3.54 + (i as f32 * 0.0002)))); + let val_c_array = + Float32Array::from_iter((0..TOTAL_MEASUREMENTS).map(|i| Some(2.94 + (i as f32 * 0.0001)))); + let val_d_array = + Float32Array::from_iter((0..TOTAL_MEASUREMENTS).map(|i| Some(((i % 50) + 10) as f32))); + let seq_high_array = Int32Array::from_iter((0..TOTAL_MEASUREMENTS).map(|_| Some(1736962329))); + let seq_low_array = + Int32Array::from_iter((0..TOTAL_MEASUREMENTS).map(|i| Some(304403000 + (i * 1000) as i32))); + + // Create the struct array with all measurements + let struct_array = StructArray::from(vec![ + ( + Arc::new(ArrowField::new("val_a", DataType::Float32, true)), + Arc::new(val_a_array) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("val_b", DataType::Float32, true)), + Arc::new(val_b_array) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("val_c", DataType::Float32, true)), + Arc::new(val_c_array) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("val_d", DataType::Float32, true)), + Arc::new(val_d_array) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("seq_high", DataType::Int32, true)), + Arc::new(seq_high_array) as ArrayRef, + ), + ( + Arc::new(ArrowField::new("seq_low", DataType::Int32, true)), + Arc::new(seq_low_array) as ArrayRef, + ), + ]); + + // Create list array with pattern: [70520 items, 0 items, 70520 items] + // This pattern triggers the issue with V2.1 encoding + let offsets = vec![ + 0i32, + NUM_MEASUREMENTS as i32, // End of row 0 + NUM_MEASUREMENTS as i32, // End of row 1 (empty) + (NUM_MEASUREMENTS * 2) as i32, // End of row 2 + ]; + let list_array = ListArray::try_new( + Arc::new(ArrowField::new("item", measurement_type, true)), + arrow_buffer::OffsetBuffer::new(arrow_buffer::ScalarBuffer::from(offsets)), + Arc::new(struct_array) as ArrayRef, + None, + ) + .unwrap(); + + // Create the outer struct wrapping the list + let data_struct = StructArray::from(vec![( + Arc::new(ArrowField::new( + "measurements", + DataType::List(Arc::new(ArrowField::new( + "item", + DataType::Struct(measurement_fields.into()), + true, + ))), + true, + )), + Arc::new(list_array) as ArrayRef, + )]); + + // Create the final record batch with 3 rows + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(data_struct) as ArrayRef]).unwrap(); + + assert_eq!(batch.num_rows(), 3, "Should have exactly 3 rows"); + + let test_uri = TempStrDir::default(); + + // Test with V2.1 format which has different encoding behavior + let batches = vec![batch]; + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema.clone()); + + // V2.1 format triggers miniblock encoding for narrow structs + let write_params = WriteParams { + data_storage_version: Some(lance_file::version::LanceFileVersion::V2_1), + ..Default::default() + }; + + // Write dataset - this will panic with miniblock 16KB assertion + let dataset = Dataset::write(reader, &test_uri, Some(write_params)) + .await + .unwrap(); + + dataset.validate().await.unwrap(); + assert_eq!(dataset.count_rows(None).await.unwrap(), 3); +} diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs new file mode 100644 index 00000000000..68ada995e8c --- /dev/null +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -0,0 +1,358 @@ +#![allow(clippy::redundant_pub_crate)] +use super::dataset_common::*; + +pub(crate) async fn scan_dataset(uri: &str) -> Result> { + let results = Dataset::open(uri) + .await? + .scan() + .try_into_stream() + .await? + .try_collect::>() + .await?; + Ok(results) +} + +#[rstest] +#[tokio::test] +async fn test_v0_7_5_migration() { + // We migrate to add Fragment.physical_rows and DeletionFile.num_deletions + // after this version. + + // Copy over table + let test_dir = copy_test_data_to_tmp("v0.7.5/with_deletions").unwrap(); + let test_uri = test_dir.path_str(); + + // Assert num rows, deletions, and physical rows are all correct. + let dataset = Dataset::open(&test_uri).await.unwrap(); + assert_eq!(dataset.count_rows(None).await.unwrap(), 90); + assert_eq!(dataset.count_deleted_rows().await.unwrap(), 10); + let total_physical_rows = futures::stream::iter(dataset.get_fragments()) + .then(|f| async move { f.physical_rows().await }) + .try_fold(0, |acc, x| async move { Ok(acc + x) }) + .await + .unwrap(); + assert_eq!(total_physical_rows, 100); + + // Append 5 rows + let schema = Arc::new(ArrowSchema::from(dataset.schema())); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from_iter_values(100..105))], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let write_params = WriteParams { + mode: WriteMode::Append, + ..Default::default() + }; + let dataset = Dataset::write(batches, &test_uri, Some(write_params)) + .await + .unwrap(); + + // Assert num rows, deletions, and physical rows are all correct. + assert_eq!(dataset.count_rows(None).await.unwrap(), 95); + assert_eq!(dataset.count_deleted_rows().await.unwrap(), 10); + let total_physical_rows = futures::stream::iter(dataset.get_fragments()) + .then(|f| async move { f.physical_rows().await }) + .try_fold(0, |acc, x| async move { Ok(acc + x) }) + .await + .unwrap(); + assert_eq!(total_physical_rows, 105); + + dataset.validate().await.unwrap(); + + // Scan data and assert it is as expected. + let expected = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from_iter_values( + (0..10).chain(20..105), + ))], + ) + .unwrap(); + let actual_batches = dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let actual = concat_batches(&actual_batches[0].schema(), &actual_batches).unwrap(); + assert_eq!(actual, expected); +} + +#[rstest] +#[tokio::test] +async fn test_fix_v0_8_0_broken_migration() { + // The migration from v0.7.5 was broken in 0.8.0. This validates we can + // automatically fix tables that have this problem. + + // Copy over table + let test_dir = copy_test_data_to_tmp("v0.8.0/migrated_from_v0.7.5").unwrap(); + let test_uri = test_dir.path_str(); + let test_uri = &test_uri; + + // Assert num rows, deletions, and physical rows are all correct, even + // though stats are bad. + let dataset = Dataset::open(test_uri).await.unwrap(); + assert_eq!(dataset.count_rows(None).await.unwrap(), 92); + assert_eq!(dataset.count_deleted_rows().await.unwrap(), 10); + let total_physical_rows = futures::stream::iter(dataset.get_fragments()) + .then(|f| async move { f.physical_rows().await }) + .try_fold(0, |acc, x| async move { Ok(acc + x) }) + .await + .unwrap(); + assert_eq!(total_physical_rows, 102); + + // Append 5 rows to table. + let schema = Arc::new(ArrowSchema::from(dataset.schema())); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from_iter_values(100..105))], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let write_params = WriteParams { + mode: WriteMode::Append, + data_storage_version: Some(LanceFileVersion::Legacy), + ..Default::default() + }; + let dataset = Dataset::write(batches, test_uri, Some(write_params)) + .await + .unwrap(); + + // Assert statistics are all now correct. + let physical_rows: Vec<_> = dataset + .get_fragments() + .iter() + .map(|f| f.metadata.physical_rows) + .collect(); + assert_eq!(physical_rows, vec![Some(100), Some(2), Some(5)]); + let num_deletions: Vec<_> = dataset + .get_fragments() + .iter() + .map(|f| { + f.metadata + .deletion_file + .as_ref() + .and_then(|df| df.num_deleted_rows) + }) + .collect(); + assert_eq!(num_deletions, vec![Some(10), None, None]); + assert_eq!(dataset.count_rows(None).await.unwrap(), 97); + + // Scan data and assert it is as expected. + let expected = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from_iter_values( + (0..10).chain(20..100).chain(0..2).chain(100..105), + ))], + ) + .unwrap(); + let actual_batches = dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let actual = concat_batches(&actual_batches[0].schema(), &actual_batches).unwrap(); + assert_eq!(actual, expected); +} + +#[rstest] +#[tokio::test] +async fn test_v0_8_14_invalid_index_fragment_bitmap( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + // Old versions of lance could create an index whose fragment bitmap was + // invalid because it did not include fragments that were part of the index + // + // We need to make sure we do not rely on the fragment bitmap in these older + // versions and instead fall back to a slower legacy behavior + let test_dir = copy_test_data_to_tmp("v0.8.14/corrupt_index").unwrap(); + let test_uri = test_dir.path_str(); + let test_uri = &test_uri; + + let mut dataset = Dataset::open(test_uri).await.unwrap(); + + // Uncomment to reproduce the issue. The below query will panic + // let mut scan = dataset.scan(); + // let query_vec = Float32Array::from(vec![0_f32; 128]); + // let scan_fut = scan + // .nearest("vector", &query_vec, 2000) + // .unwrap() + // .nprobes(4) + // .prefilter(true) + // .try_into_stream() + // .await + // .unwrap() + // .try_collect::>() + // .await + // .unwrap(); + + // Add some data and recalculate the index, forcing a migration + let mut scan = dataset.scan(); + let data = scan + .limit(Some(10), None) + .unwrap() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let schema = data[0].schema(); + let data = RecordBatchIterator::new(data.into_iter().map(arrow::error::Result::Ok), schema); + + let broken_version = dataset.version().version; + + // Any transaction, no matter how simple, should trigger the fragment bitmap to be recalculated + dataset + .append( + data, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await + .unwrap(); + + for idx in dataset.load_indices().await.unwrap().iter() { + // The corrupt fragment_bitmap does not contain 0 but the + // restored one should + assert!(idx.fragment_bitmap.as_ref().unwrap().contains(0)); + } + + let mut dataset = dataset.checkout_version(broken_version).await.unwrap(); + dataset.restore().await.unwrap(); + + // Running compaction right away should work (this is verifying compaction + // is not broken by the potentially malformed fragment bitmaps) + compact_files(&mut dataset, CompactionOptions::default(), None) + .await + .unwrap(); + + for idx in dataset.load_indices().await.unwrap().iter() { + assert!(idx.fragment_bitmap.as_ref().unwrap().contains(0)); + } + + let mut scan = dataset.scan(); + let query_vec = Float32Array::from(vec![0_f32; 128]); + let batches = scan + .nearest("vector", &query_vec, 2000) + .unwrap() + .nprobes(4) + .prefilter(true) + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + + let row_count = batches.iter().map(|batch| batch.num_rows()).sum::(); + assert_eq!(row_count, 1900); +} + +#[tokio::test] +async fn test_fix_v0_10_5_corrupt_schema() { + // Schemas could be corrupted by successive calls to `add_columns` and + // `drop_columns`. We should be able to detect this by checking for + // duplicate field ids. We should be able to fix this in new commits + // by dropping unused data files and re-writing the schema. + + // Copy over table + let test_dir = copy_test_data_to_tmp("v0.10.5/corrupt_schema").unwrap(); + let test_uri = test_dir.path_str(); + let test_uri = &test_uri; + + let mut dataset = Dataset::open(test_uri).await.unwrap(); + + let validate_res = dataset.validate().await; + assert!(validate_res.is_err()); + + // Force a migration. + dataset.delete("false").await.unwrap(); + dataset.validate().await.unwrap(); + + let data = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!( + data["b"] + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[0, 4, 8, 12] + ); + assert_eq!( + data["c"] + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[0, 5, 10, 15] + ); +} + +#[tokio::test] +async fn test_fix_v0_21_0_corrupt_fragment_bitmap() { + // In v0.21.0 and earlier, delta indices had a bug where the fragment bitmap + // could contain fragments that are part of other index deltas. + + // Copy over table + let test_dir = copy_test_data_to_tmp("v0.21.0/bad_index_fragment_bitmap").unwrap(); + let test_uri = test_dir.path_str(); + let test_uri = &test_uri; + + let mut dataset = Dataset::open(test_uri).await.unwrap(); + + let validate_res = dataset.validate().await; + assert!(validate_res.is_err()); + assert_eq!(dataset.load_indices().await.unwrap()[0].name, "vector_idx"); + + // Calling index statistics will force a migration + let stats = dataset.index_statistics("vector_idx").await.unwrap(); + let stats: serde_json::Value = serde_json::from_str(&stats).unwrap(); + assert_eq!(stats["num_indexed_fragments"], 2); + + dataset.checkout_latest().await.unwrap(); + dataset.validate().await.unwrap(); + + let indices = dataset.load_indices().await.unwrap(); + assert_eq!(indices.len(), 2); + fn get_bitmap(meta: &IndexMetadata) -> Vec { + meta.fragment_bitmap.as_ref().unwrap().iter().collect() + } + assert_eq!(get_bitmap(&indices[0]), vec![0]); + assert_eq!(get_bitmap(&indices[1]), vec![1]); +} + +#[tokio::test] +async fn test_max_fragment_id_migration() { + // v0.5.9 and earlier did not store the max fragment id in the manifest. + // This test ensures that we can read such datasets and migrate them to + // the latest version, which requires the max fragment id to be present. + { + let test_dir = copy_test_data_to_tmp("v0.5.9/no_fragments").unwrap(); + let test_uri = test_dir.path_str(); + let test_uri = &test_uri; + let dataset = Dataset::open(test_uri).await.unwrap(); + + assert_eq!(dataset.manifest.max_fragment_id, None); + assert_eq!(dataset.manifest.max_fragment_id(), None); + } + + { + let test_dir = copy_test_data_to_tmp("v0.5.9/dataset_with_fragments").unwrap(); + let test_uri = test_dir.path_str(); + let test_uri = &test_uri; + let dataset = Dataset::open(test_uri).await.unwrap(); + + assert_eq!(dataset.manifest.max_fragment_id, None); + assert_eq!(dataset.manifest.max_fragment_id(), Some(2)); + } +} diff --git a/rust/lance/src/dataset/tests/dataset_transactions.rs b/rust/lance/src/dataset/tests/dataset_transactions.rs new file mode 100644 index 00000000000..552f3ee2df3 --- /dev/null +++ b/rust/lance/src/dataset/tests/dataset_transactions.rs @@ -0,0 +1,344 @@ +#![allow(clippy::redundant_pub_crate)] +use super::dataset_common::*; + +#[tokio::test] +async fn test_read_transaction_properties() { + const LANCE_COMMIT_MESSAGE_KEY: &str = "__lance_commit_message"; + // Create a test dataset + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("value", DataType::Utf8, false), + ])); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + ) + .unwrap(); + + let test_uri = TempStrDir::default(); + + // Create WriteParams with properties + let mut properties1 = HashMap::new(); + properties1.insert( + LANCE_COMMIT_MESSAGE_KEY.to_string(), + "First commit".to_string(), + ); + properties1.insert("custom_prop".to_string(), "custom_value".to_string()); + + let write_params = WriteParams { + transaction_properties: Some(Arc::new(properties1)), + ..Default::default() + }; + + let dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch.clone())], schema.clone()), + &test_uri, + Some(write_params), + ) + .await + .unwrap(); + + let transaction = dataset.read_transaction_by_version(1).await.unwrap(); + assert!(transaction.is_some()); + let props = transaction.unwrap().transaction_properties.unwrap(); + assert_eq!(props.len(), 2); + assert_eq!( + props.get(LANCE_COMMIT_MESSAGE_KEY), + Some(&"First commit".to_string()) + ); + assert_eq!(props.get("custom_prop"), Some(&"custom_value".to_string())); + + let mut properties2 = HashMap::new(); + properties2.insert( + LANCE_COMMIT_MESSAGE_KEY.to_string(), + "Second commit".to_string(), + ); + properties2.insert("another_prop".to_string(), "another_value".to_string()); + + let write_params = WriteParams { + transaction_properties: Some(Arc::new(properties2)), + mode: WriteMode::Append, + ..Default::default() + }; + + let batch2 = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4, 5])), + Arc::new(StringArray::from(vec!["d", "e"])), + ], + ) + .unwrap(); + + let mut dataset = dataset; + dataset + .append( + RecordBatchIterator::new([Ok(batch2)], schema.clone()), + Some(write_params), + ) + .await + .unwrap(); + + let transaction = dataset.read_transaction_by_version(2).await.unwrap(); + assert!(transaction.is_some()); + let props = transaction.unwrap().transaction_properties.unwrap(); + assert_eq!(props.len(), 2); + assert_eq!( + props.get(LANCE_COMMIT_MESSAGE_KEY), + Some(&"Second commit".to_string()) + ); + assert_eq!( + props.get("another_prop"), + Some(&"another_value".to_string()) + ); + + let transaction = dataset.read_transaction_by_version(1).await.unwrap(); + assert!(transaction.is_some()); + let props = transaction.unwrap().transaction_properties.unwrap(); + assert_eq!(props.len(), 2); + assert_eq!( + props.get(LANCE_COMMIT_MESSAGE_KEY), + Some(&"First commit".to_string()) + ); + assert_eq!(props.get("custom_prop"), Some(&"custom_value".to_string())); + + let result = dataset.read_transaction_by_version(999).await; + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_session_store_registry() { + // Create a session + let session = Arc::new(Session::default()); + let registry = session.store_registry(); + assert!(registry.active_stores().is_empty()); + + // Create a dataset with memory store + let write_params = WriteParams { + session: Some(session.clone()), + ..Default::default() + }; + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ArrowField::new( + "a", + DataType::Int32, + false, + )])), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let dataset = InsertBuilder::new("memory://test") + .with_params(&write_params) + .execute(vec![batch.clone()]) + .await + .unwrap(); + + // Assert there is one active store. + assert_eq!(registry.active_stores().len(), 1); + + // If we create another dataset also in memory, it should re-use the + // existing store. + let dataset2 = InsertBuilder::new("memory://test2") + .with_params(&write_params) + .execute(vec![batch.clone()]) + .await + .unwrap(); + assert_eq!(registry.active_stores().len(), 1); + assert_eq!( + Arc::as_ptr(&dataset.object_store().inner), + Arc::as_ptr(&dataset2.object_store().inner) + ); + + // If we create another with **different parameters**, it should create a new store. + let write_params2 = WriteParams { + session: Some(session.clone()), + store_params: Some(ObjectStoreParams { + block_size: Some(10_000), + ..Default::default() + }), + ..Default::default() + }; + let dataset3 = InsertBuilder::new("memory://test3") + .with_params(&write_params2) + .execute(vec![batch.clone()]) + .await + .unwrap(); + assert_eq!(registry.active_stores().len(), 2); + assert_ne!( + Arc::as_ptr(&dataset.object_store().inner), + Arc::as_ptr(&dataset3.object_store().inner) + ); + + // Remove both datasets + drop(dataset3); + assert_eq!(registry.active_stores().len(), 1); + drop(dataset2); + drop(dataset); + assert_eq!(registry.active_stores().len(), 0); +} + +#[tokio::test] +async fn test_migrate_v2_manifest_paths() { + let test_uri = TempStrDir::default(); + + let data = lance_datagen::gen_batch() + .col("key", array::step::()) + .into_reader_rows(RowCount::from(10), BatchCount::from(1)); + let mut dataset = Dataset::write(data, &test_uri, None).await.unwrap(); + assert_eq!( + dataset.manifest_location().naming_scheme, + ManifestNamingScheme::V1 + ); + + dataset.migrate_manifest_paths_v2().await.unwrap(); + assert_eq!( + dataset.manifest_location().naming_scheme, + ManifestNamingScheme::V2 + ); +} + +pub(crate) async fn execute_sql( + sql: &str, + table: String, + dataset: Arc, +) -> Result> { + let ctx = SessionContext::new(); + ctx.register_table( + table, + Arc::new(LanceTableProvider::new(dataset, false, false)), + )?; + register_functions(&ctx); + + let df = ctx.sql(sql).await?; + Ok(df + .execute_stream() + .await + .unwrap() + .try_collect::>() + .await?) +} + +pub(crate) fn assert_results( + results: Vec, + values: &T, +) { + assert_eq!(results.len(), 1); + let results = results.into_iter().next().unwrap(); + assert_eq!(results.num_columns(), 1); + + assert_eq!( + results.column(0).as_any().downcast_ref::().unwrap(), + values + ) +} + +#[tokio::test] +async fn test_inline_transaction() { + use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator}; + use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; + use std::sync::Arc; + + async fn create_dataset(rows: i32) -> Arc { + let dir = TempDir::default(); + let uri = dir.path_str(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..rows))], + ) + .unwrap(); + let ds = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + uri.as_str(), + None, + ) + .await + .unwrap(); + Arc::new(ds) + } + + fn make_tx(read_version: u64) -> Transaction { + Transaction::new(read_version, Operation::Append { fragments: vec![] }, None) + } + + async fn delete_external_tx_file(ds: &Dataset) { + if let Some(tx_file) = ds.manifest.transaction_file.as_ref() { + let tx_path = ds.base.child("_transactions").child(tx_file.as_str()); + let _ = ds.object_store.inner.delete(&tx_path).await; // ignore errors + } + } + + let session = Arc::new(Session::default()); + + // Case 1: Default write_flag=true, delete external transaction file, read should use inline transaction + let ds = create_dataset(5).await; + let read_version = ds.manifest().version; + let tx = make_tx(read_version); + let ds2 = CommitBuilder::new(ds.clone()) + .execute(tx.clone()) + .await + .unwrap(); + delete_external_tx_file(&ds2).await; + let read_tx = ds2.read_transaction().await.unwrap().unwrap(); + assert_eq!(read_tx, tx.clone()); + + // Case 2: reading small manifest caches transaction data, eliminating transaction reading IO. + let read_ds2 = DatasetBuilder::from_uri(ds2.uri.clone()) + .with_session(session.clone()) + .load() + .await + .unwrap(); + let stats = read_ds2.object_store().io_stats_incremental(); // Reset + assert!(stats.read_bytes < 64 * 1024); + // Because the manifest is so small, we should have opportunistically + // cached the transaction in memory already. + let inline_tx = read_ds2.read_transaction().await.unwrap().unwrap(); + let stats = read_ds2.object_store().io_stats_incremental(); + assert_eq!(stats.read_iops, 0); + assert_eq!(stats.read_bytes, 0); + assert_eq!(inline_tx, tx); + + // Case 3: manifest does not contain inline transaction, read should fall back to external transaction file + let ds = create_dataset(2).await; + let tx = make_tx(ds.manifest().version); + let tx_file = crate::io::commit::write_transaction_file(ds.object_store(), &ds.base, &tx) + .await + .unwrap(); + let (mut manifest, indices) = tx + .build_manifest( + Some(ds.manifest.as_ref()), + ds.load_indices().await.unwrap().as_ref().clone(), + &tx_file, + &ManifestWriteConfig::default(), + ) + .unwrap(); + let location = write_manifest_file( + ds.object_store(), + ds.commit_handler.as_ref(), + &ds.base, + &mut manifest, + if indices.is_empty() { + None + } else { + Some(indices.clone()) + }, + &ManifestWriteConfig::default(), + ds.manifest_location.naming_scheme, + None, + ) + .await + .unwrap(); + let ds_new = ds.checkout_version(location.version).await.unwrap(); + assert!(ds_new.manifest.transaction_section.is_none()); + assert!(ds_new.manifest.transaction_file.is_some()); + let read_tx = ds_new.read_transaction().await.unwrap().unwrap(); + assert_eq!(read_tx, tx); +} diff --git a/rust/lance/src/dataset/tests/dataset_versioning.rs b/rust/lance/src/dataset/tests/dataset_versioning.rs new file mode 100644 index 00000000000..3d8e00f0435 --- /dev/null +++ b/rust/lance/src/dataset/tests/dataset_versioning.rs @@ -0,0 +1,738 @@ +use super::dataset_common::*; + +fn assert_all_manifests_use_scheme(test_dir: &TempStdDir, scheme: ManifestNamingScheme) { + let entries_names = test_dir + .join("_versions") + .read_dir() + .unwrap() + .map(|entry| entry.unwrap().file_name().into_string().unwrap()) + .collect::>(); + assert!( + entries_names + .iter() + .all(|name| ManifestNamingScheme::detect_scheme(name) == Some(scheme)), + "Entries: {:?}", + entries_names + ); +} + +#[tokio::test] +async fn test_v2_manifest_path_create() { + // Can create a dataset, using V2 paths + let data = lance_datagen::gen_batch() + .col("key", array::step::()) + .into_batch_rows(RowCount::from(10)) + .unwrap(); + let test_dir = TempStdDir::default(); + let test_uri = test_dir.to_str().unwrap(); + Dataset::write( + RecordBatchIterator::new([Ok(data.clone())], data.schema().clone()), + test_uri, + Some(WriteParams { + enable_v2_manifest_paths: true, + ..Default::default() + }), + ) + .await + .unwrap(); + + assert_all_manifests_use_scheme(&test_dir, ManifestNamingScheme::V2); + + // Appending to it will continue to use those paths + let dataset = Dataset::write( + RecordBatchIterator::new([Ok(data.clone())], data.schema().clone()), + test_uri, + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap(); + + assert_all_manifests_use_scheme(&test_dir, ManifestNamingScheme::V2); + + UpdateBuilder::new(Arc::new(dataset)) + .update_where("key = 5") + .unwrap() + .set("key", "200") + .unwrap() + .build() + .unwrap() + .execute() + .await + .unwrap(); + + assert_all_manifests_use_scheme(&test_dir, ManifestNamingScheme::V2); +} + +#[tokio::test] +async fn test_v2_manifest_path_commit() { + let schema = Schema::try_from(&ArrowSchema::new(vec![ArrowField::new( + "x", + DataType::Int32, + false, + )])) + .unwrap(); + let operation = Operation::Overwrite { + fragments: vec![], + schema, + config_upsert_values: None, + initial_bases: None, + }; + let test_dir = TempStdDir::default(); + let test_uri = test_dir.to_str().unwrap(); + let dataset = Dataset::commit( + test_uri, + operation, + None, + None, + None, + Default::default(), + true, // enable_v2_manifest_paths + ) + .await + .unwrap(); + + assert!(dataset.manifest_location.naming_scheme == ManifestNamingScheme::V2); + + assert_all_manifests_use_scheme(&test_dir, ManifestNamingScheme::V2); +} + +#[tokio::test] +async fn test_strict_overwrite() { + let schema = Schema::try_from(&ArrowSchema::new(vec![ArrowField::new( + "x", + DataType::Int32, + false, + )])) + .unwrap(); + let operation = Operation::Overwrite { + fragments: vec![], + schema, + config_upsert_values: None, + initial_bases: None, + }; + let test_uri = TempStrDir::default(); + let read_version_0_transaction = Transaction::new(0, operation, None); + let strict_builder = CommitBuilder::new(&test_uri).with_max_retries(0); + let unstrict_builder = CommitBuilder::new(&test_uri).with_max_retries(1); + strict_builder + .clone() + .execute(read_version_0_transaction.clone()) + .await + .expect("Strict overwrite should succeed when writing a new dataset"); + strict_builder + .clone() + .execute(read_version_0_transaction.clone()) + .await + .expect_err("Strict overwrite should fail when committing to a stale version"); + unstrict_builder + .clone() + .execute(read_version_0_transaction.clone()) + .await + .expect("Unstrict overwrite should succeed when committing to a stale version"); +} + +#[rstest] +#[tokio::test] +async fn test_restore( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + // Create a table + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::UInt32, + false, + )])); + + let test_uri = TempStrDir::default(); + + let data = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(UInt32Array::from_iter_values(0..100))], + ); + let reader = RecordBatchIterator::new(vec![data.unwrap()].into_iter().map(Ok), schema); + let mut dataset = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.manifest.version, 1); + let original_manifest = dataset.manifest.clone(); + + // Delete some rows + dataset.delete("i > 50").await.unwrap(); + assert_eq!(dataset.manifest.version, 2); + + // Checkout a previous version + let mut dataset = dataset.checkout_version(1).await.unwrap(); + assert_eq!(dataset.manifest.version, 1); + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1); + assert_eq!(dataset.count_fragments(), 1); + assert_eq!(fragments[0].metadata.deletion_file, None); + assert_eq!(dataset.manifest, original_manifest); + + // Checkout latest and then go back. + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.manifest.version, 2); + let mut dataset = dataset.checkout_version(1).await.unwrap(); + + // Restore to a previous version + dataset.restore().await.unwrap(); + assert_eq!(dataset.manifest.version, 3); + assert_eq!(dataset.manifest.fragments, original_manifest.fragments); + assert_eq!(dataset.manifest.schema, original_manifest.schema); + + // Delete some rows again (make sure we can still write as usual) + dataset.delete("i > 30").await.unwrap(); + assert_eq!(dataset.manifest.version, 4); + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1); + assert_eq!(dataset.count_fragments(), 1); + assert!(fragments[0].metadata.deletion_file.is_some()); +} + +#[rstest] +#[tokio::test] +async fn test_tag( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, +) { + // Create a table + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::UInt32, + false, + )])); + + let test_uri = TempStrDir::default(); + + let data = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(UInt32Array::from_iter_values(0..100))], + ); + let reader = RecordBatchIterator::new(vec![data.unwrap()].into_iter().map(Ok), schema); + let mut dataset = Dataset::write( + reader, + &test_uri, + Some(WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.manifest.version, 1); + + // delete some rows + dataset.delete("i > 50").await.unwrap(); + assert_eq!(dataset.manifest.version, 2); + + assert_eq!(dataset.tags().list().await.unwrap().len(), 0); + + let bad_tag_creation = dataset.tags().create("tag1", 3).await; + assert_eq!( + bad_tag_creation.err().unwrap().to_string(), + "Version not found error: version Main::3 does not exist" + ); + + let bad_tag_deletion = dataset.tags().delete("tag1").await; + assert_eq!( + bad_tag_deletion.err().unwrap().to_string(), + "Ref not found error: tag tag1 does not exist" + ); + + dataset.tags().create("tag1", 1).await.unwrap(); + + assert_eq!(dataset.tags().list().await.unwrap().len(), 1); + + let another_bad_tag_creation = dataset.tags().create("tag1", 1).await; + assert_eq!( + another_bad_tag_creation.err().unwrap().to_string(), + "Ref conflict error: tag tag1 already exists" + ); + + dataset.tags().delete("tag1").await.unwrap(); + + assert_eq!(dataset.tags().list().await.unwrap().len(), 0); + + dataset.tags().create("tag1", 1).await.unwrap(); + dataset.tags().create("tag2", 1).await.unwrap(); + dataset.tags().create("v1.0.0-rc1", 2).await.unwrap(); + + let default_order = dataset.tags().list_tags_ordered(None).await.unwrap(); + let default_names: Vec<_> = default_order.iter().map(|t| &t.0).collect(); + assert_eq!( + default_names, + ["v1.0.0-rc1", "tag1", "tag2"], + "Default ordering mismatch" + ); + + let asc_order = dataset + .tags() + .list_tags_ordered(Some(Ordering::Less)) + .await + .unwrap(); + let asc_names: Vec<_> = asc_order.iter().map(|t| &t.0).collect(); + assert_eq!( + asc_names, + ["tag1", "tag2", "v1.0.0-rc1"], + "Ascending ordering mismatch" + ); + + let desc_order = dataset + .tags() + .list_tags_ordered(Some(Ordering::Greater)) + .await + .unwrap(); + let desc_names: Vec<_> = desc_order.iter().map(|t| &t.0).collect(); + assert_eq!( + desc_names, + ["v1.0.0-rc1", "tag1", "tag2"], + "Descending ordering mismatch" + ); + + assert_eq!(dataset.tags().list().await.unwrap().len(), 3); + + let bad_checkout = dataset.checkout_version("tag3").await; + assert_eq!( + bad_checkout.err().unwrap().to_string(), + "Ref not found error: tag tag3 does not exist" + ); + + dataset = dataset.checkout_version("tag1").await.unwrap(); + assert_eq!(dataset.manifest.version, 1); + + let first_ver = DatasetBuilder::from_uri(&test_uri) + .with_tag("tag1") + .load() + .await + .unwrap(); + assert_eq!(first_ver.version().version, 1); + + // test update tag + let bad_tag_update = dataset.tags().update("tag3", 1).await; + assert_eq!( + bad_tag_update.err().unwrap().to_string(), + "Ref not found error: tag tag3 does not exist" + ); + + let another_bad_tag_update = dataset.tags().update("tag1", 3).await; + assert_eq!( + another_bad_tag_update.err().unwrap().to_string(), + "Version not found error: version 3 does not exist" + ); + + dataset.tags().update("tag1", 2).await.unwrap(); + dataset = dataset.checkout_version("tag1").await.unwrap(); + assert_eq!(dataset.manifest.version, 2); + + dataset.tags().update("tag1", 1).await.unwrap(); + dataset = dataset.checkout_version("tag1").await.unwrap(); + assert_eq!(dataset.manifest.version, 1); +} + +#[rstest] +#[tokio::test] +async fn test_fragment_id_zero_not_reused() { + // Test case 1: Fragment id zero isn't re-used + // 1. Create a dataset with 1 fragment + // 2. Delete all rows + // 3. Append another fragment + // 4. Assert new fragment has id 1 not 0 + + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::UInt32, + false, + )])); + + // Create dataset with 1 fragment + let data = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(UInt32Array::from_iter_values(0..10))], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![data].into_iter().map(Ok), schema.clone()); + let mut dataset = Dataset::write(batches, &test_uri, None).await.unwrap(); + + // Verify we have 1 fragment with id 0 + assert_eq!(dataset.get_fragments().len(), 1); + assert_eq!(dataset.get_fragments()[0].id(), 0); + assert_eq!(dataset.manifest.max_fragment_id(), Some(0)); + + // Delete all rows + dataset.delete("true").await.unwrap(); + + // After deletion, dataset should be empty but max_fragment_id preserved + assert_eq!(dataset.get_fragments().len(), 0); + assert_eq!(dataset.count_rows(None).await.unwrap(), 0); + assert_eq!(dataset.manifest.max_fragment_id(), Some(0)); + + // Append another fragment + let data = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(UInt32Array::from_iter_values(20..30))], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![data].into_iter().map(Ok), schema.clone()); + let write_params = WriteParams { + mode: WriteMode::Append, + ..Default::default() + }; + let dataset = Dataset::write(batches, &test_uri, Some(write_params)) + .await + .unwrap(); + + // Assert new fragment has id 1, not 0 + assert_eq!(dataset.get_fragments().len(), 1); + assert_eq!(dataset.get_fragments()[0].id(), 1); + assert_eq!(dataset.manifest.max_fragment_id(), Some(1)); +} + +#[rstest] +#[tokio::test] +async fn test_fragment_id_never_reset() { + // Test case 2: Fragment id is never reset, even if all rows are deleted + // 1. Create dataset with N fragments + // 2. Delete all rows + // 3. Append more fragments + // 4. Assert new fragments have ids >= N + + let test_uri = TempStrDir::default(); + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::UInt32, + false, + )])); + + // Create dataset with 3 fragments (N=3) + let data = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(UInt32Array::from_iter_values(0..30))], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(data)], schema.clone()); + let write_params = WriteParams { + max_rows_per_file: 10, // Force multiple fragments + ..Default::default() + }; + let mut dataset = Dataset::write(batches, &test_uri, Some(write_params)) + .await + .unwrap(); + + // Verify we have 3 fragments with ids 0, 1, 2 + assert_eq!(dataset.get_fragments().len(), 3); + assert_eq!(dataset.get_fragments()[0].id(), 0); + assert_eq!(dataset.get_fragments()[1].id(), 1); + assert_eq!(dataset.get_fragments()[2].id(), 2); + assert_eq!(dataset.manifest.max_fragment_id(), Some(2)); + + // Delete all rows + dataset.delete("true").await.unwrap(); + + // After deletion, dataset should be empty but max_fragment_id preserved + assert_eq!(dataset.get_fragments().len(), 0); + assert_eq!(dataset.count_rows(None).await.unwrap(), 0); + assert_eq!(dataset.manifest.max_fragment_id(), Some(2)); + + // Append more fragments (2 new fragments) + let data = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(UInt32Array::from_iter_values(100..120))], + ) + .unwrap(); + let batches = RecordBatchIterator::new(vec![Ok(data)], schema.clone()); + let write_params = WriteParams { + mode: WriteMode::Append, + max_rows_per_file: 10, // Force multiple fragments + ..Default::default() + }; + let dataset = Dataset::write(batches, &test_uri, Some(write_params)) + .await + .unwrap(); + + // Assert new fragments have ids >= N (3, 4) + assert_eq!(dataset.get_fragments().len(), 2); + assert_eq!(dataset.get_fragments()[0].id(), 3); + assert_eq!(dataset.get_fragments()[1].id(), 4); + assert_eq!(dataset.manifest.max_fragment_id(), Some(4)); +} + +#[tokio::test] +async fn test_branch() { + let tempdir = TempDir::default(); + let test_uri = tempdir.path_str(); + let data_storage_version = LanceFileVersion::Stable; + + // Generate consistent test data batches + let generate_data = |prefix: &str, start_id: i32, row_count: u64| { + gen_batch() + .col("id", array::step_custom::(start_id, 1)) + .col("value", array::fill_utf8(format!("{prefix}_data"))) + .into_reader_rows(RowCount::from(row_count), BatchCount::from(1)) + }; + + // Reusable dataset writer with configurable mode + async fn write_dataset( + uri: &str, + data_reader: impl RecordBatchReader + Send + 'static, + mode: WriteMode, + version: LanceFileVersion, + ) -> Dataset { + let params = WriteParams { + max_rows_per_file: 100, + max_rows_per_group: 20, + data_storage_version: Some(version), + mode, + ..Default::default() + }; + Dataset::write(data_reader, uri, Some(params)) + .await + .unwrap() + } + + // Unified dataset scanning and row counting + async fn collect_rows(dataset: &Dataset) -> (usize, Vec) { + let batches = dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + (batches.iter().map(|b| b.num_rows()).sum(), batches) + } + + // Phase 1: Create empty dataset, write data batch 1, create branch1 based on version_number, write data batch 2 + let mut dataset = write_dataset( + &test_uri, + generate_data("batch1", 0, 50), + WriteMode::Create, + data_storage_version, + ) + .await; + + let original_version = dataset.version().version; + assert_eq!(original_version, 1); + + // Create branch1 on the latest version and write data batch 2 + let mut branch1_dataset = dataset + .create_branch("branch1", original_version, None) + .await + .unwrap(); + assert_eq!(branch1_dataset.uri, format!("{}/tree/branch1", test_uri)); + + branch1_dataset = write_dataset( + branch1_dataset.uri(), + generate_data("batch2", 50, 30), + WriteMode::Append, + data_storage_version, + ) + .await; + + // Phase 2: Create branch2 based on branch1's latest version_number, write data batch 3 + let mut branch2_dataset = branch1_dataset + .create_branch( + "dev/branch2", + ("branch1", branch1_dataset.version().version), + None, + ) + .await + .unwrap(); + assert_eq!( + branch2_dataset.uri, + format!("{}/tree/dev/branch2", test_uri) + ); + + branch2_dataset = write_dataset( + branch2_dataset.uri(), + generate_data("batch3", 80, 20), + WriteMode::Append, + data_storage_version, + ) + .await; + + // Phase 3: Create a tag on branch2, the actual tag content is under root dataset + // create branch3 based on that tag, write data batch 4 + branch2_dataset + .tags() + .create_on_branch( + "tag1", + branch2_dataset.version().version, + Some("dev/branch2"), + ) + .await + .unwrap(); + + let mut branch3_dataset = branch2_dataset + .create_branch("feature/nathan/branch3", "tag1", None) + .await + .unwrap(); + assert_eq!( + branch3_dataset.uri, + format!("{}/tree/feature/nathan/branch3", test_uri) + ); + + branch3_dataset = write_dataset( + branch3_dataset.uri(), + generate_data("batch4", 100, 25), + WriteMode::Append, + data_storage_version, + ) + .await; + + // Verify data correctness and independence of each branch + // Main branch only has data 1 (50 rows) + let main_dataset = Dataset::open(&test_uri).await.unwrap(); + let (main_rows, _) = collect_rows(&main_dataset).await; + assert_eq!(main_rows, 50); // only batch1 + assert_eq!(main_dataset.version().version, 1); + + // branch1 has data 1 + 2 (80 rows) + let updated_branch1 = Dataset::open(branch1_dataset.uri()).await.unwrap(); + let (branch1_rows, _) = collect_rows(&updated_branch1).await; + assert_eq!(branch1_rows, 80); // batch1+batch2 + assert_eq!(updated_branch1.version().version, 2); + + // branch2 has data 1 + 2 + 3 (100 rows) + let updated_branch2 = Dataset::open(branch2_dataset.uri()).await.unwrap(); + let (branch2_rows, _) = collect_rows(&updated_branch2).await; + assert_eq!(branch2_rows, 100); // batch1+batch2+batch3 + assert_eq!(updated_branch2.version().version, 3); + + // branch3 has data 1 + 2 + 3 + 4 (125 rows) + let updated_branch3 = Dataset::open(branch3_dataset.uri()).await.unwrap(); + let (branch3_rows, _) = collect_rows(&updated_branch3).await; + assert_eq!(branch3_rows, 125); // batch1+batch2+batch3+batch4 + assert_eq!(updated_branch3.version().version, 4); + + // Use list_branches to get branch list and verify each field of branch_content + let branches = dataset.list_branches().await.unwrap(); + assert_eq!(branches.len(), 3); + assert!(branches.contains_key("branch1")); + assert!(branches.contains_key("dev/branch2")); + assert!(branches.contains_key("feature/nathan/branch3")); + + // Verify branch1 content + let branch1_content = branches.get("branch1").unwrap(); + assert_eq!(branch1_content.parent_branch, None); // Created based on main branch + assert_eq!(branch1_content.parent_version, 1); + assert!(branch1_content.create_at > 0); + assert!(branch1_content.manifest_size > 0); + + // Verify branch2 content + let branch2_content = branches.get("dev/branch2").unwrap(); + assert_eq!(branch2_content.parent_branch.as_deref().unwrap(), "branch1"); + assert_eq!(branch2_content.parent_version, 2); + assert!(branch2_content.create_at > 0); + assert!(branch2_content.manifest_size > 0); + assert!(branch2_content.create_at >= branch1_content.create_at); + + // Verify branch3 content + let branch3_content = branches.get("feature/nathan/branch3").unwrap(); + // Created based on tag pointed to branch2 + assert_eq!( + branch3_content.parent_branch.as_deref().unwrap(), + "dev/branch2" + ); + assert_eq!(branch3_content.parent_version, 3); + assert!(branch3_content.create_at > 0); + assert!(branch3_content.manifest_size > 0); + assert!(branch3_content.create_at >= branch2_content.create_at); + + // Verify checkout_branch + let checkout_branch1 = main_dataset.checkout_branch("branch1").await.unwrap(); + let checkout_branch2 = checkout_branch1 + .checkout_branch("dev/branch2") + .await + .unwrap(); + let checkout_branch2_tag = checkout_branch1.checkout_version("tag1").await.unwrap(); + let checkout_branch3 = checkout_branch2_tag + .checkout_branch("feature/nathan/branch3") + .await + .unwrap(); + let checkout_branch3_at_version3 = checkout_branch2 + .checkout_version(("feature/nathan/branch3", 3)) + .await + .unwrap(); + assert_eq!(checkout_branch3.version().version, 4); + assert_eq!(checkout_branch3_at_version3.version().version, 3); + assert_eq!(checkout_branch2.version().version, 3); + assert_eq!(checkout_branch2_tag.version().version, 3); + assert_eq!(checkout_branch1.version().version, 2); + assert_eq!(checkout_branch3.count_rows(None).await.unwrap(), 125); + assert_eq!( + checkout_branch3_at_version3.count_rows(None).await.unwrap(), + 100 + ); + assert_eq!(checkout_branch2.count_rows(None).await.unwrap(), 100); + assert_eq!(checkout_branch2_tag.count_rows(None).await.unwrap(), 100); + assert_eq!(checkout_branch1.count_rows(None).await.unwrap(), 80); + assert_eq!( + checkout_branch3.manifest.branch.as_deref().unwrap(), + "feature/nathan/branch3" + ); + assert_eq!( + checkout_branch3_at_version3 + .manifest + .branch + .as_deref() + .unwrap(), + "feature/nathan/branch3" + ); + assert_eq!( + checkout_branch2.manifest.branch.as_deref().unwrap(), + "dev/branch2" + ); + assert_eq!( + checkout_branch2_tag.manifest.branch.as_deref().unwrap(), + "dev/branch2" + ); + assert_eq!( + checkout_branch1.manifest.branch.as_deref().unwrap(), + "branch1" + ); + + let mut dataset = main_dataset; + // Finally delete all branches + dataset.delete_branch("branch1").await.unwrap(); + dataset.delete_branch("dev/branch2").await.unwrap(); + // Test deleting zombie branch + let root_location = dataset.refs.root().unwrap(); + let branch_file = branch_contents_path(&root_location.path, "feature/nathan/branch3"); + dataset.object_store.delete(&branch_file).await.unwrap(); + // Now "feature/nathan/branch3" is a zombie branch + // Use delete_branch to verify if the directory is cleaned up + dataset + .force_delete_branch("feature/nathan/branch3") + .await + .unwrap(); + let cleaned_path = Path::parse(format!("{}/tree/feature", test_uri)).unwrap(); + assert!(!dataset.object_store.exists(&cleaned_path).await.unwrap()); + + // Verify list_branches is empty + let branches_after_delete = dataset.list_branches().await.unwrap(); + assert!(branches_after_delete.is_empty()); + + // Verify branch directories are all deleted cleanly + let test_path = tempdir.obj_path(); + let branches = dataset + .object_store + .read_dir(test_path.child("tree")) + .await + .unwrap(); + assert!(branches.is_empty()); +} diff --git a/rust/lance/src/dataset/tests/mod.rs b/rust/lance/src/dataset/tests/mod.rs new file mode 100644 index 00000000000..e0657c8b5de --- /dev/null +++ b/rust/lance/src/dataset/tests/mod.rs @@ -0,0 +1,10 @@ +#![allow(clippy::redundant_pub_crate)] +pub(crate) mod dataset_common; +mod dataset_concurrency_store; +mod dataset_geo; +mod dataset_index; +mod dataset_io; +mod dataset_merge_update; +mod dataset_migrations; +mod dataset_transactions; +mod dataset_versioning; From 7fc3736b74fb3bd9290c498b22add1eb386df6ce Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 3 Dec 2025 02:05:39 +0800 Subject: [PATCH 2/3] Fix build Signed-off-by: Xuanwo --- .../lance/src/dataset/tests/dataset_common.rs | 111 ++++-------------- .../tests/dataset_concurrency_store.rs | 15 ++- rust/lance/src/dataset/tests/dataset_geo.rs | 11 +- rust/lance/src/dataset/tests/dataset_index.rs | 46 +++++++- rust/lance/src/dataset/tests/dataset_io.rs | 39 +++++- .../src/dataset/tests/dataset_merge_update.rs | 36 +++++- .../src/dataset/tests/dataset_migrations.rs | 24 +++- .../src/dataset/tests/dataset_transactions.rs | 31 ++++- .../src/dataset/tests/dataset_versioning.rs | 25 +++- rust/lance/src/dataset/tests/mod.rs | 3 +- 10 files changed, 234 insertions(+), 107 deletions(-) diff --git a/rust/lance/src/dataset/tests/dataset_common.rs b/rust/lance/src/dataset/tests/dataset_common.rs index 20adaf5dff8..b8e46afa221 100644 --- a/rust/lance/src/dataset/tests/dataset_common.rs +++ b/rust/lance/src/dataset/tests/dataset_common.rs @@ -1,99 +1,28 @@ -#![allow(clippy::redundant_pub_crate)] -pub(crate) use std::collections::{HashMap, HashSet}; -pub(crate) use std::sync::Arc; -pub(crate) use std::vec; +use std::sync::Arc; -pub(crate) use crate::dataset::builder::DatasetBuilder; -pub(crate) use crate::dataset::tests::dataset_migrations::scan_dataset; -pub(crate) use crate::dataset::tests::dataset_transactions::{assert_results, execute_sql}; -pub(crate) use crate::dataset::{ - AutoCleanupParams, - ManifestWriteConfig, - ProjectionRequest, - write_manifest_file, +use arrow::array::as_struct_array; +use arrow::compute::concat_batches; +use arrow_array::{ + ArrayRef, DictionaryArray, Int32Array, RecordBatch, RecordBatchIterator, StringArray, + StructArray, UInt16Array, }; -pub(crate) use crate::{Dataset, Error, Result}; -pub(crate) use crate::session::Session; -pub(crate) use crate::dataset::optimize::{compact_files, CompactionOptions}; -pub(crate) use crate::dataset::transaction::{DataReplacementGroup, Operation, Transaction}; -pub(crate) use crate::dataset::WriteMode::Overwrite; -pub(crate) use crate::dataset::ROW_ID; -pub(crate) use crate::datatypes::Schema; -pub(crate) use crate::io::ObjectStoreParams; -pub(crate) use lance_core::ROW_ADDR; -pub(crate) use lance_table::format::{DataStorageFormat, IndexMetadata}; -pub(crate) use lance_table::io::commit::ManifestNamingScheme; -pub(crate) use crate::dataset::WriteDestination; -pub(crate) use crate::dataset::UpdateBuilder; -pub(crate) use crate::index::vector::VectorIndexParams; -pub(crate) use crate::utils::test::copy_test_data_to_tmp; -pub(crate) use lance_arrow::FixedSizeListArrayExt; -pub(crate) use mock_instant::thread_local::MockClock; +use arrow_ord::sort::sort_to_indices; +use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; +use arrow_select::take::take; +use futures::TryStreamExt; +use lance_file::version::LanceFileVersion; +use lance_table::format::WriterVersion; -pub(crate) use crate::dataset::write::{CommitBuilder, InsertBuilder, WriteMode, WriteParams}; -pub(crate) use arrow::array::{as_struct_array, AsArray, GenericListBuilder, GenericStringBuilder}; -pub(crate) use arrow::compute::concat_batches; -pub(crate) use arrow::datatypes::UInt64Type; -pub(crate) use arrow_array::{ - builder::StringDictionaryBuilder, - cast::as_string_array, - types::{Float32Type, Int32Type}, - ArrayRef, DictionaryArray, Float32Array, Int32Array, Int64Array, Int8Array, - Int8DictionaryArray, ListArray, RecordBatchIterator, StringArray, UInt16Array, UInt32Array, -}; -pub(crate) use arrow_array::{ - Array, FixedSizeListArray, GenericStringArray, Int16Array, Int16DictionaryArray, - LargeBinaryArray, StructArray, UInt64Array, -}; -pub(crate) use arrow_array::RecordBatch; -pub(crate) use arrow_array::RecordBatchReader; -pub(crate) use arrow_ord::sort::sort_to_indices; -pub(crate) use arrow_schema::{ - DataType, Field as ArrowField, Field, Fields as ArrowFields, Schema as ArrowSchema, -}; -pub(crate) use lance_arrow::bfloat16::{self, BFLOAT16_EXT_NAME}; -pub(crate) use lance_arrow::{ARROW_EXT_META_KEY, ARROW_EXT_NAME_KEY, BLOB_META_KEY}; -pub(crate) use lance_core::utils::tempfile::{TempDir, TempStdDir, TempStrDir}; -pub(crate) use lance_datagen::{array, gen_batch, BatchCount, Dimension, RowCount}; -pub(crate) use lance_file::version::LanceFileVersion; -pub(crate) use lance_file::writer::FileWriter; -pub(crate) use lance_index::scalar::inverted::{ - query::{BooleanQuery, MatchQuery, Occur, Operator, PhraseQuery}, - tokenizer::InvertedIndexParams, -}; -pub(crate) use lance_index::DatasetIndexExt; -pub(crate) use lance_index::scalar::FullTextSearchQuery; -pub(crate) use lance_index::{scalar::ScalarIndexParams, vector::DIST_COL, IndexType}; -pub(crate) use lance_io::assert_io_eq; -pub(crate) use lance_io::utils::CachedFileSize; -pub(crate) use lance_linalg::distance::MetricType; -pub(crate) use lance_table::feature_flags; -pub(crate) use lance_table::format::{DataFile, WriterVersion}; - -pub(crate) use crate::datafusion::LanceTableProvider; -pub(crate) use crate::dataset::refs::branch_contents_path; -pub(crate) use datafusion::common::{assert_contains, assert_not_contains}; -pub(crate) use datafusion::prelude::SessionContext; -pub(crate) use lance_arrow::json::ARROW_JSON_EXT_NAME; -pub(crate) use lance_datafusion::datagen::DatafusionDatagenExt; -pub(crate) use lance_datafusion::udf::register_functions; -pub(crate) use lance_index::scalar::inverted::query::{FtsQuery, MultiMatchQuery}; -pub(crate) use lance_testing::datagen::generate_random_array; -pub(crate) use itertools::Itertools; -pub(crate) use rand::seq::SliceRandom; -pub(crate) use rand::Rng; -pub(crate) use rstest::rstest; -pub(crate) use futures::{StreamExt, TryStreamExt}; -pub(crate) use std::cmp::Ordering; -pub(crate) use object_store::path::Path; -pub(crate) use lance_table::io::manifest::read_manifest; +use crate::dataset::write::WriteParams; +use crate::dataset::WriteMode; +use crate::Dataset; // Used to validate that futures returned are Send. -pub(crate) fn require_send(t: T) -> T { +pub(super) fn require_send(t: T) -> T { t } -pub(crate) async fn create_file( +pub(super) async fn create_file( path: &std::path::Path, mode: WriteMode, data_storage_version: LanceFileVersion, @@ -168,13 +97,13 @@ pub(crate) async fn create_file( let idx_arr = actual_batch.column_by_name("i").unwrap(); let sorted_indices = sort_to_indices(idx_arr, None, None).unwrap(); let struct_arr: StructArray = actual_batch.into(); - let sorted_arr = arrow_select::take::take(&struct_arr, &sorted_indices, None).unwrap(); + let sorted_arr = take(&struct_arr, &sorted_indices, None).unwrap(); let expected_struct_arr: StructArray = concat_batches(&schema, &expected_batches).unwrap().into(); assert_eq!(&expected_struct_arr, as_struct_array(sorted_arr.as_ref())); - // Each fragments has different fragment ID + // Each fragment has different fragment ID assert_eq!( actual_ds .fragments() @@ -182,5 +111,5 @@ pub(crate) async fn create_file( .map(|f| f.id) .collect::>(), (0..10).collect::>() - ) + ); } diff --git a/rust/lance/src/dataset/tests/dataset_concurrency_store.rs b/rust/lance/src/dataset/tests/dataset_concurrency_store.rs index 182f9d0e4e3..f2b3eeabea0 100644 --- a/rust/lance/src/dataset/tests/dataset_concurrency_store.rs +++ b/rust/lance/src/dataset/tests/dataset_concurrency_store.rs @@ -1,4 +1,17 @@ -use super::dataset_common::*; +use std::sync::Arc; +use std::vec; + +use crate::dataset::WriteDestination; +use crate::{Dataset, Error, Result}; + +use crate::dataset::write::{WriteMode, WriteParams}; +use arrow_array::RecordBatch; +use arrow_array::{Int32Array, RecordBatchIterator}; +use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; +use futures::TryStreamExt; +use lance_core::utils::tempfile::TempStrDir; +use lance_index::DatasetIndexExt; +use lance_index::{scalar::ScalarIndexParams, IndexType}; #[tokio::test] async fn concurrent_create() { diff --git a/rust/lance/src/dataset/tests/dataset_geo.rs b/rust/lance/src/dataset/tests/dataset_geo.rs index 379fb2ae036..aab867b02dd 100644 --- a/rust/lance/src/dataset/tests/dataset_geo.rs +++ b/rust/lance/src/dataset/tests/dataset_geo.rs @@ -1,4 +1,13 @@ -use super::dataset_common::*; +use std::sync::Arc; +use std::vec; + +use crate::dataset::tests::dataset_transactions::execute_sql; +use crate::Dataset; + +use arrow_array::cast::AsArray; +use arrow_array::RecordBatch; +use arrow_array::RecordBatchIterator; +use lance_core::utils::tempfile::TempStrDir; #[tokio::test] async fn test_geo_types() { diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 921cd8360aa..55a14a3f5f9 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -1,4 +1,48 @@ -use super::dataset_common::*; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::vec; + +use crate::dataset::tests::dataset_migrations::scan_dataset; +use crate::dataset::tests::dataset_transactions::{assert_results, execute_sql}; +use crate::dataset::ROW_ID; +use crate::index::vector::VectorIndexParams; +use crate::{Dataset, Error, Result}; +use lance_arrow::FixedSizeListArrayExt; + +use crate::dataset::write::{WriteMode, WriteParams}; +use arrow::array::{AsArray, GenericListBuilder, GenericStringBuilder}; +use arrow::datatypes::UInt64Type; +use arrow_array::RecordBatch; +use arrow_array::{ + builder::StringDictionaryBuilder, + types::{Float32Type, Int32Type}, + ArrayRef, Float32Array, Int32Array, RecordBatchIterator, StringArray, +}; +use arrow_array::{Array, GenericStringArray, StructArray, UInt64Array}; +use arrow_schema::{ + DataType, Field as ArrowField, Field, Fields as ArrowFields, Schema as ArrowSchema, +}; +use lance_arrow::ARROW_EXT_NAME_KEY; +use lance_core::utils::tempfile::TempStrDir; +use lance_datagen::{array, gen_batch, BatchCount, Dimension, RowCount}; +use lance_file::version::LanceFileVersion; +use lance_index::scalar::inverted::{ + query::{BooleanQuery, MatchQuery, Occur, Operator, PhraseQuery}, + tokenizer::InvertedIndexParams, +}; +use lance_index::scalar::FullTextSearchQuery; +use lance_index::DatasetIndexExt; +use lance_index::{scalar::ScalarIndexParams, vector::DIST_COL, IndexType}; +use lance_linalg::distance::MetricType; + +use datafusion::common::{assert_contains, assert_not_contains}; +use futures::{StreamExt, TryStreamExt}; +use itertools::Itertools; +use lance_arrow::json::ARROW_JSON_EXT_NAME; +use lance_index::scalar::inverted::query::{FtsQuery, MultiMatchQuery}; +use lance_testing::datagen::generate_random_array; +use rand::Rng; +use rstest::rstest; #[rstest] #[tokio::test] diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index 6377a9cb016..b62922beb89 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -1,4 +1,41 @@ -use super::dataset_common::*; +use std::sync::Arc; +use std::vec; + +use super::dataset_common::{create_file, require_send}; + +use crate::dataset::builder::DatasetBuilder; +use crate::dataset::WriteDestination; +use crate::dataset::WriteMode::Overwrite; +use crate::dataset::{write_manifest_file, ManifestWriteConfig}; +use crate::session::Session; +use crate::{Dataset, Error, Result}; +use lance_table::format::DataStorageFormat; + +use crate::dataset::write::{WriteMode, WriteParams}; +use arrow::array::as_struct_array; +use arrow::compute::concat_batches; +use arrow_array::RecordBatch; +use arrow_array::RecordBatchReader; +use arrow_array::{ + cast::as_string_array, + types::{Float32Type, Int32Type}, + ArrayRef, Int32Array, Int64Array, Int8Array, Int8DictionaryArray, RecordBatchIterator, + StringArray, +}; +use arrow_array::{Array, FixedSizeListArray, Int16Array, Int16DictionaryArray, StructArray}; +use arrow_ord::sort::sort_to_indices; +use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; +use lance_arrow::bfloat16::{self, BFLOAT16_EXT_NAME}; +use lance_arrow::{ARROW_EXT_META_KEY, ARROW_EXT_NAME_KEY}; +use lance_core::utils::tempfile::{TempStdDir, TempStrDir}; +use lance_datagen::{array, gen_batch, BatchCount, RowCount}; +use lance_file::version::LanceFileVersion; +use lance_io::assert_io_eq; +use lance_table::feature_flags; + +use futures::TryStreamExt; +use lance_table::io::manifest::read_manifest; +use rstest::rstest; #[rstest] #[lance_test_macros::test(tokio::test)] diff --git a/rust/lance/src/dataset/tests/dataset_merge_update.rs b/rust/lance/src/dataset/tests/dataset_merge_update.rs index b4dab03341f..a7e960c0501 100644 --- a/rust/lance/src/dataset/tests/dataset_merge_update.rs +++ b/rust/lance/src/dataset/tests/dataset_merge_update.rs @@ -1,4 +1,38 @@ -use super::dataset_common::*; +use std::sync::Arc; +use std::vec; + +use crate::dataset::optimize::{compact_files, CompactionOptions}; +use crate::dataset::transaction::{DataReplacementGroup, Operation}; +use crate::dataset::WriteDestination; +use crate::dataset::ROW_ID; +use crate::dataset::{AutoCleanupParams, ProjectionRequest}; +use crate::{Dataset, Error}; +use lance_core::ROW_ADDR; +use mock_instant::thread_local::MockClock; + +use crate::dataset::write::{InsertBuilder, WriteMode, WriteParams}; +use arrow::array::AsArray; +use arrow::compute::concat_batches; +use arrow_array::RecordBatch; +use arrow_array::{ + types::Int32Type, ArrayRef, Float32Array, Int32Array, ListArray, RecordBatchIterator, + StringArray, +}; +use arrow_array::{Array, LargeBinaryArray, StructArray}; +use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; +use lance_arrow::BLOB_META_KEY; +use lance_core::utils::tempfile::{TempDir, TempStrDir}; +use lance_datagen::{array, gen_batch, BatchCount, RowCount}; +use lance_file::version::LanceFileVersion; +use lance_file::writer::FileWriter; +use lance_io::utils::CachedFileSize; +use lance_table::format::DataFile; + +use futures::TryStreamExt; +use lance_datafusion::datagen::DatafusionDatagenExt; +use object_store::path::Path; +use rand::seq::SliceRandom; +use rstest::rstest; #[rstest] #[tokio::test] diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index 68ada995e8c..29255c08d26 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -1,7 +1,23 @@ -#![allow(clippy::redundant_pub_crate)] -use super::dataset_common::*; - -pub(crate) async fn scan_dataset(uri: &str) -> Result> { +use std::sync::Arc; +use std::vec; + +use crate::dataset::optimize::{compact_files, CompactionOptions}; +use crate::utils::test::copy_test_data_to_tmp; +use crate::{Dataset, Result}; +use lance_table::format::IndexMetadata; + +use crate::dataset::write::{WriteMode, WriteParams}; +use arrow::compute::concat_batches; +use arrow_array::RecordBatch; +use arrow_array::{Float32Array, Int64Array, RecordBatchIterator}; +use arrow_schema::Schema as ArrowSchema; +use lance_file::version::LanceFileVersion; +use lance_index::DatasetIndexExt; + +use futures::{StreamExt, TryStreamExt}; +use rstest::rstest; + +pub(super) async fn scan_dataset(uri: &str) -> Result> { let results = Dataset::open(uri) .await? .scan() diff --git a/rust/lance/src/dataset/tests/dataset_transactions.rs b/rust/lance/src/dataset/tests/dataset_transactions.rs index 552f3ee2df3..4250c2fed44 100644 --- a/rust/lance/src/dataset/tests/dataset_transactions.rs +++ b/rust/lance/src/dataset/tests/dataset_transactions.rs @@ -1,5 +1,28 @@ -#![allow(clippy::redundant_pub_crate)] -use super::dataset_common::*; +use std::collections::HashMap; +use std::sync::Arc; +use std::vec; + +use crate::dataset::builder::DatasetBuilder; +use crate::dataset::transaction::{Operation, Transaction}; +use crate::dataset::{write_manifest_file, ManifestWriteConfig}; +use crate::io::ObjectStoreParams; +use crate::session::Session; +use crate::{Dataset, Result}; +use lance_table::io::commit::ManifestNamingScheme; + +use crate::dataset::write::{CommitBuilder, InsertBuilder, WriteMode, WriteParams}; +use arrow_array::Array; +use arrow_array::RecordBatch; +use arrow_array::{types::Int32Type, Int32Array, RecordBatchIterator, StringArray}; +use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; +use lance_core::utils::tempfile::{TempDir, TempStrDir}; +use lance_datagen::{array, BatchCount, RowCount}; +use lance_index::DatasetIndexExt; + +use crate::datafusion::LanceTableProvider; +use datafusion::prelude::SessionContext; +use futures::TryStreamExt; +use lance_datafusion::udf::register_functions; #[tokio::test] async fn test_read_transaction_properties() { @@ -201,7 +224,7 @@ async fn test_migrate_v2_manifest_paths() { ); } -pub(crate) async fn execute_sql( +pub(super) async fn execute_sql( sql: &str, table: String, dataset: Arc, @@ -222,7 +245,7 @@ pub(crate) async fn execute_sql( .await?) } -pub(crate) fn assert_results( +pub(super) fn assert_results( results: Vec, values: &T, ) { diff --git a/rust/lance/src/dataset/tests/dataset_versioning.rs b/rust/lance/src/dataset/tests/dataset_versioning.rs index 3d8e00f0435..1bac0073426 100644 --- a/rust/lance/src/dataset/tests/dataset_versioning.rs +++ b/rust/lance/src/dataset/tests/dataset_versioning.rs @@ -1,4 +1,27 @@ -use super::dataset_common::*; +use std::sync::Arc; +use std::vec; + +use crate::dataset::builder::DatasetBuilder; +use crate::dataset::transaction::{Operation, Transaction}; +use crate::dataset::UpdateBuilder; +use crate::datatypes::Schema; +use crate::Dataset; +use lance_table::io::commit::ManifestNamingScheme; + +use crate::dataset::write::{CommitBuilder, WriteMode, WriteParams}; +use arrow_array::RecordBatch; +use arrow_array::RecordBatchReader; +use arrow_array::{types::Int32Type, RecordBatchIterator, UInt32Array}; +use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; +use lance_core::utils::tempfile::{TempDir, TempStdDir, TempStrDir}; +use lance_datagen::{array, gen_batch, BatchCount, RowCount}; +use lance_file::version::LanceFileVersion; + +use crate::dataset::refs::branch_contents_path; +use futures::TryStreamExt; +use object_store::path::Path; +use rstest::rstest; +use std::cmp::Ordering; fn assert_all_manifests_use_scheme(test_dir: &TempStdDir, scheme: ManifestNamingScheme) { let entries_names = test_dir diff --git a/rust/lance/src/dataset/tests/mod.rs b/rust/lance/src/dataset/tests/mod.rs index e0657c8b5de..c0e4f7d9ee1 100644 --- a/rust/lance/src/dataset/tests/mod.rs +++ b/rust/lance/src/dataset/tests/mod.rs @@ -1,5 +1,4 @@ -#![allow(clippy::redundant_pub_crate)] -pub(crate) mod dataset_common; +mod dataset_common; mod dataset_concurrency_store; mod dataset_geo; mod dataset_index; From eef7da6c6dec555dedc04b64afac14c12cf1d426 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 3 Dec 2025 02:09:44 +0800 Subject: [PATCH 3/3] Fix licenses Signed-off-by: Xuanwo --- rust/lance/src/dataset/tests/dataset_common.rs | 3 +++ rust/lance/src/dataset/tests/dataset_concurrency_store.rs | 3 +++ rust/lance/src/dataset/tests/dataset_geo.rs | 3 +++ rust/lance/src/dataset/tests/dataset_index.rs | 3 +++ rust/lance/src/dataset/tests/dataset_io.rs | 3 +++ rust/lance/src/dataset/tests/dataset_merge_update.rs | 3 +++ rust/lance/src/dataset/tests/dataset_migrations.rs | 3 +++ rust/lance/src/dataset/tests/dataset_transactions.rs | 3 +++ rust/lance/src/dataset/tests/dataset_versioning.rs | 3 +++ rust/lance/src/dataset/tests/mod.rs | 3 +++ 10 files changed, 30 insertions(+) diff --git a/rust/lance/src/dataset/tests/dataset_common.rs b/rust/lance/src/dataset/tests/dataset_common.rs index b8e46afa221..88fc419067d 100644 --- a/rust/lance/src/dataset/tests/dataset_common.rs +++ b/rust/lance/src/dataset/tests/dataset_common.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + use std::sync::Arc; use arrow::array::as_struct_array; diff --git a/rust/lance/src/dataset/tests/dataset_concurrency_store.rs b/rust/lance/src/dataset/tests/dataset_concurrency_store.rs index f2b3eeabea0..cb445347ab3 100644 --- a/rust/lance/src/dataset/tests/dataset_concurrency_store.rs +++ b/rust/lance/src/dataset/tests/dataset_concurrency_store.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + use std::sync::Arc; use std::vec; diff --git a/rust/lance/src/dataset/tests/dataset_geo.rs b/rust/lance/src/dataset/tests/dataset_geo.rs index aab867b02dd..2c8bcccdd13 100644 --- a/rust/lance/src/dataset/tests/dataset_geo.rs +++ b/rust/lance/src/dataset/tests/dataset_geo.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + use std::sync::Arc; use std::vec; diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 55a14a3f5f9..ce8254cbd52 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::vec; diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index b62922beb89..0c163de3f96 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + use std::sync::Arc; use std::vec; diff --git a/rust/lance/src/dataset/tests/dataset_merge_update.rs b/rust/lance/src/dataset/tests/dataset_merge_update.rs index a7e960c0501..aa35f1b6408 100644 --- a/rust/lance/src/dataset/tests/dataset_merge_update.rs +++ b/rust/lance/src/dataset/tests/dataset_merge_update.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + use std::sync::Arc; use std::vec; diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index 29255c08d26..abccd20edd3 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + use std::sync::Arc; use std::vec; diff --git a/rust/lance/src/dataset/tests/dataset_transactions.rs b/rust/lance/src/dataset/tests/dataset_transactions.rs index 4250c2fed44..387512497bc 100644 --- a/rust/lance/src/dataset/tests/dataset_transactions.rs +++ b/rust/lance/src/dataset/tests/dataset_transactions.rs @@ -1,4 +1,7 @@ use std::collections::HashMap; +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + use std::sync::Arc; use std::vec; diff --git a/rust/lance/src/dataset/tests/dataset_versioning.rs b/rust/lance/src/dataset/tests/dataset_versioning.rs index 1bac0073426..cfefa23ab3b 100644 --- a/rust/lance/src/dataset/tests/dataset_versioning.rs +++ b/rust/lance/src/dataset/tests/dataset_versioning.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + use std::sync::Arc; use std::vec; diff --git a/rust/lance/src/dataset/tests/mod.rs b/rust/lance/src/dataset/tests/mod.rs index c0e4f7d9ee1..95c82c8c732 100644 --- a/rust/lance/src/dataset/tests/mod.rs +++ b/rust/lance/src/dataset/tests/mod.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + mod dataset_common; mod dataset_concurrency_store; mod dataset_geo;