This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-3260-219571a4185df2bc6e144dbbcb72cf016714b61b in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git
commit b9d285567526b384585b1613c76b46ab5ef063dd Author: Dhruv Arya <[email protected]> AuthorDate: Wed Oct 21 07:58:40 2026 +0000 feat(scan): support timestamp-based time travel (#3260) * feat(scan): support timestamp-based time travel * test(scan): verify timestamp scans read historical rows * docs: shorten timestamp time travel docstrings * fix: refine timestamp selection API and edge-case coverage * docs: add timestamp time travel changelog entry * test: consolidate timestamp time travel coverage * refactor: use a single snapshot selector with last-call-wins semantics * docs: clarify timestamp history selection and PyIceberg differences * test: share the Parquet fixture row count * docs: note Java timestamp ordering compatibility * docs: clarify minimum timestamp difference from Java * docs: describe snapshot selection enum * test: adapt timestamp time travel tests to relocated scan fixture Co-authored-by: Isaac <[email protected]> * fix: name the timestamp when an as-of-time snapshot is missing Co-authored-by: Isaac <[email protected]> * refactor: spell out the as-of-time tie-break and i64::MIN edge Co-authored-by: Isaac <[email protected]> --------- Co-authored-by: Dhruv Arya <[email protected]> Co-authored-by: Isaac <[email protected]> --- CHANGELOG.md | 4 + crates/iceberg/public-api.txt | 2 + crates/iceberg/src/scan/mod.rs | 456 ++++++++++++++++++++++++++++++++-- crates/iceberg/src/test_utils/scan.rs | 7 +- crates/iceberg/src/util/snapshot.rs | 125 ++++++++++ 5 files changed, 576 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3076335c9..728b057e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/). ## [Unreleased] +### Added + +* Support timestamp-based time travel through `TableScanBuilder::as_of_time` and `snapshot_id_as_of_time`. + ### Breaking Changes * `FileWrite::close`, `EncryptedOutputFile::write`, `ManifestListWriter::close`, and `PuffinWriter::close` now return `Result<io::FileMetadata>` instead of `Result<()>`. The returned size is the on-disk size, including encryption overhead. Custom `FileWrite` implementations must return the number of bytes stored. diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 9d65474ad..87fd3943c 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1438,6 +1438,7 @@ impl core::fmt::Debug for iceberg::scan::TableScan pub fn iceberg::scan::TableScan::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::scan::TableScanBuilder<'a> impl<'a> iceberg::scan::TableScanBuilder<'a> +pub fn iceberg::scan::TableScanBuilder<'a>::as_of_time(self, timestamp_ms: i64) -> Self pub fn iceberg::scan::TableScanBuilder<'a>::build(self) -> iceberg::Result<iceberg::scan::TableScan> pub fn iceberg::scan::TableScanBuilder<'a>::select(self, column_names: impl core::iter::traits::collect::IntoIterator<Item = impl alloc::string::ToString>) -> Self pub fn iceberg::scan::TableScanBuilder<'a>::select_all(self) -> Self @@ -3514,6 +3515,7 @@ pub mod iceberg::util::location pub mod iceberg::util::snapshot pub fn iceberg::util::snapshot::ancestors_between(table_metadata: &iceberg::spec::TableMetadataRef, latest_snapshot_id: i64, oldest_snapshot_id: core::option::Option<i64>) -> impl core::iter::traits::iterator::Iterator<Item = iceberg::spec::SnapshotRef> + core::marker::Send pub fn iceberg::util::snapshot::ancestors_of(table_metadata: &iceberg::spec::TableMetadataRef, snapshot_id: i64) -> impl core::iter::traits::iterator::Iterator<Item = iceberg::spec::SnapshotRef> + core::marker::Send +pub fn iceberg::util::snapshot::snapshot_id_as_of_time(table_metadata: &iceberg::spec::TableMetadataRef, timestamp_ms: i64) -> iceberg::Result<i64> pub mod iceberg::writer pub mod iceberg::writer::base_writer pub mod iceberg::writer::base_writer::data_file_writer diff --git a/crates/iceberg/src/scan/mod.rs b/crates/iceberg/src/scan/mod.rs index 6d23e9470..59a353c5e 100644 --- a/crates/iceberg/src/scan/mod.rs +++ b/crates/iceberg/src/scan/mod.rs @@ -47,6 +47,7 @@ use crate::runtime::Runtime; use crate::spec::{DataContentType, Schema, SchemaRef, SnapshotRef, SortOrderRef, StructType}; use crate::table::Table; use crate::util::available_parallelism; +use crate::util::snapshot::snapshot_id_as_of_time; use crate::{Error, ErrorKind, Result}; /// A stream of arrow [`RecordBatch`]es. @@ -130,12 +131,18 @@ fn projected_partition_type( .map(Some) } +/// Selects a snapshot by ID or timestamp. +enum SnapshotSelection { + SnapshotId(i64), + AsOfTime(i64), +} + /// Builder to create table scan. pub struct TableScanBuilder<'a> { table: &'a Table, // Defaults to none which means select all columns column_names: Option<Vec<String>>, - snapshot_id: Option<i64>, + snapshot_selection: Option<SnapshotSelection>, batch_size: Option<usize>, case_sensitive: bool, filter: Option<Predicate>, @@ -154,7 +161,7 @@ impl<'a> TableScanBuilder<'a> { Self { table, column_names: None, - snapshot_id: None, + snapshot_selection: None, batch_size: None, case_sensitive: true, filter: None, @@ -212,10 +219,38 @@ impl<'a> TableScanBuilder<'a> { } /// Set the snapshot to scan. The scan projects the schema that snapshot was - /// written with. When not set, the scan reads the current snapshot with the - /// table's current schema. + /// written with. + /// + /// The last call to this or [`Self::as_of_time`] takes precedence. + /// If neither is set, the scan reads the current snapshot with the table's + /// current schema. pub fn snapshot_id(mut self, snapshot_id: i64) -> Self { - self.snapshot_id = Some(snapshot_id); + self.snapshot_selection = Some(SnapshotSelection::SnapshotId(snapshot_id)); + self + } + + /// Select a snapshot at or before a timestamp in milliseconds since the Unix epoch. + /// + /// Uses the table's main history, and projects the schema the selected snapshot + /// was written with. The last call to this or [`Self::snapshot_id`] takes precedence. + /// + /// ``` + /// # use iceberg::TableIdent; + /// # use iceberg::io::FileIO; + /// # use iceberg::table::StaticTable; + /// # #[tokio::main] + /// # async fn main() -> iceberg::Result<()> { + /// # let location = concat!(env!("CARGO_MANIFEST_DIR"), "/testdata/example_table_metadata_v2.json"); + /// # let ident = TableIdent::from_strs(["ns", "t"])?; + /// # let table = StaticTable::from_metadata_file(location, ident, FileIO::new_with_fs()).await?; + /// // The snapshot log has entries at 1515100955770 and 1555100955770. + /// let scan = table.scan().as_of_time(1_555_100_955_769).build()?; + /// assert_eq!(scan.snapshot().unwrap().snapshot_id(), 3051729675574597004); + /// # Ok(()) + /// # } + /// ``` + pub fn as_of_time(mut self, timestamp_ms: i64) -> Self { + self.snapshot_selection = Some(SnapshotSelection::AsOfTime(timestamp_ms)); self } @@ -288,12 +323,29 @@ impl<'a> TableScanBuilder<'a> { /// Build the table scan. pub fn build(self) -> Result<TableScan> { - let snapshot = match self.snapshot_id { - Some(snapshot_id) => self + let selected_snapshot = match self.snapshot_selection { + Some(SnapshotSelection::SnapshotId(id)) => { + Some((id, format!("Snapshot with id {id} not found"))) + } + Some(SnapshotSelection::AsOfTime(timestamp_ms)) => { + let id = snapshot_id_as_of_time(&self.table.metadata_ref(), timestamp_ms)?; + // History outlives expired snapshots; name the timestamp that resolved to `id`. + Some(( + id, + format!( + "Snapshot with id {id} resolved from timestamp {timestamp_ms} ms not found; it may have expired" + ), + )) + } + None => None, + }; + let is_time_travel = selected_snapshot.is_some(); + let snapshot = match selected_snapshot { + Some((snapshot_id, not_found_message)) => self .table .metadata() .snapshot_by_id(snapshot_id) - .ok_or_else(|| invalid_data!("Snapshot with id {snapshot_id} not found"))? + .ok_or_else(|| invalid_data!(not_found_message))? .clone(), None => { let Some(current_snapshot_id) = self.table.metadata().current_snapshot() else { @@ -317,9 +369,10 @@ impl<'a> TableScanBuilder<'a> { // A scan of the current state uses the table's current schema; time travel uses the // schema its snapshot was written with. - let schema = match self.snapshot_id { - Some(_) => snapshot.schema(self.table.metadata())?, - None => self.table.metadata().current_schema().clone(), + let schema = if is_time_travel { + snapshot.schema(self.table.metadata())? + } else { + self.table.metadata().current_schema().clone() }; let field_ids = collect_scan_field_ids(&schema, self.column_names.as_deref(), self.case_sensitive)?; @@ -692,14 +745,18 @@ mod tests { //! shared tests for the table scan API use std::collections::HashMap; + use std::fs; + use std::fs::File; use std::sync::Arc; use arrow_array::cast::AsArray; use arrow_array::types::Int32Type; use arrow_array::{ - Array, ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RunArray, StringArray, + Array, ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, RunArray, + StringArray, }; use futures::{TryStreamExt, stream}; + use parquet::arrow::ArrowWriter; use uuid::Uuid; use super::resolve_field_id; @@ -712,14 +769,18 @@ mod tests { }; use crate::scan::{FileScanTask, FileScanTaskDeleteFile}; use crate::spec::{ - DataContentType, DataFileFormat, Datum, Literal, MAIN_BRANCH, MappedField, NameMapping, - NestedField, NullOrder, Operation, PartitionSpec, PrimitiveType, Schema, Snapshot, - SortDirection, SortField, SortOrder, Struct, Summary, TableMetadataBuilder, + DataContentType, DataFileBuilder, DataFileFormat, Datum, Literal, MAIN_BRANCH, + ManifestEntry, ManifestListWriter, ManifestStatus, ManifestWriterBuilder, MappedField, + NameMapping, NestedField, NullOrder, Operation, PartitionSpec, PrimitiveType, Schema, + Snapshot, SortDirection, SortField, SortOrder, Struct, Summary, TableMetadataBuilder, TableProperties, Transform, Type, UnboundPartitionSpec, }; use crate::table::Table; - use crate::test_utils::scan::{TableTestFixture, assert_last_updated_seq_all}; + use crate::test_utils::scan::{ + ROWS_PER_FIXTURE_FILE, TableTestFixture, assert_last_updated_seq_all, + }; use crate::test_utils::test_runtime; + use crate::util::snapshot::snapshot_id_as_of_time; use crate::{ErrorKind, TableIdent}; #[tokio::test] @@ -906,6 +967,369 @@ mod tests { ); } + #[test] + fn test_table_scan_as_of_time_after_rollback() { + for elapsed_ms in [0, 1000] { + let table = TableTestFixture::new().table; + let mut metadata = table.metadata().clone(); + let original = metadata.history()[0].snapshot_id; + let previous = metadata.history()[1].clone(); + let rollback_time = previous.timestamp_ms + elapsed_ms; + metadata.snapshot_log.push(crate::spec::SnapshotLog { + timestamp_ms: rollback_time, + snapshot_id: original, + }); + metadata.current_snapshot_id = Some(original); + metadata.refs.get_mut(MAIN_BRANCH).unwrap().snapshot_id = original; + let table = table.with_metadata(Arc::new(metadata)); + let scan = table.scan().as_of_time(rollback_time).build().unwrap(); + // A rollback at the same timestamp keeps the earlier history entry. + let expected = if elapsed_ms == 0 { + previous.snapshot_id + } else { + original + }; + assert_eq!(scan.snapshot().unwrap().snapshot_id(), expected); + } + } + + #[test] + fn test_table_scan_as_of_time_rejects_empty_history() { + let table = TableTestFixture::new_empty().table; + assert!(table.scan().build().unwrap().snapshot().is_none()); + let err = table.scan().as_of_time(0).build().unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!(err.message().contains("No snapshot history")); + } + + #[test] + fn test_table_scan_as_of_time_rejects_before_history() { + let table = TableTestFixture::new().table; + let before = table.metadata().history()[0].timestamp_ms - 1; + let err = table.scan().as_of_time(before).build().unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!(err.message().contains(&before.to_string())); + } + + #[test] + fn test_table_scan_as_of_time_rejects_all_qualifying_snapshots_expired() { + let table = TableTestFixture::new_with_deep_history().table; + let mut metadata = table.metadata().clone(); + let entry = metadata.history()[3].clone(); + let expired_ids: Vec<_> = metadata + .history() + .iter() + .filter(|log| log.timestamp_ms <= entry.timestamp_ms) + .map(|log| log.snapshot_id) + .collect(); + assert_eq!(expired_ids.len(), 4); + for snapshot_id in expired_ids { + metadata.snapshots.remove(&snapshot_id); + } + let table = table.with_metadata(Arc::new(metadata)); + // History resolution succeeds even though every qualifying snapshot expired. + assert_eq!( + snapshot_id_as_of_time(&table.metadata_ref(), entry.timestamp_ms).unwrap(), + entry.snapshot_id + ); + assert!(table.scan().build().unwrap().snapshot().is_some()); + let err = table + .scan() + .as_of_time(entry.timestamp_ms) + .build() + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert_eq!( + err.message(), + format!( + "Snapshot with id {} resolved from timestamp {} ms not found; it may have expired", + entry.snapshot_id, entry.timestamp_ms + ) + ); + } + + #[test] + fn test_table_scan_last_snapshot_selection_wins() { + let table = TableTestFixture::new().table; + let first = &table.metadata().history()[0]; + let last = &table.metadata().history()[1]; + let cases = [ + ( + table + .scan() + .snapshot_id(last.snapshot_id) + .as_of_time(first.timestamp_ms), + first.snapshot_id, + ), + ( + table + .scan() + .as_of_time(first.timestamp_ms) + .snapshot_id(last.snapshot_id), + last.snapshot_id, + ), + // An overwritten selector is not resolved or validated. + ( + table.scan().snapshot_id(-1).as_of_time(first.timestamp_ms), + first.snapshot_id, + ), + ( + table + .scan() + .as_of_time(i64::MIN) + .snapshot_id(last.snapshot_id), + last.snapshot_id, + ), + ( + table + .scan() + .snapshot_id(first.snapshot_id) + .snapshot_id(last.snapshot_id), + last.snapshot_id, + ), + ( + table + .scan() + .as_of_time(i64::MAX) + .as_of_time(first.timestamp_ms), + first.snapshot_id, + ), + ]; + for (builder, expected) in cases { + let scan = builder.build().unwrap(); + assert_eq!(scan.snapshot().unwrap().snapshot_id(), expected); + } + } + + #[test] + fn test_table_scan_invalid_last_snapshot_selection() { + let table = TableTestFixture::new().table; + let first = &table.metadata().history()[0]; + let cases = [ + ( + table + .scan() + .snapshot_id(first.snapshot_id) + .as_of_time(first.timestamp_ms - 1), + "No snapshot history", + ), + ( + table.scan().as_of_time(first.timestamp_ms).snapshot_id(-1), + "Snapshot with id -1 not found", + ), + ]; + for (builder, expected) in cases { + let err = builder.build().unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!(err.message().contains(expected)); + } + } + + #[test] + fn test_table_scan_as_of_time_uses_snapshot_schema() { + for select_current in [false, true] { + let table = TableTestFixture::new().table; + let mut metadata = table.metadata().clone(); + let entry = metadata.history()[usize::from(select_current)].clone(); + // Simulate a schema-only update after this snapshot: schema 0 has x only, + // while the current table schema also contains y and other columns. + Arc::make_mut(metadata.snapshots.get_mut(&entry.snapshot_id).unwrap()).schema_id = + Some(0); + let table = table.with_metadata(Arc::new(metadata)); + let scan = table + .scan() + .as_of_time(entry.timestamp_ms) + .select(["x"]) + .with_filter(Reference::new("x").greater_than(Datum::long(0))) + .build() + .unwrap(); + let context = scan.plan_context.as_ref().unwrap(); + assert_eq!(context.snapshot_schema.schema_id(), 0); + assert!(context.snapshot_bound_predicate.is_some()); + assert!( + table + .scan() + .as_of_time(entry.timestamp_ms) + .select(["y"]) + .build() + .is_err() + ); + assert!( + table + .scan() + .as_of_time(entry.timestamp_ms) + .with_filter(Reference::new("y").greater_than(Datum::long(0))) + .build() + .is_err() + ); + } + } + + #[test] + fn test_table_scan_as_of_time_schema_compatibility() { + let table = TableTestFixture::new().table; + let entry = table.metadata().history()[0].clone(); + // Older snapshots may omit schema-id and use the current-schema fallback. + assert!( + table + .metadata() + .snapshot_by_id(entry.snapshot_id) + .unwrap() + .schema_id() + .is_none() + ); + let scan = table.scan().as_of_time(entry.timestamp_ms).build().unwrap(); + assert_eq!( + scan.plan_context.as_ref().unwrap().snapshot_schema, + *table.metadata().current_schema() + ); + + let mut metadata = table.metadata().clone(); + Arc::make_mut(metadata.snapshots.get_mut(&entry.snapshot_id).unwrap()).schema_id = + Some(1234); + let table = table.with_metadata(Arc::new(metadata)); + let err = table + .scan() + .as_of_time(entry.timestamp_ms) + .build() + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!(err.message().contains("Schema with id 1234 not found")); + } + + #[tokio::test] + async fn test_table_scan_as_of_time_reads_historical_rows() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + let entry = fixture.table.metadata().history()[0].clone(); + + // The older snapshot has only x, while the current snapshot has eight columns. + let mut metadata = fixture.table.metadata().clone(); + Arc::make_mut(metadata.snapshots.get_mut(&entry.snapshot_id).unwrap()).schema_id = Some(0); + fixture.table = fixture.table.with_metadata(Arc::new(metadata)); + let snapshot = fixture + .table + .metadata() + .snapshot_by_id(entry.snapshot_id) + .unwrap(); + let schema = snapshot.schema(fixture.table.metadata()).unwrap(); + let arrow_schema = Arc::new(crate::arrow::schema_to_arrow_schema(&schema).unwrap()); + + // Write distinct historical data: x = 100, versus x = 1 in the current files. + let path = format!("{}/historical.parquet", fixture.table_location); + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(Int64Array::from( + vec![100], + ))]) + .unwrap(); + let mut writer = + ArrowWriter::try_new(File::create(&path).unwrap(), arrow_schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let mut manifest_writer = ManifestWriterBuilder::new( + fixture.next_manifest_file(), + Some(snapshot.snapshot_id()), + schema, + fixture + .table + .metadata() + .default_partition_spec() + .as_ref() + .clone(), + ) + .build_v2_data(); + manifest_writer + .add_entry( + ManifestEntry::builder() + .status(ManifestStatus::Added) + .data_file( + DataFileBuilder::default() + .partition_spec_id(0) + .content(DataContentType::Data) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(fs::metadata(&path).unwrap().len()) + .file_path(path) + .record_count(1) + .partition(Struct::from_iter([Some(Literal::long(100))])) + .build() + .unwrap(), + ) + .build(), + ) + .unwrap(); + let manifest = manifest_writer.write_manifest_file().await.unwrap(); + let output = fixture + .table + .file_io() + .new_output(snapshot.manifest_list()) + .unwrap() + .writer() + .await + .unwrap(); + let mut list_writer = ManifestListWriter::v2( + output, + snapshot.snapshot_id(), + snapshot.parent_snapshot_id(), + snapshot.sequence_number(), + ); + list_writer.add_manifests([manifest].into_iter()).unwrap(); + list_writer.close().await.unwrap(); + + let table = fixture.table; + let [by_time, by_id] = [ + table.scan().as_of_time(entry.timestamp_ms), + table.scan().snapshot_id(entry.snapshot_id), + ] + .map(|scan| { + scan.with_filter(Reference::new("x").greater_than_or_equal_to(Datum::long(100))) + .build() + .unwrap() + }); + let current = table.scan().build().unwrap(); + let mut planned_files = Vec::new(); + let mut rows = Vec::new(); + let mut schemas = Vec::new(); + for scan in [by_time, by_id, current] { + let mut tasks: Vec<_> = scan + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + tasks.sort_by_key(|task| task.data_file_path().to_string()); + planned_files.push(tasks); + let batches: Vec<_> = scan.to_arrow().await.unwrap().try_collect().await.unwrap(); + assert!(!batches.is_empty()); + schemas.push(batches[0].schema()); + let mut values: Vec<i64> = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("x") + .unwrap() + .as_any() + .downcast_ref::<Int64Array>() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + values.sort_unstable(); + rows.push(values); + } + assert!(!planned_files[0].is_empty()); + assert_eq!(planned_files[0], planned_files[1]); + assert_eq!(rows[0], vec![100]); + assert_eq!(rows[0], rows[1]); + // The current snapshot contains two live fixture files. + assert_eq!(rows[2], vec![1; 2 * ROWS_PER_FIXTURE_FILE]); + assert_eq!(schemas[0], schemas[1]); + assert_eq!(schemas[0].fields().len(), 1); + assert_eq!(schemas[2].fields().len(), 8); + } + /// A scan of the current state reads the current snapshot with the table's /// current schema, not the older one that snapshot was written with. #[test] diff --git a/crates/iceberg/src/test_utils/scan.rs b/crates/iceberg/src/test_utils/scan.rs index 63f6418ef..2f96e61ad 100644 --- a/crates/iceberg/src/test_utils/scan.rs +++ b/crates/iceberg/src/test_utils/scan.rs @@ -49,6 +49,8 @@ use crate::spec::{ use crate::table::Table; use crate::test_utils::test_runtime; +pub(crate) const ROWS_PER_FIXTURE_FILE: usize = 1024; + fn render_template(template: &str, ctx: Value) -> String { let mut env = Environment::new(); env.set_auto_escape_callback(|_| AutoEscape::None); @@ -293,7 +295,7 @@ impl TableTestFixture { } } - fn next_manifest_file(&self) -> OutputFile { + pub(crate) fn next_manifest_file(&self) -> OutputFile { self.table .file_io() .new_output(format!( @@ -647,7 +649,8 @@ impl TableTestFixture { Arc::new(arrow_schema::Schema::new(fields)) }; // x: [1, 1, 1, 1, ...] - let col1 = Arc::new(Int64Array::from_iter_values(vec![1; 1024])) as ArrayRef; + let col1 = + Arc::new(Int64Array::from_iter_values(vec![1; ROWS_PER_FIXTURE_FILE])) as ArrayRef; let mut values = vec![2; 512]; values.append(vec![3; 200].as_mut()); diff --git a/crates/iceberg/src/util/snapshot.rs b/crates/iceberg/src/util/snapshot.rs index deac53ef8..c5f307d61 100644 --- a/crates/iceberg/src/util/snapshot.rs +++ b/crates/iceberg/src/util/snapshot.rs @@ -18,6 +18,7 @@ use std::collections::HashSet; use crate::spec::{SnapshotRef, TableMetadataRef}; +use crate::{Error, ErrorKind, Result}; struct Ancestors { next: Option<SnapshotRef>, @@ -76,9 +77,62 @@ pub fn ancestors_between( }) } +/// Resolve the snapshot ID from the main-history entry with the greatest timestamp +/// at or before `timestamp_ms` (milliseconds since the Unix epoch), taking the first +/// entry on ties. +/// +/// This matches Java's ordering, except that an entry at `i64::MIN` can match here +/// but never matches in Java, which seeds its search with `Long.MIN_VALUE`. +/// PyIceberg takes the last qualifying log entry, so ties or out-of-order +/// timestamps can produce different results. +/// +/// Returns [`ErrorKind::DataInvalid`] if no matching history exists. +/// The returned snapshot may have expired, so +/// [`snapshot_by_id`](crate::spec::TableMetadata::snapshot_by_id) can still return `None`. +/// +/// ``` +/// # use iceberg::TableIdent; +/// # use iceberg::io::FileIO; +/// # use iceberg::table::StaticTable; +/// use iceberg::util::snapshot::snapshot_id_as_of_time; +/// # #[tokio::main] +/// # async fn main() -> iceberg::Result<()> { +/// # let location = concat!(env!("CARGO_MANIFEST_DIR"), "/testdata/example_table_metadata_v2.json"); +/// # let ident = TableIdent::from_strs(["ns", "t"])?; +/// # let table = StaticTable::from_metadata_file(location, ident, FileIO::new_with_fs()).await?; +/// let metadata = table.metadata(); +/// // The snapshot log has entries at 1515100955770 and 1555100955770. +/// let snapshot_id = snapshot_id_as_of_time(&metadata, 1_555_100_955_770)?; +/// assert_eq!(snapshot_id, 3055729675574597004); +/// assert!(snapshot_id_as_of_time(&metadata, 1_515_100_955_769).is_err()); +/// # Ok(()) +/// # } +/// ``` +pub fn snapshot_id_as_of_time(table_metadata: &TableMetadataRef, timestamp_ms: i64) -> Result<i64> { + table_metadata + .history() + .iter() + .enumerate() + .filter(|(_, entry)| entry.timestamp_ms() <= timestamp_ms) + .max_by(|(idx_a, a), (idx_b, b)| { + a.timestamp_ms() + .cmp(&b.timestamp_ms()) + // Earlier entries win ties, matching Java's timestamp selection. + .then(idx_b.cmp(idx_a)) + }) + .map(|(_, entry)| entry.snapshot_id) + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!("No snapshot history at or before timestamp {timestamp_ms} ms"), + ) + }) +} + #[cfg(test)] mod tests { use super::*; + use crate::spec::SnapshotLog; use crate::test_utils::scan::TableTestFixture; // Five snapshots chained as: S1 (root) -> S2 -> S3 -> S4 -> S5 (current) @@ -93,6 +147,77 @@ mod tests { std::sync::Arc::new(fixture.table.metadata().clone()) } + type History = [(i64, i64)]; + + fn metadata_with_history(history: &History) -> TableMetadataRef { + let mut metadata = metadata().as_ref().clone(); + metadata.snapshot_log = history + .iter() + .map(|&(timestamp_ms, snapshot_id)| SnapshotLog { + timestamp_ms, + snapshot_id, + }) + .collect(); + metadata.into() + } + + #[test] + fn test_snapshot_id_as_of_time() { + let cases: &[(&History, i64, i64)] = &[ + (&[(1000, S1), (2000, S2)], 1000, S1), + (&[(1000, S1), (2000, S2)], 1500, S1), + (&[(1000, S1), (2000, S2)], 2000, S2), + (&[(1000, S1), (2000, S2)], 3000, S2), + // Rollback records when an existing snapshot becomes current again. + (&[(1000, S1), (2000, S2), (3000, S1)], 3500, S1), + // The first entry wins on ties. + (&[(1000, S1), (2000, S2), (2000, S3)], 2000, S2), + // Clock skew means the log need not be sorted by timestamp. + (&[(1000, S1), (3000, S2), (2500, S3)], 3500, S2), + (&[(1000, S1), (3000, S2), (2500, S3)], 2600, S3), + (&[(-2000, S1), (-1000, S2)], -1500, S1), + // Java returns no match here; see the docstring. + (&[(i64::MIN, S1), (0, S2)], i64::MIN, S1), + (&[(0, S1), (i64::MAX, S2)], i64::MAX, S2), + ]; + for &(history, timestamp_ms, expected) in cases { + let metadata = metadata_with_history(history); + assert_eq!( + snapshot_id_as_of_time(&metadata, timestamp_ms).unwrap(), + expected, + "timestamp {timestamp_ms}, history {history:?}" + ); + } + } + + #[test] + fn test_snapshot_id_as_of_time_without_history_at_timestamp() { + let cases: &[(&History, i64)] = &[ + (&[], 1000), + (&[(1000, S1)], 999), + (&[(3000, S1)], 1500), // Expired prefix; S1 is still retained. + ]; + for &(history, timestamp_ms) in cases { + let metadata = metadata_with_history(history); + assert!(metadata.snapshot_by_id(S1).is_some()); + let err = snapshot_id_as_of_time(&metadata, timestamp_ms).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!(err.message().contains(×tamp_ms.to_string())); + } + } + + #[test] + fn test_snapshot_id_as_of_time_ignores_snapshots_outside_main_history() { + let mut metadata = metadata().as_ref().clone(); + metadata.snapshot_log.truncate(2); + // Newer retained snapshots may be staged or belong only to another branch. + assert!(metadata.snapshot_by_id(S5).is_some()); + assert_eq!( + snapshot_id_as_of_time(&metadata.into(), i64::MAX).unwrap(), + S2 + ); + } + // --- ancestors_of --- #[test]
