From c0bb59accc1f42fb9250cecf60c7d8d6dd85092b Mon Sep 17 00:00:00 2001 From: Noah Date: Sat, 26 Sep 2026 16:41:06 -0400 Subject: [PATCH 1/5] init --- crates/datafusion/src/physical_plan/commit.rs | 112 +++++++++++++- .../src/physical_plan/metadata_scan.rs | 7 + crates/datafusion/src/physical_plan/mod.rs | 3 + crates/datafusion/src/physical_plan/scan.rs | 140 +++++++++++++++--- crates/datafusion/src/physical_plan/write.rs | 13 +- crates/datafusion/src/table/metadata_table.rs | 25 +++- crates/datafusion/src/table/mod.rs | 50 ++++++- .../tests/integration_datafusion_test.rs | 129 +++++++++++++++- 8 files changed, 449 insertions(+), 30 deletions(-) diff --git a/crates/datafusion/src/physical_plan/commit.rs b/crates/datafusion/src/physical_plan/commit.rs index dcbee3a..3981f57 100644 --- a/crates/datafusion/src/physical_plan/commit.rs +++ b/crates/datafusion/src/physical_plan/commit.rs @@ -23,7 +23,7 @@ use datafusion::arrow::datatypes::{ DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, }; use datafusion::common::{ - internal_datafusion_err, internal_err, tree_node::TreeNodeRecursion, + exec_err, internal_datafusion_err, internal_err, tree_node::TreeNodeRecursion, }; use datafusion::error::{DataFusionError, Result}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; @@ -44,8 +44,12 @@ use crate::to_datafusion_error; /// IcebergCommitExec is responsible for collecting the files written and use /// [`Transaction::fast_append`] to commit the data files written. +/// +/// Its input produces the data files to commit, in the form +/// [`IcebergWriteExec`](super::IcebergWriteExec) outputs them, in a single +/// partition. Its output is one row holding the number of rows committed. #[derive(Debug)] -pub(crate) struct IcebergCommitExec { +pub struct IcebergCommitExec { table: Table, catalog: Arc, input: Arc, @@ -55,6 +59,13 @@ pub(crate) struct IcebergCommitExec { } impl IcebergCommitExec { + /// Commits the data files `input` produces to `table` through `catalog`. + /// + /// `input` must have a single partition, such as a + /// [`CoalescePartitionsExec`](datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec) + /// over an [`IcebergWriteExec`](super::IcebergWriteExec); executing the node + /// fails otherwise. `schema` is the table's Arrow schema, shown in the + /// verbose plan display. pub fn new( table: Table, catalog: Arc, @@ -75,6 +86,16 @@ impl IcebergCommitExec { } } + /// The catalog this node commits through. + pub fn catalog(&self) -> &Arc { + &self.catalog + } + + /// The table this node commits to, as loaded when the node was planned. + pub fn table(&self) -> &Table { + &self.table + } + // Compute the plan properties for this execution plan fn compute_properties(schema: ArrowSchemaRef) -> Arc { Arc::new(PlanProperties::new( @@ -193,6 +214,16 @@ impl ExecutionPlan for IcebergCommitExec { ); } + // Only partition 0 of the input is read below, so the files of any + // other partition would silently go uncommitted. + let input_partitions = self.input.properties().partitioning.partition_count(); + if input_partitions != 1 { + return exec_err!( + "IcebergCommitExec requires an input with one partition, but it has \ + {input_partitions}; coalesce the input first" + ); + } + let table = self.table.clone(); let input_plan = self.input.clone(); @@ -297,6 +328,7 @@ mod tests { use datafusion::physical_plan::common::collect; use datafusion::physical_plan::execution_plan::Boundedness; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, }; @@ -559,6 +591,82 @@ mod tests { Ok(()) } + /// The commit reads a single input partition, so an input with more is + /// refused rather than committing the files of its first partition alone. + #[tokio::test] + async fn test_iceberg_commit_exec_rejects_multiple_input_partitions() + -> Result<(), Box> { + let catalog = Arc::new( + MemoryCatalogBuilder::default() + .load( + "memory", + HashMap::from([( + MEMORY_CATALOG_WAREHOUSE.to_string(), + "memory://root".to_string(), + )]), + ) + .await?, + ); + let namespace = NamespaceIdent::new("test_namespace".to_string()); + catalog.create_namespace(&namespace, HashMap::new()).await?; + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)) + .into(), + ]) + .build()?; + let table_creation = TableCreation::builder() + .name("test_table".to_string()) + .schema(schema) + .location("memory://root/test_table".to_string()) + .build(); + let table = catalog.create_table(&namespace, table_creation).await?; + + // One data file in each of two input partitions. + let partition_type = table.metadata().default_partition_type().clone(); + let mut partitions: Vec> = vec![]; + for path in ["path/to/file1.parquet", "path/to/file2.parquet"] { + let data_file = DataFileBuilder::default() + .content(DataContentType::Data) + .file_path(path.to_string()) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(1024) + .record_count(100) + .partition_spec_id(table.metadata().default_partition_spec_id()) + .partition(Struct::empty()) + .build()?; + let json = iceberg::spec::serialize_data_file_to_json( + data_file, + &partition_type, + table.metadata().format_version(), + )?; + partitions.push(Arc::new(MockWriteExec::new(vec![json]))); + } + let input = UnionExec::try_new(partitions)?; + assert_eq!(input.properties().partitioning.partition_count(), 2); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( + DATA_FILES_COL_NAME, + DataType::Utf8, + false, + )])); + let commit_exec = + IcebergCommitExec::new(table.clone(), catalog.clone(), input, arrow_schema); + let err = match commit_exec.execute(0, Arc::new(TaskContext::default())) { + Ok(_) => panic!("a commit over two input partitions must not execute"), + Err(err) => err.to_string(), + }; + assert!( + err.contains("requires an input with one partition"), + "{err}" + ); + + let table = catalog.load_table(table.identifier()).await?; + assert!(table.metadata().current_snapshot().is_none()); + + Ok(()) + } + #[tokio::test] async fn test_iceberg_commit_exec_empty_insert() -> Result<(), Box> { diff --git a/crates/datafusion/src/physical_plan/metadata_scan.rs b/crates/datafusion/src/physical_plan/metadata_scan.rs index 92e0671..1394865 100644 --- a/crates/datafusion/src/physical_plan/metadata_scan.rs +++ b/crates/datafusion/src/physical_plan/metadata_scan.rs @@ -27,6 +27,7 @@ use futures::TryStreamExt; use crate::metadata_table::IcebergMetadataTableProvider; +/// Scans an Iceberg metadata table, such as `$snapshots`. #[derive(Debug)] pub struct IcebergMetadataScan { provider: IcebergMetadataTableProvider, @@ -34,6 +35,7 @@ pub struct IcebergMetadataScan { } impl IcebergMetadataScan { + /// Creates a scan of the metadata table `provider` reads. pub fn new(provider: IcebergMetadataTableProvider) -> Self { let properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(provider.schema()), @@ -46,6 +48,11 @@ impl IcebergMetadataScan { properties, } } + + /// The provider this node scans. + pub fn provider(&self) -> &IcebergMetadataTableProvider { + &self.provider + } } impl DisplayAs for IcebergMetadataScan { diff --git a/crates/datafusion/src/physical_plan/mod.rs b/crates/datafusion/src/physical_plan/mod.rs index 578dcb5..024266a 100644 --- a/crates/datafusion/src/physical_plan/mod.rs +++ b/crates/datafusion/src/physical_plan/mod.rs @@ -26,6 +26,9 @@ pub(crate) mod write; pub(crate) const DATA_FILES_COL_NAME: &str = "data_files"; +pub use commit::IcebergCommitExec; pub use expr_to_predicate::convert_filters_to_predicate; +pub use metadata_scan::IcebergMetadataScan; pub use project::{PartitionExpr, project_with_partition}; pub use scan::IcebergTableScan; +pub use write::IcebergWriteExec; diff --git a/crates/datafusion/src/physical_plan/scan.rs b/crates/datafusion/src/physical_plan/scan.rs index b9b15c5..e3c9ed7 100644 --- a/crates/datafusion/src/physical_plan/scan.rs +++ b/crates/datafusion/src/physical_plan/scan.rs @@ -64,23 +64,134 @@ impl IcebergTableScan { projection: Option<&Vec>, filters: &[Expr], limit: Option, - ) -> Self { + ) -> Result { + Self::new_with_predicate( + table, + snapshot_id, + schema, + projection.map(Vec::as_slice), + convert_filters_to_predicate(filters), + limit, + ) + } + + /// Creates a scan of `table` from an already-converted Iceberg + /// [`Predicate`] rather than DataFusion filters, for rebuilding a scan from + /// its parts, such as after sending them to another process. A predicate + /// cannot be converted back to the filters it came from. + /// + /// The arguments mean what the matching accessors return: + /// + /// - `snapshot_id`: the snapshot to read, or `None` for the table's current + /// snapshot. + /// - `schema`: the Arrow schema of the table the scan reads, as its + /// provider reports it. + /// - `projection`: indices into `schema` of the columns to read, or `None` + /// for all. The columns are read from the table by name. + /// - `predicate`: pushed down to Iceberg to skip data files and rows. The + /// table providers report their filters as + /// [`Inexact`](datafusion::logical_expr::TableProviderFilterPushDown::Inexact), + /// so DataFusion still applies them above the scan. + /// + /// # Errors + /// + /// Returns an error if `projection` holds an index outside `schema`. + /// + /// # Example + /// + /// ``` + /// use std::collections::HashMap; + /// + /// use datafusion::catalog::TableProvider; + /// use datafusion::physical_plan::ExecutionPlan; + /// use datafusion::prelude::{SessionContext, col, lit}; + /// use datafusion_iceberg::IcebergStaticTableProvider; + /// use datafusion_iceberg::physical_plan::IcebergTableScan; + /// use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; + /// use iceberg::spec::{NestedField, PrimitiveType, Schema, Type}; + /// use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation}; + /// + /// # tokio::runtime::Runtime::new()?.block_on(async { + /// # let warehouse = tempfile::tempdir()?; + /// # let props = HashMap::from([( + /// # MEMORY_CATALOG_WAREHOUSE.to_string(), + /// # warehouse.path().display().to_string(), + /// # )]); + /// # let catalog = MemoryCatalogBuilder::default().load("memory", props).await?; + /// # let namespace = NamespaceIdent::new("ns".to_string()); + /// # catalog.create_namespace(&namespace, HashMap::new()).await?; + /// # let schema = Schema::builder() + /// # .with_fields(vec![ + /// # NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + /// # NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String)) + /// # .into(), + /// # ]) + /// # .build()?; + /// # let creation = TableCreation::builder().name("t".to_string()).schema(schema).build(); + /// # let table = catalog.create_table(&namespace, creation).await?; + /// let provider = IcebergStaticTableProvider::try_new_from_table(table).await?; + /// let ctx = SessionContext::new(); + /// let filters = [col("id").gt(lit(1))]; + /// let plan = provider + /// .scan(&ctx.state(), Some(&vec![1]), &filters, None) + /// .await?; + /// let scan = plan.downcast_ref::().unwrap(); + /// + /// // Rebuild an equivalent scan from the original's parts. + /// let schema = provider.schema(); + /// let projection = scan + /// .projection() + /// .map(|names| { + /// names + /// .iter() + /// .map(|name| schema.index_of(name)) + /// .collect::, _>>() + /// }) + /// .transpose()?; + /// let rebuilt = IcebergTableScan::new_with_predicate( + /// scan.table().clone(), + /// scan.snapshot_id(), + /// schema, + /// projection.as_deref(), + /// scan.predicates().cloned(), + /// scan.limit(), + /// )?; + /// assert_eq!(rebuilt.schema(), scan.schema()); + /// assert_eq!(rebuilt.predicates(), scan.predicates()); + /// # Ok::<(), Box>(()) + /// # })?; + /// # Ok::<(), Box>(()) + /// ``` + pub fn new_with_predicate( + table: Table, + snapshot_id: Option, + schema: ArrowSchemaRef, + projection: Option<&[usize]>, + predicate: Option, + limit: Option, + ) -> Result { let output_schema = match projection { - None => schema.clone(), - Some(projection) => Arc::new(schema.project(projection).unwrap()), + None => schema, + Some(projection) => Arc::new(schema.project(projection)?), }; - let plan_properties = Self::compute_properties(output_schema.clone()); - let projection = get_column_names(schema.clone(), projection); - let predicates = convert_filters_to_predicate(filters); + // The columns to read, by name; `None` reads them all. + let projection = projection.map(|_| { + output_schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect() + }); + let plan_properties = Self::compute_properties(output_schema); - Self { + Ok(Self { table, snapshot_id, plan_properties, projection, - predicates, + predicates: predicate, limit, - } + }) } pub fn table(&self) -> &Table { @@ -239,14 +350,3 @@ async fn get_batch_stream( .map_err(to_datafusion_error); Ok(Box::pin(stream)) } - -fn get_column_names( - schema: ArrowSchemaRef, - projection: Option<&Vec>, -) -> Option> { - projection.map(|v| { - v.iter() - .map(|p| schema.field(*p).name().clone()) - .collect::>() - }) -} diff --git a/crates/datafusion/src/physical_plan/write.rs b/crates/datafusion/src/physical_plan/write.rs index 96c0ac0..86f364c 100644 --- a/crates/datafusion/src/physical_plan/write.rs +++ b/crates/datafusion/src/physical_plan/write.rs @@ -58,7 +58,7 @@ use crate::to_datafusion_error; /// The output of this execution plan is a record batch containing a single column with serialized /// data file information that can be used for committing the write operation to the table. #[derive(Debug)] -pub(crate) struct IcebergWriteExec { +pub struct IcebergWriteExec { table: Table, input: Arc, result_schema: ArrowSchemaRef, @@ -66,6 +66,12 @@ pub(crate) struct IcebergWriteExec { } impl IcebergWriteExec { + /// Writes the rows of `input` to new data files of `table`, without + /// committing them; see [`IcebergCommitExec`](super::IcebergCommitExec). + /// + /// `input`'s columns are matched to the table's by name. For a partitioned + /// table, `input` must also carry the partition values that + /// [`project_with_partition`](super::project_with_partition) adds. pub fn new(table: Table, input: Arc) -> Self { let plan_properties = Self::compute_properties(&input, Self::make_result_schema()); @@ -78,6 +84,11 @@ impl IcebergWriteExec { } } + /// The table this node writes to, as loaded when the node was planned. + pub fn table(&self) -> &Table { + &self.table + } + fn compute_properties( input: &Arc, schema: ArrowSchemaRef, diff --git a/crates/datafusion/src/table/metadata_table.rs b/crates/datafusion/src/table/metadata_table.rs index e9ef9fa..bcde88a 100644 --- a/crates/datafusion/src/table/metadata_table.rs +++ b/crates/datafusion/src/table/metadata_table.rs @@ -38,8 +38,29 @@ use crate::to_datafusion_error; /// managing access to a [`iceberg::inspect::MetadataTable`]. #[derive(Debug, Clone)] pub struct IcebergMetadataTableProvider { - pub(crate) table: Table, - pub(crate) r#type: MetadataTableType, + table: Table, + r#type: MetadataTableType, +} + +impl IcebergMetadataTableProvider { + /// Creates a provider for the `metadata_type` metadata table of an + /// already loaded table. + pub fn new(table: Table, metadata_type: MetadataTableType) -> Self { + Self { + table, + r#type: metadata_type, + } + } + + /// The table whose metadata this provider reads. + pub fn table(&self) -> &Table { + &self.table + } + + /// Which metadata table of [`Self::table`] this provider reads. + pub fn metadata_type(&self) -> &MetadataTableType { + &self.r#type + } } #[async_trait] diff --git a/crates/datafusion/src/table/mod.rs b/crates/datafusion/src/table/mod.rs index 69f4b3a..eaf6479 100644 --- a/crates/datafusion/src/table/mod.rs +++ b/crates/datafusion/src/table/mod.rs @@ -46,7 +46,7 @@ use iceberg::inspect::MetadataTableType; use iceberg::spec::TableProperties; use iceberg::table::Table; use iceberg::{Catalog, NamespaceIdent, TableIdent}; -use metadata_table::IcebergMetadataTableProvider; +pub use metadata_table::IcebergMetadataTableProvider; use crate::error::to_datafusion_error; use crate::physical_plan::commit::IcebergCommitExec; @@ -113,7 +113,17 @@ impl IcebergTableProvider { .load_table(&self.table_ident) .await .map_err(to_datafusion_error)?; - Ok(IcebergMetadataTableProvider { table, r#type }) + Ok(IcebergMetadataTableProvider::new(table, r#type)) + } + + /// The catalog this provider loads its table from and commits through. + pub fn catalog(&self) -> &Arc { + &self.catalog + } + + /// The identifier of the table this provider reads and writes. + pub fn table_ident(&self) -> &TableIdent { + &self.table_ident } } @@ -149,7 +159,7 @@ impl TableProvider for IcebergTableProvider { projection, filters, limit, - ))) + )?)) } fn supports_filters_pushdown( @@ -295,6 +305,17 @@ impl IcebergStaticTableProvider { schema, }) } + + /// The table as loaded when this provider was built. + pub fn table(&self) -> &Table { + &self.table + } + + /// The snapshot this provider reads, or `None` for the current snapshot of + /// [`Self::table`]. + pub fn snapshot_id(&self) -> Option { + self.snapshot_id + } } #[async_trait] @@ -322,7 +343,7 @@ impl TableProvider for IcebergStaticTableProvider { projection, filters, limit, - ))) + )?)) } fn supports_filters_pushdown( @@ -950,4 +971,25 @@ mod tests { "Limit should be None when not specified" ); } + + #[tokio::test] + async fn test_scan_rejects_out_of_range_projection() { + let table = get_test_table_from_metadata_file().await; + let provider = IcebergStaticTableProvider::try_new_from_table(table.clone()) + .await + .unwrap(); + let schema = provider.schema(); + let out_of_range = schema.fields().len(); + + let err = IcebergTableScan::new_with_predicate( + table, + None, + schema, + Some(&[0, out_of_range]), + None, + None, + ) + .unwrap_err(); + assert!(err.to_string().contains("project index"), "{err}"); + } } diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index dcea6a5..d4a912e 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -24,9 +24,18 @@ use std::vec; use datafusion::arrow::array::{Array, StringArray, UInt64Array}; use datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; +use datafusion::arrow::util::pretty::pretty_format_batches; +use datafusion::datasource::TableProvider; use datafusion::execution::context::SessionContext; use datafusion::parquet::arrow::PARQUET_FIELD_ID_META_KEY; -use datafusion_iceberg::IcebergCatalogProvider; +use datafusion::physical_plan::{ExecutionPlan, collect}; +use datafusion_iceberg::physical_plan::{ + IcebergCommitExec, IcebergMetadataScan, IcebergTableScan, IcebergWriteExec, +}; +use datafusion_iceberg::{ + IcebergCatalogProvider, IcebergMetadataTableProvider, IcebergStaticTableProvider, + IcebergTableProvider, +}; use expect_test::expect; use iceberg::io::LocalFsStorageFactory; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; @@ -977,3 +986,121 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { Ok(()) } + +/// Executes the single partition of `plan` and renders its rows as a table. +async fn run( + plan: &dyn ExecutionPlan, + ctx: &SessionContext, +) -> Result> { + let stream = plan.execute(0, ctx.task_ctx())?; + let batches = datafusion::physical_plan::common::collect(stream).await?; + Ok(pretty_format_batches(&batches)?.to_string()) +} + +/// Returns the first node of type `T` in `plan`, depth first. +fn find_node(plan: &Arc) -> Option<&T> { + plan.downcast_ref::() + .or_else(|| plan.children().into_iter().find_map(find_node::)) +} + +/// The plan nodes and providers can be named and inspected from outside this +/// crate, and rebuilt from their parts, as a codec that serializes them does. +#[tokio::test] +async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { + let iceberg_catalog = get_iceberg_catalog().await; + let namespace = NamespaceIdent::new("test_plan_nodes".to_string()); + set_test_namespace(&iceberg_catalog, &namespace).await?; + let creation = get_table_creation(temp_path(), "my_table", None)?; + iceberg_catalog.create_table(&namespace, creation).await?; + let ident = TableIdent::new(namespace.clone(), "my_table".to_string()); + let client: Arc = Arc::new(iceberg_catalog); + + let ctx = SessionContext::new(); + let catalog = IcebergCatalogProvider::try_new(client.clone()).await?; + ctx.register_catalog("catalog", Arc::new(catalog)); + let provider = ctx + .table_provider("catalog.test_plan_nodes.my_table") + .await?; + let provider = provider + .downcast_ref::() + .expect("a catalog-backed provider"); + assert_eq!(provider.table_ident(), &ident); + assert!(Arc::ptr_eq(provider.catalog(), &client)); + + // Write path: a commit above a write, both holding the table, and the + // commit going through the provider's catalog. + let insert = ctx + .sql("INSERT INTO catalog.test_plan_nodes.my_table VALUES (1, 'alan'), (2, 'turing')") + .await? + .create_physical_plan() + .await?; + let commit = insert + .downcast_ref::() + .expect("the insert plan is rooted at a commit"); + assert_eq!(commit.table().identifier(), &ident); + assert!(Arc::ptr_eq(commit.catalog(), &client)); + let write = find_node::(&insert).expect("a write below the commit"); + assert_eq!(write.table().identifier(), &ident); + collect(insert, ctx.task_ctx()).await?; + + // Read path: a scan pinned to a snapshot, rebuilt from its parts, + // returns the same rows. + let table = client.load_table(&ident).await?; + let snapshot_id = table.metadata().current_snapshot_id().unwrap(); + let pinned = + IcebergStaticTableProvider::try_new_from_table_snapshot(table, snapshot_id) + .await?; + assert_eq!(pinned.snapshot_id(), Some(snapshot_id)); + ctx.register_table("pinned", Arc::new(pinned.clone()))?; + let plan = ctx + .sql("SELECT foo2 FROM pinned WHERE foo1 = 1") + .await? + .create_physical_plan() + .await?; + let scan = find_node::(&plan).expect("a scan"); + assert!(scan.predicates().is_some(), "the filter is pushed down"); + let schema = pinned.schema(); + let projection = scan + .projection() + .map(|names| { + names + .iter() + .map(|name| schema.index_of(name)) + .collect::, _>>() + }) + .transpose()?; + let rebuilt = IcebergTableScan::new_with_predicate( + pinned.table().clone(), + scan.snapshot_id(), + schema, + projection.as_deref(), + scan.predicates().cloned(), + scan.limit(), + )?; + let expected = run(scan, &ctx).await?; + assert!( + expected.contains("alan") && !expected.contains("turing"), + "{expected}" + ); + assert_eq!(run(&rebuilt, &ctx).await?, expected); + + // Metadata tables: a scan rebuilt from a metadata scan's parts reads the + // same rows. + let plan = ctx + .sql("SELECT * FROM catalog.test_plan_nodes.\"my_table$snapshots\"") + .await? + .create_physical_plan() + .await?; + let metadata_scan = find_node::(&plan).expect("a metadata scan"); + let provider = metadata_scan.provider(); + assert_eq!(provider.table().identifier(), &ident); + let rebuilt = IcebergMetadataScan::new(IcebergMetadataTableProvider::new( + provider.table().clone(), + provider.metadata_type().clone(), + )); + let expected = run(metadata_scan, &ctx).await?; + assert!(expected.contains(&snapshot_id.to_string()), "{expected}"); + assert_eq!(run(&rebuilt, &ctx).await?, expected); + + Ok(()) +} From 27d6c7cca7f4c2c14125f4da02f9799d213ec791 Mon Sep 17 00:00:00 2001 From: Noah Date: Sun, 27 Sep 2026 13:06:18 -0400 Subject: [PATCH 2/5] feat: make IcebergTableProvider::try_new public A plan codec rebuilds a catalog-backed provider from its catalog and table identifier, which needs the constructor. Co-Authored-By: Claude Opus 5.5 --- crates/datafusion/src/table/mod.rs | 2 +- crates/datafusion/tests/integration_datafusion_test.rs | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/crates/datafusion/src/table/mod.rs b/crates/datafusion/src/table/mod.rs index eaf6479..190da1c 100644 --- a/crates/datafusion/src/table/mod.rs +++ b/crates/datafusion/src/table/mod.rs @@ -79,7 +79,7 @@ impl IcebergTableProvider { /// /// Loads the table once to get the initial schema, then stores the catalog /// reference for future metadata refreshes on each operation. - pub(crate) async fn try_new( + pub async fn try_new( catalog: Arc, namespace: NamespaceIdent, name: impl Into, diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index d4a912e..3160231 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -1026,6 +1026,14 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { .expect("a catalog-backed provider"); assert_eq!(provider.table_ident(), &ident); assert!(Arc::ptr_eq(provider.catalog(), &client)); + let rebuilt = IcebergTableProvider::try_new( + provider.catalog().clone(), + provider.table_ident().namespace().clone(), + provider.table_ident().name(), + ) + .await?; + assert_eq!(rebuilt.table_ident(), &ident); + assert_eq!(rebuilt.schema(), provider.schema()); // Write path: a commit above a write, both holding the table, and the // commit going through the provider's catalog. From 46409190f11c7d59fce944fcfc5bcb19073d0a18 Mon Sep 17 00:00:00 2001 From: Noah Date: Tue, 29 Sep 2026 07:25:03 -0400 Subject: [PATCH 3/5] fix: rebuild the commit and scan from what they expose A codec decoding these nodes in another process has only their accessors to work from, and two constructors asked for more than that. - IcebergCommitExec::new no longer takes a schema. It was only shown in the verbose display and could not be read back from the node; the display now shows the table's current schema on one line. - IcebergTableScan::new_with_predicate takes the output schema and the projected column names, as schema() and projection() return them, instead of the full table schema and indices, which the scan does not expose. It returns an error when the names are not the schema's fields in order, rather than building a scan whose batches do not match its schema. The integration test now executes a write and commit rebuilt from their accessors, and rebuilds scans from their accessors alone, with and without a projection. --- crates/datafusion/src/physical_plan/commit.rs | 72 ++++++++--------- crates/datafusion/src/physical_plan/scan.rs | 78 ++++++++++--------- crates/datafusion/src/table/mod.rs | 45 ++++++++++- .../tests/integration_datafusion_test.rs | 56 +++++++++---- 4 files changed, 155 insertions(+), 96 deletions(-) diff --git a/crates/datafusion/src/physical_plan/commit.rs b/crates/datafusion/src/physical_plan/commit.rs index 3981f57..8169570 100644 --- a/crates/datafusion/src/physical_plan/commit.rs +++ b/crates/datafusion/src/physical_plan/commit.rs @@ -53,7 +53,6 @@ pub struct IcebergCommitExec { table: Table, catalog: Arc, input: Arc, - schema: ArrowSchemaRef, count_schema: ArrowSchemaRef, plan_properties: Arc, } @@ -64,13 +63,11 @@ impl IcebergCommitExec { /// `input` must have a single partition, such as a /// [`CoalescePartitionsExec`](datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec) /// over an [`IcebergWriteExec`](super::IcebergWriteExec); executing the node - /// fails otherwise. `schema` is the table's Arrow schema, shown in the - /// verbose plan display. + /// fails otherwise. pub fn new( table: Table, catalog: Arc, input: Arc, - schema: ArrowSchemaRef, ) -> Self { let count_schema = Self::make_count_schema(); @@ -80,7 +77,6 @@ impl IcebergCommitExec { table, catalog, input, - schema, count_schema, plan_properties, } @@ -136,11 +132,22 @@ impl DisplayAs for IcebergCommitExec { write!(f, "IcebergCommitExec: table={}", self.table.identifier()) } DisplayFormatType::Verbose => { + // The fields on one line, as `Schema`'s own display spans + // several. + let fields = self + .table + .metadata() + .current_schema() + .as_struct() + .fields() + .iter() + .map(|field| field.to_string().trim_end().to_string()) + .collect::>() + .join(", "); write!( f, - "IcebergCommitExec: table={}, schema={:?}", - self.table.identifier(), - self.schema + "IcebergCommitExec: table={}, schema=[{fields}]", + self.table.identifier() ) } DisplayFormatType::TreeRender => { @@ -198,7 +205,6 @@ impl ExecutionPlan for IcebergCommitExec { self.table.clone(), self.catalog.clone(), children[0].clone(), - self.schema.clone(), ))) } @@ -519,19 +525,8 @@ mod tests { let input_exec = Arc::new(MockWriteExec::new(vec![data_file1_json, data_file2_json])); - // Create the IcebergCommitExec - let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( - DATA_FILES_COL_NAME, - DataType::Utf8, - false, - )])); - - let commit_exec = IcebergCommitExec::new( - table.clone(), - catalog.clone(), - input_exec, - arrow_schema, - ); + let commit_exec = + IcebergCommitExec::new(table.clone(), catalog.clone(), input_exec); // Verify Execution Plan schema matches the count schema assert_eq!(commit_exec.schema(), IcebergCommitExec::make_count_schema()); @@ -644,14 +639,7 @@ mod tests { } let input = UnionExec::try_new(partitions)?; assert_eq!(input.properties().partitioning.partition_count(), 2); - - let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( - DATA_FILES_COL_NAME, - DataType::Utf8, - false, - )])); - let commit_exec = - IcebergCommitExec::new(table.clone(), catalog.clone(), input, arrow_schema); + let commit_exec = IcebergCommitExec::new(table.clone(), catalog.clone(), input); let err = match commit_exec.execute(0, Arc::new(TaskContext::default())) { Ok(_) => panic!("a commit over two input partitions must not execute"), Err(err) => err.to_string(), @@ -706,16 +694,20 @@ mod tests { // Mock write plan produces no data files let input_exec = Arc::new(MockWriteExec::new(vec![])); - let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( - DATA_FILES_COL_NAME, - DataType::Utf8, - false, - )])); - let commit_exec = IcebergCommitExec::new( - table.clone(), - catalog.clone(), - input_exec, - arrow_schema, + let commit_exec = + IcebergCommitExec::new(table.clone(), catalog.clone(), input_exec); + + // The verbose display shows the table's schema on one line. + struct Verbose<'a>(&'a IcebergCommitExec); + impl fmt::Display for Verbose<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + self.0.fmt_as(DisplayFormatType::Verbose, f) + } + } + assert_eq!( + Verbose(&commit_exec).to_string(), + "IcebergCommitExec: table=test_empty_insert.empty_insert_table, \ + schema=[1: id: required int]" ); let task_ctx = Arc::new(TaskContext::default()); diff --git a/crates/datafusion/src/physical_plan/scan.rs b/crates/datafusion/src/physical_plan/scan.rs index e3c9ed7..3816918 100644 --- a/crates/datafusion/src/physical_plan/scan.rs +++ b/crates/datafusion/src/physical_plan/scan.rs @@ -21,6 +21,7 @@ use std::vec; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datafusion::common::plan_err; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::error::Result; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; @@ -65,11 +66,23 @@ impl IcebergTableScan { filters: &[Expr], limit: Option, ) -> Result { + let (output_schema, projection) = match projection { + None => (schema, None), + Some(projection) => { + let output_schema = Arc::new(schema.project(projection)?); + let names = output_schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(); + (output_schema, Some(names)) + } + }; Self::new_with_predicate( table, snapshot_id, - schema, - projection.map(Vec::as_slice), + output_schema, + projection, convert_filters_to_predicate(filters), limit, ) @@ -80,14 +93,17 @@ impl IcebergTableScan { /// its parts, such as after sending them to another process. A predicate /// cannot be converted back to the filters it came from. /// - /// The arguments mean what the matching accessors return: + /// Each argument takes what the matching accessor returns (`predicate` + /// what [`Self::predicates`] does), so a scan is rebuilt from + /// [`ExecutionPlan::schema`] and [`Self::projection`]: /// /// - `snapshot_id`: the snapshot to read, or `None` for the table's current /// snapshot. - /// - `schema`: the Arrow schema of the table the scan reads, as its - /// provider reports it. - /// - `projection`: indices into `schema` of the columns to read, or `None` - /// for all. The columns are read from the table by name. + /// - `schema`: the Arrow schema the scan outputs. With `projection` set, + /// its fields must be the projected columns, in order. Without one, it + /// must be the full schema of the table as read, which is not checked. + /// - `projection`: the names of the columns to read from the table, or + /// `None` for all of them. /// - `predicate`: pushed down to Iceberg to skip data files and rows. The /// table providers report their filters as /// [`Inexact`](datafusion::logical_expr::TableProviderFilterPushDown::Inexact), @@ -95,7 +111,8 @@ impl IcebergTableScan { /// /// # Errors /// - /// Returns an error if `projection` holds an index outside `schema`. + /// Returns an error if `projection` does not name the fields of `schema`, + /// in order. /// /// # Example /// @@ -137,26 +154,17 @@ impl IcebergTableScan { /// .await?; /// let scan = plan.downcast_ref::().unwrap(); /// - /// // Rebuild an equivalent scan from the original's parts. - /// let schema = provider.schema(); - /// let projection = scan - /// .projection() - /// .map(|names| { - /// names - /// .iter() - /// .map(|name| schema.index_of(name)) - /// .collect::, _>>() - /// }) - /// .transpose()?; + /// // Rebuild an equivalent scan from the original's accessors alone. /// let rebuilt = IcebergTableScan::new_with_predicate( /// scan.table().clone(), /// scan.snapshot_id(), - /// schema, - /// projection.as_deref(), + /// scan.schema(), + /// scan.projection().map(<[String]>::to_vec), /// scan.predicates().cloned(), /// scan.limit(), /// )?; /// assert_eq!(rebuilt.schema(), scan.schema()); + /// assert_eq!(rebuilt.projection(), scan.projection()); /// assert_eq!(rebuilt.predicates(), scan.predicates()); /// # Ok::<(), Box>(()) /// # })?; @@ -166,23 +174,23 @@ impl IcebergTableScan { table: Table, snapshot_id: Option, schema: ArrowSchemaRef, - projection: Option<&[usize]>, + projection: Option>, predicate: Option, limit: Option, ) -> Result { - let output_schema = match projection { - None => schema, - Some(projection) => Arc::new(schema.project(projection)?), - }; - // The columns to read, by name; `None` reads them all. - let projection = projection.map(|_| { - output_schema - .fields() - .iter() - .map(|field| field.name().clone()) - .collect() - }); - let plan_properties = Self::compute_properties(output_schema); + // The scan reads the named columns and reports `schema`, so the two + // must agree, or its batches would not match its schema. + if let Some(projection) = &projection { + let fields: Vec<&String> = + schema.fields().iter().map(|field| field.name()).collect(); + if !projection.iter().eq(fields.iter().copied()) { + return plan_err!( + "IcebergTableScan projection {projection:?} does not match the \ + fields of its schema {fields:?}" + ); + } + } + let plan_properties = Self::compute_properties(schema); Ok(Self { table, diff --git a/crates/datafusion/src/table/mod.rs b/crates/datafusion/src/table/mod.rs index 190da1c..bf78248 100644 --- a/crates/datafusion/src/table/mod.rs +++ b/crates/datafusion/src/table/mod.rs @@ -238,7 +238,6 @@ impl TableProvider for IcebergTableProvider { table, self.catalog.clone(), coalesce_partitions, - self.schema.clone(), ))) } } @@ -981,15 +980,53 @@ mod tests { let schema = provider.schema(); let out_of_range = schema.fields().len(); - let err = IcebergTableScan::new_with_predicate( + let err = IcebergTableScan::new( table, None, schema, - Some(&[0, out_of_range]), - None, + Some(&vec![0, out_of_range]), + &[], None, ) .unwrap_err(); assert!(err.to_string().contains("project index"), "{err}"); } + + /// A projection that does not name the schema's fields would make the scan + /// return batches that do not match its schema, so it is refused. + #[tokio::test] + async fn test_scan_rejects_projection_not_matching_schema() { + let table = get_test_table_from_metadata_file().await; + let provider = IcebergStaticTableProvider::try_new_from_table(table.clone()) + .await + .unwrap(); + let schema = provider.schema(); + assert!(schema.fields().len() >= 2); + let first = schema.field(0).name().clone(); + let second = schema.field(1).name().clone(); + let projected = Arc::new(schema.project(&[1]).unwrap()); + + let mismatched = [ + // Names in the wrong order for the schema. + (schema.clone(), vec![second.clone(), first.clone()]), + // A column the schema does not have. + (projected.clone(), vec![first.clone()]), + // More columns than the schema has. + (projected, vec![first.clone(), second]), + // Fewer columns than the schema has. + (schema.clone(), vec![first]), + ]; + for (schema, projection) in mismatched { + let err = IcebergTableScan::new_with_predicate( + table.clone(), + None, + schema, + Some(projection), + None, + None, + ) + .unwrap_err(); + assert!(err.to_string().contains("does not match"), "{err}"); + } + } } diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 3160231..44a52d1 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -28,7 +28,8 @@ use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::datasource::TableProvider; use datafusion::execution::context::SessionContext; use datafusion::parquet::arrow::PARQUET_FIELD_ID_META_KEY; -use datafusion::physical_plan::{ExecutionPlan, collect}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_iceberg::physical_plan::{ IcebergCommitExec, IcebergMetadataScan, IcebergTableScan, IcebergWriteExec, }; @@ -1036,7 +1037,10 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { assert_eq!(rebuilt.schema(), provider.schema()); // Write path: a commit above a write, both holding the table, and the - // commit going through the provider's catalog. + // commit going through the provider's catalog. The plan that runs is + // rebuilt from their accessors and children alone. The optimizer drops + // the coalesce above a single-partition write, so the rebuilt commit + // always gets one, as a codec would. let insert = ctx .sql("INSERT INTO catalog.test_plan_nodes.my_table VALUES (1, 'alan'), (2, 'turing')") .await? @@ -1049,9 +1053,19 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { assert!(Arc::ptr_eq(commit.catalog(), &client)); let write = find_node::(&insert).expect("a write below the commit"); assert_eq!(write.table().identifier(), &ident); - collect(insert, ctx.task_ctx()).await?; + let rebuilt_write: Arc = Arc::new(IcebergWriteExec::new( + write.table().clone(), + write.children()[0].clone(), + )); + let rebuilt_commit = IcebergCommitExec::new( + commit.table().clone(), + commit.catalog().clone(), + Arc::new(CoalescePartitionsExec::new(rebuilt_write)), + ); + let inserted = run(&rebuilt_commit, &ctx).await?; + assert!(inserted.contains("| 2 |"), "{inserted}"); - // Read path: a scan pinned to a snapshot, rebuilt from its parts, + // Read path: a scan pinned to a snapshot, rebuilt from its accessors, // returns the same rows. let table = client.load_table(&ident).await?; let snapshot_id = table.metadata().current_snapshot_id().unwrap(); @@ -1067,24 +1081,15 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { .await?; let scan = find_node::(&plan).expect("a scan"); assert!(scan.predicates().is_some(), "the filter is pushed down"); - let schema = pinned.schema(); - let projection = scan - .projection() - .map(|names| { - names - .iter() - .map(|name| schema.index_of(name)) - .collect::, _>>() - }) - .transpose()?; let rebuilt = IcebergTableScan::new_with_predicate( - pinned.table().clone(), + scan.table().clone(), scan.snapshot_id(), - schema, - projection.as_deref(), + scan.schema(), + scan.projection().map(<[String]>::to_vec), scan.predicates().cloned(), scan.limit(), )?; + assert_eq!(rebuilt.schema(), scan.schema()); let expected = run(scan, &ctx).await?; assert!( expected.contains("alan") && !expected.contains("turing"), @@ -1092,6 +1097,23 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { ); assert_eq!(run(&rebuilt, &ctx).await?, expected); + // Without a projection the scan reads every column, and its limit is kept. + let plan = pinned.scan(&ctx.state(), None, &[], Some(1)).await?; + let scan = plan.downcast_ref::().expect("a scan"); + assert_eq!(scan.projection(), None); + let rebuilt = IcebergTableScan::new_with_predicate( + scan.table().clone(), + scan.snapshot_id(), + scan.schema(), + scan.projection().map(<[String]>::to_vec), + scan.predicates().cloned(), + scan.limit(), + )?; + assert_eq!(rebuilt.limit(), Some(1)); + let expected = run(scan, &ctx).await?; + assert_eq!(expected.lines().count(), 5, "one row:\n{expected}"); + assert_eq!(run(&rebuilt, &ctx).await?, expected); + // Metadata tables: a scan rebuilt from a metadata scan's parts reads the // same rows. let plan = ctx From 14aaa2487b6d1688178bb496f728dcd567718511 Mon Sep 17 00:00:00 2001 From: Noah Date: Tue, 29 Sep 2026 13:36:24 -0400 Subject: [PATCH 4/5] fix: read the scan's columns by name, and assert full test outputs new_with_predicate no longer takes a projection. When given, it had to repeat the schema's field names; when not, the schema had to be the full table schema, which was not checked, so a one-column schema with no projection built a scan that reported one column and returned two. The scan now always selects the names of its schema's fields, which reads the same columns as select_all for a full schema, so its batches match its schema by construction and the constructor cannot fail. This also fixes the catalog-backed provider, which pairs the schema it cached when built with the table as loaded on each scan: after a column was added and written to, a scan with no projection returned that column too. projection() keeps its signature and returns the schema's field names. The tests now compare full outputs rather than parts of them, check the pinned snapshot is read rather than the current one, and check each $snapshots row against the table's metadata. --- crates/datafusion/src/physical_plan/commit.rs | 7 +- crates/datafusion/src/physical_plan/scan.rs | 88 ++--- crates/datafusion/src/table/mod.rs | 44 +-- .../tests/integration_datafusion_test.rs | 315 ++++++++++++++++-- 4 files changed, 333 insertions(+), 121 deletions(-) diff --git a/crates/datafusion/src/physical_plan/commit.rs b/crates/datafusion/src/physical_plan/commit.rs index 8169570..468bf5a 100644 --- a/crates/datafusion/src/physical_plan/commit.rs +++ b/crates/datafusion/src/physical_plan/commit.rs @@ -644,9 +644,10 @@ mod tests { Ok(_) => panic!("a commit over two input partitions must not execute"), Err(err) => err.to_string(), }; - assert!( - err.contains("requires an input with one partition"), - "{err}" + assert_eq!( + err, + "Execution error: IcebergCommitExec requires an input with one \ + partition, but it has 2; coalesce the input first" ); let table = catalog.load_table(table.identifier()).await?; diff --git a/crates/datafusion/src/physical_plan/scan.rs b/crates/datafusion/src/physical_plan/scan.rs index 3816918..70dfdcd 100644 --- a/crates/datafusion/src/physical_plan/scan.rs +++ b/crates/datafusion/src/physical_plan/scan.rs @@ -21,7 +21,6 @@ use std::vec; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; -use datafusion::common::plan_err; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::error::Result; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; @@ -48,8 +47,8 @@ pub struct IcebergTableScan { /// Stores certain, often expensive to compute, /// plan properties used in query optimization. plan_properties: Arc, - /// Projection column names, None means all columns - projection: Option>, + /// The columns to read, by name: the fields of the output schema + projection: Vec, /// Filters to apply to the table scan predicates: Option, /// Optional limit on the number of rows to return @@ -66,26 +65,17 @@ impl IcebergTableScan { filters: &[Expr], limit: Option, ) -> Result { - let (output_schema, projection) = match projection { - None => (schema, None), - Some(projection) => { - let output_schema = Arc::new(schema.project(projection)?); - let names = output_schema - .fields() - .iter() - .map(|field| field.name().clone()) - .collect(); - (output_schema, Some(names)) - } + let output_schema = match projection { + None => schema, + Some(projection) => Arc::new(schema.project(projection)?), }; - Self::new_with_predicate( + Ok(Self::new_with_predicate( table, snapshot_id, output_schema, - projection, convert_filters_to_predicate(filters), limit, - ) + )) } /// Creates a scan of `table` from an already-converted Iceberg @@ -93,26 +83,20 @@ impl IcebergTableScan { /// its parts, such as after sending them to another process. A predicate /// cannot be converted back to the filters it came from. /// - /// Each argument takes what the matching accessor returns (`predicate` - /// what [`Self::predicates`] does), so a scan is rebuilt from - /// [`ExecutionPlan::schema`] and [`Self::projection`]: + /// Each argument takes what the matching accessor returns (`schema` what + /// [`ExecutionPlan::schema`] does, and `predicate` what + /// [`Self::predicates`] does): /// /// - `snapshot_id`: the snapshot to read, or `None` for the table's current /// snapshot. - /// - `schema`: the Arrow schema the scan outputs. With `projection` set, - /// its fields must be the projected columns, in order. Without one, it - /// must be the full schema of the table as read, which is not checked. - /// - `projection`: the names of the columns to read from the table, or - /// `None` for all of them. + /// - `schema`: the Arrow schema the scan outputs. The scan reads the + /// columns of the same names from the snapshot it scans, and no others. + /// A name that snapshot's schema lacks fails the scan when it runs. /// - `predicate`: pushed down to Iceberg to skip data files and rows. The /// table providers report their filters as /// [`Inexact`](datafusion::logical_expr::TableProviderFilterPushDown::Inexact), /// so DataFusion still applies them above the scan. - /// - /// # Errors - /// - /// Returns an error if `projection` does not name the fields of `schema`, - /// in order. + /// - `limit`: the most rows the scan returns, or `None` for all of them. /// /// # Example /// @@ -159,10 +143,9 @@ impl IcebergTableScan { /// scan.table().clone(), /// scan.snapshot_id(), /// scan.schema(), - /// scan.projection().map(<[String]>::to_vec), /// scan.predicates().cloned(), /// scan.limit(), - /// )?; + /// ); /// assert_eq!(rebuilt.schema(), scan.schema()); /// assert_eq!(rebuilt.projection(), scan.projection()); /// assert_eq!(rebuilt.predicates(), scan.predicates()); @@ -174,32 +157,26 @@ impl IcebergTableScan { table: Table, snapshot_id: Option, schema: ArrowSchemaRef, - projection: Option>, predicate: Option, limit: Option, - ) -> Result { - // The scan reads the named columns and reports `schema`, so the two - // must agree, or its batches would not match its schema. - if let Some(projection) = &projection { - let fields: Vec<&String> = - schema.fields().iter().map(|field| field.name()).collect(); - if !projection.iter().eq(fields.iter().copied()) { - return plan_err!( - "IcebergTableScan projection {projection:?} does not match the \ - fields of its schema {fields:?}" - ); - } - } + ) -> Self { + // Reading the columns by name, rather than all of them, keeps the + // batches matching `schema` even when the table has columns it lacks. + let projection = schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(); let plan_properties = Self::compute_properties(schema); - Ok(Self { + Self { table, snapshot_id, plan_properties, projection, predicates: predicate, limit, - }) + } } pub fn table(&self) -> &Table { @@ -210,8 +187,10 @@ impl IcebergTableScan { self.snapshot_id } + /// The names of the columns the scan reads, which are the fields of its + /// schema. Always `Some`. pub fn projection(&self) -> Option<&[String]> { - self.projection.as_deref() + Some(&self.projection) } pub fn predicates(&self) -> Option<&Predicate> { @@ -312,9 +291,7 @@ impl DisplayAs for IcebergTableScan { write!( f, "IcebergTableScan projection:[{}] predicate:[{}]", - self.projection - .clone() - .map_or(String::new(), |v| v.join(",")), + self.projection.join(","), self.predicates .clone() .map_or(String::from(""), |p| format!("{p}")) @@ -334,7 +311,7 @@ impl DisplayAs for IcebergTableScan { async fn get_batch_stream( table: Table, snapshot_id: Option, - column_names: Option>, + column_names: Vec, predicates: Option, ) -> Result> + Send>>> { let scan_builder = match snapshot_id { @@ -342,10 +319,7 @@ async fn get_batch_stream( None => table.scan(), }; - let mut scan_builder = match column_names { - Some(column_names) => scan_builder.select(column_names), - None => scan_builder.select_all(), - }; + let mut scan_builder = scan_builder.select(column_names); if let Some(pred) = predicates { scan_builder = scan_builder.with_filter(pred); } diff --git a/crates/datafusion/src/table/mod.rs b/crates/datafusion/src/table/mod.rs index bf78248..dc20a4a 100644 --- a/crates/datafusion/src/table/mod.rs +++ b/crates/datafusion/src/table/mod.rs @@ -989,44 +989,12 @@ mod tests { None, ) .unwrap_err(); - assert!(err.to_string().contains("project index"), "{err}"); - } - - /// A projection that does not name the schema's fields would make the scan - /// return batches that do not match its schema, so it is refused. - #[tokio::test] - async fn test_scan_rejects_projection_not_matching_schema() { - let table = get_test_table_from_metadata_file().await; - let provider = IcebergStaticTableProvider::try_new_from_table(table.clone()) - .await - .unwrap(); - let schema = provider.schema(); - assert!(schema.fields().len() >= 2); - let first = schema.field(0).name().clone(); - let second = schema.field(1).name().clone(); - let projected = Arc::new(schema.project(&[1]).unwrap()); - - let mismatched = [ - // Names in the wrong order for the schema. - (schema.clone(), vec![second.clone(), first.clone()]), - // A column the schema does not have. - (projected.clone(), vec![first.clone()]), - // More columns than the schema has. - (projected, vec![first.clone(), second]), - // Fewer columns than the schema has. - (schema.clone(), vec![first]), - ]; - for (schema, projection) in mismatched { - let err = IcebergTableScan::new_with_predicate( - table.clone(), - None, - schema, - Some(projection), - None, - None, + assert_eq!( + err.to_string(), + format!( + "Arrow error: Schema error: project index {out_of_range} out of \ + bounds, max field {out_of_range}" ) - .unwrap_err(); - assert!(err.to_string().contains("does not match"), "{err}"); - } + ); } } diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 44a52d1..54c6fa8 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -17,12 +17,15 @@ //! Integration tests for Iceberg Datafusion with Hive Metastore. -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap}; use std::error::Error; use std::sync::Arc; use std::vec; -use datafusion::arrow::array::{Array, StringArray, UInt64Array}; +use datafusion::arrow::array::{Array, AsArray, RecordBatch, StringArray, UInt64Array}; +use datafusion::arrow::compute::cast; +use datafusion::arrow::compute::{concat_batches, sort_to_indices, take_record_batch}; +use datafusion::arrow::datatypes::Int64Type; use datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::datasource::TableProvider; @@ -44,6 +47,7 @@ use iceberg::spec::{ NestedField, PrimitiveType, Schema, StructType, Transform, Type, UnboundPartitionSpec, }; use iceberg::test_utils::check_record_batches; +use iceberg::transaction::{AddColumn, ApplyTransactionAction, Transaction}; use iceberg::{ Catalog, CatalogBuilder, MemoryCatalog, NamespaceIdent, Result as IcebergResult, TableCreation, TableIdent, @@ -988,14 +992,23 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { Ok(()) } -/// Executes the single partition of `plan` and renders its rows as a table. +/// Executes `plan`, which must have a single partition, and returns its rows. +async fn run_batches( + plan: &dyn ExecutionPlan, + ctx: &SessionContext, +) -> Result, Box> { + assert_eq!(plan.properties().partitioning.partition_count(), 1); + let stream = plan.execute(0, ctx.task_ctx())?; + Ok(datafusion::physical_plan::common::collect(stream).await?) +} + +/// Executes `plan`, which must have a single partition, and renders its rows +/// as a table. async fn run( plan: &dyn ExecutionPlan, ctx: &SessionContext, ) -> Result> { - let stream = plan.execute(0, ctx.task_ctx())?; - let batches = datafusion::physical_plan::common::collect(stream).await?; - Ok(pretty_format_batches(&batches)?.to_string()) + Ok(pretty_format_batches(&run_batches(plan, ctx).await?)?.to_string()) } /// Returns the first node of type `T` in `plan`, depth first. @@ -1062,8 +1075,13 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { commit.catalog().clone(), Arc::new(CoalescePartitionsExec::new(rebuilt_write)), ); - let inserted = run(&rebuilt_commit, &ctx).await?; - assert!(inserted.contains("| 2 |"), "{inserted}"); + expect![[r#" + +-------+ + | count | + +-------+ + | 2 | + +-------+"#]] + .assert_eq(&run(&rebuilt_commit, &ctx).await?); // Read path: a scan pinned to a snapshot, rebuilt from its accessors, // returns the same rows. @@ -1074,46 +1092,93 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { .await?; assert_eq!(pinned.snapshot_id(), Some(snapshot_id)); ctx.register_table("pinned", Arc::new(pinned.clone()))?; + // A later write, so a scan reading the current snapshot rather than the + // pinned one would return its row too. + ctx.sql("INSERT INTO catalog.test_plan_nodes.my_table VALUES (3, 'hopper')") + .await? + .collect() + .await?; + let latest_snapshot_id = client + .load_table(&ident) + .await? + .metadata() + .current_snapshot_id() + .unwrap(); + assert_ne!(latest_snapshot_id, snapshot_id); let plan = ctx .sql("SELECT foo2 FROM pinned WHERE foo1 = 1") .await? .create_physical_plan() .await?; let scan = find_node::(&plan).expect("a scan"); - assert!(scan.predicates().is_some(), "the filter is pushed down"); + assert_eq!( + scan.predicates().map(ToString::to_string).as_deref(), + Some("foo1 = 1") + ); let rebuilt = IcebergTableScan::new_with_predicate( scan.table().clone(), scan.snapshot_id(), scan.schema(), - scan.projection().map(<[String]>::to_vec), scan.predicates().cloned(), scan.limit(), - )?; + ); assert_eq!(rebuilt.schema(), scan.schema()); + assert_eq!(rebuilt.projection(), scan.projection()); let expected = run(scan, &ctx).await?; - assert!( - expected.contains("alan") && !expected.contains("turing"), - "{expected}" - ); + expect![[r#" + +------+------+ + | foo1 | foo2 | + +------+------+ + | 1 | alan | + +------+------+"#]] + .assert_eq(&expected); assert_eq!(run(&rebuilt, &ctx).await?, expected); - // Without a projection the scan reads every column, and its limit is kept. + // Without a projection the scan reads every column by name, and its limit + // is kept. let plan = pinned.scan(&ctx.state(), None, &[], Some(1)).await?; let scan = plan.downcast_ref::().expect("a scan"); - assert_eq!(scan.projection(), None); + assert_eq!( + scan.projection(), + Some(&["foo1".to_string(), "foo2".to_string()][..]) + ); let rebuilt = IcebergTableScan::new_with_predicate( scan.table().clone(), scan.snapshot_id(), scan.schema(), - scan.projection().map(<[String]>::to_vec), scan.predicates().cloned(), scan.limit(), - )?; + ); assert_eq!(rebuilt.limit(), Some(1)); let expected = run(scan, &ctx).await?; - assert_eq!(expected.lines().count(), 5, "one row:\n{expected}"); + expect![[r#" + +------+------+ + | foo1 | foo2 | + +------+------+ + | 1 | alan | + +------+------+"#]] + .assert_eq(&expected); assert_eq!(run(&rebuilt, &ctx).await?, expected); + // A scan reads the columns of its schema and no others, so one built over + // part of the table returns only those columns, from the pinned snapshot. + let foo2_only = Arc::new(pinned.schema().project(&[1])?); + let partial = IcebergTableScan::new_with_predicate( + pinned.table().clone(), + Some(snapshot_id), + foo2_only, + None, + None, + ); + expect![[r#" + +--------+ + | foo2 | + +--------+ + | alan | + | turing | + +--------+"#]] + .assert_eq(&run(&partial, &ctx).await?); + // Metadata tables: a scan rebuilt from a metadata scan's parts reads the // same rows. let plan = ctx @@ -1128,9 +1193,213 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { provider.table().clone(), provider.metadata_type().clone(), )); - let expected = run(metadata_scan, &ctx).await?; - assert!(expected.contains(&snapshot_id.to_string()), "{expected}"); - assert_eq!(run(&rebuilt, &ctx).await?, expected); + let batches = run_batches(metadata_scan, &ctx).await?; + // Snapshots come back in no set order, so put the first, which has no + // parent, first. Their ids, times and paths differ on every run, so each + // row is checked against its snapshot's metadata. + let batch = concat_batches(&batches[0].schema(), &batches)?; + let order = sort_to_indices(batch.column_by_name("parent_id").unwrap(), None, None)?; + let batch = take_record_batch(&batch, &order)?; + let column = |name: &str| batch.column_by_name(name).unwrap().clone(); + let longs = |name: &str| -> Result>, Box> { + Ok(cast(&column(name), &DataType::Int64)? + .as_primitive::() + .iter() + .collect()) + }; + let strings = |name: &str| -> Vec> { + column(name) + .as_string::() + .iter() + .map(|value| value.map(str::to_string)) + .collect() + }; + let metadata = provider.table().metadata(); + let snapshots = [snapshot_id, latest_snapshot_id].map(|id| { + metadata + .snapshot_by_id(id) + .expect("a snapshot of the table") + }); + assert_eq!( + batch + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + [ + "committed_at", + "snapshot_id", + "parent_id", + "operation", + "manifest_list", + "summary" + ] + ); + assert_eq!( + longs("committed_at")?, + snapshots.map(|snapshot| Some(snapshot.timestamp_ms() * 1000)) + ); + assert_eq!( + longs("snapshot_id")?, + [Some(snapshot_id), Some(latest_snapshot_id)] + ); + assert_eq!(longs("parent_id")?, [None, Some(snapshot_id)]); + assert_eq!( + strings("operation"), + [Some("append".to_string()), Some("append".to_string())] + ); + assert_eq!( + strings("manifest_list"), + snapshots.map(|snapshot| Some(snapshot.manifest_list().to_string())) + ); + // The summaries print their keys in no set order, so sort them. + let summaries = column("summary"); + let summaries = summaries.as_map(); + let summaries = (0..summaries.len()) + .map(|row| { + let entries = summaries.value(row); + let keys = entries.column(0).as_string::(); + let values = entries.column(1).as_string::(); + (0..entries.len()) + .map(|entry| format!("{}: {}", keys.value(entry), values.value(entry))) + .collect::>() + }) + .collect::>(); + expect![[r#" + [ + { + "added-data-files: 1", + "added-files-size: 913", + "added-records: 2", + "total-data-files: 1", + "total-delete-files: 0", + "total-equality-deletes: 0", + "total-files-size: 913", + "total-position-deletes: 0", + "total-records: 2", + }, + { + "added-data-files: 1", + "added-files-size: 901", + "added-records: 1", + "total-data-files: 2", + "total-delete-files: 0", + "total-equality-deletes: 0", + "total-files-size: 1814", + "total-position-deletes: 0", + "total-records: 3", + }, + ] + "#]] + .assert_debug_eq(&summaries); + // The rebuilt scan holds the same table, so lists its snapshots in the + // same order. + assert_eq!( + pretty_format_batches(&run_batches(&rebuilt, &ctx).await?)?.to_string(), + pretty_format_batches(&batches)?.to_string() + ); + + Ok(()) +} + +/// A catalog-backed provider keeps the schema it was built with, but scans the +/// table as it is now. After a column is added to the table and written to, a +/// scan with no projection still returns only the provider's columns, matching +/// its schema. +#[tokio::test] +async fn test_scan_after_schema_evolution_reads_provider_columns() +-> Result<(), Box> { + let iceberg_catalog = get_iceberg_catalog().await; + let namespace = NamespaceIdent::new("test_schema_evolution".to_string()); + set_test_namespace(&iceberg_catalog, &namespace).await?; + let creation = get_table_creation(temp_path(), "my_table", None)?; + iceberg_catalog.create_table(&namespace, creation).await?; + let ident = TableIdent::new(namespace.clone(), "my_table".to_string()); + let client: Arc = Arc::new(iceberg_catalog); + + let provider = Arc::new( + IcebergTableProvider::try_new(client.clone(), namespace, "my_table").await?, + ); + let ctx = SessionContext::new(); + ctx.register_table("t", provider.clone())?; + ctx.sql("INSERT INTO t VALUES (1, 'alan')") + .await? + .collect() + .await?; + + let table = client.load_table(&ident).await?; + let tx = Transaction::new(&table); + tx.update_schema() + .add_column(AddColumn::optional( + "foo3", + Type::Primitive(PrimitiveType::Int), + )) + .apply(tx)? + .commit(client.as_ref()) + .await?; + + // A scan reads the schema of the snapshot it reads, so write a snapshot + // with the new column, through a provider that sees it. + let evolved = IcebergTableProvider::try_new( + client.clone(), + ident.namespace().clone(), + ident.name(), + ) + .await?; + expect![[r#" + Schema { + fields: [ + Field { + name: "foo1", + data_type: Int32, + metadata: { + "PARQUET:field_id": "1", + }, + }, + Field { + name: "foo2", + data_type: Utf8, + metadata: { + "PARQUET:field_id": "2", + }, + }, + Field { + name: "foo3", + data_type: Int32, + nullable: true, + metadata: { + "PARQUET:field_id": "3", + }, + }, + ], + metadata: {}, + } + "#]] + .assert_debug_eq(&evolved.schema()); + ctx.register_table("evolved", Arc::new(evolved))?; + ctx.sql("INSERT INTO evolved VALUES (2, 'turing', 3)") + .await? + .collect() + .await?; + + let plan = provider.scan(&ctx.state(), None, &[], None).await?; + assert_eq!(plan.schema(), provider.schema()); + // The rows come from two data files, which may be read in either order. + // They are joined under their own schema, not the plan's, so that a column + // the plan does not report would show. + let batches = run_batches(plan.as_ref(), &ctx).await?; + let batch = concat_batches(&batches[0].schema(), &batches)?; + let order = sort_to_indices(batch.column_by_name("foo1").unwrap(), None, None)?; + let sorted = take_record_batch(&batch, &order)?; + expect![[r#" + +------+--------+ + | foo1 | foo2 | + +------+--------+ + | 1 | alan | + | 2 | turing | + +------+--------+"#]] + .assert_eq(&pretty_format_batches(&[sorted])?.to_string()); Ok(()) } From 342145911a849d4d62982b32b369d125506459e7 Mon Sep 17 00:00:00 2001 From: Noah Date: Tue, 29 Sep 2026 13:51:36 -0400 Subject: [PATCH 5/5] refactor: keep IcebergMetadataTableProvider's methods in one impl block The constructor and accessors went into a second inherent impl block; move them into the existing one, beside scan. --- crates/datafusion/src/table/metadata_table.rs | 40 +++++++++---------- 1 file changed, 19 insertions(+), 21 deletions(-) diff --git a/crates/datafusion/src/table/metadata_table.rs b/crates/datafusion/src/table/metadata_table.rs index bcde88a..b2f3d19 100644 --- a/crates/datafusion/src/table/metadata_table.rs +++ b/crates/datafusion/src/table/metadata_table.rs @@ -42,27 +42,6 @@ pub struct IcebergMetadataTableProvider { r#type: MetadataTableType, } -impl IcebergMetadataTableProvider { - /// Creates a provider for the `metadata_type` metadata table of an - /// already loaded table. - pub fn new(table: Table, metadata_type: MetadataTableType) -> Self { - Self { - table, - r#type: metadata_type, - } - } - - /// The table whose metadata this provider reads. - pub fn table(&self) -> &Table { - &self.table - } - - /// Which metadata table of [`Self::table`] this provider reads. - pub fn metadata_type(&self) -> &MetadataTableType { - &self.r#type - } -} - #[async_trait] impl TableProvider for IcebergMetadataTableProvider { fn schema(&self) -> ArrowSchemaRef { @@ -91,6 +70,25 @@ impl TableProvider for IcebergMetadataTableProvider { } impl IcebergMetadataTableProvider { + /// Creates a provider for the `metadata_type` metadata table of an + /// already loaded table. + pub fn new(table: Table, metadata_type: MetadataTableType) -> Self { + Self { + table, + r#type: metadata_type, + } + } + + /// The table whose metadata this provider reads. + pub fn table(&self) -> &Table { + &self.table + } + + /// Which metadata table of [`Self::table`] this provider reads. + pub fn metadata_type(&self) -> &MetadataTableType { + &self.r#type + } + pub async fn scan(self) -> Result>> { let metadata_table = self.table.inspect(); let stream = match self.r#type {