From a43741b094f621c8b1a0e16ddff4492f9e3fcb12 Mon Sep 17 00:00:00 2001 From: Noah Date: Mon, 20 Jul 2026 22:05:44 -0400 Subject: [PATCH 01/10] init --- crates/integrations/datafusion/public-api.txt | 6 + .../integrations/datafusion/src/table/mod.rs | 145 +++++++++++++++++- 2 files changed, 149 insertions(+), 2 deletions(-) diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index d24bd9fc9e..4aabf77294 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -74,6 +74,9 @@ pub fn iceberg_datafusion::IcebergStaticTableProvider::schema(&self) -> arrow_sc pub fn iceberg_datafusion::IcebergStaticTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergStaticTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType pub struct iceberg_datafusion::table::IcebergTableProvider +impl iceberg_datafusion::IcebergTableProvider +pub fn iceberg_datafusion::IcebergTableProvider::snapshot_id(&self) -> core::option::Option +pub fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> Self impl core::clone::Clone for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafusion::IcebergTableProvider impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider @@ -121,6 +124,9 @@ pub fn iceberg_datafusion::IcebergStaticTableProvider::schema(&self) -> arrow_sc pub fn iceberg_datafusion::IcebergStaticTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergStaticTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType pub struct iceberg_datafusion::IcebergTableProvider +impl iceberg_datafusion::IcebergTableProvider +pub fn iceberg_datafusion::IcebergTableProvider::snapshot_id(&self) -> core::option::Option +pub fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> Self impl core::clone::Clone for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafusion::IcebergTableProvider impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 75b7988d8d..a2e5b837e8 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -73,6 +73,10 @@ pub struct IcebergTableProvider { table_ident: TableIdent, /// A reference-counted arrow `Schema` (cached at construction) schema: ArrowSchemaRef, + /// Optional snapshot to read. `None` reads the current snapshot (refreshed + /// from the catalog on each scan); `Some` pins reads to that snapshot for + /// time-travel. Writes always target the current table state. + snapshot_id: Option, } impl IcebergTableProvider { @@ -95,9 +99,23 @@ impl IcebergTableProvider { catalog, table_ident, schema, + snapshot_id: None, }) } + /// Pins reads to a specific snapshot for time-travel. `None` (the default) + /// reads the current snapshot. The snapshot id is threaded into the scan + /// node, so it is serialized and honored by a distributed engine as well. + pub fn with_snapshot_id(mut self, snapshot_id: Option) -> Self { + self.snapshot_id = snapshot_id; + self + } + + /// Returns the snapshot this provider reads, if pinned for time-travel. + pub fn snapshot_id(&self) -> Option { + self.snapshot_id + } + pub(crate) async fn metadata_table( &self, r#type: MetadataTableType, @@ -136,10 +154,10 @@ impl TableProvider for IcebergTableProvider { .await .map_err(to_datafusion_error)?; - // Create scan with fresh metadata (always use current snapshot) + // Create scan with fresh metadata, honoring a pinned snapshot if set. Ok(Arc::new(IcebergTableScan::new( table, - None, // Always use current snapshot for catalog-backed provider + self.snapshot_id, self.schema.clone(), projection, filters, @@ -865,4 +883,127 @@ mod tests { "Limit should be None when not specified" ); } + + /// Counts the rows a provider returns for `SELECT * FROM t`. + async fn scan_row_count(provider: IcebergTableProvider) -> usize { + let ctx = SessionContext::new(); + ctx.register_table("t", Arc::new(provider)).unwrap(); + let batches = ctx + .sql("SELECT * FROM t") + .await + .unwrap() + .collect() + .await + .unwrap(); + batches.iter().map(|b| b.num_rows()).sum() + } + + #[tokio::test] + async fn test_pinned_snapshot_reads_historical_data() { + let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; + + // First append -> snapshot with a single row. + let writer = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) + .await + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table("t", Arc::new(writer)).unwrap(); + ctx.sql("INSERT INTO t VALUES (1, 'a')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + // Capture the snapshot produced by the first append. + let table_ident = TableIdent::new(namespace.clone(), table_name.clone()); + let first_snapshot = catalog + .load_table(&table_ident) + .await + .unwrap() + .metadata() + .current_snapshot() + .unwrap() + .snapshot_id(); + + // Second append -> current snapshot now has two rows. + ctx.sql("INSERT INTO t VALUES (2, 'b')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + // Default (unpinned) provider sees the latest state: both rows. + let current = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) + .await + .unwrap(); + assert_eq!( + scan_row_count(current).await, + 2, + "unpinned provider should read the current snapshot" + ); + + // Pinning the first snapshot time-travels: only the first row is visible, + // even though a newer snapshot exists. + let pinned = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) + .await + .unwrap() + .with_snapshot_id(Some(first_snapshot)); + assert_eq!(pinned.snapshot_id(), Some(first_snapshot)); + + // The pin is threaded onto the scan node itself — this is the value a + // distributed engine's codec serializes to reproduce the scan remotely. + let scan_plan = pinned + .scan(&SessionContext::new().state(), None, &[], None) + .await + .unwrap(); + let iceberg_scan = scan_plan + .as_any() + .downcast_ref::() + .expect("Expected IcebergTableScan"); + assert_eq!( + iceberg_scan.snapshot_id(), + Some(first_snapshot), + "pinned snapshot should propagate to the scan node" + ); + + // And it changes what is actually read: only the historical row. + assert_eq!( + scan_row_count(pinned).await, + 1, + "provider pinned to the first snapshot should read only historical data" + ); + } + + #[tokio::test] + async fn test_with_snapshot_id_none_reads_current() { + let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; + + let writer = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) + .await + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table("t", Arc::new(writer)).unwrap(); + ctx.sql("INSERT INTO t VALUES (1, 'a')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + // Explicitly unpinning (None) is equivalent to the default: read current. + let provider = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) + .await + .unwrap() + .with_snapshot_id(Some(1)) + .with_snapshot_id(None); + assert_eq!(provider.snapshot_id(), None); + assert_eq!(scan_row_count(provider).await, 1); + } } From 63e9b614a215974d31073218551782d02fce7449 Mon Sep 17 00:00:00 2001 From: Noah Date: Mon, 20 Jul 2026 22:12:30 -0400 Subject: [PATCH 02/10] merge tests --- .../integrations/datafusion/src/table/mod.rs | 33 ++++++------------- 1 file changed, 10 insertions(+), 23 deletions(-) diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index a2e5b837e8..85b03735f4 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -977,33 +977,20 @@ mod tests { 1, "provider pinned to the first snapshot should read only historical data" ); - } - - #[tokio::test] - async fn test_with_snapshot_id_none_reads_current() { - let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; - let writer = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap(); - let ctx = SessionContext::new(); - ctx.register_table("t", Arc::new(writer)).unwrap(); - ctx.sql("INSERT INTO t VALUES (1, 'a')") - .await - .unwrap() - .collect() - .await - .unwrap(); - - // Explicitly unpinning (None) is equivalent to the default: read current. - let provider = + // Clearing the pin (Some -> None) unpins back to the current snapshot, + // so the newer row becomes visible again. + let unpinned = IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) .await .unwrap() - .with_snapshot_id(Some(1)) + .with_snapshot_id(Some(first_snapshot)) .with_snapshot_id(None); - assert_eq!(provider.snapshot_id(), None); - assert_eq!(scan_row_count(provider).await, 1); + assert_eq!(unpinned.snapshot_id(), None); + assert_eq!( + scan_row_count(unpinned).await, + 2, + "clearing the pin should read the current snapshot again" + ); } } From 60b5c1a041dea035f3ca524b51db84285a17ccd4 Mon Sep 17 00:00:00 2001 From: Noah Date: Mon, 20 Jul 2026 22:33:36 -0400 Subject: [PATCH 03/10] cleanup --- crates/integrations/datafusion/src/table/mod.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 76cbbae892..3c0cfafca7 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -991,7 +991,6 @@ mod tests { .await .unwrap(); let iceberg_scan = scan_plan - .as_any() .downcast_ref::() .expect("Expected IcebergTableScan"); assert_eq!( From 4fe1ebb41752c23e735b4f6b9d47fe4c901d21bb Mon Sep 17 00:00:00 2001 From: Noah Date: Mon, 20 Jul 2026 22:49:59 -0400 Subject: [PATCH 04/10] make tests assert batch --- .../integrations/datafusion/src/table/mod.rs | 55 +++++++++++++------ 1 file changed, 37 insertions(+), 18 deletions(-) diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 3c0cfafca7..50e1dead71 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -913,22 +913,24 @@ mod tests { ); } - /// Counts the rows a provider returns for `SELECT * FROM t`. - async fn scan_row_count(provider: IcebergTableProvider) -> usize { + /// Runs `SELECT * FROM t` against `provider` and returns the result batches. + async fn scan_rows( + provider: IcebergTableProvider, + ) -> Vec { let ctx = SessionContext::new(); ctx.register_table("t", Arc::new(provider)).unwrap(); - let batches = ctx - .sql("SELECT * FROM t") + ctx.sql("SELECT * FROM t") .await .unwrap() .collect() .await - .unwrap(); - batches.iter().map(|b| b.num_rows()).sum() + .unwrap() } #[tokio::test] async fn test_pinned_snapshot_reads_historical_data() { + use datafusion::assert_batches_sorted_eq; + let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; // First append -> snapshot with a single row. @@ -969,10 +971,16 @@ mod tests { IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) .await .unwrap(); - assert_eq!( - scan_row_count(current).await, - 2, - "unpinned provider should read the current snapshot" + assert_batches_sorted_eq!( + [ + "+----+------+", + "| id | name |", + "+----+------+", + "| 1 | a |", + "| 2 | b |", + "+----+------+", + ], + &scan_rows(current).await ); // Pinning the first snapshot time-travels: only the first row is visible, @@ -1000,10 +1008,15 @@ mod tests { ); // And it changes what is actually read: only the historical row. - assert_eq!( - scan_row_count(pinned).await, - 1, - "provider pinned to the first snapshot should read only historical data" + assert_batches_sorted_eq!( + [ + "+----+------+", + "| id | name |", + "+----+------+", + "| 1 | a |", + "+----+------+", + ], + &scan_rows(pinned).await ); // Clearing the pin (Some -> None) unpins back to the current snapshot, @@ -1015,10 +1028,16 @@ mod tests { .with_snapshot_id(Some(first_snapshot)) .with_snapshot_id(None); assert_eq!(unpinned.snapshot_id(), None); - assert_eq!( - scan_row_count(unpinned).await, - 2, - "clearing the pin should read the current snapshot again" + assert_batches_sorted_eq!( + [ + "+----+------+", + "| id | name |", + "+----+------+", + "| 1 | a |", + "| 2 | b |", + "+----+------+", + ], + &scan_rows(unpinned).await ); } } From b97ce72169666941e8317899e453fe28cf263240 Mon Sep 17 00:00:00 2001 From: Noah Date: Thu, 23 Jul 2026 16:12:26 -0400 Subject: [PATCH 05/10] fix read-insert-read snapshot --- .../integrations/datafusion/src/table/mod.rs | 42 ++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 50e1dead71..73fa063ca5 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -74,7 +74,7 @@ pub struct IcebergTableProvider { schema: ArrowSchemaRef, /// Optional snapshot to read. `None` reads the current snapshot (refreshed /// from the catalog on each scan); `Some` pins reads to that snapshot for - /// time-travel. Writes always target the current table state. + /// time-travel and rejects writes. snapshot_id: Option, } @@ -105,6 +105,9 @@ impl IcebergTableProvider { /// Pins reads to a specific snapshot for time-travel. `None` (the default) /// reads the current snapshot. The snapshot id is threaded into the scan /// node, so it is serialized and honored by a distributed engine as well. + /// A pinned provider is read-only: `insert_into` returns an error, since a + /// write would commit to the current table state and be invisible to this + /// provider's pinned reads. pub fn with_snapshot_id(mut self, snapshot_id: Option) -> Self { self.snapshot_id = snapshot_id; self @@ -180,6 +183,16 @@ impl TableProvider for IcebergTableProvider { ))); } + // A pinned provider reads a fixed snapshot, but writes would commit to + // the table's current state — the inserted rows would be invisible to + // this provider's own reads. Reject the write instead. + if let Some(snapshot_id) = self.snapshot_id { + return Err(DataFusionError::NotImplemented(format!( + "IcebergTableProvider is pinned to snapshot {snapshot_id} and cannot be \ + written to; use an unpinned IcebergTableProvider for writes" + ))); + } + // Load fresh table metadata from catalog let table = self .catalog @@ -760,6 +773,33 @@ mod tests { } } + #[tokio::test] + async fn test_pinned_provider_rejects_writes() { + use datafusion::physical_plan::empty::EmptyExec; + + let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; + let provider = IcebergTableProvider::try_new(catalog, namespace, table_name) + .await + .unwrap() + .with_snapshot_id(Some(42)); + let ctx = SessionContext::new(); + + let input = Arc::new(EmptyExec::new(provider.schema())) as Arc; + let error = provider + .insert_into(&ctx.state(), input, InsertOp::Append) + .await + .expect_err("writes to a pinned provider should be rejected"); + + assert!( + matches!( + error, + DataFusionError::NotImplemented(ref message) + if message.contains("pinned to snapshot 42") + ), + "unexpected error: {error}" + ); + } + #[tokio::test] async fn test_insert_plan_fanout_enabled_no_sort() { use datafusion::datasource::TableProvider; From 3e42a9eca45df2d2b14f5245156c338b00a4962a Mon Sep 17 00:00:00 2001 From: Noah Date: Thu, 23 Jul 2026 23:40:41 -0400 Subject: [PATCH 06/10] fix schema version error --- .../integrations/datafusion/src/table/mod.rs | 223 ++++++++++++++++-- 1 file changed, 199 insertions(+), 24 deletions(-) diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 73fa063ca5..c54bb92cfd 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -56,6 +56,34 @@ use crate::physical_plan::scan::IcebergTableScan; use crate::physical_plan::sort::sort_by_partition; use crate::physical_plan::write::IcebergWriteExec; +/// Computes the arrow schema DataFusion should expose for reading `table` at the +/// given snapshot. +/// +/// `None` uses the table's current schema. `Some` validates the snapshot exists +/// and uses that snapshot's schema, so time-travel reads stay consistent with the +/// schema in effect at that snapshot even after the table's schema has evolved. +fn snapshot_arrow_schema(table: &Table, snapshot_id: Option) -> Result { + let iceberg_schema = match snapshot_id { + None => table.metadata().current_schema().clone(), + Some(snapshot_id) => { + let snapshot = table + .metadata() + .snapshot_by_id(snapshot_id) + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "snapshot id {snapshot_id} not found in table {}", + table.identifier().name() + ), + ) + })?; + snapshot.schema(table.metadata())? + } + }; + Ok(Arc::new(schema_to_arrow_schema(&iceberg_schema)?)) +} + /// Catalog-backed table provider with automatic metadata refresh. /// /// This provider loads fresh table metadata from the catalog on every scan and write @@ -70,7 +98,12 @@ pub struct IcebergTableProvider { catalog: Arc, /// The table identifier (namespace + name) table_ident: TableIdent, - /// A reference-counted arrow `Schema` (cached at construction) + /// The table as loaded at construction. Used only to resolve schemas at + /// construction time and when pinning a snapshot; scans and writes always + /// reload fresh metadata from the catalog. + table: Table, + /// A reference-counted arrow `Schema` (cached at construction, recomputed + /// when a snapshot is pinned so it always matches the schema being read) schema: ArrowSchemaRef, /// Optional snapshot to read. `None` reads the current snapshot (refreshed /// from the catalog on each scan); `Some` pins reads to that snapshot for @@ -92,11 +125,12 @@ impl IcebergTableProvider { // Load table once to get initial schema let table = catalog.load_table(&table_ident).await?; - let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?); + let schema = snapshot_arrow_schema(&table, None)?; Ok(IcebergTableProvider { catalog, table_ident, + table, schema, snapshot_id: None, }) @@ -108,9 +142,18 @@ impl IcebergTableProvider { /// A pinned provider is read-only: `insert_into` returns an error, since a /// write would commit to the current table state and be invisible to this /// provider's pinned reads. - pub fn with_snapshot_id(mut self, snapshot_id: Option) -> Self { + /// + /// Validates that the snapshot exists and re-derives the exposed schema from + /// it, so `schema()` and the columns pushed down to the scan match the schema + /// actually being read even when the table's schema has since evolved. + /// + /// The snapshot is resolved against the metadata loaded when this provider + /// was constructed, so a snapshot committed afterwards is not visible here; + /// rebuild the provider to pin one. + pub fn with_snapshot_id(mut self, snapshot_id: Option) -> Result { + self.schema = snapshot_arrow_schema(&self.table, snapshot_id)?; self.snapshot_id = snapshot_id; - self + Ok(self) } /// Returns the snapshot this provider reads, if pinned for time-travel. @@ -290,7 +333,7 @@ impl IcebergStaticTableProvider { /// /// Uses the table's current snapshot for all queries. Does not support write operations. pub async fn try_new_from_table(table: Table) -> Result { - let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?); + let schema = snapshot_arrow_schema(&table, None)?; Ok(IcebergStaticTableProvider { table, snapshot_id: None, @@ -303,20 +346,7 @@ impl IcebergStaticTableProvider { /// Queries the specified snapshot for all operations. Useful for time-travel queries. /// Does not support write operations. pub async fn try_new_from_table_snapshot(table: Table, snapshot_id: i64) -> Result { - let snapshot = table - .metadata() - .snapshot_by_id(snapshot_id) - .ok_or_else(|| { - Error::new( - ErrorKind::Unexpected, - format!( - "snapshot id {snapshot_id} not found in table {}", - table.identifier().name() - ), - ) - })?; - let table_schema = snapshot.schema(table.metadata())?; - let schema = Arc::new(schema_to_arrow_schema(&table_schema)?); + let schema = snapshot_arrow_schema(&table, Some(snapshot_id))?; Ok(IcebergStaticTableProvider { table, snapshot_id: Some(snapshot_id), @@ -778,11 +808,36 @@ mod tests { use datafusion::physical_plan::empty::EmptyExec; let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; + + // An append so the table has a snapshot to pin. + let writer = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) + .await + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table("t", Arc::new(writer)).unwrap(); + ctx.sql("INSERT INTO t VALUES (1, 'a')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let table_ident = TableIdent::new(namespace.clone(), table_name.clone()); + let snapshot = catalog + .load_table(&table_ident) + .await + .unwrap() + .metadata() + .current_snapshot() + .unwrap() + .snapshot_id(); + let provider = IcebergTableProvider::try_new(catalog, namespace, table_name) .await .unwrap() - .with_snapshot_id(Some(42)); - let ctx = SessionContext::new(); + .with_snapshot_id(Some(snapshot)) + .unwrap(); let input = Arc::new(EmptyExec::new(provider.schema())) as Arc; let error = provider @@ -794,7 +849,7 @@ mod tests { matches!( error, DataFusionError::NotImplemented(ref message) - if message.contains("pinned to snapshot 42") + if message.contains(&format!("pinned to snapshot {snapshot}")) ), "unexpected error: {error}" ); @@ -1029,7 +1084,8 @@ mod tests { IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) .await .unwrap() - .with_snapshot_id(Some(first_snapshot)); + .with_snapshot_id(Some(first_snapshot)) + .unwrap(); assert_eq!(pinned.snapshot_id(), Some(first_snapshot)); // The pin is threaded onto the scan node itself — this is the value a @@ -1066,7 +1122,9 @@ mod tests { .await .unwrap() .with_snapshot_id(Some(first_snapshot)) - .with_snapshot_id(None); + .unwrap() + .with_snapshot_id(None) + .unwrap(); assert_eq!(unpinned.snapshot_id(), None); assert_batches_sorted_eq!( [ @@ -1080,4 +1138,121 @@ mod tests { &scan_rows(unpinned).await ); } + + #[tokio::test] + async fn test_with_snapshot_id_recomputes_schema_on_evolution() { + use datafusion::assert_batches_sorted_eq; + use iceberg::transaction::{AddColumn, ApplyTransactionAction, Transaction}; + + // A table with a single `id` column. + let temp_dir = TempDir::new().unwrap(); + let warehouse_path = temp_dir.path().to_str().unwrap().to_string(); + let catalog: Arc = Arc::new( + MemoryCatalogBuilder::default() + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse_path.clone())]), + ) + .await + .unwrap(), + ); + let namespace = NamespaceIdent::new("test_ns".to_string()); + catalog + .create_namespace(&namespace, HashMap::new()) + .await + .unwrap(); + let schema = Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build() + .unwrap(); + catalog + .create_table( + &namespace, + TableCreation::builder() + .name("t".to_string()) + .location(format!("{warehouse_path}/t")) + .schema(schema) + .properties(HashMap::new()) + .build(), + ) + .await + .unwrap(); + + // Append a row while the schema is still {id} -> snapshot S1 (schema id 0). + let writer = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), "t".to_string()) + .await + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table("t", Arc::new(writer)).unwrap(); + ctx.sql("INSERT INTO t VALUES (1)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let table_ident = TableIdent::new(namespace.clone(), "t".to_string()); + let s1 = catalog + .load_table(&table_ident) + .await + .unwrap() + .metadata() + .current_snapshot() + .unwrap() + .snapshot_id(); + + // Evolve the schema: add a `name` column. Current schema becomes + // {id, name}, but S1 still references the {id}-only schema. + let table = catalog.load_table(&table_ident).await.unwrap(); + let tx = Transaction::new(&table); + let tx = tx + .update_schema() + .add_column(AddColumn::optional( + "name", + Type::Primitive(PrimitiveType::String), + )) + .apply(tx) + .unwrap(); + tx.commit(catalog.as_ref()).await.unwrap(); + + // Unpinned provider exposes the evolved schema {id, name}. + let current = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), "t".to_string()) + .await + .unwrap(); + assert_eq!(current.schema().fields().len(), 2); + + // Pinning S1 re-derives the historical schema {id}. Before the fix, + // schema() returned the current {id, name}, and the scan below failed with + // "Column name not found in table". + let pinned = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), "t".to_string()) + .await + .unwrap() + .with_snapshot_id(Some(s1)) + .unwrap(); + assert_eq!(pinned.schema().fields().len(), 1); + assert_eq!(pinned.schema().field(0).name(), "id"); + + // A SELECT * against the pinned snapshot reads the historical row cleanly. + assert_batches_sorted_eq!( + ["+----+", "| id |", "+----+", "| 1 |", "+----+",], + &scan_rows(pinned).await + ); + + // A nonexistent snapshot id is rejected up front by with_snapshot_id. + let err = IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), "t".to_string()) + .await + .unwrap() + .with_snapshot_id(Some(9_999_999)) + .unwrap_err(); + assert!( + err.to_string().contains("snapshot id"), + "unexpected error: {err}" + ); + } } From a9215c7d013e3ce52a7db69a28060443b2c7a4f3 Mon Sep 17 00:00:00 2001 From: Noah Date: Fri, 24 Jul 2026 00:45:57 -0400 Subject: [PATCH 07/10] fix test --- crates/integrations/datafusion/public-api.txt | 4 ++-- crates/integrations/datafusion/src/table/mod.rs | 11 ++++++----- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index be095bbef2..1321e1a0c9 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -72,7 +72,7 @@ pub fn iceberg_datafusion::IcebergStaticTableProvider::table_type(&self) -> data pub struct iceberg_datafusion::table::IcebergTableProvider impl iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::snapshot_id(&self) -> core::option::Option -pub fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> Self +pub fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> iceberg::error::Result impl core::clone::Clone for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafusion::IcebergTableProvider impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider @@ -119,7 +119,7 @@ pub fn iceberg_datafusion::IcebergStaticTableProvider::table_type(&self) -> data pub struct iceberg_datafusion::IcebergTableProvider impl iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::snapshot_id(&self) -> core::option::Option -pub fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> Self +pub fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> iceberg::error::Result impl core::clone::Clone for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafusion::IcebergTableProvider impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index c54bb92cfd..44ae0af9d8 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -1245,11 +1245,12 @@ mod tests { ); // A nonexistent snapshot id is rejected up front by with_snapshot_id. - let err = IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), "t".to_string()) - .await - .unwrap() - .with_snapshot_id(Some(9_999_999)) - .unwrap_err(); + let err = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), "t".to_string()) + .await + .unwrap() + .with_snapshot_id(Some(9_999_999)) + .unwrap_err(); assert!( err.to_string().contains("snapshot id"), "unexpected error: {err}" From cb218e5df80d611203b80b4a3b35eef644e1120e Mon Sep 17 00:00:00 2001 From: Noah Date: Sat, 25 Jul 2026 17:19:13 -0400 Subject: [PATCH 08/10] make snapshot_arrow_schema public --- crates/integrations/datafusion/public-api.txt | 2 ++ crates/integrations/datafusion/src/table/mod.rs | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index 1321e1a0c9..9acdba7292 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -83,6 +83,7 @@ pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'l pub fn iceberg_datafusion::IcebergTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::IcebergTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType +pub fn iceberg_datafusion::table::snapshot_arrow_schema(table: &iceberg::table::Table, snapshot_id: core::option::Option) -> iceberg::error::Result pub mod iceberg_datafusion::table_provider_factory pub struct iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory impl iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory @@ -141,4 +142,5 @@ pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory:: impl datafusion_catalog::table::TableProviderFactory for iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory::create<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, cmd: &'life2 datafusion_expr::logical_plan::ddl::CreateExternalTable) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait pub fn iceberg_datafusion::from_datafusion_error(error: datafusion_common::error::DataFusionError) -> iceberg::error::Error +pub fn iceberg_datafusion::snapshot_arrow_schema(table: &iceberg::table::Table, snapshot_id: core::option::Option) -> iceberg::error::Result pub fn iceberg_datafusion::to_datafusion_error(error: iceberg::error::Error) -> datafusion_common::error::DataFusionError diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 44ae0af9d8..c7ece34434 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -62,7 +62,7 @@ use crate::physical_plan::write::IcebergWriteExec; /// `None` uses the table's current schema. `Some` validates the snapshot exists /// and uses that snapshot's schema, so time-travel reads stay consistent with the /// schema in effect at that snapshot even after the table's schema has evolved. -fn snapshot_arrow_schema(table: &Table, snapshot_id: Option) -> Result { +pub fn snapshot_arrow_schema(table: &Table, snapshot_id: Option) -> Result { let iceberg_schema = match snapshot_id { None => table.metadata().current_schema().clone(), Some(snapshot_id) => { @@ -74,7 +74,7 @@ fn snapshot_arrow_schema(table: &Table, snapshot_id: Option) -> Result Date: Sun, 26 Jul 2026 00:11:03 -0400 Subject: [PATCH 09/10] reload table in with_snapshot_id, get rid of Table Parameter --- crates/integrations/datafusion/public-api.txt | 4 +- .../integrations/datafusion/src/table/mod.rs | 227 +++++++++++++++--- 2 files changed, 189 insertions(+), 42 deletions(-) diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index 9acdba7292..8dc130f16c 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -72,7 +72,7 @@ pub fn iceberg_datafusion::IcebergStaticTableProvider::table_type(&self) -> data pub struct iceberg_datafusion::table::IcebergTableProvider impl iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::snapshot_id(&self) -> core::option::Option -pub fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> iceberg::error::Result +pub async fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> iceberg::error::Result impl core::clone::Clone for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafusion::IcebergTableProvider impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider @@ -120,7 +120,7 @@ pub fn iceberg_datafusion::IcebergStaticTableProvider::table_type(&self) -> data pub struct iceberg_datafusion::IcebergTableProvider impl iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::snapshot_id(&self) -> core::option::Option -pub fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> iceberg::error::Result +pub async fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> iceberg::error::Result impl core::clone::Clone for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafusion::IcebergTableProvider impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index c7ece34434..1a42b63c9a 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -56,12 +56,9 @@ use crate::physical_plan::scan::IcebergTableScan; use crate::physical_plan::sort::sort_by_partition; use crate::physical_plan::write::IcebergWriteExec; -/// Computes the arrow schema DataFusion should expose for reading `table` at the -/// given snapshot. -/// -/// `None` uses the table's current schema. `Some` validates the snapshot exists -/// and uses that snapshot's schema, so time-travel reads stay consistent with the -/// schema in effect at that snapshot even after the table's schema has evolved. +/// Computes the arrow schema to expose for reading `table` at the given snapshot: +/// the table's current schema for `None`, that snapshot's schema for `Some`, so +/// a time-travel read plans against the schema in effect at that snapshot. pub fn snapshot_arrow_schema(table: &Table, snapshot_id: Option) -> Result { let iceberg_schema = match snapshot_id { None => table.metadata().current_schema().clone(), @@ -98,16 +95,11 @@ pub struct IcebergTableProvider { catalog: Arc, /// The table identifier (namespace + name) table_ident: TableIdent, - /// The table as loaded at construction. Used only to resolve schemas at - /// construction time and when pinning a snapshot; scans and writes always - /// reload fresh metadata from the catalog. - table: Table, /// A reference-counted arrow `Schema` (cached at construction, recomputed - /// when a snapshot is pinned so it always matches the schema being read) + /// when a snapshot is pinned) schema: ArrowSchemaRef, - /// Optional snapshot to read. `None` reads the current snapshot (refreshed - /// from the catalog on each scan); `Some` pins reads to that snapshot for - /// time-travel and rejects writes. + /// Snapshot to read. `None` reads the current snapshot; `Some` pins reads to + /// that snapshot and rejects writes. snapshot_id: Option, } @@ -130,28 +122,23 @@ impl IcebergTableProvider { Ok(IcebergTableProvider { catalog, table_ident, - table, schema, snapshot_id: None, }) } - /// Pins reads to a specific snapshot for time-travel. `None` (the default) - /// reads the current snapshot. The snapshot id is threaded into the scan - /// node, so it is serialized and honored by a distributed engine as well. - /// A pinned provider is read-only: `insert_into` returns an error, since a - /// write would commit to the current table state and be invisible to this - /// provider's pinned reads. + /// Pins reads to a snapshot for time travel; `None` (the default) reads the + /// current snapshot. A pinned provider is read-only: `insert_into` would + /// commit to the current table state, invisible to its own reads, so it + /// errors instead. /// - /// Validates that the snapshot exists and re-derives the exposed schema from - /// it, so `schema()` and the columns pushed down to the scan match the schema - /// actually being read even when the table's schema has since evolved. - /// - /// The snapshot is resolved against the metadata loaded when this provider - /// was constructed, so a snapshot committed afterwards is not visible here; - /// rebuild the provider to pin one. - pub fn with_snapshot_id(mut self, snapshot_id: Option) -> Result { - self.schema = snapshot_arrow_schema(&self.table, snapshot_id)?; + /// Reloads metadata from the catalog to validate the snapshot and re-derive + /// `schema()` from it. The reload matters because providers handed out by + /// the catalog are cached, so a snapshot committed after construction would + /// otherwise look nonexistent. + pub async fn with_snapshot_id(mut self, snapshot_id: Option) -> Result { + let table = self.catalog.load_table(&self.table_ident).await?; + self.schema = snapshot_arrow_schema(&table, snapshot_id)?; self.snapshot_id = snapshot_id; Ok(self) } @@ -226,9 +213,8 @@ impl TableProvider for IcebergTableProvider { ))); } - // A pinned provider reads a fixed snapshot, but writes would commit to - // the table's current state — the inserted rows would be invisible to - // this provider's own reads. Reject the write instead. + // A write would commit to the table's current state, invisible to this + // provider's pinned reads. if let Some(snapshot_id) = self.snapshot_id { return Err(DataFusionError::NotImplemented(format!( "IcebergTableProvider is pinned to snapshot {snapshot_id} and cannot be \ @@ -837,6 +823,7 @@ mod tests { .await .unwrap() .with_snapshot_id(Some(snapshot)) + .await .unwrap(); let input = Arc::new(EmptyExec::new(provider.schema())) as Arc; @@ -1078,18 +1065,18 @@ mod tests { &scan_rows(current).await ); - // Pinning the first snapshot time-travels: only the first row is visible, - // even though a newer snapshot exists. + // Pinning the first snapshot hides the newer row. let pinned = IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) .await .unwrap() .with_snapshot_id(Some(first_snapshot)) + .await .unwrap(); assert_eq!(pinned.snapshot_id(), Some(first_snapshot)); - // The pin is threaded onto the scan node itself — this is the value a - // distributed engine's codec serializes to reproduce the scan remotely. + // The pin is carried on the scan node itself, so a distributed engine's + // codec can serialize it. let scan_plan = pinned .scan(&SessionContext::new().state(), None, &[], None) .await @@ -1103,7 +1090,7 @@ mod tests { "pinned snapshot should propagate to the scan node" ); - // And it changes what is actually read: only the historical row. + // And it changes what is read. assert_batches_sorted_eq!( [ "+----+------+", @@ -1115,15 +1102,16 @@ mod tests { &scan_rows(pinned).await ); - // Clearing the pin (Some -> None) unpins back to the current snapshot, - // so the newer row becomes visible again. + // Clearing the pin brings the newer row back. let unpinned = IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) .await .unwrap() .with_snapshot_id(Some(first_snapshot)) + .await .unwrap() .with_snapshot_id(None) + .await .unwrap(); assert_eq!(unpinned.snapshot_id(), None); assert_batches_sorted_eq!( @@ -1139,6 +1127,58 @@ mod tests { ); } + #[tokio::test] + async fn test_with_snapshot_id_sees_snapshot_created_after_construction() { + use datafusion::assert_batches_sorted_eq; + + let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; + + // Built before the table has any snapshot, then left alone — the shape of + // the providers the catalog builds once and caches. + let stale = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) + .await + .unwrap(); + + // Someone else writes, producing a snapshot `stale` has never seen. + let writer = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) + .await + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table("t", Arc::new(writer)).unwrap(); + ctx.sql("INSERT INTO t VALUES (1, 'a')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let table_ident = TableIdent::new(namespace.clone(), table_name.clone()); + let new_snapshot = catalog + .load_table(&table_ident) + .await + .unwrap() + .metadata() + .current_snapshot() + .unwrap() + .snapshot_id(); + + // Pinning reloads, so the newer snapshot resolves. + let pinned = stale.with_snapshot_id(Some(new_snapshot)).await.unwrap(); + assert_eq!(pinned.snapshot_id(), Some(new_snapshot)); + assert_batches_sorted_eq!( + [ + "+----+------+", + "| id | name |", + "+----+------+", + "| 1 | a |", + "+----+------+", + ], + &scan_rows(pinned).await + ); + } + #[tokio::test] async fn test_with_snapshot_id_recomputes_schema_on_evolution() { use datafusion::assert_batches_sorted_eq; @@ -1234,6 +1274,7 @@ mod tests { .await .unwrap() .with_snapshot_id(Some(s1)) + .await .unwrap(); assert_eq!(pinned.schema().fields().len(), 1); assert_eq!(pinned.schema().field(0).name(), "id"); @@ -1250,10 +1291,116 @@ mod tests { .await .unwrap() .with_snapshot_id(Some(9_999_999)) + .await .unwrap_err(); assert!( err.to_string().contains("snapshot id"), "unexpected error: {err}" ); } + + const PINNED_SNAPSHOT: i64 = 3_051_729_675_574_597_004; + + /// A table whose `PINNED_SNAPSHOT` was written under `{id, name}`, with + /// `email` added after it. Built from metadata rather than a catalog, so the + /// tests below stay synchronous and do no I/O. + fn evolved_table() -> Table { + use iceberg::spec::{ + FormatVersion, Operation, PartitionSpec, Snapshot, SortOrder, Summary, + TableMetadataBuilder, + }; + use iceberg::test_utils::test_runtime; + + let v1 = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(); + let v2 = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String)).into(), + NestedField::optional(3, "email", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(); + + // Two stages, so the snapshot gets the schema id the builder assigned to + // `v1` rather than a guessed one. + let base = TableMetadataBuilder::new( + v1, + PartitionSpec::unpartition_spec(), + SortOrder::unsorted_order(), + "/test/table".to_string(), + FormatVersion::V2, + HashMap::new(), + ) + .unwrap() + .build() + .unwrap() + .metadata; + let v1_schema_id = base.current_schema().schema_id(); + + let snapshot = Snapshot::builder() + .with_snapshot_id(PINNED_SNAPSHOT) + .with_sequence_number(1) + .with_timestamp_ms(base.last_updated_ms()) + .with_manifest_list("/test/snap-1.avro") + .with_schema_id(v1_schema_id) + .with_summary(Summary { + operation: Operation::Append, + additional_properties: HashMap::new(), + }) + .build(); + + let metadata = TableMetadataBuilder::new_from_metadata(base, None) + .add_snapshot(snapshot) + .unwrap() + .add_schema(v2) + .unwrap() + .set_current_schema(-1) + .unwrap() + .build() + .unwrap() + .metadata; + + Table::builder() + .metadata(metadata) + .identifier(TableIdent::from_strs(["ns", "tbl"]).unwrap()) + .file_io(FileIO::new_with_fs()) + .metadata_location("/test/metadata.json") + .runtime(test_runtime()) + .build() + .unwrap() + } + + fn field_names(schema: &ArrowSchemaRef) -> Vec<&str> { + schema.fields().iter().map(|f| f.name().as_str()).collect() + } + + #[test] + fn test_snapshot_arrow_schema_pinned_uses_that_snapshot_schema() { + // `email` came later, so a pinned read must not advertise it. + let schema = snapshot_arrow_schema(&evolved_table(), Some(PINNED_SNAPSHOT)).unwrap(); + assert_eq!(field_names(&schema), vec!["id", "name"]); + } + + #[test] + fn test_snapshot_arrow_schema_unpinned_uses_current_schema() { + let schema = snapshot_arrow_schema(&evolved_table(), None).unwrap(); + assert_eq!(field_names(&schema), vec!["id", "name", "email"]); + } + + #[test] + fn test_snapshot_arrow_schema_unknown_snapshot_errors() { + // Falling back to the current schema would turn a stale pin into wrong + // output instead of an error. + let err = snapshot_arrow_schema(&evolved_table(), Some(-1)).unwrap_err(); + let msg = err.to_string(); + assert!(msg.contains("snapshot id -1"), "names the snapshot: {msg}"); + // Namespace-qualified: a bare table name is ambiguous across namespaces. + assert!(msg.contains("ns.tbl"), "names the table: {msg}"); + } } From fe72134f335289e1432f0e132a9d40b32aa7d892 Mon Sep 17 00:00:00 2001 From: Noah Date: Sun, 26 Jul 2026 00:17:51 -0400 Subject: [PATCH 10/10] simplify tests --- .../integrations/datafusion/src/table/mod.rs | 395 ++++-------------- 1 file changed, 91 insertions(+), 304 deletions(-) diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 1a42b63c9a..45d4cf18f7 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -809,16 +809,7 @@ mod tests { .await .unwrap(); - let table_ident = TableIdent::new(namespace.clone(), table_name.clone()); - let snapshot = catalog - .load_table(&table_ident) - .await - .unwrap() - .metadata() - .current_snapshot() - .unwrap() - .snapshot_id(); - + let snapshot = current_snapshot_id(&catalog, &namespace, &table_name).await; let provider = IcebergTableProvider::try_new(catalog, namespace, table_name) .await .unwrap() @@ -1009,38 +1000,43 @@ mod tests { .unwrap() } + async fn current_snapshot_id( + catalog: &Arc, + namespace: &NamespaceIdent, + table_name: &str, + ) -> i64 { + catalog + .load_table(&TableIdent::new(namespace.clone(), table_name.to_string())) + .await + .unwrap() + .metadata() + .current_snapshot() + .unwrap() + .snapshot_id() + } + #[tokio::test] async fn test_pinned_snapshot_reads_historical_data() { use datafusion::assert_batches_sorted_eq; let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; - // First append -> snapshot with a single row. - let writer = + // Built while the table still has no snapshot at all, then left alone — + // the shape of the providers the catalog builds once and caches. + let provider = IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) .await .unwrap(); + let ctx = SessionContext::new(); - ctx.register_table("t", Arc::new(writer)).unwrap(); + ctx.register_table("t", Arc::new(provider.clone())).unwrap(); ctx.sql("INSERT INTO t VALUES (1, 'a')") .await .unwrap() .collect() .await .unwrap(); - - // Capture the snapshot produced by the first append. - let table_ident = TableIdent::new(namespace.clone(), table_name.clone()); - let first_snapshot = catalog - .load_table(&table_ident) - .await - .unwrap() - .metadata() - .current_snapshot() - .unwrap() - .snapshot_id(); - - // Second append -> current snapshot now has two rows. + let first_snapshot = current_snapshot_id(&catalog, &namespace, &table_name).await; ctx.sql("INSERT INTO t VALUES (2, 'b')") .await .unwrap() @@ -1048,11 +1044,8 @@ mod tests { .await .unwrap(); - // Default (unpinned) provider sees the latest state: both rows. - let current = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap(); + // Unpinned, it reloads on every scan and sees the latest state. + assert_eq!(provider.snapshot_id(), None); assert_batches_sorted_eq!( [ "+----+------+", @@ -1062,18 +1055,27 @@ mod tests { "| 2 | b |", "+----+------+", ], - &scan_rows(current).await + &scan_rows(provider.clone()).await ); - // Pinning the first snapshot hides the newer row. - let pinned = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap() - .with_snapshot_id(Some(first_snapshot)) - .await - .unwrap(); + // Pinning resolves `first_snapshot` even though the provider predates it, + // and hides the newer row. + let pinned = provider + .clone() + .with_snapshot_id(Some(first_snapshot)) + .await + .unwrap(); assert_eq!(pinned.snapshot_id(), Some(first_snapshot)); + assert_batches_sorted_eq!( + [ + "+----+------+", + "| id | name |", + "+----+------+", + "| 1 | a |", + "+----+------+", + ], + &scan_rows(pinned.clone()).await + ); // The pin is carried on the scan node itself, so a distributed engine's // codec can serialize it. @@ -1081,38 +1083,17 @@ mod tests { .scan(&SessionContext::new().state(), None, &[], None) .await .unwrap(); - let iceberg_scan = scan_plan - .downcast_ref::() - .expect("Expected IcebergTableScan"); assert_eq!( - iceberg_scan.snapshot_id(), + scan_plan + .downcast_ref::() + .expect("Expected IcebergTableScan") + .snapshot_id(), Some(first_snapshot), "pinned snapshot should propagate to the scan node" ); - // And it changes what is read. - assert_batches_sorted_eq!( - [ - "+----+------+", - "| id | name |", - "+----+------+", - "| 1 | a |", - "+----+------+", - ], - &scan_rows(pinned).await - ); - // Clearing the pin brings the newer row back. - let unpinned = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap() - .with_snapshot_id(Some(first_snapshot)) - .await - .unwrap() - .with_snapshot_id(None) - .await - .unwrap(); + let unpinned = pinned.with_snapshot_id(None).await.unwrap(); assert_eq!(unpinned.snapshot_id(), None); assert_batches_sorted_eq!( [ @@ -1128,19 +1109,14 @@ mod tests { } #[tokio::test] - async fn test_with_snapshot_id_sees_snapshot_created_after_construction() { + async fn test_with_snapshot_id_recomputes_schema_on_evolution() { use datafusion::assert_batches_sorted_eq; + use iceberg::transaction::{AddColumn, ApplyTransactionAction, Transaction}; let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; + let table_ident = TableIdent::new(namespace.clone(), table_name.clone()); - // Built before the table has any snapshot, then left alone — the shape of - // the providers the catalog builds once and caches. - let stale = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap(); - - // Someone else writes, producing a snapshot `stale` has never seen. + // Append under the original {id, name} schema. let writer = IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) .await @@ -1153,254 +1129,65 @@ mod tests { .collect() .await .unwrap(); + let snapshot = current_snapshot_id(&catalog, &namespace, &table_name).await; - let table_ident = TableIdent::new(namespace.clone(), table_name.clone()); - let new_snapshot = catalog - .load_table(&table_ident) - .await - .unwrap() - .metadata() - .current_snapshot() - .unwrap() - .snapshot_id(); - - // Pinning reloads, so the newer snapshot resolves. - let pinned = stale.with_snapshot_id(Some(new_snapshot)).await.unwrap(); - assert_eq!(pinned.snapshot_id(), Some(new_snapshot)); - assert_batches_sorted_eq!( - [ - "+----+------+", - "| id | name |", - "+----+------+", - "| 1 | a |", - "+----+------+", - ], - &scan_rows(pinned).await - ); - } - - #[tokio::test] - async fn test_with_snapshot_id_recomputes_schema_on_evolution() { - use datafusion::assert_batches_sorted_eq; - use iceberg::transaction::{AddColumn, ApplyTransactionAction, Transaction}; - - // A table with a single `id` column. - let temp_dir = TempDir::new().unwrap(); - let warehouse_path = temp_dir.path().to_str().unwrap().to_string(); - let catalog: Arc = Arc::new( - MemoryCatalogBuilder::default() - .load( - "memory", - HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse_path.clone())]), - ) - .await - .unwrap(), - ); - let namespace = NamespaceIdent::new("test_ns".to_string()); - catalog - .create_namespace(&namespace, HashMap::new()) - .await - .unwrap(); - let schema = Schema::builder() - .with_schema_id(0) - .with_fields(vec![ - NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), - ]) - .build() - .unwrap(); - catalog - .create_table( - &namespace, - TableCreation::builder() - .name("t".to_string()) - .location(format!("{warehouse_path}/t")) - .schema(schema) - .properties(HashMap::new()) - .build(), - ) - .await - .unwrap(); - - // Append a row while the schema is still {id} -> snapshot S1 (schema id 0). - let writer = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), "t".to_string()) - .await - .unwrap(); - let ctx = SessionContext::new(); - ctx.register_table("t", Arc::new(writer)).unwrap(); - ctx.sql("INSERT INTO t VALUES (1)") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let table_ident = TableIdent::new(namespace.clone(), "t".to_string()); - let s1 = catalog - .load_table(&table_ident) - .await - .unwrap() - .metadata() - .current_snapshot() - .unwrap() - .snapshot_id(); - - // Evolve the schema: add a `name` column. Current schema becomes - // {id, name}, but S1 still references the {id}-only schema. + // Evolve the schema. The snapshot above still references {id, name}. let table = catalog.load_table(&table_ident).await.unwrap(); let tx = Transaction::new(&table); let tx = tx .update_schema() .add_column(AddColumn::optional( - "name", + "email", Type::Primitive(PrimitiveType::String), )) .apply(tx) .unwrap(); tx.commit(catalog.as_ref()).await.unwrap(); - // Unpinned provider exposes the evolved schema {id, name}. - let current = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), "t".to_string()) - .await - .unwrap(); - assert_eq!(current.schema().fields().len(), 2); - - // Pinning S1 re-derives the historical schema {id}. Before the fix, - // schema() returned the current {id, name}, and the scan below failed with - // "Column name not found in table". - let pinned = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), "t".to_string()) - .await - .unwrap() - .with_snapshot_id(Some(s1)) + let provider = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) .await .unwrap(); - assert_eq!(pinned.schema().fields().len(), 1); - assert_eq!(pinned.schema().field(0).name(), "id"); + assert_eq!(provider.schema().fields().len(), 3); - // A SELECT * against the pinned snapshot reads the historical row cleanly. + // Pinning re-derives the historical schema. Before the fix, `schema()` + // kept advertising `email` and the scan below failed with "Column email + // not found in table". + let pinned = provider + .clone() + .with_snapshot_id(Some(snapshot)) + .await + .unwrap(); + assert_eq!( + pinned + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(), + ["id", "name"] + ); assert_batches_sorted_eq!( - ["+----+", "| id |", "+----+", "| 1 |", "+----+",], + [ + "+----+------+", + "| id | name |", + "+----+------+", + "| 1 | a |", + "+----+------+", + ], &scan_rows(pinned).await ); - // A nonexistent snapshot id is rejected up front by with_snapshot_id. - let err = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), "t".to_string()) - .await - .unwrap() - .with_snapshot_id(Some(9_999_999)) - .await - .unwrap_err(); - assert!( - err.to_string().contains("snapshot id"), - "unexpected error: {err}" - ); - } - - const PINNED_SNAPSHOT: i64 = 3_051_729_675_574_597_004; - - /// A table whose `PINNED_SNAPSHOT` was written under `{id, name}`, with - /// `email` added after it. Built from metadata rather than a catalog, so the - /// tests below stay synchronous and do no I/O. - fn evolved_table() -> Table { - use iceberg::spec::{ - FormatVersion, Operation, PartitionSpec, Snapshot, SortOrder, Summary, - TableMetadataBuilder, - }; - use iceberg::test_utils::test_runtime; - - let v1 = Schema::builder() - .with_fields(vec![ - NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), - NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String)).into(), - ]) - .build() - .unwrap(); - let v2 = Schema::builder() - .with_fields(vec![ - NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), - NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String)).into(), - NestedField::optional(3, "email", Type::Primitive(PrimitiveType::String)).into(), - ]) - .build() - .unwrap(); - - // Two stages, so the snapshot gets the schema id the builder assigned to - // `v1` rather than a guessed one. - let base = TableMetadataBuilder::new( - v1, - PartitionSpec::unpartition_spec(), - SortOrder::unsorted_order(), - "/test/table".to_string(), - FormatVersion::V2, - HashMap::new(), - ) - .unwrap() - .build() - .unwrap() - .metadata; - let v1_schema_id = base.current_schema().schema_id(); - - let snapshot = Snapshot::builder() - .with_snapshot_id(PINNED_SNAPSHOT) - .with_sequence_number(1) - .with_timestamp_ms(base.last_updated_ms()) - .with_manifest_list("/test/snap-1.avro") - .with_schema_id(v1_schema_id) - .with_summary(Summary { - operation: Operation::Append, - additional_properties: HashMap::new(), - }) - .build(); - - let metadata = TableMetadataBuilder::new_from_metadata(base, None) - .add_snapshot(snapshot) - .unwrap() - .add_schema(v2) - .unwrap() - .set_current_schema(-1) - .unwrap() - .build() - .unwrap() - .metadata; - - Table::builder() - .metadata(metadata) - .identifier(TableIdent::from_strs(["ns", "tbl"]).unwrap()) - .file_io(FileIO::new_with_fs()) - .metadata_location("/test/metadata.json") - .runtime(test_runtime()) - .build() - .unwrap() - } - - fn field_names(schema: &ArrowSchemaRef) -> Vec<&str> { - schema.fields().iter().map(|f| f.name().as_str()).collect() - } - - #[test] - fn test_snapshot_arrow_schema_pinned_uses_that_snapshot_schema() { - // `email` came later, so a pinned read must not advertise it. - let schema = snapshot_arrow_schema(&evolved_table(), Some(PINNED_SNAPSHOT)).unwrap(); - assert_eq!(field_names(&schema), vec!["id", "name"]); - } - - #[test] - fn test_snapshot_arrow_schema_unpinned_uses_current_schema() { - let schema = snapshot_arrow_schema(&evolved_table(), None).unwrap(); - assert_eq!(field_names(&schema), vec!["id", "name", "email"]); - } - - #[test] - fn test_snapshot_arrow_schema_unknown_snapshot_errors() { - // Falling back to the current schema would turn a stale pin into wrong - // output instead of an error. - let err = snapshot_arrow_schema(&evolved_table(), Some(-1)).unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("snapshot id -1"), "names the snapshot: {msg}"); + // A snapshot the table doesn't have is rejected up front. Falling back to + // the current schema would turn a stale pin into wrong output instead of + // an error. + let error = provider + .with_snapshot_id(Some(9_999_999)) + .await + .unwrap_err(); + let message = error.to_string(); + assert!(message.contains("snapshot id 9999999"), "{message}"); // Namespace-qualified: a bare table name is ambiguous across namespaces. - assert!(msg.contains("ns.tbl"), "names the table: {msg}"); + assert!(message.contains("test_ns.test_table"), "{message}"); } }