diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index e197c2057d..8dc130f16c 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -70,6 +70,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 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 @@ -80,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 @@ -114,6 +118,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 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 @@ -135,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 2fd958dff4..45d4cf18f7 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -56,6 +56,31 @@ use crate::physical_plan::scan::IcebergTableScan; use crate::physical_plan::sort::sort_by_partition; use crate::physical_plan::write::IcebergWriteExec; +/// 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(), + 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() + ), + ) + })?; + 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,8 +95,12 @@ pub struct IcebergTableProvider { catalog: Arc, /// The table identifier (namespace + name) table_ident: TableIdent, - /// A reference-counted arrow `Schema` (cached at construction) + /// A reference-counted arrow `Schema` (cached at construction, recomputed + /// when a snapshot is pinned) schema: ArrowSchemaRef, + /// Snapshot to read. `None` reads the current snapshot; `Some` pins reads to + /// that snapshot and rejects writes. + snapshot_id: Option, } impl IcebergTableProvider { @@ -88,15 +117,37 @@ 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, schema, + snapshot_id: None, }) } + /// 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. + /// + /// 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) + } + + /// 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, @@ -131,10 +182,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, @@ -162,6 +213,15 @@ impl TableProvider for IcebergTableProvider { ))); } + // 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 \ + written to; use an unpinned IcebergTableProvider for writes" + ))); + } + // Load fresh table metadata from catalog let table = self .catalog @@ -259,7 +319,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, @@ -272,20 +332,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), @@ -742,6 +789,50 @@ 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; + + // 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 snapshot = current_snapshot_id(&catalog, &namespace, &table_name).await; + let provider = IcebergTableProvider::try_new(catalog, namespace, table_name) + .await + .unwrap() + .with_snapshot_id(Some(snapshot)) + .await + .unwrap(); + + 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(&format!("pinned to snapshot {snapshot}")) + ), + "unexpected error: {error}" + ); + } + #[tokio::test] async fn test_insert_plan_fanout_enabled_no_sort() { use datafusion::datasource::TableProvider; @@ -894,4 +985,209 @@ mod tests { "Limit should be None when not specified" ); } + + /// 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(); + ctx.sql("SELECT * FROM t") + .await + .unwrap() + .collect() + .await + .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; + + // 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(provider.clone())).unwrap(); + ctx.sql("INSERT INTO t VALUES (1, 'a')") + .await + .unwrap() + .collect() + .await + .unwrap(); + let first_snapshot = current_snapshot_id(&catalog, &namespace, &table_name).await; + ctx.sql("INSERT INTO t VALUES (2, 'b')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + // Unpinned, it reloads on every scan and sees the latest state. + assert_eq!(provider.snapshot_id(), None); + assert_batches_sorted_eq!( + [ + "+----+------+", + "| id | name |", + "+----+------+", + "| 1 | a |", + "| 2 | b |", + "+----+------+", + ], + &scan_rows(provider.clone()).await + ); + + // 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. + let scan_plan = pinned + .scan(&SessionContext::new().state(), None, &[], None) + .await + .unwrap(); + assert_eq!( + scan_plan + .downcast_ref::() + .expect("Expected IcebergTableScan") + .snapshot_id(), + Some(first_snapshot), + "pinned snapshot should propagate to the scan node" + ); + + // Clearing the pin brings the newer row back. + let unpinned = pinned.with_snapshot_id(None).await.unwrap(); + assert_eq!(unpinned.snapshot_id(), None); + assert_batches_sorted_eq!( + [ + "+----+------+", + "| id | name |", + "+----+------+", + "| 1 | a |", + "| 2 | b |", + "+----+------+", + ], + &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}; + + let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; + let table_ident = TableIdent::new(namespace.clone(), table_name.clone()); + + // Append under the original {id, name} schema. + 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 snapshot = current_snapshot_id(&catalog, &namespace, &table_name).await; + + // 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( + "email", + Type::Primitive(PrimitiveType::String), + )) + .apply(tx) + .unwrap(); + tx.commit(catalog.as_ref()).await.unwrap(); + + let provider = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) + .await + .unwrap(); + assert_eq!(provider.schema().fields().len(), 3); + + // 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 | name |", + "+----+------+", + "| 1 | a |", + "+----+------+", + ], + &scan_rows(pinned).await + ); + + // 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!(message.contains("test_ns.test_table"), "{message}"); + } }