diff --git a/crates/datafusion/src/physical_plan/commit.rs b/crates/datafusion/src/physical_plan/commit.rs index dcbee3a..468bf5a 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,22 +44,30 @@ 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, - schema: ArrowSchemaRef, count_schema: ArrowSchemaRef, plan_properties: Arc, } 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. pub fn new( table: Table, catalog: Arc, input: Arc, - schema: ArrowSchemaRef, ) -> Self { let count_schema = Self::make_count_schema(); @@ -69,12 +77,21 @@ impl IcebergCommitExec { table, catalog, input, - schema, count_schema, plan_properties, } } + /// 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( @@ -115,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 => { @@ -177,7 +205,6 @@ impl ExecutionPlan for IcebergCommitExec { self.table.clone(), self.catalog.clone(), children[0].clone(), - self.schema.clone(), ))) } @@ -193,6 +220,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 +334,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, }; @@ -487,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()); @@ -559,6 +586,76 @@ 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 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(), + }; + 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?; + assert!(table.metadata().current_snapshot().is_none()); + + Ok(()) + } + #[tokio::test] async fn test_iceberg_commit_exec_empty_insert() -> Result<(), Box> { @@ -598,16 +695,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/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..70dfdcd 100644 --- a/crates/datafusion/src/physical_plan/scan.rs +++ b/crates/datafusion/src/physical_plan/scan.rs @@ -47,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 @@ -64,21 +64,117 @@ impl IcebergTableScan { projection: Option<&Vec>, filters: &[Expr], limit: Option, - ) -> Self { + ) -> 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); + Ok(Self::new_with_predicate( + table, + snapshot_id, + output_schema, + 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. + /// + /// 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. 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. + /// - `limit`: the most rows the scan returns, or `None` for all of them. + /// + /// # 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 accessors alone. + /// let rebuilt = IcebergTableScan::new_with_predicate( + /// scan.table().clone(), + /// scan.snapshot_id(), + /// scan.schema(), + /// 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>(()) + /// # })?; + /// # Ok::<(), Box>(()) + /// ``` + pub fn new_with_predicate( + table: Table, + snapshot_id: Option, + schema: ArrowSchemaRef, + predicate: Option, + limit: Option, + ) -> 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); Self { table, snapshot_id, plan_properties, projection, - predicates, + predicates: predicate, limit, } } @@ -91,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> { @@ -193,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}")) @@ -215,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 { @@ -223,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); } @@ -239,14 +332,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..b2f3d19 100644 --- a/crates/datafusion/src/table/metadata_table.rs +++ b/crates/datafusion/src/table/metadata_table.rs @@ -38,8 +38,8 @@ 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, } #[async_trait] @@ -70,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 { diff --git a/crates/datafusion/src/table/mod.rs b/crates/datafusion/src/table/mod.rs index 69f4b3a..dc20a4a 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; @@ -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, @@ -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( @@ -228,7 +238,6 @@ impl TableProvider for IcebergTableProvider { table, self.catalog.clone(), coalesce_partitions, - self.schema.clone(), ))) } } @@ -295,6 +304,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 +342,7 @@ impl TableProvider for IcebergStaticTableProvider { projection, filters, limit, - ))) + )?)) } fn supports_filters_pushdown( @@ -950,4 +970,31 @@ 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( + table, + None, + schema, + Some(&vec![0, out_of_range]), + &[], + None, + ) + .unwrap_err(); + assert_eq!( + err.to_string(), + format!( + "Arrow error: Schema error: project index {out_of_range} out of \ + bounds, max field {out_of_range}" + ) + ); + } } diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index dcea6a5..54c6fa8 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -17,16 +17,29 @@ //! 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; use datafusion::execution::context::SessionContext; use datafusion::parquet::arrow::PARQUET_FIELD_ID_META_KEY; -use datafusion_iceberg::IcebergCatalogProvider; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +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}; @@ -34,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, @@ -977,3 +991,415 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { Ok(()) } + +/// 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> { + Ok(pretty_format_batches(&run_batches(plan, ctx).await?)?.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)); + 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. 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? + .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); + 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)), + ); + 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. + 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()))?; + // 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_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.predicates().cloned(), + scan.limit(), + ); + assert_eq!(rebuilt.schema(), scan.schema()); + assert_eq!(rebuilt.projection(), scan.projection()); + let expected = run(scan, &ctx).await?; + expect![[r#" + +------+------+ + | foo1 | foo2 | + +------+------+ + | 1 | alan | + +------+------+"#]] + .assert_eq(&expected); + assert_eq!(run(&rebuilt, &ctx).await?, expected); + + // 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(), + Some(&["foo1".to_string(), "foo2".to_string()][..]) + ); + let rebuilt = IcebergTableScan::new_with_predicate( + scan.table().clone(), + scan.snapshot_id(), + scan.schema(), + scan.predicates().cloned(), + scan.limit(), + ); + assert_eq!(rebuilt.limit(), Some(1)); + let expected = run(scan, &ctx).await?; + 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 + .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 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(()) +}