diff --git a/rust/lance/Cargo.toml b/rust/lance/Cargo.toml index 417d41a5fe1..bf19b7659b0 100644 --- a/rust/lance/Cargo.toml +++ b/rust/lance/Cargo.toml @@ -198,6 +198,10 @@ harness = false name = "scan" harness = false +[[bench]] +name = "cdf" +harness = false + [[bench]] name = "s3_file_reader_diagnostics" harness = false diff --git a/rust/lance/benches/cdf.rs b/rust/lance/benches/cdf.rs new file mode 100644 index 00000000000..ce751cf1cc0 --- /dev/null +++ b/rust/lance/benches/cdf.rs @@ -0,0 +1,186 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Compare CDF fragment pruning with equivalent full scans. +//! +//! Run with `cargo bench -p lance --bench cdf --profile release-with-debug`. + +use std::sync::Arc; +use std::time::Duration; + +use arrow_array::types::Int32Type; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use futures::TryStreamExt; +use lance::Dataset; +use lance::dataset::{UpdateBuilder, WriteParams}; +use lance_core::utils::tempfile::TempStrDir; +use lance_core::{ROW_CREATED_AT_VERSION, ROW_ID, ROW_LAST_UPDATED_AT_VERSION, WILDCARD}; +use lance_datagen::{BatchCount, RowCount, array, gen_batch}; + +// Small fragments isolate planning and file-open overhead. +const ROWS_PER_FRAGMENT: usize = 16; + +#[derive(Clone, Copy, Debug)] +enum Query { + Inserted, + Updated, + Upserted, + AllInserted, +} + +#[derive(Clone, Copy, Debug)] +enum ScanPath { + Full, + Pruned, +} + +struct Fixture { + _dir: TempStrDir, + dataset: Dataset, +} + +impl Fixture { + async fn new(fragments: u32) -> Self { + let dir = TempStrDir::default(); + let reader = gen_batch() + .col("key", array::step::()) + .col("value", array::fill_utf8("initial".into())) + .into_reader_rows( + RowCount::from(ROWS_PER_FRAGMENT as u64), + BatchCount::from(fragments), + ); + let mut dataset = Dataset::write( + reader, + dir.as_str(), + Some(WriteParams { + max_rows_per_file: ROWS_PER_FRAGMENT, + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), fragments as usize); + + let reader = gen_batch() + .col( + "key", + array::step_custom::(fragments as i32 * ROWS_PER_FRAGMENT as i32, 1), + ) + .col("value", array::fill_utf8("inserted".into())) + .into_reader_rows( + RowCount::from(ROWS_PER_FRAGMENT as u64), + BatchCount::from(1), + ); + dataset.append(reader, None).await.unwrap(); + let result = UpdateBuilder::new(Arc::new(dataset)) + .update_where("key = 0") + .unwrap() + .set("value", "'updated'") + .unwrap() + .build() + .unwrap() + .execute() + .await + .unwrap(); + Self { + _dir: dir, + dataset: result.new_dataset.as_ref().clone(), + } + } +} + +async fn scan(dataset: &Dataset, query: Query, path: ScanPath) -> usize { + let begin = if matches!(query, Query::AllInserted) { + 0 + } else { + 1 + }; + let mut stream = match path { + ScanPath::Full => { + let mut scanner = dataset.scan(); + scanner + .project(&[ + WILDCARD, + ROW_ID, + ROW_CREATED_AT_VERSION, + ROW_LAST_UPDATED_AT_VERSION, + ]) + .unwrap(); + let inserted = + format!("_row_created_at_version > {begin} AND _row_created_at_version <= 3"); + let updated = "_row_created_at_version <= 1 AND _row_last_updated_at_version > 1 AND _row_last_updated_at_version <= 3"; + let filter = match query { + Query::Inserted | Query::AllInserted => inserted, + Query::Updated => updated.to_string(), + Query::Upserted => format!("({inserted}) OR ({updated})"), + }; + scanner.filter(&filter).unwrap(); + scanner.try_into_stream().await.unwrap() + } + ScanPath::Pruned => { + let delta = dataset + .delta() + .with_begin_version(begin) + .with_end_version(3) + .build() + .unwrap(); + match query { + Query::Inserted | Query::AllInserted => delta.get_inserted_rows().await.unwrap(), + Query::Updated => delta.get_updated_rows().await.unwrap(), + Query::Upserted => delta.get_upserted_rows().await.unwrap(), + } + } + }; + let mut rows = 0; + while let Some(batch) = stream.try_next().await.unwrap() { + rows += batch.num_rows(); + } + // Dev-dependencies enable lance-io's `test-util`, which keeps a record of every + // I/O request until drained. Left to grow, that log slows each later scan, so + // sequential Criterion runs would misattribute the drift to whichever path ran last. + dataset + .object_store(None) + .await + .unwrap() + .io_stats_incremental(); + rows +} + +fn bench_cdf(c: &mut Criterion) { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("cdf"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(2)); + for fragments in [100, 1_000, 20_000] { + let fixture = runtime.block_on(Fixture::new(fragments)); + for query in [ + Query::Inserted, + Query::Updated, + Query::Upserted, + Query::AllInserted, + ] { + let expected = match query { + Query::Inserted => ROWS_PER_FRAGMENT, + Query::Updated => 1, + Query::Upserted => ROWS_PER_FRAGMENT + 1, + Query::AllInserted => (fragments as usize + 1) * ROWS_PER_FRAGMENT, + }; + for path in [ScanPath::Full, ScanPath::Pruned] { + assert_eq!( + runtime.block_on(scan(&fixture.dataset, query, path)), + expected + ); + group.bench_function( + BenchmarkId::new(format!("{query:?}/{path:?}"), fragments), + |b| b.iter(|| runtime.block_on(scan(&fixture.dataset, query, path))), + ); + } + } + } + group.finish(); +} + +criterion_group!(benches, bench_cdf); +criterion_main!(benches); diff --git a/rust/lance/src/dataset/delta.rs b/rust/lance/src/dataset/delta.rs index b9e12da2482..833617c63f3 100644 --- a/rust/lance/src/dataset/delta.rs +++ b/rust/lance/src/dataset/delta.rs @@ -37,8 +37,10 @@ use lance_datafusion::exec::{LanceExecutionOptions, OneShotExec, execute_plan}; use lance_table::format::Fragment; use lance_table::rowids::RowIdSequence; use lance_table::rowids::segment::U64Segment; +use lance_table::rowids::version::{RowDatasetVersionMeta, read_dataset_versions}; use std::collections::HashMap; use std::sync::Arc; +use tracing::warn; /// Rows per batch of [`DatasetDelta::get_deleted_row_ids`], taken from the /// scanner so it matches the sibling readers. @@ -241,6 +243,13 @@ pub struct DatasetDelta { pub(crate) end_timestamp: Option>, } +#[derive(Clone, Copy, Debug)] +enum RowChange { + Inserted, + Updated, + Upserted, +} + impl DatasetDelta { /// Resolve the effective version range for this delta. /// @@ -421,6 +430,7 @@ impl DatasetDelta { /// /// This returns rows where `_row_created_at_version` is greater than `begin_version` /// and less than or equal to `end_version`. + /// Available inline row-version metadata is used to prune fragments before reading data files. /// /// The result always includes: /// - `_row_created_at_version`: Version when the row was created @@ -449,29 +459,7 @@ impl DatasetDelta { /// # } /// ``` pub async fn get_inserted_rows(&self) -> Result { - let mut scanner = self.base_dataset.scan(); - - // Enable version columns - scanner.project(&[ - WILDCARD, - ROW_ID, - ROW_CREATED_AT_VERSION, - ROW_LAST_UPDATED_AT_VERSION, - ])?; - - // Filter for rows created in the version range - let filter = self.build_inserted_rows_filter().await?; - scanner.filter(&filter)?; - - scanner.try_into_stream().await - } - - async fn build_inserted_rows_filter(&self) -> Result { - let (begin_version, end_version) = self.resolve_range().await?; - Ok(format!( - "_row_created_at_version > {} AND _row_created_at_version <= {}", - begin_version, end_version - )) + self.scan_rows(RowChange::Inserted).await } /// Get updated rows between the two versions. @@ -479,6 +467,7 @@ impl DatasetDelta { /// This returns rows where `_row_last_updated_at_version` is greater than `begin_version` /// and less than or equal to `end_version`, but `_row_created_at_version` is less than /// or equal to `begin_version` (to exclude newly inserted rows). + /// Available inline row-version metadata is used to prune fragments before reading data files. /// /// The result always includes: /// - `_row_created_at_version`: Version when the row was created @@ -507,29 +496,7 @@ impl DatasetDelta { /// # } /// ``` pub async fn get_updated_rows(&self) -> Result { - let mut scanner = self.base_dataset.scan(); - - // Enable version columns - scanner.project(&[ - WILDCARD, - ROW_ID, - ROW_CREATED_AT_VERSION, - ROW_LAST_UPDATED_AT_VERSION, - ])?; - - // Filter for rows that were updated (not inserted) in the version range - let filter = self.build_updated_rows_batch_filter().await?; - scanner.filter(&filter)?; - - scanner.try_into_stream().await - } - - async fn build_updated_rows_batch_filter(&self) -> Result { - let (begin_version, end_version) = self.resolve_range().await?; - Ok(format!( - "_row_created_at_version <= {} AND _row_last_updated_at_version > {} AND _row_last_updated_at_version <= {}", - begin_version, begin_version, end_version - )) + self.scan_rows(RowChange::Updated).await } /// Get upserted rows between the two versions. @@ -542,6 +509,7 @@ impl DatasetDelta { /// Condition 2: /// This returns rows where `_row_created_at_version` is greater than `begin_version` /// and less than or equal to `end_version`. + /// Available inline row-version metadata is used to prune fragments before reading data files. /// /// The result always includes: /// - `_row_created_at_version`: Version when the row was created @@ -570,9 +538,61 @@ impl DatasetDelta { /// # } /// ``` pub async fn get_upserted_rows(&self) -> Result { - let mut scanner = self.base_dataset.scan(); + self.scan_rows(RowChange::Upserted).await + } - // Enable version columns + async fn scan_rows(&self, change: RowChange) -> Result { + let (begin_version, end_version) = self.resolve_range().await?; + let inserted_filter = || { + format!( + "{ROW_CREATED_AT_VERSION} > {begin_version} AND {ROW_CREATED_AT_VERSION} <= {end_version}" + ) + }; + let updated_filter = || { + format!( + "{ROW_CREATED_AT_VERSION} <= {begin_version} AND {ROW_LAST_UPDATED_AT_VERSION} > {begin_version} AND {ROW_LAST_UPDATED_AT_VERSION} <= {end_version}" + ) + }; + let filter = match change { + RowChange::Inserted => inserted_filter(), + RowChange::Updated => updated_filter(), + RowChange::Upserted => format!("({}) OR ({})", inserted_filter(), updated_filter()), + }; + + let fragments = pruned_fragments(&self.base_dataset.manifest.fragments, |fragment| { + let in_range = |version| version > begin_version && version <= end_version; + let may_have_inserts = || { + version_meta_may_match( + fragment.id, + ROW_CREATED_AT_VERSION, + fragment.created_at_version_meta.as_ref(), + in_range, + ) + }; + let may_have_updates = || { + version_meta_may_match( + fragment.id, + ROW_LAST_UPDATED_AT_VERSION, + fragment.last_updated_at_version_meta.as_ref(), + in_range, + ) && version_meta_may_match( + fragment.id, + ROW_CREATED_AT_VERSION, + fragment.created_at_version_meta.as_ref(), + |version| version <= begin_version, + ) + }; + match change { + RowChange::Inserted => may_have_inserts(), + RowChange::Updated => may_have_updates(), + RowChange::Upserted => may_have_inserts() || may_have_updates(), + } + }); + + let mut scanner = self.base_dataset.scan(); + if let Some(fragments) = fragments { + scanner.with_fragments(fragments); + } scanner.project(&[ WILDCARD, ROW_ID, @@ -580,21 +600,49 @@ impl DatasetDelta { ROW_LAST_UPDATED_AT_VERSION, ])?; - // Filter for rows that were updated or inserted in the version range - let filter = self.build_upserted_rows_filter().await?; + // Fragment pruning is conservative; the row filter determines the exact result. scanner.filter(&filter)?; - scanner.try_into_stream().await } +} - async fn build_upserted_rows_filter(&self) -> Result { - let inserted_row_filter = self.build_inserted_rows_filter().await?; - let updated_rows_filter = self.build_updated_rows_batch_filter().await?; - Ok(format!( - "({}) OR ({})", - inserted_row_filter, updated_rows_filter - )) - } +fn pruned_fragments( + fragments: &[Fragment], + may_match: impl Fn(&Fragment) -> bool, +) -> Option> { + // Reuse the manifest's shared list unless pruning excludes a fragment. + let mut remaining = fragments.iter(); + let first_excluded = remaining.position(|fragment| !may_match(fragment))?; + let mut surviving = Vec::with_capacity(first_excluded); + surviving.extend(fragments[..first_excluded].iter().cloned()); + surviving.extend(remaining.filter(|fragment| may_match(fragment)).cloned()); + Some(surviving) +} + +fn version_meta_may_match( + fragment_id: u64, + column: &str, + meta: Option<&RowDatasetVersionMeta>, + predicate: impl Fn(u64) -> bool, +) -> bool { + // Missing or non-inline metadata cannot safely rule out a match. + let Some(RowDatasetVersionMeta::Inline(data)) = meta else { + return true; + }; + let sequence = match read_dataset_versions(data) { + Ok(sequence) => sequence, + Err(error) => { + warn!( + fragment_id, + column, + metadata_bytes = data.len(), + error = %error, + "Cannot decode row-version metadata; skipping pruning for this column" + ); + return true; + } + }; + sequence.runs.is_empty() || sequence.runs.iter().any(|run| predicate(run.version)) } /// A fragment's deletion vector at this version, empty where it has none. @@ -956,10 +1004,22 @@ mod tests { use arrow_array::types::UInt64Type; use chrono::Duration; use futures::TryStreamExt; - use lance_core::{ROW_CREATED_AT_VERSION, ROW_ID, ROW_LAST_UPDATED_AT_VERSION}; + use lance_core::{ROW_CREATED_AT_VERSION, ROW_ID, ROW_LAST_UPDATED_AT_VERSION, WILDCARD}; use lance_datagen::{BatchCount, RowCount, array}; + use lance_io::stream::RecordBatchStream; + use lance_table::rowids::version::{ + RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, + }; use mock_instant::thread_local::MockClock; + use rstest::rstest; + use std::cell::RefCell; + use std::collections::BTreeSet; use std::sync::Arc; + use tempfile::NamedTempFile; + use tracing::Level; + + use super::{RowChange, pruned_fragments, version_meta_may_match}; + use crate::dataset::optimize::{CompactionOptions, compact_files}; async fn create_test_dataset( rows: usize, @@ -1044,8 +1104,333 @@ mod tests { async fn collect_stream( stream: crate::dataset::scanner::DatasetRecordBatchStream, ) -> arrow_array::RecordBatch { + let schema = stream.schema(); let batches: Vec<_> = stream.try_collect().await.unwrap(); - arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap() + arrow_select::concat::concat_batches(&schema, &batches).unwrap() + } + + #[rstest] + #[case::empty(vec![], None)] + #[case::all_kept(vec![true, true, true], None)] + #[case::first_excluded(vec![false, true, true], Some(vec![1, 2]))] + #[case::middle_excluded(vec![true, false, true], Some(vec![0, 2]))] + #[case::last_excluded(vec![true, true, false], Some(vec![0, 1]))] + #[case::all_excluded(vec![false, false, false], Some(vec![]))] + fn test_pruned_fragments(#[case] matches: Vec, #[case] expected: Option>) { + let fragments: Vec<_> = (0..matches.len()) + .map(|id| super::Fragment::new(id as u64)) + .collect(); + let visited = RefCell::new(Vec::new()); + let actual = pruned_fragments(&fragments, |fragment| { + visited.borrow_mut().push(fragment.id); + matches[fragment.id as usize] + }); + assert_eq!( + actual.map(|fragments| fragments.into_iter().map(|fragment| fragment.id).collect()), + expected + ); + assert_eq!( + visited.into_inner(), + fragments + .iter() + .map(|fragment| fragment.id) + .collect::>() + ); + } + + #[rstest] + #[case::before(vec![1], 2, 4, false)] + #[case::at_begin(vec![2], 2, 4, false)] + #[case::at_end(vec![4], 2, 4, true)] + #[case::after(vec![5], 2, 4, false)] + #[case::mixed(vec![5, 1, 3], 2, 4, true)] + #[case::gap(vec![1, 5], 2, 4, false)] + #[case::empty_range(vec![2], 2, 2, false)] + #[case::reversed_range(vec![2], 3, 1, false)] + #[case::max_version(vec![u64::MAX], u64::MAX - 1, u64::MAX, true)] + fn test_version_meta_pruning( + #[case] versions: Vec, + #[case] begin: u64, + #[case] end: u64, + #[case] expected: bool, + ) { + let sequence = RowDatasetVersionSequence { + runs: versions + .into_iter() + .enumerate() + .map(|(offset, version)| RowDatasetVersionRun { + span: super::U64Segment::Range(offset as u64..offset as u64 + 1), + version, + }) + .collect(), + }; + let meta = RowDatasetVersionMeta::from_sequence(&sequence).unwrap(); + assert_eq!( + version_meta_may_match(42, ROW_CREATED_AT_VERSION, Some(&meta), |version| { + version > begin && version <= end + }), + expected + ); + } + + #[rstest] + #[case::missing(None, false)] + #[case::undecodable(Some(RowDatasetVersionMeta::Inline(Arc::from([0xff]))), true)] + #[case::external(Some(RowDatasetVersionMeta::from_external_file("versions".into(), 0, 10)), false)] + #[case::empty(Some(RowDatasetVersionMeta::from_sequence(&RowDatasetVersionSequence::new()).unwrap()), false)] + fn test_version_meta_pruning_unknown( + #[case] meta: Option, + #[case] has_warning: bool, + #[values(ROW_CREATED_AT_VERSION, ROW_LAST_UPDATED_AT_VERSION)] column: &str, + ) { + let log = NamedTempFile::new().unwrap(); + let subscriber = tracing_subscriber::fmt() + .with_max_level(Level::WARN) + .with_ansi(false) + .without_time() + .with_writer(log.reopen().unwrap()) + .finish(); + tracing::subscriber::with_default(subscriber, || { + assert!(version_meta_may_match(42, column, meta.as_ref(), |_| false)); + }); + let output = std::fs::read_to_string(log.path()).unwrap(); + if has_warning { + assert!(output.contains("WARN")); + assert!(output.contains("fragment_id=42")); + assert!(output.contains(column)); + assert!(output.contains("metadata_bytes=1")); + assert!(output.contains("Failed to decode RowDatasetVersionSequence")); + assert_eq!( + output + .matches("Cannot decode row-version metadata; skipping pruning for this column") + .count(), + 1 + ); + } else { + assert!(output.is_empty(), "Unexpected warning: {output}"); + } + } + + #[rstest] + #[case::inserted( + RowChange::Inserted, + "_row_created_at_version > 1 AND _row_created_at_version <= 3", + vec![13, 14, 15] + )] + #[case::updated( + RowChange::Updated, + "_row_created_at_version <= 1 AND _row_last_updated_at_version > 1 AND _row_last_updated_at_version <= 3", + vec![1] + )] + #[case::upserted( + RowChange::Upserted, + "(_row_created_at_version > 1 AND _row_created_at_version <= 3) OR (_row_created_at_version <= 1 AND _row_last_updated_at_version > 1 AND _row_last_updated_at_version <= 3)", + vec![1, 13, 14, 15] + )] + #[tokio::test] + async fn test_delta_prunes_version_fragments( + #[case] change: RowChange, + #[case] filter: &str, + #[case] expected_keys: Vec, + ) { + let dir = lance_core::utils::tempfile::TempStrDir::default(); + let data = lance_datagen::gen_batch() + .col("key", array::step::()) + .col("value", array::fill_utf8("initial".into())) + .into_reader_rows(RowCount::from(4), BatchCount::from(3)); + let mut dataset = Dataset::write( + data, + dir.as_str(), + Some(WriteParams { + max_rows_per_file: 4, + enable_stable_row_ids: true, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 3); + + dataset = write_dataset_temp(&dir, 12, 4, 1, "inserted", true, true).await; + let inserted_file = dataset.manifest.fragments.last().unwrap().files[0] + .path + .clone(); + dataset = update_where(dataset, "key = 1", "updated").await; + let updated_file = dataset.manifest.fragments.last().unwrap().files[0] + .path + .clone(); + dataset = write_dataset_temp(&dir, 16, 4, 1, "later", true, true).await; + dataset.delete("key = 12").await.unwrap(); + + let columns = &[ + WILDCARD, + ROW_ID, + ROW_CREATED_AT_VERSION, + ROW_LAST_UPDATED_AT_VERSION, + ]; + let collect_changes = async |dataset: &Dataset| { + let delta = dataset + .delta() + .with_begin_version(1) + .with_end_version(3) + .build() + .unwrap(); + let stream = match change { + RowChange::Inserted => delta.get_inserted_rows().await, + RowChange::Updated => delta.get_updated_rows().await, + RowChange::Upserted => delta.get_upserted_rows().await, + } + .unwrap(); + collect_stream(stream).await + }; + let take_read_data_files = |dataset: &Dataset| -> BTreeSet { + dataset + .object_store + .io_stats_incremental() + .requests + .iter() + .filter_map(|request| request.path.filename()) + .filter(|filename| filename.ends_with(".lance")) + .map(str::to_owned) + .collect() + }; + dataset.session().file_metadata_cache().clear().await; + dataset.object_store.io_stats_incremental(); + let expected = scan_project_filter(&dataset, columns, Some(filter)).await; + let full_scan_files = take_read_data_files(&dataset); + assert_eq!(full_scan_files.len(), dataset.get_fragments().len()); + // Tiny files can cache inline pages with their metadata, hiding file reads. + dataset.session().file_metadata_cache().clear().await; + let actual = collect_changes(&dataset).await; + assert_eq!(actual, expected); + let mut keys = actual["key"].as_primitive::().values().to_vec(); + keys.sort_unstable(); + assert_eq!(keys, expected_keys); + + let read_files = take_read_data_files(&dataset); + let expected_files = match change { + RowChange::Inserted => BTreeSet::from([inserted_file]), + RowChange::Updated => BTreeSet::from([updated_file]), + RowChange::Upserted => BTreeSet::from([inserted_file, updated_file]), + }; + assert_eq!(read_files, expected_files); + assert!(read_files.len() < full_scan_files.len()); + + for column in [ROW_CREATED_AT_VERSION, ROW_LAST_UPDATED_AT_VERSION] { + for meta in [None, Some(RowDatasetVersionMeta::Inline(Arc::from([0xff])))] { + let mut unknown_dataset = dataset.clone(); + let manifest = Arc::make_mut(&mut unknown_dataset.manifest); + for fragment in Arc::make_mut(&mut manifest.fragments) { + if column == ROW_CREATED_AT_VERSION { + fragment.created_at_version_meta = meta.clone(); + } else { + fragment.last_updated_at_version_meta = meta.clone(); + } + } + let expected = scan_project_filter(&unknown_dataset, columns, Some(filter)).await; + assert_eq!( + collect_changes(&unknown_dataset).await, + expected, + "Pruning changed results for {change:?} with {column} metadata {meta:?}" + ); + } + } + + let metrics = compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 100, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + assert!(metrics.fragments_removed > 0); + assert_eq!(dataset.get_fragments().len(), 1); + let sequence = dataset.manifest.fragments[0] + .created_at_version_meta + .as_ref() + .unwrap() + .load_sequence() + .unwrap(); + assert!(sequence.runs.len() > 1); + let expected = scan_project_filter(&dataset, columns, Some(filter)).await; + assert_eq!(collect_changes(&dataset).await, expected); + } + + #[rstest] + #[case::same_version(1, 1)] + #[case::future_range(4, 5)] + #[case::reversed_range(3, 1)] + #[tokio::test] + async fn test_delta_pruning_empty(#[case] begin: u64, #[case] end: u64) { + let dataset = create_test_dataset(4, 1, "value", true).await; + dataset.object_store.io_stats_incremental(); + let delta = dataset + .delta() + .with_begin_version(begin) + .with_end_version(end) + .build() + .unwrap(); + for stream in [ + delta.get_inserted_rows().await.unwrap(), + delta.get_updated_rows().await.unwrap(), + delta.get_upserted_rows().await.unwrap(), + ] { + let batch = collect_stream(stream).await; + assert_eq!(batch.num_rows(), 0); + for column in [ + "key", + "value", + ROW_ID, + ROW_CREATED_AT_VERSION, + ROW_LAST_UPDATED_AT_VERSION, + ] { + assert!(batch.column_by_name(column).is_some()); + } + } + assert_eq!(dataset.object_store.io_stats_incremental().read_iops, 0); + } + + #[rstest] + #[case::missing(None)] + #[case::undecodable(Some(RowDatasetVersionMeta::Inline(Arc::from([0xff]))))] + #[tokio::test] + async fn test_delta_pruning_preserves_default_versions( + #[case] meta: Option, + ) { + let mut dataset = create_test_dataset(4, 1, "value", true).await; + let manifest = Arc::make_mut(&mut dataset.manifest); + for fragment in Arc::make_mut(&mut manifest.fragments) { + fragment.created_at_version_meta = meta.clone(); + fragment.last_updated_at_version_meta = meta.clone(); + } + let delta = dataset + .delta() + .with_begin_version(0) + .with_end_version(1) + .build() + .unwrap(); + for stream in [ + delta.get_inserted_rows().await.unwrap(), + delta.get_upserted_rows().await.unwrap(), + ] { + let batch = collect_stream(stream).await; + assert_eq!(batch.num_rows(), 4); + for column in [ROW_CREATED_AT_VERSION, ROW_LAST_UPDATED_AT_VERSION] { + assert_eq!( + batch[column].as_primitive::().values().as_ref(), + &[1; 4] + ); + } + } + assert_eq!( + collect_stream(delta.get_updated_rows().await.unwrap()) + .await + .num_rows(), + 0 + ); } #[tokio::test] @@ -2525,6 +2910,10 @@ mod tests { let ds = update_where(ds, "key < 10", "updated_v3").await; assert_eq!(ds.version().version, 3); + // An insertion still belongs to the range if it was updated after the range. + let ds = update_where(ds, "key = 50", "updated_v4").await; + assert_eq!(ds.version().version, 4); + // Get upserted rows between version 1 and 3 let delta = ds .delta() @@ -2560,7 +2949,7 @@ mod tests { // Inserted rows from version 2 assert!((50..70).contains(&key)); assert_eq!(created_at[i], 2); - assert_eq!(updated_at[i], 2); + assert_eq!(updated_at[i], if key == 50 { 4 } else { 2 }); } } } @@ -2591,6 +2980,23 @@ mod tests { let txs = delta.list_transactions().await.unwrap(); assert_eq!(txs.len(), 1); + + assert_eq!( + collect_stream(delta.get_inserted_rows().await.unwrap()) + .await + .num_rows(), + 0 + ); + for stream in [ + delta.get_updated_rows().await.unwrap(), + delta.get_upserted_rows().await.unwrap(), + ] { + let batch = collect_stream(stream).await; + assert_eq!(batch.num_rows(), 10); + let mut keys = batch["key"].as_primitive::().values().to_vec(); + keys.sort_unstable(); + assert_eq!(keys, (0..10).collect::>()); + } } #[tokio::test]