From 5917fd3a298ea0571b5f27d4c5d7c3934c5402c6 Mon Sep 17 00:00:00 2001 From: Noah Date: Sat, 27 Jun 2026 21:43:55 -0400 Subject: [PATCH 01/11] init --- crates/integrations/datafusion/public-api.txt | 156 ++++++++++++ crates/integrations/datafusion/src/catalog.rs | 20 ++ .../datafusion/src/catalog_config.rs | 54 +++++ crates/integrations/datafusion/src/lib.rs | 3 + .../datafusion/src/physical_plan/commit.rs | 44 +++- .../src/physical_plan/metadata_scan.rs | 7 + .../datafusion/src/physical_plan/mod.rs | 5 +- .../datafusion/src/physical_plan/project.rs | 51 ++-- .../datafusion/src/physical_plan/scan.rs | 33 ++- .../datafusion/src/physical_plan/write.rs | 36 ++- crates/integrations/datafusion/src/schema.rs | 107 ++++++++- .../datafusion/src/table/metadata_table.rs | 41 ++++ .../integrations/datafusion/src/table/mod.rs | 222 ++++++++++++++---- 13 files changed, 695 insertions(+), 84 deletions(-) create mode 100644 crates/integrations/datafusion/src/catalog_config.rs diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index d24bd9fc9e..6aac904627 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -2,6 +2,12 @@ pub mod iceberg_datafusion pub mod iceberg_datafusion::metadata_table pub struct iceberg_datafusion::metadata_table::IcebergMetadataTableProvider impl iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::catalog_config(&self) -> core::option::Option<&iceberg_datafusion::IcebergCatalogConfig> +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::metadata_type(&self) -> &iceberg::inspect::metadata_table::MetadataTableType +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::new(table: iceberg::table::Table, type: iceberg::inspect::metadata_table::MetadataTableType) -> Self +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::table(&self) -> &iceberg::table::Table +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::with_catalog_config(self, catalog_config: core::option::Option) -> Self +impl iceberg_datafusion::metadata_table::IcebergMetadataTableProvider pub async fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan(self) -> datafusion_common::error::Result>> impl core::clone::Clone for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::clone(&self) -> iceberg_datafusion::metadata_table::IcebergMetadataTableProvider @@ -13,13 +19,51 @@ pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan<'l pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType pub mod iceberg_datafusion::physical_plan +pub struct iceberg_datafusion::physical_plan::IcebergCommitExec +impl iceberg_datafusion::physical_plan::IcebergCommitExec +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::catalog_config(&self) -> core::option::Option<&iceberg_datafusion::IcebergCatalogConfig> +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::new(table: iceberg::table::Table, catalog: alloc::sync::Arc, input: alloc::sync::Arc, schema: arrow_schema::schema::SchemaRef) -> Self +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::table(&self) -> &iceberg::table::Table +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::with_catalog_config(self, catalog_config: core::option::Option) -> Self +impl core::fmt::Debug for iceberg_datafusion::physical_plan::IcebergCommitExec +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl datafusion_physical_plan::display::DisplayAs for iceberg_datafusion::physical_plan::IcebergCommitExec +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::fmt_as(&self, t: datafusion_physical_plan::display::DisplayFormatType, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl datafusion_physical_plan::execution_plan::ExecutionPlan for iceberg_datafusion::physical_plan::IcebergCommitExec +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::as_any(&self) -> &dyn core::any::Any +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::benefits_from_input_partitioning(&self) -> alloc::vec::Vec +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::children(&self) -> alloc::vec::Vec<&alloc::sync::Arc> +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::execute(&self, partition: usize, context: alloc::sync::Arc) -> datafusion_common::error::Result +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::name(&self) -> &str +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::properties(&self) -> &alloc::sync::Arc +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::required_input_distribution(&self) -> alloc::vec::Vec +pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::with_new_children(self: alloc::sync::Arc, children: alloc::vec::Vec>) -> datafusion_common::error::Result> +pub struct iceberg_datafusion::physical_plan::IcebergMetadataScan +impl iceberg_datafusion::physical_plan::IcebergMetadataScan +pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::new(provider: iceberg_datafusion::metadata_table::IcebergMetadataTableProvider) -> Self +pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::provider(&self) -> &iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +impl core::fmt::Debug for iceberg_datafusion::physical_plan::IcebergMetadataScan +pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl datafusion_physical_plan::display::DisplayAs for iceberg_datafusion::physical_plan::IcebergMetadataScan +pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::fmt_as(&self, _t: datafusion_physical_plan::display::DisplayFormatType, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl datafusion_physical_plan::execution_plan::ExecutionPlan for iceberg_datafusion::physical_plan::IcebergMetadataScan +pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::as_any(&self) -> &dyn core::any::Any +pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::children(&self) -> alloc::vec::Vec<&alloc::sync::Arc> +pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::execute(&self, _partition: usize, _context: alloc::sync::Arc) -> datafusion_common::error::Result +pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::name(&self) -> &str +pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::properties(&self) -> &alloc::sync::Arc +pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::with_new_children(self: alloc::sync::Arc, _children: alloc::vec::Vec>) -> datafusion_common::error::Result> pub struct iceberg_datafusion::physical_plan::IcebergTableScan impl iceberg_datafusion::physical_plan::IcebergTableScan +pub fn iceberg_datafusion::physical_plan::IcebergTableScan::catalog_config(&self) -> core::option::Option<&iceberg_datafusion::IcebergCatalogConfig> pub fn iceberg_datafusion::physical_plan::IcebergTableScan::limit(&self) -> core::option::Option +pub fn iceberg_datafusion::physical_plan::IcebergTableScan::new(table: iceberg::table::Table, snapshot_id: core::option::Option, schema: arrow_schema::schema::SchemaRef, projection: core::option::Option<&alloc::vec::Vec>, filters: &[datafusion_expr::expr::Expr], limit: core::option::Option) -> Self pub fn iceberg_datafusion::physical_plan::IcebergTableScan::predicates(&self) -> core::option::Option<&iceberg::expr::predicate::Predicate> pub fn iceberg_datafusion::physical_plan::IcebergTableScan::projection(&self) -> core::option::Option<&[alloc::string::String]> pub fn iceberg_datafusion::physical_plan::IcebergTableScan::snapshot_id(&self) -> core::option::Option pub fn iceberg_datafusion::physical_plan::IcebergTableScan::table(&self) -> &iceberg::table::Table +pub fn iceberg_datafusion::physical_plan::IcebergTableScan::with_catalog_config(self, catalog_config: core::option::Option) -> Self +pub fn iceberg_datafusion::physical_plan::IcebergTableScan::with_predicates(self, predicates: core::option::Option) -> Self impl core::fmt::Debug for iceberg_datafusion::physical_plan::IcebergTableScan pub fn iceberg_datafusion::physical_plan::IcebergTableScan::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_physical_plan::display::DisplayAs for iceberg_datafusion::physical_plan::IcebergTableScan @@ -31,12 +75,61 @@ pub fn iceberg_datafusion::physical_plan::IcebergTableScan::execute(&self, _part pub fn iceberg_datafusion::physical_plan::IcebergTableScan::name(&self) -> &str pub fn iceberg_datafusion::physical_plan::IcebergTableScan::properties(&self) -> &alloc::sync::Arc pub fn iceberg_datafusion::physical_plan::IcebergTableScan::with_new_children(self: alloc::sync::Arc, _children: alloc::vec::Vec>) -> datafusion_common::error::Result> +pub struct iceberg_datafusion::physical_plan::IcebergWriteExec +impl iceberg_datafusion::physical_plan::IcebergWriteExec +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::catalog_config(&self) -> core::option::Option<&iceberg_datafusion::IcebergCatalogConfig> +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::new(table: iceberg::table::Table, input: alloc::sync::Arc, schema: arrow_schema::schema::SchemaRef) -> Self +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::table(&self) -> &iceberg::table::Table +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::with_catalog_config(self, catalog_config: core::option::Option) -> Self +impl core::fmt::Debug for iceberg_datafusion::physical_plan::IcebergWriteExec +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl datafusion_physical_plan::display::DisplayAs for iceberg_datafusion::physical_plan::IcebergWriteExec +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::fmt_as(&self, t: datafusion_physical_plan::display::DisplayFormatType, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl datafusion_physical_plan::execution_plan::ExecutionPlan for iceberg_datafusion::physical_plan::IcebergWriteExec +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::as_any(&self) -> &dyn core::any::Any +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::benefits_from_input_partitioning(&self) -> alloc::vec::Vec +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::children(&self) -> alloc::vec::Vec<&alloc::sync::Arc> +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::execute(&self, partition: usize, context: alloc::sync::Arc) -> datafusion_common::error::Result +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::maintains_input_order(&self) -> alloc::vec::Vec +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::name(&self) -> &str +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::properties(&self) -> &alloc::sync::Arc +pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::with_new_children(self: alloc::sync::Arc, children: alloc::vec::Vec>) -> datafusion_common::error::Result> +pub struct iceberg_datafusion::physical_plan::PartitionExpr +impl iceberg_datafusion::physical_plan::PartitionExpr +pub fn iceberg_datafusion::physical_plan::PartitionExpr::partition_spec(&self) -> &alloc::sync::Arc +pub fn iceberg_datafusion::physical_plan::PartitionExpr::table_schema(&self) -> &iceberg::spec::schema::SchemaRef +pub fn iceberg_datafusion::physical_plan::PartitionExpr::try_new(partition_spec: alloc::sync::Arc, table_schema: iceberg::spec::schema::SchemaRef) -> datafusion_common::error::Result +impl core::clone::Clone for iceberg_datafusion::physical_plan::PartitionExpr +pub fn iceberg_datafusion::physical_plan::PartitionExpr::clone(&self) -> iceberg_datafusion::physical_plan::PartitionExpr +impl core::cmp::Eq for iceberg_datafusion::physical_plan::PartitionExpr +impl core::cmp::PartialEq for iceberg_datafusion::physical_plan::PartitionExpr +pub fn iceberg_datafusion::physical_plan::PartitionExpr::eq(&self, other: &Self) -> bool +impl core::fmt::Debug for iceberg_datafusion::physical_plan::PartitionExpr +pub fn iceberg_datafusion::physical_plan::PartitionExpr::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::fmt::Display for iceberg_datafusion::physical_plan::PartitionExpr +pub fn iceberg_datafusion::physical_plan::PartitionExpr::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::hash::Hash for iceberg_datafusion::physical_plan::PartitionExpr +pub fn iceberg_datafusion::physical_plan::PartitionExpr::hash(&self, state: &mut H) +impl datafusion_physical_expr_common::physical_expr::PhysicalExpr for iceberg_datafusion::physical_plan::PartitionExpr +pub fn iceberg_datafusion::physical_plan::PartitionExpr::as_any(&self) -> &dyn core::any::Any +pub fn iceberg_datafusion::physical_plan::PartitionExpr::children(&self) -> alloc::vec::Vec<&alloc::sync::Arc> +pub fn iceberg_datafusion::physical_plan::PartitionExpr::data_type(&self, _input_schema: &arrow_schema::schema::Schema) -> datafusion_common::error::Result +pub fn iceberg_datafusion::physical_plan::PartitionExpr::evaluate(&self, batch: &arrow_array::record_batch::RecordBatch) -> datafusion_common::error::Result +pub fn iceberg_datafusion::physical_plan::PartitionExpr::fmt_sql(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +pub fn iceberg_datafusion::physical_plan::PartitionExpr::nullable(&self, _input_schema: &arrow_schema::schema::Schema) -> datafusion_common::error::Result +pub fn iceberg_datafusion::physical_plan::PartitionExpr::with_new_children(self: alloc::sync::Arc, _children: alloc::vec::Vec>) -> datafusion_common::error::Result> pub fn iceberg_datafusion::physical_plan::convert_filters_to_predicate(filters: &[datafusion_expr::expr::Expr]) -> core::option::Option pub fn iceberg_datafusion::physical_plan::project_with_partition(input: alloc::sync::Arc, table: &iceberg::table::Table) -> datafusion_common::error::Result> pub mod iceberg_datafusion::table pub mod iceberg_datafusion::table::metadata_table pub struct iceberg_datafusion::table::metadata_table::IcebergMetadataTableProvider impl iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::catalog_config(&self) -> core::option::Option<&iceberg_datafusion::IcebergCatalogConfig> +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::metadata_type(&self) -> &iceberg::inspect::metadata_table::MetadataTableType +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::new(table: iceberg::table::Table, type: iceberg::inspect::metadata_table::MetadataTableType) -> Self +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::table(&self) -> &iceberg::table::Table +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::with_catalog_config(self, catalog_config: core::option::Option) -> Self +impl iceberg_datafusion::metadata_table::IcebergMetadataTableProvider pub async fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan(self) -> datafusion_common::error::Result>> impl core::clone::Clone for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::clone(&self) -> iceberg_datafusion::metadata_table::IcebergMetadataTableProvider @@ -58,6 +151,24 @@ impl core::fmt::Debug for iceberg_datafusion::table_provider_factory::IcebergTab pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProviderFactory for iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory::create<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, cmd: &'life2 datafusion_expr::logical_plan::ddl::CreateExternalTable) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub struct iceberg_datafusion::table::IcebergMetadataTableProvider +impl iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::catalog_config(&self) -> core::option::Option<&iceberg_datafusion::IcebergCatalogConfig> +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::metadata_type(&self) -> &iceberg::inspect::metadata_table::MetadataTableType +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::new(table: iceberg::table::Table, type: iceberg::inspect::metadata_table::MetadataTableType) -> Self +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::table(&self) -> &iceberg::table::Table +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::with_catalog_config(self, catalog_config: core::option::Option) -> Self +impl iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub async fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan(self) -> datafusion_common::error::Result>> +impl core::clone::Clone for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::clone(&self) -> iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +impl core::fmt::Debug for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl datafusion_catalog::table::TableProvider for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::as_any(&self) -> &dyn core::any::Any +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, _projection: core::option::Option<&'life2 alloc::vec::Vec>, _filters: &'life3 [datafusion_expr::expr::Expr], _limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType pub struct iceberg_datafusion::table::IcebergStaticTableProvider impl iceberg_datafusion::IcebergStaticTableProvider pub async fn iceberg_datafusion::IcebergStaticTableProvider::try_new_from_table(table: iceberg::table::Table) -> iceberg::error::Result @@ -74,6 +185,12 @@ pub fn iceberg_datafusion::IcebergStaticTableProvider::schema(&self) -> arrow_sc pub fn iceberg_datafusion::IcebergStaticTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergStaticTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType pub struct iceberg_datafusion::table::IcebergTableProvider +impl iceberg_datafusion::IcebergTableProvider +pub fn iceberg_datafusion::IcebergTableProvider::config(&self) -> core::option::Option<&iceberg_datafusion::IcebergCatalogConfig> +pub fn iceberg_datafusion::IcebergTableProvider::snapshot_id(&self) -> core::option::Option +pub fn iceberg_datafusion::IcebergTableProvider::table_ident(&self) -> &iceberg::catalog::TableIdent +pub async fn iceberg_datafusion::IcebergTableProvider::try_new_with_config(catalog: alloc::sync::Arc, config: iceberg_datafusion::IcebergCatalogConfig, namespace: iceberg::catalog::NamespaceIdent, name: impl core::convert::Into) -> iceberg::error::Result +pub fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> Self impl core::clone::Clone for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafusion::IcebergTableProvider impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider @@ -96,15 +213,48 @@ impl core::fmt::Debug for iceberg_datafusion::table_provider_factory::IcebergTab pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProviderFactory for iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory::create<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, cmd: &'life2 datafusion_expr::logical_plan::ddl::CreateExternalTable) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub struct iceberg_datafusion::IcebergCatalogConfig +pub iceberg_datafusion::IcebergCatalogConfig::name: alloc::string::String +pub iceberg_datafusion::IcebergCatalogConfig::props: std::collections::hash::map::HashMap +pub iceberg_datafusion::IcebergCatalogConfig::type: alloc::string::String +impl iceberg_datafusion::IcebergCatalogConfig +pub fn iceberg_datafusion::IcebergCatalogConfig::new(type: impl core::convert::Into, name: impl core::convert::Into, props: std::collections::hash::map::HashMap) -> Self +impl core::clone::Clone for iceberg_datafusion::IcebergCatalogConfig +pub fn iceberg_datafusion::IcebergCatalogConfig::clone(&self) -> iceberg_datafusion::IcebergCatalogConfig +impl core::cmp::Eq for iceberg_datafusion::IcebergCatalogConfig +impl core::cmp::PartialEq for iceberg_datafusion::IcebergCatalogConfig +pub fn iceberg_datafusion::IcebergCatalogConfig::eq(&self, other: &iceberg_datafusion::IcebergCatalogConfig) -> bool +impl core::fmt::Debug for iceberg_datafusion::IcebergCatalogConfig +pub fn iceberg_datafusion::IcebergCatalogConfig::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::StructuralPartialEq for iceberg_datafusion::IcebergCatalogConfig pub struct iceberg_datafusion::IcebergCatalogProvider impl iceberg_datafusion::IcebergCatalogProvider pub async fn iceberg_datafusion::IcebergCatalogProvider::try_new(client: alloc::sync::Arc) -> iceberg::error::Result +pub async fn iceberg_datafusion::IcebergCatalogProvider::try_new_with_config(client: alloc::sync::Arc, config: iceberg_datafusion::IcebergCatalogConfig) -> iceberg::error::Result impl core::fmt::Debug for iceberg_datafusion::IcebergCatalogProvider pub fn iceberg_datafusion::IcebergCatalogProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::catalog::CatalogProvider for iceberg_datafusion::IcebergCatalogProvider pub fn iceberg_datafusion::IcebergCatalogProvider::as_any(&self) -> &dyn core::any::Any pub fn iceberg_datafusion::IcebergCatalogProvider::schema(&self, name: &str) -> core::option::Option> pub fn iceberg_datafusion::IcebergCatalogProvider::schema_names(&self) -> alloc::vec::Vec +pub struct iceberg_datafusion::IcebergMetadataTableProvider +impl iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::catalog_config(&self) -> core::option::Option<&iceberg_datafusion::IcebergCatalogConfig> +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::metadata_type(&self) -> &iceberg::inspect::metadata_table::MetadataTableType +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::new(table: iceberg::table::Table, type: iceberg::inspect::metadata_table::MetadataTableType) -> Self +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::table(&self) -> &iceberg::table::Table +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::with_catalog_config(self, catalog_config: core::option::Option) -> Self +impl iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub async fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan(self) -> datafusion_common::error::Result>> +impl core::clone::Clone for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::clone(&self) -> iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +impl core::fmt::Debug for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl datafusion_catalog::table::TableProvider for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::as_any(&self) -> &dyn core::any::Any +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, _projection: core::option::Option<&'life2 alloc::vec::Vec>, _filters: &'life3 [datafusion_expr::expr::Expr], _limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType pub struct iceberg_datafusion::IcebergStaticTableProvider impl iceberg_datafusion::IcebergStaticTableProvider pub async fn iceberg_datafusion::IcebergStaticTableProvider::try_new_from_table(table: iceberg::table::Table) -> iceberg::error::Result @@ -121,6 +271,12 @@ pub fn iceberg_datafusion::IcebergStaticTableProvider::schema(&self) -> arrow_sc pub fn iceberg_datafusion::IcebergStaticTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergStaticTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType pub struct iceberg_datafusion::IcebergTableProvider +impl iceberg_datafusion::IcebergTableProvider +pub fn iceberg_datafusion::IcebergTableProvider::config(&self) -> core::option::Option<&iceberg_datafusion::IcebergCatalogConfig> +pub fn iceberg_datafusion::IcebergTableProvider::snapshot_id(&self) -> core::option::Option +pub fn iceberg_datafusion::IcebergTableProvider::table_ident(&self) -> &iceberg::catalog::TableIdent +pub async fn iceberg_datafusion::IcebergTableProvider::try_new_with_config(catalog: alloc::sync::Arc, config: iceberg_datafusion::IcebergCatalogConfig, namespace: iceberg::catalog::NamespaceIdent, name: impl core::convert::Into) -> iceberg::error::Result +pub fn iceberg_datafusion::IcebergTableProvider::with_snapshot_id(self, snapshot_id: core::option::Option) -> Self impl core::clone::Clone for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafusion::IcebergTableProvider impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs index 69ab68f93a..21ea5afbf7 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog.rs @@ -23,6 +23,7 @@ use datafusion::catalog::{CatalogProvider, SchemaProvider}; use futures::future::try_join_all; use iceberg::{Catalog, NamespaceIdent, Result}; +use crate::IcebergCatalogConfig; use crate::schema::IcebergSchemaProvider; /// Provides an interface to manage and access multiple schemas @@ -47,6 +48,24 @@ impl IcebergCatalogProvider { /// attempts to create a schema provider for each namespace, and /// collects these providers into a `HashMap`. pub async fn try_new(client: Arc) -> Result { + Self::try_new_impl(client, None).await + } + + /// Like [`try_new`](Self::try_new), but threads a serializable + /// [`IcebergCatalogConfig`] into every schema and table provider it creates, + /// so the catalog's tables can be queried by a distributed engine such as + /// Ballista. The `client` must already be built from the same `config`. + pub async fn try_new_with_config( + client: Arc, + config: IcebergCatalogConfig, + ) -> Result { + Self::try_new_impl(client, Some(config)).await + } + + async fn try_new_impl( + client: Arc, + config: Option, + ) -> Result { // TODO: // Schemas and providers should be cached and evicted based on time // As of right now; schemas might become stale. @@ -63,6 +82,7 @@ impl IcebergCatalogProvider { .map(|name| { IcebergSchemaProvider::try_new( client.clone(), + config.clone(), NamespaceIdent::new(name.clone()), ) }) diff --git a/crates/integrations/datafusion/src/catalog_config.rs b/crates/integrations/datafusion/src/catalog_config.rs new file mode 100644 index 0000000000..b62e500ffe --- /dev/null +++ b/crates/integrations/datafusion/src/catalog_config.rs @@ -0,0 +1,54 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::HashMap; + +/// A serializable description of the catalog (and storage) that backs an +/// [`IcebergTableProvider`](crate::table::IcebergTableProvider). +/// +/// This is the minimal, self-contained handle needed to *reconstruct* a catalog +/// and its associated `FileIO` on a remote node. It deliberately holds only +/// plain data (no live connections) so that distributed query engines such as +/// Ballista can serialize it, ship it to executors, and rebuild the catalog +/// there via a catalog loader (e.g. `iceberg-catalog-loader`) and the storage +/// via `FileIOBuilder::with_props`. +/// +/// The `props` map carries both the catalog connection properties (e.g. the +/// REST catalog URI) and the storage/`FileIO` properties (e.g. S3 endpoint and +/// credentials); in practice these live together in a single map. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct IcebergCatalogConfig { + /// The catalog type, e.g. `"rest"`, `"sql"`, `"glue"`. + pub r#type: String, + pub name: String, + /// Catalog connection and storage properties. + pub props: HashMap, +} + +impl IcebergCatalogConfig { + pub fn new( + r#type: impl Into, + name: impl Into, + props: HashMap, + ) -> Self { + Self { + r#type: r#type.into(), + name: name.into(), + props, + } + } +} diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index 4b0ea8606d..19475aeda2 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -18,6 +18,9 @@ mod catalog; pub use catalog::*; +mod catalog_config; +pub use catalog_config::*; + mod error; pub use error::*; diff --git a/crates/integrations/datafusion/src/physical_plan/commit.rs b/crates/integrations/datafusion/src/physical_plan/commit.rs index 835c804908..d3a43c4006 100644 --- a/crates/integrations/datafusion/src/physical_plan/commit.rs +++ b/crates/integrations/datafusion/src/physical_plan/commit.rs @@ -41,13 +41,17 @@ use crate::to_datafusion_error; /// IcebergCommitExec is responsible for collecting the files written and use /// [`Transaction::fast_append`] to commit the data files written. #[derive(Debug)] -pub(crate) struct IcebergCommitExec { +pub struct IcebergCommitExec { table: Table, catalog: Arc, input: Arc, schema: ArrowSchemaRef, count_schema: ArrowSchemaRef, plan_properties: Arc, + /// Optional serializable catalog/storage config, populated when this node is + /// built through a config-backed provider so it can be reconstructed on a + /// remote node by a distributed engine. + catalog_config: Option, } impl IcebergCommitExec { @@ -68,9 +72,31 @@ impl IcebergCommitExec { schema, count_schema, plan_properties, + catalog_config: None, } } + /// Attaches a serializable catalog/storage config to this node so that a + /// distributed engine can reconstruct it (including the catalog) on a remote + /// node. + pub fn with_catalog_config( + mut self, + catalog_config: Option, + ) -> Self { + self.catalog_config = catalog_config; + self + } + + /// Returns the serializable catalog/storage config, if any. + pub fn catalog_config(&self) -> Option<&crate::IcebergCatalogConfig> { + self.catalog_config.as_ref() + } + + /// Returns the table this node commits to. + 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( @@ -160,12 +186,15 @@ impl ExecutionPlan for IcebergCommitExec { ))); } - Ok(Arc::new(IcebergCommitExec::new( - self.table.clone(), - self.catalog.clone(), - children[0].clone(), - self.schema.clone(), - ))) + Ok(Arc::new( + IcebergCommitExec::new( + self.table.clone(), + self.catalog.clone(), + children[0].clone(), + self.schema.clone(), + ) + .with_catalog_config(self.catalog_config.clone()), + )) } fn execute( @@ -592,6 +621,7 @@ mod tests { let iceberg_table_provider = IcebergTableProvider::try_new( catalog.clone(), + None, namespace.clone(), table_name.to_string(), ) diff --git a/crates/integrations/datafusion/src/physical_plan/metadata_scan.rs b/crates/integrations/datafusion/src/physical_plan/metadata_scan.rs index a1a65dec1f..d25a897cc6 100644 --- a/crates/integrations/datafusion/src/physical_plan/metadata_scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/metadata_scan.rs @@ -45,6 +45,13 @@ impl IcebergMetadataScan { properties, } } + + /// Returns the metadata-table provider this node scans, so a distributed + /// engine can serialize the catalog config + table identifier + metadata type + /// it carries and rebuild it on a remote node. + pub fn provider(&self) -> &IcebergMetadataTableProvider { + &self.provider + } } impl DisplayAs for IcebergMetadataScan { diff --git a/crates/integrations/datafusion/src/physical_plan/mod.rs b/crates/integrations/datafusion/src/physical_plan/mod.rs index aeac30de32..024266a4e9 100644 --- a/crates/integrations/datafusion/src/physical_plan/mod.rs +++ b/crates/integrations/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 project::project_with_partition; +pub use metadata_scan::IcebergMetadataScan; +pub use project::{PartitionExpr, project_with_partition}; pub use scan::IcebergTableScan; +pub use write::IcebergWriteExec; diff --git a/crates/integrations/datafusion/src/physical_plan/project.rs b/crates/integrations/datafusion/src/physical_plan/project.rs index 670d961f91..00c3c848d6 100644 --- a/crates/integrations/datafusion/src/physical_plan/project.rs +++ b/crates/integrations/datafusion/src/physical_plan/project.rs @@ -30,7 +30,7 @@ use iceberg::arrow::{ PROJECTED_PARTITION_VALUE_COLUMN, PartitionValueCalculator, schema_to_arrow_schema, strip_metadata_from_schema, }; -use iceberg::spec::PartitionSpec; +use iceberg::spec::{PartitionSpec, SchemaRef}; use iceberg::table::Table; use crate::to_datafusion_error; @@ -79,10 +79,6 @@ pub fn project_with_partition( ))); } - let calculator = - PartitionValueCalculator::try_new(partition_spec.as_ref(), table_schema.as_ref()) - .map_err(to_datafusion_error)?; - let mut projection_exprs: Vec<(Arc, String)> = Vec::with_capacity(input_schema.fields().len() + 1); @@ -91,26 +87,51 @@ pub fn project_with_partition( projection_exprs.push((column_expr, field.name().clone())); } - let partition_expr = Arc::new(PartitionExpr::new(calculator, partition_spec.clone())); + let partition_expr = Arc::new(PartitionExpr::try_new( + partition_spec.clone(), + table_schema.clone(), + )?); projection_exprs.push((partition_expr, PROJECTED_PARTITION_VALUE_COLUMN.to_string())); let projection = ProjectionExec::try_new(projection_exprs, input)?; Ok(Arc::new(projection)) } -/// PhysicalExpr implementation for partition value calculation +/// `PhysicalExpr` that computes Iceberg partition values for each input row. +/// +/// Alongside the live (non-serializable) [`PartitionValueCalculator`], it retains +/// the [`PartitionSpec`] and table schema it was built from. A distributed engine +/// can serialize those two — both are self-contained iceberg spec types — and +/// rebuild an equivalent expression on a remote node via [`PartitionExpr::try_new`]. #[derive(Debug, Clone)] -struct PartitionExpr { +pub struct PartitionExpr { calculator: Arc, partition_spec: Arc, + table_schema: SchemaRef, } impl PartitionExpr { - fn new(calculator: PartitionValueCalculator, partition_spec: Arc) -> Self { - Self { + /// Builds a partition expression from a partition spec and the table schema + /// it is bound to, constructing the underlying [`PartitionValueCalculator`]. + pub fn try_new(partition_spec: Arc, table_schema: SchemaRef) -> DFResult { + let calculator = + PartitionValueCalculator::try_new(partition_spec.as_ref(), table_schema.as_ref()) + .map_err(to_datafusion_error)?; + Ok(Self { calculator: Arc::new(calculator), partition_spec, - } + table_schema, + }) + } + + /// The partition spec whose values this expression computes. + pub fn partition_spec(&self) -> &Arc { + &self.partition_spec + } + + /// The table schema the partition values are derived from. + pub fn table_schema(&self) -> &SchemaRef { + &self.table_schema } } @@ -248,8 +269,6 @@ mod tests { let input = Arc::new(EmptyExec::new(arrow_schema.clone())); - let calculator = PartitionValueCalculator::try_new(&partition_spec, &table_schema).unwrap(); - let mut projection_exprs: Vec<(Arc, String)> = Vec::with_capacity(arrow_schema.fields().len() + 1); for (i, field) in arrow_schema.fields().iter().enumerate() { @@ -257,7 +276,9 @@ mod tests { projection_exprs.push((column_expr, field.name().clone())); } - let partition_expr = Arc::new(PartitionExpr::new(calculator, partition_spec)); + let partition_expr = Arc::new( + PartitionExpr::try_new(partition_spec, Arc::new(table_schema.clone())).unwrap(), + ); projection_exprs.push((partition_expr, PROJECTED_PARTITION_VALUE_COLUMN.to_string())); let projection = ProjectionExec::try_new(projection_exprs, input).unwrap(); @@ -302,7 +323,7 @@ mod tests { let partition_spec = Arc::new(partition_spec); let calculator = PartitionValueCalculator::try_new(&partition_spec, &table_schema).unwrap(); let partition_type = calculator.partition_arrow_type().clone(); - let expr = PartitionExpr::new(calculator, partition_spec); + let expr = PartitionExpr::try_new(partition_spec, Arc::new(table_schema.clone())).unwrap(); assert_eq!(expr.data_type(&arrow_schema).unwrap(), partition_type); assert!(!expr.nullable(&arrow_schema).unwrap()); diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index 36539ae503..098af883e8 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -53,11 +53,15 @@ pub struct IcebergTableScan { predicates: Option, /// Optional limit on the number of rows to return limit: Option, + /// Optional serializable catalog/storage config, populated when this scan is + /// built through a config-backed provider so it can be reconstructed on a + /// remote node by a distributed engine. + catalog_config: Option, } impl IcebergTableScan { /// Creates a new [`IcebergTableScan`] object. - pub(crate) fn new( + pub fn new( table: Table, snapshot_id: Option, schema: ArrowSchemaRef, @@ -80,9 +84,36 @@ impl IcebergTableScan { projection, predicates, limit, + catalog_config: None, } } + /// Attaches a serializable catalog/storage config to this scan so that a + /// distributed engine can reconstruct it on a remote node. + pub fn with_catalog_config( + mut self, + catalog_config: Option, + ) -> Self { + self.catalog_config = catalog_config; + self + } + + /// Returns the serializable catalog/storage config, if any. + pub fn catalog_config(&self) -> Option<&crate::IcebergCatalogConfig> { + self.catalog_config.as_ref() + } + + /// Replaces the scan's pushed-down filter predicate. + /// + /// `IcebergTableScan::new` derives the predicate from DataFusion `Expr` + /// filters; this setter lets a distributed engine restore an already-built + /// [`Predicate`] directly (e.g. after deserializing it), so file pruning is + /// preserved on remote nodes. + pub fn with_predicates(mut self, predicates: Option) -> Self { + self.predicates = predicates; + self + } + pub fn table(&self) -> &Table { &self.table } diff --git a/crates/integrations/datafusion/src/physical_plan/write.rs b/crates/integrations/datafusion/src/physical_plan/write.rs index 697eeec659..b9aa80d493 100644 --- a/crates/integrations/datafusion/src/physical_plan/write.rs +++ b/crates/integrations/datafusion/src/physical_plan/write.rs @@ -60,11 +60,15 @@ 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, plan_properties: Arc, + /// Optional serializable catalog/storage config, populated when this node is + /// built through a config-backed provider so it can be reconstructed on a + /// remote node by a distributed engine. + catalog_config: Option, } impl IcebergWriteExec { @@ -76,9 +80,30 @@ impl IcebergWriteExec { input, result_schema: Self::make_result_schema(), plan_properties, + catalog_config: None, } } + /// Attaches a serializable catalog/storage config to this node so that a + /// distributed engine can reconstruct it on a remote node. + pub fn with_catalog_config( + mut self, + catalog_config: Option, + ) -> Self { + self.catalog_config = catalog_config; + self + } + + /// Returns the serializable catalog/storage config, if any. + pub fn catalog_config(&self) -> Option<&crate::IcebergCatalogConfig> { + self.catalog_config.as_ref() + } + + /// Returns the table this node writes to. + pub fn table(&self) -> &Table { + &self.table + } + fn compute_properties( input: &Arc, schema: ArrowSchemaRef, @@ -172,11 +197,10 @@ impl ExecutionPlan for IcebergWriteExec { ))); } - Ok(Arc::new(Self::new( - self.table.clone(), - Arc::clone(&children[0]), - self.schema(), - ))) + Ok(Arc::new( + Self::new(self.table.clone(), Arc::clone(&children[0]), self.schema()) + .with_catalog_config(self.catalog_config.clone()), + )) } /// Executes the write operation for the given partition. diff --git a/crates/integrations/datafusion/src/schema.rs b/crates/integrations/datafusion/src/schema.rs index 76cf599062..f5c7d32912 100644 --- a/crates/integrations/datafusion/src/schema.rs +++ b/crates/integrations/datafusion/src/schema.rs @@ -32,7 +32,7 @@ use iceberg::inspect::MetadataTableType; use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableCreation, TableIdent}; use crate::table::IcebergTableProvider; -use crate::to_datafusion_error; +use crate::{IcebergCatalogConfig, to_datafusion_error}; /// Represents a [`SchemaProvider`] for the Iceberg [`Catalog`], managing /// access to table providers within a specific namespace. @@ -42,6 +42,10 @@ pub(crate) struct IcebergSchemaProvider { catalog: Arc, /// The namespace this schema represents namespace: NamespaceIdent, + /// Optional serializable catalog/storage config. When present, every table + /// provider this schema creates carries it, so catalog-registered tables can + /// be queried by a distributed engine. + config: Option, /// A concurrent map where keys are table names /// and values are dynamic references to objects implementing the /// [`TableProvider`] trait. @@ -57,8 +61,12 @@ impl IcebergSchemaProvider { /// This method retrieves a list of table names /// attempts to create a table provider for each table name, and /// collects these providers into a `HashMap`. + /// + /// When `config` is present it is threaded into every table provider this + /// schema creates, so the tables can be queried by a distributed engine. pub(crate) async fn try_new( client: Arc, + config: Option, namespace: NamespaceIdent, ) -> Result { // TODO: @@ -75,7 +83,14 @@ impl IcebergSchemaProvider { let providers = try_join_all( table_names .iter() - .map(|name| IcebergTableProvider::try_new(client.clone(), namespace.clone(), name)) + .map(|name| { + IcebergTableProvider::try_new( + client.clone(), + config.clone(), + namespace.clone(), + name.clone(), + ) + }) .collect::>(), ) .await?; @@ -88,6 +103,7 @@ impl IcebergSchemaProvider { Ok(IcebergSchemaProvider { catalog: client, namespace, + config, tables, }) } @@ -171,6 +187,7 @@ impl SchemaProvider for IcebergSchemaProvider { let catalog = self.catalog.clone(); let namespace = self.namespace.clone(); + let config = self.config.clone(); let tables = self.tables.clone(); let name_clone = name.clone(); @@ -189,9 +206,11 @@ impl SchemaProvider for IcebergSchemaProvider { .await .map_err(to_datafusion_error)?; - // Create a new table provider using the catalog reference + // Create a new table provider using the catalog reference, + // carrying the config so it stays distributable. let table_provider = IcebergTableProvider::try_new( catalog.clone(), + config.clone(), namespace.clone(), name_clone.clone(), ) @@ -315,13 +334,93 @@ mod tests { .await .unwrap(); - let provider = IcebergSchemaProvider::try_new(Arc::new(catalog), namespace) + let provider = IcebergSchemaProvider::try_new(Arc::new(catalog), None, namespace) .await .unwrap(); (provider, temp_dir) } + #[tokio::test] + async fn test_schema_provider_with_config_propagates_to_tables() { + use iceberg::TableCreation; + use iceberg::spec::{NestedField, PrimitiveType, Schema, Type}; + + let temp_dir = TempDir::new().unwrap(); + let warehouse_path = temp_dir.path().to_str().unwrap().to_string(); + let catalog = Arc::new( + MemoryCatalogBuilder::default() + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse_path.clone())]), + ) + .await + .unwrap(), + ); + + let namespace = NamespaceIdent::new("test_ns".to_string()); + catalog + .create_namespace(&namespace, HashMap::new()) + .await + .unwrap(); + + let schema = Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build() + .unwrap(); + catalog + .create_table( + &namespace, + TableCreation::builder() + .name("t".to_string()) + .location(format!("{warehouse_path}/t")) + .schema(schema) + .properties(HashMap::new()) + .build(), + ) + .await + .unwrap(); + + // With config: the table provider carries it (and is therefore distributable). + let config = crate::IcebergCatalogConfig::new("memory", "memory", HashMap::new()); + let with_config = + IcebergSchemaProvider::try_new(catalog.clone(), Some(config), namespace.clone()) + .await + .unwrap(); + let provider = with_config + .table("t") + .await + .unwrap() + .expect("table provider"); + let iceberg = provider + .as_any() + .downcast_ref::() + .expect("IcebergTableProvider"); + assert!( + iceberg.config().is_some(), + "try_new_with_config should propagate the config to its tables" + ); + + // Without config: providers stay config-less (legacy behavior). + let without_config = + IcebergSchemaProvider::try_new(catalog.clone(), None, namespace.clone()) + .await + .unwrap(); + let provider = without_config + .table("t") + .await + .unwrap() + .expect("table provider"); + let iceberg = provider + .as_any() + .downcast_ref::() + .expect("IcebergTableProvider"); + assert!(iceberg.config().is_none()); + } + #[tokio::test] async fn test_register_table_with_data_fails() { let (schema_provider, _temp_dir) = create_test_schema_provider().await; diff --git a/crates/integrations/datafusion/src/table/metadata_table.rs b/crates/integrations/datafusion/src/table/metadata_table.rs index 38148b4084..ab3c1befa7 100644 --- a/crates/integrations/datafusion/src/table/metadata_table.rs +++ b/crates/integrations/datafusion/src/table/metadata_table.rs @@ -41,6 +41,47 @@ use crate::to_datafusion_error; pub struct IcebergMetadataTableProvider { pub(crate) table: Table, pub(crate) r#type: MetadataTableType, + /// Optional serializable catalog/storage config, populated when this provider + /// is built through a config-backed table provider so that a distributed + /// engine can reconstruct it (reload the table from the catalog) on a remote + /// node. + catalog_config: Option, +} + +impl IcebergMetadataTableProvider { + /// Creates a metadata-table provider over an already-loaded table. + pub fn new(table: Table, r#type: MetadataTableType) -> Self { + Self { + table, + r#type, + catalog_config: None, + } + } + + /// Attaches a serializable catalog/storage config so that a distributed engine + /// can reconstruct this provider on a remote node. + pub fn with_catalog_config( + mut self, + catalog_config: Option, + ) -> Self { + self.catalog_config = catalog_config; + self + } + + /// Returns the serializable catalog/storage config, if any. + pub fn catalog_config(&self) -> Option<&crate::IcebergCatalogConfig> { + self.catalog_config.as_ref() + } + + /// Returns the table this provider inspects. + pub fn table(&self) -> &Table { + &self.table + } + + /// Returns which metadata table this provider serves. + pub fn metadata_type(&self) -> &MetadataTableType { + &self.r#type + } } #[async_trait] diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 75b7988d8d..aa8a2a0872 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -47,7 +47,7 @@ use iceberg::inspect::MetadataTableType; use iceberg::spec::TableProperties; use iceberg::table::Table; use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableIdent}; -use metadata_table::IcebergMetadataTableProvider; +pub use metadata_table::IcebergMetadataTableProvider; use crate::error::to_datafusion_error; use crate::physical_plan::commit::IcebergCommitExec; @@ -73,6 +73,15 @@ pub struct IcebergTableProvider { table_ident: TableIdent, /// A reference-counted arrow `Schema` (cached at construction) schema: ArrowSchemaRef, + /// Optional serializable catalog/storage config. When present, it is + /// threaded into the execution plan nodes produced by `scan`/`insert_into` + /// so that a distributed engine can reconstruct them (and their catalog and + /// storage) on remote nodes. + config: Option, + /// Optional snapshot to read. `None` reads the current snapshot (refreshed + /// from the catalog on each scan); `Some` pins reads to that snapshot for + /// time-travel. Writes always target the current table state. + snapshot_id: Option, } impl IcebergTableProvider { @@ -82,6 +91,7 @@ impl IcebergTableProvider { /// reference for future metadata refreshes on each operation. pub(crate) async fn try_new( catalog: Arc, + config: Option, namespace: NamespaceIdent, name: impl Into, ) -> Result { @@ -95,16 +105,59 @@ impl IcebergTableProvider { catalog, table_ident, schema, + config, + snapshot_id: None, }) } + /// Creates a catalog-backed table provider that carries a serializable + /// [`IcebergCatalogConfig`](crate::IcebergCatalogConfig). + /// + /// The `catalog` must already be built from the same `config`. The config is + /// threaded into the execution plan nodes this provider produces so that a + /// distributed engine (e.g. Ballista) can serialize those nodes and rebuild + /// the catalog/storage on remote executors. + pub async fn try_new_with_config( + catalog: Arc, + config: crate::IcebergCatalogConfig, + namespace: NamespaceIdent, + name: impl Into, + ) -> Result { + Self::try_new(catalog, Some(config), namespace, name).await + } + + /// Pins reads to a specific snapshot for time-travel. `None` (the default) + /// reads the current snapshot. The snapshot id is threaded into the scan + /// node, so it is serialized and honored by a distributed engine as well. + pub fn with_snapshot_id(mut self, snapshot_id: Option) -> Self { + self.snapshot_id = snapshot_id; + self + } + + /// Returns the snapshot this provider reads, if pinned for time-travel. + pub fn snapshot_id(&self) -> Option { + self.snapshot_id + } + + /// Returns the serializable catalog/storage config, if this provider was + /// created with one. + pub fn config(&self) -> Option<&crate::IcebergCatalogConfig> { + self.config.as_ref() + } + + /// Returns the identifier of the table this provider serves. + pub fn table_ident(&self) -> &TableIdent { + &self.table_ident + } + pub(crate) async fn metadata_table( &self, r#type: MetadataTableType, ) -> Result { // Load fresh table metadata for metadata table access let table = self.catalog.load_table(&self.table_ident).await?; - Ok(IcebergMetadataTableProvider { table, r#type }) + Ok(IcebergMetadataTableProvider::new(table, r#type) + .with_catalog_config(self.config.clone())) } } @@ -136,15 +189,18 @@ impl TableProvider for IcebergTableProvider { .await .map_err(to_datafusion_error)?; - // Create scan with fresh metadata (always use current snapshot) - Ok(Arc::new(IcebergTableScan::new( - table, - None, // Always use current snapshot for catalog-backed provider - self.schema.clone(), - projection, - filters, - limit, - ))) + // Create scan with fresh metadata, honoring a pinned snapshot if set. + Ok(Arc::new( + IcebergTableScan::new( + table, + self.snapshot_id, + self.schema.clone(), + projection, + filters, + limit, + ) + .with_catalog_config(self.config.clone()), + )) } fn supports_filters_pushdown( @@ -217,21 +273,23 @@ impl TableProvider for IcebergTableProvider { sort_by_partition(repartitioned_plan)? }; - let write_plan = Arc::new(IcebergWriteExec::new( - table.clone(), - write_input, - self.schema.clone(), - )); + let write_plan = Arc::new( + IcebergWriteExec::new(table.clone(), write_input, self.schema.clone()) + .with_catalog_config(self.config.clone()), + ); // Merge the outputs of write_plan into one so we can commit all files together let coalesce_partitions = Arc::new(CoalescePartitionsExec::new(write_plan)); - Ok(Arc::new(IcebergCommitExec::new( - table, - self.catalog.clone(), - coalesce_partitions, - self.schema.clone(), - ))) + Ok(Arc::new( + IcebergCommitExec::new( + table, + self.catalog.clone(), + coalesce_partitions, + self.schema.clone(), + ) + .with_catalog_config(self.config.clone()), + )) } } @@ -526,10 +584,14 @@ mod tests { let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; // Test creating a catalog-backed provider - let provider = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap(); + let provider = IcebergTableProvider::try_new( + catalog.clone(), + None, + namespace.clone(), + table_name.clone(), + ) + .await + .unwrap(); // Verify the schema is loaded correctly let schema = provider.schema(); @@ -542,10 +604,14 @@ mod tests { async fn test_catalog_backed_provider_scan() { let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; - let provider = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap(); + let provider = IcebergTableProvider::try_new( + catalog.clone(), + None, + namespace.clone(), + table_name.clone(), + ) + .await + .unwrap(); let ctx = SessionContext::new(); ctx.register_table("test_table", Arc::new(provider)) @@ -568,10 +634,14 @@ mod tests { async fn test_catalog_backed_provider_insert() { let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; - let provider = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap(); + let provider = IcebergTableProvider::try_new( + catalog.clone(), + None, + namespace.clone(), + table_name.clone(), + ) + .await + .unwrap(); let ctx = SessionContext::new(); ctx.register_table("test_table", Arc::new(provider)) @@ -595,10 +665,14 @@ mod tests { async fn test_physical_input_schema_consistent_with_logical_input_schema() { let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; - let provider = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap(); + let provider = IcebergTableProvider::try_new( + catalog.clone(), + None, + namespace.clone(), + table_name.clone(), + ) + .await + .unwrap(); let ctx = SessionContext::new(); ctx.register_table("test_table", Arc::new(provider)) @@ -720,10 +794,14 @@ mod tests { let (catalog, namespace, table_name, _temp_dir) = get_partitioned_test_catalog_and_table(Some(true)).await; - let provider = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap(); + let provider = IcebergTableProvider::try_new( + catalog.clone(), + None, + namespace.clone(), + table_name.clone(), + ) + .await + .unwrap(); let ctx = SessionContext::new(); let input_schema = provider.schema(); @@ -752,10 +830,14 @@ mod tests { let (catalog, namespace, table_name, _temp_dir) = get_partitioned_test_catalog_and_table(Some(false)).await; - let provider = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap(); + let provider = IcebergTableProvider::try_new( + catalog.clone(), + None, + namespace.clone(), + table_name.clone(), + ) + .await + .unwrap(); let ctx = SessionContext::new(); let input_schema = provider.schema(); @@ -812,10 +894,14 @@ mod tests { let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; - let provider = - IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) - .await - .unwrap(); + let provider = IcebergTableProvider::try_new( + catalog.clone(), + None, + namespace.clone(), + table_name.clone(), + ) + .await + .unwrap(); let ctx = SessionContext::new(); let state = ctx.state(); @@ -865,4 +951,40 @@ mod tests { "Limit should be None when not specified" ); } + + #[tokio::test] + async fn test_with_snapshot_id_pins_scan() { + use datafusion::datasource::TableProvider; + + let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; + + // Default provider reads the current snapshot (None in the scan). + let provider = IcebergTableProvider::try_new( + catalog.clone(), + None, + namespace.clone(), + table_name.clone(), + ) + .await + .unwrap(); + assert_eq!(provider.snapshot_id(), None); + + // Pinning a snapshot threads it into the scan node, where the codec reads + // it — so time-travel is honored locally and distributed. + let pinned = provider.with_snapshot_id(Some(123)); + assert_eq!(pinned.snapshot_id(), Some(123)); + + let ctx = SessionContext::new(); + let state = ctx.state(); + let scan_plan = pinned.scan(&state, None, &[], None).await.unwrap(); + let iceberg_scan = scan_plan + .as_any() + .downcast_ref::() + .expect("Expected IcebergTableScan"); + assert_eq!( + iceberg_scan.snapshot_id(), + Some(123), + "pinned snapshot should propagate to the scan node" + ); + } } From a4e8657f4b8027e7f1e8830679fa6cd18a2596ad Mon Sep 17 00:00:00 2001 From: Noah Date: Wed, 15 Jul 2026 15:13:05 -0400 Subject: [PATCH 02/11] update --- crates/integrations/datafusion/src/schema.rs | 4 ---- crates/integrations/datafusion/src/table/mod.rs | 4 +--- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/crates/integrations/datafusion/src/schema.rs b/crates/integrations/datafusion/src/schema.rs index f739d9250f..91526935c2 100644 --- a/crates/integrations/datafusion/src/schema.rs +++ b/crates/integrations/datafusion/src/schema.rs @@ -201,8 +201,6 @@ impl SchemaProvider for IcebergSchemaProvider { .await .map_err(to_datafusion_error)?; - // Create a new table provider using the catalog reference, - // carrying the config so it stays distributable. let table_provider = IcebergTableProvider::try_new( catalog.clone(), config.clone(), @@ -391,7 +389,6 @@ mod tests { .unwrap() .expect("table provider"); let iceberg = provider - .as_any() .downcast_ref::() .expect("IcebergTableProvider"); assert!( @@ -410,7 +407,6 @@ mod tests { .unwrap() .expect("table provider"); let iceberg = provider - .as_any() .downcast_ref::() .expect("IcebergTableProvider"); assert!(iceberg.config().is_none()); diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 6b23da2ef5..9f688491d7 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -184,7 +184,6 @@ impl TableProvider for IcebergTableProvider { .await .map_err(to_datafusion_error)?; - // Create scan with fresh metadata, honoring a pinned snapshot if set. Ok(Arc::new( IcebergTableScan::new( table, @@ -786,7 +785,7 @@ mod tests { use datafusion::physical_plan::empty::EmptyExec; let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; - let provider = IcebergTableProvider::try_new(catalog, namespace, table_name) + let provider = IcebergTableProvider::try_new(catalog, None, namespace, table_name) .await .unwrap(); let ctx = SessionContext::new(); @@ -1008,7 +1007,6 @@ mod tests { let state = ctx.state(); let scan_plan = pinned.scan(&state, None, &[], None).await.unwrap(); let iceberg_scan = scan_plan - .as_any() .downcast_ref::() .expect("Expected IcebergTableScan"); assert_eq!( From b4d13f061c70127b768988029c2ed3f4e84b5af2 Mon Sep 17 00:00:00 2001 From: Noah Date: Wed, 15 Jul 2026 15:26:29 -0400 Subject: [PATCH 03/11] cleanup --- crates/integrations/datafusion/src/physical_plan/commit.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/crates/integrations/datafusion/src/physical_plan/commit.rs b/crates/integrations/datafusion/src/physical_plan/commit.rs index ce74b1103b..a9cb294122 100644 --- a/crates/integrations/datafusion/src/physical_plan/commit.rs +++ b/crates/integrations/datafusion/src/physical_plan/commit.rs @@ -86,12 +86,10 @@ impl IcebergCommitExec { self } - /// Returns the serializable catalog/storage config, if any. pub fn catalog_config(&self) -> Option<&crate::IcebergCatalogConfig> { self.catalog_config.as_ref() } - /// Returns the table this node commits to. pub fn table(&self) -> &Table { &self.table } From c9621d1720aadaa8132d001cfd912b0a0306c04c Mon Sep 17 00:00:00 2001 From: hsiang-c <137842490+hsiang-c@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:50:22 -0700 Subject: [PATCH 04/11] feat: add support for `_spec_id` metadata column (#2695) ## Which issue does this PR close? - Partially closes https://github.com/apache/iceberg-rust/issues/2607 ## What changes are included in this PR? - If a projection includes `_spec_id`, which is a constant like the `_file` metadata column for all rows, add it to `RecordBatchTransformerBuilder` ## Are these changes tested? - Yes, unit tests. - Iceberg Spark: https://github.com/apache/datafusion-comet/pull/4752 Sample plan ```shell == Physical Plan == AdaptiveSparkPlan isFinalPlan=false +- CometSort [id#1247L, data#1248, _spec_id#1249], [id#1247L ASC NULLS FIRST] +- CometExchange rangepartitioning(id#1247L ASC NULLS FIRST, 200), ENSURE_REQUIREMENTS, CometNativeShuffle, [plan_id=3533] +- CometFilter [id#1247L, data#1248, _spec_id#1249], (id#1247L >= 10) +- CometIcebergNativeScan [id#1247L, data#1248, _spec_id#1249], /var/folders/d2/b93h6k7174ddqxltrxgb51040000gn/T/iceberg_warehouse16657344237870771906/iceberg_data/default/table/metadata/00004-70b616c5-1bbe-4e1d-bdd6-e9ac4910674a.metadata.json, testrest.default.table (branch=null) [filters=id IS NOT NULL, id >= 10, groupedBy=], 1 ``` --- crates/iceberg/src/arrow/reader/pipeline.rs | 18 +- .../src/arrow/record_batch_transformer.rs | 105 +++--- crates/iceberg/src/arrow/value.rs | 5 +- crates/iceberg/src/scan/context.rs | 15 +- crates/iceberg/src/scan/mod.rs | 325 +++++++++++++++++- ...table_metadata_v2_partition_evolution.json | 128 +++++++ 6 files changed, 534 insertions(+), 62 deletions(-) create mode 100644 crates/iceberg/testdata/example_table_metadata_v2_partition_evolution.json diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 2b7d68aa43..a4072854e4 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -37,7 +37,9 @@ use crate::arrow::record_batch_transformer::RecordBatchTransformerBuilder; use crate::arrow::scan_metrics::{CountingFileRead, ScanMetrics, ScanResult}; use crate::error::Result; use crate::io::{FileIO, FileMetadata, FileRead}; -use crate::metadata_columns::{RESERVED_FIELD_ID_FILE, is_metadata_field}; +use crate::metadata_columns::{ + RESERVED_FIELD_ID_FILE, RESERVED_FIELD_ID_SPEC_ID, is_metadata_field, +}; use crate::scan::{ArrowRecordBatchStream, FileScanTask, FileScanTaskStream}; use crate::spec::Datum; use crate::{Error, ErrorKind}; @@ -248,6 +250,20 @@ impl FileScanTaskReader { record_batch_transformer_builder.with_constant(RESERVED_FIELD_ID_FILE, file_datum); } + if task + .project_field_ids() + .contains(&RESERVED_FIELD_ID_SPEC_ID) + { + let partition_spec = task + .partition_spec + .as_ref() + .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Partition spec is missing"))?; + + let spec_id_datum = Datum::int(partition_spec.spec_id()); + record_batch_transformer_builder = record_batch_transformer_builder + .with_constant(RESERVED_FIELD_ID_SPEC_ID, spec_id_datum); + } + if let (Some(partition_spec), Some(partition_data)) = (task.partition_spec.clone(), task.partition.clone()) { diff --git a/crates/iceberg/src/arrow/record_batch_transformer.rs b/crates/iceberg/src/arrow/record_batch_transformer.rs index c780221d4a..b7b8d609a7 100644 --- a/crates/iceberg/src/arrow/record_batch_transformer.rs +++ b/crates/iceberg/src/arrow/record_batch_transformer.rs @@ -359,49 +359,29 @@ impl RecordBatchTransformer { let fields: Result> = projected_iceberg_field_ids .iter() .map(|field_id| { - // Check if this is a constant field - if constant_fields.contains_key(field_id) { - // For metadata/virtual fields (like _file), get name from metadata_columns - // For partition fields, get name from schema (they exist in schema) - if let Ok(iceberg_field) = get_metadata_field(*field_id) { - // This is a metadata/virtual field - convert Iceberg field to Arrow - let datum = constant_fields.get(field_id).ok_or(Error::new( - ErrorKind::Unexpected, - "constant field not found", - ))?; - let arrow_type = datum_to_arrow_type_with_ree(datum); - let arrow_field = - Field::new(&iceberg_field.name, arrow_type, !iceberg_field.required) - .with_metadata(HashMap::from([( - PARQUET_FIELD_ID_META_KEY.to_string(), - iceberg_field.id.to_string(), - )])); - Ok(Arc::new(arrow_field)) - } else { - // This is a partition constant field (exists in schema but uses constant value) - let field = &field_id_to_mapped_schema_map - .get(field_id) - .ok_or(Error::new(ErrorKind::Unexpected, "field not found"))? - .0; - let datum = constant_fields.get(field_id).ok_or(Error::new( - ErrorKind::Unexpected, - "constant field not found", - ))?; - let arrow_type = datum_to_arrow_type_with_ree(datum); - // Use the type from constant_fields (REE for constants) - let constant_field = - Field::new(field.name(), arrow_type, field.is_nullable()) - .with_metadata(field.metadata().clone()); - Ok(Arc::new(constant_field)) - } - } else { - // Regular field - use schema as-is - Ok(field_id_to_mapped_schema_map - .get(field_id) - .ok_or(Error::new(ErrorKind::Unexpected, "field not found"))? - .0 - .clone()) + // Metadata/virtual fields (like _file, _spec_id) don't exist in the table + // schema, so build their Arrow field from the metadata column definition and + // the pre-computed constant's type. + if let Some(datum) = constant_fields.get(field_id) + && let Ok(iceberg_field) = get_metadata_field(*field_id) + { + let arrow_type = datum_to_arrow_type_with_ree(datum); + let arrow_field = + Field::new(&iceberg_field.name, arrow_type, !iceberg_field.required) + .with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + iceberg_field.id.to_string(), + )])); + return Ok(Arc::new(arrow_field)); } + + // Regular fields and identity-partitioned constant fields both exist in the + // table schema, so use the mapped Arrow field as-is. + Ok(field_id_to_mapped_schema_map + .get(field_id) + .ok_or(Error::new(ErrorKind::Unexpected, "field not found"))? + .0 + .clone()) }) .collect(); @@ -481,16 +461,39 @@ impl RecordBatchTransformer { projected_iceberg_field_ids .iter() .map(|field_id| { - // Check if this is a constant field (metadata/virtual or identity-partitioned) - // Constant fields always use their pre-computed constant values, regardless of whether - // they exist in the Parquet file. This is per Iceberg spec rule #1: partition metadata - // is authoritative and should be preferred over file data. + // Check if this is a constant field (metadata/virtual or identity-partitioned). + // + // Metadata/virtual fields (like _file, _spec_id) never exist in the data file, + // so they always use their pre-computed constant value. + // + // For identity-partitioned fields, the Iceberg spec's "Column Projection" rules + // only apply to "field ids which are not present in a data file". When the column + // IS present in the Parquet file, it must be read from the file; the partition + // metadata constant is only a fallback for when the column is absent (e.g. add_files). if let Some(datum) = constant_fields.get(field_id) { - let arrow_type = datum_to_arrow_type_with_ree(datum); - return Ok(ColumnSource::Add { - value: Some(datum.literal().clone()), - target_type: arrow_type, - }); + let is_metadata_field = get_metadata_field(*field_id).is_ok(); + let present_in_file = field_id_to_source_schema_map.contains_key(field_id); + + if is_metadata_field || !present_in_file { + let arrow_type = if is_metadata_field { + datum_to_arrow_type_with_ree(datum) + } else { + field_id_to_mapped_schema_map + .get(field_id) + .ok_or(Error::new( + ErrorKind::Unexpected, + "could not find field in schema", + ))? + .0 + .data_type() + .clone() + }; + + return Ok(ColumnSource::Add { + value: Some(datum.literal().clone()), + target_type: arrow_type, + }); + } } let (target_field, _) = diff --git a/crates/iceberg/src/arrow/value.rs b/crates/iceberg/src/arrow/value.rs index d07233c420..c2c0d650d8 100644 --- a/crates/iceberg/src/arrow/value.rs +++ b/crates/iceberg/src/arrow/value.rs @@ -834,6 +834,9 @@ pub(crate) fn create_primitive_array_repeated( let vals: Vec> = vec![None; num_rows]; Arc::new(Date32Array::from(vals)) } + (DataType::Int64, Some(PrimitiveLiteral::Int(value))) => { + Arc::new(Int64Array::from(vec![i64::from(*value); num_rows])) + } (DataType::Int64, Some(PrimitiveLiteral::Long(value))) => { Arc::new(Int64Array::from(vec![*value; num_rows])) } @@ -969,7 +972,7 @@ pub(crate) fn create_primitive_array_repeated( (dt, _) => { return Err(Error::new( ErrorKind::Unexpected, - format!("unexpected target column type {dt}"), + format!("unexpected target column type {dt}, prim_lit {prim_lit:?}"), )); } }) diff --git a/crates/iceberg/src/scan/context.rs b/crates/iceberg/src/scan/context.rs index 75672d9cbb..6a6b2ebd65 100644 --- a/crates/iceberg/src/scan/context.rs +++ b/crates/iceberg/src/scan/context.rs @@ -28,8 +28,8 @@ use crate::scan::{ PartitionFilterCache, }; use crate::spec::{ - ManifestContentType, ManifestEntryRef, ManifestFile, ManifestList, NameMapping, SchemaRef, - SnapshotRef, TableMetadataRef, + ManifestContentType, ManifestEntryRef, ManifestFile, ManifestList, NameMapping, + PartitionSpecRef, SchemaRef, SnapshotRef, TableMetadataRef, }; use crate::{Error, ErrorKind, Result}; @@ -48,6 +48,7 @@ pub(crate) struct ManifestFileContext { delete_file_index: DeleteFileIndex, name_mapping: Option>, case_sensitive: bool, + partition_spec: Option, } /// Wraps a [`ManifestEntryRef`] alongside the objects that are needed @@ -63,6 +64,7 @@ pub(crate) struct ManifestEntryContext { pub delete_file_index: DeleteFileIndex, pub name_mapping: Option>, pub case_sensitive: bool, + pub partition_spec: Option, } impl ManifestFileContext { @@ -80,6 +82,7 @@ impl ManifestFileContext { delete_file_index, name_mapping, case_sensitive, + partition_spec, } = self; let manifest = object_cache.get_manifest(&manifest_file).await?; @@ -96,6 +99,7 @@ impl ManifestFileContext { delete_file_index: delete_file_index.clone(), name_mapping: name_mapping.clone(), case_sensitive, + partition_spec: partition_spec.clone(), }; sender @@ -135,8 +139,7 @@ impl ManifestEntryContext { ) .with_deletes(deletes) .with_partition(Some(self.manifest_entry.data_file.partition.clone())) - // TODO: Pass actual PartitionSpec through context chain for native flow - .with_partition_spec(None) + .with_partition_spec(self.partition_spec.clone()) .with_name_mapping(self.name_mapping) .with_case_sensitive(self.case_sensitive) .build()) @@ -284,6 +287,10 @@ impl PlanContext { delete_file_index, name_mapping: self.name_mapping.clone(), case_sensitive: self.case_sensitive, + partition_spec: self + .table_metadata + .partition_spec_by_id(manifest_file.partition_spec_id) + .cloned(), } } } diff --git a/crates/iceberg/src/scan/mod.rs b/crates/iceberg/src/scan/mod.rs index 9e12ffa4a9..e5f2bf9979 100644 --- a/crates/iceberg/src/scan/mod.rs +++ b/crates/iceberg/src/scan/mod.rs @@ -625,8 +625,9 @@ pub mod tests { use std::sync::Arc; use arrow_array::cast::AsArray; + use arrow_array::types::Int32Type; use arrow_array::{ - Array, ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, + Array, ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, RunArray, StringArray, }; use futures::{TryStreamExt, stream}; @@ -641,7 +642,7 @@ pub mod tests { use crate::arrow::ArrowReaderBuilder; use crate::expr::{BoundPredicate, Reference}; use crate::io::{FileIO, OutputFile}; - use crate::metadata_columns::RESERVED_COL_NAME_FILE; + use crate::metadata_columns::{RESERVED_COL_NAME_FILE, RESERVED_COL_NAME_SPEC_ID}; use crate::scan::FileScanTask; use crate::spec::{ DEFAULT_SCHEMA_NAME_MAPPING, DataContentType, DataFileBuilder, DataFileFormat, Datum, @@ -820,6 +821,40 @@ pub mod tests { } } + pub fn new_with_partition_evolution() -> Self { + let table = Self::new().table; + let table_location = table.metadata().location.clone(); + + let manifest_list1_location = + format!("{}/metadata/manifests_list_1.avro", table_location); + let manifest_list2_location = + format!("{}/metadata/manifests_list_2.avro", table_location); + let manifest_list3_location = + format!("{}/metadata/manifests_list_3.avro", table_location); + let table_metadata1_location = format!("{}/metadata/v1.json", table_location); + + let new_table_metadata = { + let template_json_str = fs::read_to_string(format!( + "{}/testdata/example_table_metadata_v2_partition_evolution.json", + env!("CARGO_MANIFEST_DIR") + )) + .unwrap(); + let metadata_json = render_template(&template_json_str, context! { + table_location => &table_location, + manifest_list_1_location => &manifest_list1_location, + manifest_list_2_location => &manifest_list2_location, + manifest_list_3_location => &manifest_list3_location, + table_metadata_1_location => &table_metadata1_location, + }); + Arc::new(serde_json::from_str::(&metadata_json).unwrap()) + }; + + Self { + table_location, + table: table.with_metadata(new_table_metadata), + } + } + fn next_manifest_file(&self) -> OutputFile { self.table .file_io() @@ -936,6 +971,123 @@ pub mod tests { manifest_list_write.close().await.unwrap(); } + pub async fn setup_manifest_files_with_partition_evolution(&mut self) { + let current_snapshot = self.table.metadata().current_snapshot().unwrap(); + let parent_snapshot = current_snapshot + .parent_snapshot(self.table.metadata()) + .unwrap(); + let current_schema = current_snapshot.schema(self.table.metadata()).unwrap(); + let current_partition_spec = self.table.metadata().default_partition_spec(); + + // Write the data files first, then use the file size in the manifest entries + let parquet_file_size = self.write_parquet_data_files(); + + let mut writer = ManifestWriterBuilder::new( + self.next_manifest_file(), + Some(current_snapshot.snapshot_id()), + current_schema.clone(), + current_partition_spec.as_ref().clone(), + ) + .build_v2_data(); + writer + .add_entry( + ManifestEntry::builder() + .status(ManifestStatus::Added) + .data_file( + DataFileBuilder::default() + .partition_spec_id(1) + .content(DataContentType::Data) + .file_path(format!("{}/1.parquet", &self.table_location)) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(parquet_file_size) + .record_count(1) + .partition(Struct::from_iter([ + Some(Literal::long(100)), + Some(Literal::string("apa")), + Some(Literal::int(27)), + ])) + .key_metadata(None) + .build() + .unwrap(), + ) + .build(), + ) + .unwrap(); + writer + .add_delete_entry( + ManifestEntry::builder() + .status(ManifestStatus::Deleted) + .snapshot_id(parent_snapshot.snapshot_id()) + .sequence_number(parent_snapshot.sequence_number()) + .file_sequence_number(parent_snapshot.sequence_number()) + .data_file( + DataFileBuilder::default() + .partition_spec_id(1) + .content(DataContentType::Data) + .file_path(format!("{}/2.parquet", &self.table_location)) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(parquet_file_size) + .record_count(1) + .partition(Struct::from_iter([ + Some(Literal::long(200)), + Some(Literal::string("ice")), + Some(Literal::int(5)), + ])) + .build() + .unwrap(), + ) + .build(), + ) + .unwrap(); + writer + .add_existing_entry( + ManifestEntry::builder() + .status(ManifestStatus::Existing) + .snapshot_id(parent_snapshot.snapshot_id()) + .sequence_number(parent_snapshot.sequence_number()) + .file_sequence_number(parent_snapshot.sequence_number()) + .data_file( + DataFileBuilder::default() + .partition_spec_id(1) + .content(DataContentType::Data) + .file_path(format!("{}/3.parquet", &self.table_location)) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(parquet_file_size) + .record_count(1) + .partition(Struct::from_iter([ + Some(Literal::long(300)), + Some(Literal::string("apa")), + Some(Literal::int(19)), + ])) + .build() + .unwrap(), + ) + .build(), + ) + .unwrap(); + let data_file_manifest = writer.write_manifest_file().await.unwrap(); + + // Write to manifest list + let manifest_list_writer = self + .table + .file_io() + .new_output(current_snapshot.manifest_list()) + .unwrap() + .writer() + .await + .unwrap(); + let mut manifest_list_write = ManifestListWriter::v2( + manifest_list_writer, + current_snapshot.snapshot_id(), + current_snapshot.parent_snapshot_id(), + current_snapshot.sequence_number(), + ); + manifest_list_write + .add_manifests(vec![data_file_manifest].into_iter()) + .unwrap(); + manifest_list_write.close().await.unwrap(); + } + /// Writes identical Parquet data files (1.parquet, 2.parquet, 3.parquet) /// and returns the file size in bytes. fn write_parquet_data_files(&self) -> u64 { @@ -2021,8 +2173,6 @@ pub mod tests { #[tokio::test] async fn test_select_with_file_column() { - use arrow_array::cast::AsArray; - let mut fixture = TableTestFixture::new(); fixture.setup_manifest_files().await; @@ -2066,7 +2216,7 @@ pub mod tests { // Decode the RunArray to verify it contains the file path let run_array = file_col .as_any() - .downcast_ref::>() + .downcast_ref::>() .expect("_file column should be a RunArray"); let values = run_array.values(); @@ -2365,4 +2515,169 @@ pub mod tests { // Assert it finished (didn't timeout) assert!(result.is_ok(), "Scan timed out - deadlock detected"); } + + #[tokio::test] + async fn test_select_with_spec_id_column() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + + // Select regular columns plus the _spec_id column + let table_scan = fixture + .table + .scan() + .select(["x", RESERVED_COL_NAME_SPEC_ID, "z"]) + .with_row_selection_enabled(true) + .build() + .unwrap(); + + let batch_stream = table_scan.to_arrow().await.unwrap(); + let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); + + // Verify we have 3 columns: x, _spec_id, and z + assert_eq!(batches[0].num_columns(), 3); + + // Verify the x column exists and has correct data + let col1 = batches[0].column_by_name("x").unwrap(); + let int64_arr = col1.as_any().downcast_ref::().unwrap(); + assert_eq!(int64_arr.value(0), 1); + + // Verify the _spec_id column exists + let spec_id_col = batches[0].column_by_name(RESERVED_COL_NAME_SPEC_ID); + assert!( + spec_id_col.is_some(), + "_spec_id column should be present in the batch" + ); + + // Verify the _spec_id data type + let spec_id_col = spec_id_col.unwrap(); + assert!( + matches!( + spec_id_col.data_type(), + arrow_schema::DataType::RunEndEncoded(_, _) + ), + "_spec_id column should use RunEndEncoded type" + ); + + // Decode the RunArray to verify it contains the spec id + let run_array = spec_id_col + .as_any() + .downcast_ref::>() + .expect("_spec_id column should be a RunArray"); + + let values = run_array.values(); + let int_values = values.as_primitive::(); + assert_eq!(int_values.len(), 1, "Should have a single _spec_id"); + + let spec_id = int_values.value(0); + assert_eq!(spec_id, 0, "_spec_id should be 0, got: {spec_id}"); + + // Verify 'z' column exists + assert!(batches[0].column_by_name("z").is_some()); + } + + #[tokio::test] + async fn test_select_with_spec_id_column_from_unpartitioned_table() { + let mut fixture = TableTestFixture::new_unpartitioned(); + fixture.setup_unpartitioned_manifest_files().await; + + // Select regular columns plus the _spec_id column + let table_scan = fixture + .table + .scan() + .select(["x", RESERVED_COL_NAME_SPEC_ID]) + .with_row_selection_enabled(true) + .build() + .unwrap(); + + let batch_stream = table_scan.to_arrow().await.unwrap(); + let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); + + // Verify we have 2 columns: x and _spec_id + assert_eq!(batches[0].num_columns(), 2); + + // Verify the _spec_id column exists + let spec_id_col = batches[0].column_by_name(RESERVED_COL_NAME_SPEC_ID); + assert!( + spec_id_col.is_some(), + "_spec_id column should be present in the batch" + ); + + // Verify the _spec_id data type + let spec_id_col = spec_id_col.unwrap(); + assert!( + matches!( + spec_id_col.data_type(), + arrow_schema::DataType::RunEndEncoded(_, _) + ), + "_spec_id column should use RunEndEncoded type" + ); + + // Decode the RunArray to verify it contains the spec id + let run_array = spec_id_col + .as_any() + .downcast_ref::>() + .expect("_spec_id column should be a RunArray"); + + let values = run_array.values(); + let int_values = values.as_primitive::(); + assert_eq!(int_values.len(), 1, "Should have a single _spec_id"); + + let spec_id = int_values.value(0); + assert_eq!(spec_id, 0, "_spec_id should be 0, got: {spec_id}"); + } + + #[tokio::test] + async fn test_select_with_spec_id_column_with_partition_evolution() { + let mut fixture = TableTestFixture::new_with_partition_evolution(); + fixture + .setup_manifest_files_with_partition_evolution() + .await; + + // Select regular columns plus the _spec_id column + let table_scan = fixture + .table + .scan() + .select(["x", RESERVED_COL_NAME_SPEC_ID, "z"]) + .with_row_selection_enabled(true) + .build() + .unwrap(); + + let batch_stream = table_scan.to_arrow().await.unwrap(); + let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); + + // Verify the x column exists and has correct data + let col1 = batches[0].column_by_name("x").unwrap(); + let int64_arr = col1.as_any().downcast_ref::().unwrap(); + assert_eq!(int64_arr.value(0), 1); + + // Verify the _spec_id column exists + let spec_id_col = batches[0].column_by_name(RESERVED_COL_NAME_SPEC_ID); + assert!( + spec_id_col.is_some(), + "_spec_id column should be present in the batch" + ); + + // Verify the _spec_id data type + let spec_id_col = spec_id_col.unwrap(); + assert!( + matches!( + spec_id_col.data_type(), + arrow_schema::DataType::RunEndEncoded(_, _) + ), + "_spec_id column should use RunEndEncoded type" + ); + + // Decode the RunArray to verify it contains the spec id + let run_array = spec_id_col + .as_any() + .downcast_ref::>() + .expect("_spec_id column should be a RunArray"); + + let values = run_array.values(); + let int_values = values.as_primitive::(); + assert_eq!(int_values.len(), 1, "Should have a single _spec_id"); + + let spec_id = int_values.value(0); + assert_eq!(spec_id, 2, "_spec_id should be 2, got: {spec_id}"); + } } diff --git a/crates/iceberg/testdata/example_table_metadata_v2_partition_evolution.json b/crates/iceberg/testdata/example_table_metadata_v2_partition_evolution.json new file mode 100644 index 0000000000..70743c394f --- /dev/null +++ b/crates/iceberg/testdata/example_table_metadata_v2_partition_evolution.json @@ -0,0 +1,128 @@ +{ + "format-version": 2, + "table-uuid": "9c12d441-03fe-4693-9a96-a0705ddf69c1", + "location": "{{ table_location }}", + "last-sequence-number": 34, + "last-updated-ms": 1602638573590, + "last-column-id": 3, + "current-schema-id": 1, + "schemas": [ + { + "type": "struct", + "schema-id": 0, + "fields": [ + {"id": 1, "name": "x", "required": true, "type": "long"} + ]}, + { + "type": "struct", + "schema-id": 1, + "identifier-field-ids": [1, 2], + "fields": [ + {"id": 1, "name": "x", "required": true, "type": "long"}, + {"id": 2, "name": "y", "required": true, "type": "long", "doc": "comment"}, + {"id": 3, "name": "z", "required": true, "type": "long"}, + {"id": 4, "name": "a", "required": true, "type": "string"}, + {"id": 5, "name": "dbl", "required": true, "type": "double"}, + {"id": 6, "name": "i32", "required": true, "type": "int"}, + {"id": 7, "name": "i64", "required": true, "type": "long"}, + {"id": 8, "name": "bool", "required": true, "type": "boolean"} + ] + } + ], + "default-spec-id": 2, + "partition-specs": [ + { + "spec-id": 0, + "fields": [ + {"name": "x", "transform": "identity", "source-id": 1, "field-id": 1000} + ] + }, + { + "spec-id": 1, + "fields": [ + { + "name": "x", + "transform": "identity", + "source-id": 1, + "field-id": 1000 + }, + { + "name": "a_trunc_3", + "transform": "truncate[3]", + "source-id": 4, + "field-id": 1001 + } + ] + }, + { + "spec-id": 2, + "fields": [ + { + "name": "x", + "transform": "identity", + "source-id": 1, + "field-id": 1000 + }, + { + "name": "a_trunc_3", + "transform": "truncate[3]", + "source-id": 4, + "field-id": 1001 + }, + { + "name": "z_bucket_32", + "transform": "bucket[32]", + "source-id": 3, + "field-id": 1002 + } + ] + } + ], + "last-partition-id": 1002, + "default-sort-order-id": 3, + "sort-orders": [ + { + "order-id": 3, + "fields": [ + {"transform": "identity", "source-id": 2, "direction": "asc", "null-order": "nulls-first"}, + {"transform": "bucket[4]", "source-id": 3, "direction": "desc", "null-order": "nulls-last"} + ] + } + ], + "properties": {"read.split.target.size": "134217728"}, + "current-snapshot-id": 3059729675574597004, + "snapshots": [ + { + "snapshot-id": 3051729675574597004, + "timestamp-ms": 1515100955770, + "sequence-number": 0, + "summary": {"operation": "append"}, + "manifest-list": "{{ manifest_list_1_location }}" + }, + { + "snapshot-id": 3055729675574597004, + "parent-snapshot-id": 3051729675574597004, + "timestamp-ms": 1555100955770, + "sequence-number": 1, + "summary": {"operation": "append"}, + "manifest-list": "{{ manifest_list_2_location }}", + "schema-id": 1 + }, + { + "snapshot-id": 3059729675574597004, + "parent-snapshot-id": 3055729675574597004, + "timestamp-ms": 1555100955771, + "sequence-number": 2, + "summary": {"operation": "append"}, + "manifest-list": "{{ manifest_list_3_location }}", + "schema-id": 1 + } + ], + "snapshot-log": [ + {"snapshot-id": 3051729675574597004, "timestamp-ms": 1515100955770}, + {"snapshot-id": 3055729675574597004, "timestamp-ms": 1555100955770}, + {"snapshot-id": 3059729675574597004, "timestamp-ms": 1555100955770} + ], + "metadata-log": [{"metadata-file": "{{ table_metadata_1_location }}", "timestamp-ms": 1515100}], + "refs": {"test": {"snapshot-id": 3051729675574597004, "type": "tag", "max-ref-age-ms": 10000000}} +} \ No newline at end of file From 5c759bd59ed779ea9e6ca5f9ce32b98c3dca061c Mon Sep 17 00:00:00 2001 From: Xander Date: Thu, 16 Jul 2026 02:56:47 +0100 Subject: [PATCH 05/11] feat(encryption) [9/N] Read encrypted parquet data-files (#2584) ## Which issue does this PR close? - Working towards https://github.com/apache/iceberg-rust/issues/2034 We're currently not able to decrypt PME encrypted files, all that's needed for this PR is to wire through the key_metadata from the manifest entries into the scan tasks and then configure the parquet reader with the correct encryption key and aad. It's important to note that no encryption manager is needed here since the key_metadata on the manifest files contain the plaintext keys. ## What changes are included in this PR? ## Are these changes tested? Tests follow the pattern on manually writing an encrypted parquet file using the raw arrow-rs writer and then using the FileScanTask and iceberg reader to read that tmp file. I looked at using the [parquet testing repo](https://github.com/apache/parquet-testing) but the encrypted files there all use column based encryption which isn't supported by the spec. I'll separately track adding some files there that we can use here but the testing strategy here is consistent with how Java does things so I think this should be sufficient for us to continue. --- crates/iceberg/public-api.txt | 6 +- .../src/arrow/caching_delete_file_loader.rs | 16 +- .../iceberg/src/arrow/delete_file_loader.rs | 160 +++++++++++- crates/iceberg/src/arrow/mod.rs | 3 + crates/iceberg/src/arrow/reader/pipeline.rs | 234 +++++++++++++++++- crates/iceberg/src/arrow/reader/row_filter.rs | 3 + crates/iceberg/src/arrow/test_utils.rs | 48 ++++ crates/iceberg/src/encryption/key_metadata.rs | 40 ++- crates/iceberg/src/scan/context.rs | 1 + crates/iceberg/src/scan/task.rs | 31 +++ 10 files changed, 530 insertions(+), 12 deletions(-) create mode 100644 crates/iceberg/src/arrow/test_utils.rs diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index c2013294e3..2de70bf265 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1244,6 +1244,7 @@ pub iceberg::scan::FileScanTask::data_file_format: iceberg::spec::DataFileFormat pub iceberg::scan::FileScanTask::data_file_path: alloc::string::String pub iceberg::scan::FileScanTask::deletes: alloc::vec::Vec pub iceberg::scan::FileScanTask::file_size_in_bytes: u64 +pub iceberg::scan::FileScanTask::key_metadata: core::option::Option> pub iceberg::scan::FileScanTask::length: u64 pub iceberg::scan::FileScanTask::name_mapping: core::option::Option> pub iceberg::scan::FileScanTask::partition: core::option::Option @@ -1267,7 +1268,7 @@ impl core::fmt::Debug for iceberg::scan::FileScanTask pub fn iceberg::scan::FileScanTask::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl core::marker::StructuralPartialEq for iceberg::scan::FileScanTask impl iceberg::scan::FileScanTask -pub fn iceberg::scan::FileScanTask::builder() -> FileScanTaskBuilder<((), (), (), (), (), (), (), (), (), (), (), (), (), ())> +pub fn iceberg::scan::FileScanTask::builder() -> FileScanTaskBuilder<((), (), (), (), (), (), (), (), (), (), (), (), (), (), ())> impl serde_core::ser::Serialize for iceberg::scan::FileScanTask pub fn iceberg::scan::FileScanTask::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg::scan::FileScanTask @@ -1277,6 +1278,7 @@ pub iceberg::scan::FileScanTaskDeleteFile::equality_ids: core::option::Option> pub iceberg::scan::FileScanTaskDeleteFile::partition_spec_id: i32 impl core::clone::Clone for iceberg::scan::FileScanTaskDeleteFile pub fn iceberg::scan::FileScanTaskDeleteFile::clone(&self) -> iceberg::scan::FileScanTaskDeleteFile @@ -1286,7 +1288,7 @@ impl core::fmt::Debug for iceberg::scan::FileScanTaskDeleteFile pub fn iceberg::scan::FileScanTaskDeleteFile::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl core::marker::StructuralPartialEq for iceberg::scan::FileScanTaskDeleteFile impl iceberg::scan::FileScanTaskDeleteFile -pub fn iceberg::scan::FileScanTaskDeleteFile::builder() -> FileScanTaskDeleteFileBuilder<((), (), (), (), ())> +pub fn iceberg::scan::FileScanTaskDeleteFile::builder() -> FileScanTaskDeleteFileBuilder<((), (), (), (), (), ())> impl serde_core::ser::Serialize for iceberg::scan::FileScanTaskDeleteFile pub fn iceberg::scan::FileScanTaskDeleteFile::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg::scan::FileScanTaskDeleteFile diff --git a/crates/iceberg/src/arrow/caching_delete_file_loader.rs b/crates/iceberg/src/arrow/caching_delete_file_loader.rs index de31d678a5..25704e5bc4 100644 --- a/crates/iceberg/src/arrow/caching_delete_file_loader.rs +++ b/crates/iceberg/src/arrow/caching_delete_file_loader.rs @@ -251,7 +251,11 @@ impl CachingDeleteFileLoader { PosDelLoadAction::Load => Ok(DeleteFileContext::PosDels { file_path: task.file_path.clone(), stream: basic_delete_file_loader - .parquet_to_batch_stream(&task.file_path, task.file_size_in_bytes) + .parquet_to_batch_stream( + &task.file_path, + task.file_size_in_bytes, + task.key_metadata.as_deref(), + ) .await?, }), } @@ -270,7 +274,11 @@ impl CachingDeleteFileLoader { let equality_ids_vec = task.equality_ids.clone().unwrap(); let evolved_stream = BasicDeleteFileLoader::evolve_schema( basic_delete_file_loader - .parquet_to_batch_stream(&task.file_path, task.file_size_in_bytes) + .parquet_to_batch_stream( + &task.file_path, + task.file_size_in_bytes, + task.key_metadata.as_deref(), + ) .await?, schema, &equality_ids_vec, @@ -647,6 +655,7 @@ mod tests { .parquet_to_batch_stream( &eq_delete_file_path, std::fs::metadata(&eq_delete_file_path).unwrap().len(), + None, ) .await .expect("could not get batch stream"); @@ -977,6 +986,7 @@ mod tests { .parquet_to_batch_stream( &delete_file_path, std::fs::metadata(&delete_file_path).unwrap().len(), + None, ) .await .unwrap(); @@ -1158,7 +1168,7 @@ mod tests { let basic_delete_file_loader = BasicDeleteFileLoader::new(file_io.clone(), ScanMetrics::new()); let record_batch_stream = basic_delete_file_loader - .parquet_to_batch_stream(&path, std::fs::metadata(&path).unwrap().len()) + .parquet_to_batch_stream(&path, std::fs::metadata(&path).unwrap().len(), None) .await .expect("could not get batch stream"); diff --git a/crates/iceberg/src/arrow/delete_file_loader.rs b/crates/iceberg/src/arrow/delete_file_loader.rs index 134b029613..9bcc965197 100644 --- a/crates/iceberg/src/arrow/delete_file_loader.rs +++ b/crates/iceberg/src/arrow/delete_file_loader.rs @@ -67,6 +67,7 @@ impl BasicDeleteFileLoader { &self, data_file_path: &str, file_size_in_bytes: u64, + key_metadata: Option<&[u8]>, ) -> Result { /* Essentially a super-cut-down ArrowReader. We can't use ArrowReader directly @@ -80,6 +81,7 @@ impl BasicDeleteFileLoader { file_size_in_bytes, parquet_read_options, self.scan_metrics.bytes_read_counter(), + key_metadata, ) .await?; @@ -121,7 +123,11 @@ impl DeleteFileLoader for BasicDeleteFileLoader { schema: SchemaRef, ) -> Result { let raw_batch_stream = self - .parquet_to_batch_stream(&task.file_path, task.file_size_in_bytes) + .parquet_to_batch_stream( + &task.file_path, + task.file_size_in_bytes, + task.key_metadata.as_deref(), + ) .await?; // For equality deletes, only evolve the equality_ids columns. @@ -141,6 +147,7 @@ mod tests { use super::*; use crate::arrow::delete_filter::tests::setup; + use crate::arrow::test_utils::write_encrypted_parquet; #[tokio::test] async fn test_basic_delete_file_loader_read_delete_file() { @@ -165,4 +172,155 @@ mod tests { assert_eq!(result.len(), 1); } + + #[tokio::test] + async fn test_read_encrypted_positional_delete_file() { + use std::sync::Arc; + + use arrow_array::{Int64Array, RecordBatch, StringArray}; + + use crate::arrow::delete_filter::tests::create_pos_del_schema; + use crate::encryption::StandardKeyMetadata; + use crate::scan::FileScanTaskDeleteFile; + use crate::spec::DataContentType; + + let encryption_key = b"0123456789abcdef"; + let aad_prefix = b"aad_prefix"; + + let tmp_dir = TempDir::new().unwrap(); + let table_location = tmp_dir.path().to_str().unwrap(); + let file_io = FileIO::new_with_fs(); + + let positional_delete_schema = create_pos_del_schema(); + let file_path_col = Arc::new(StringArray::from_iter_values(vec!["data.parquet"; 4])); + let pos_col = Arc::new(Int64Array::from(vec![0i64, 1, 5, 10])); + let batch = RecordBatch::try_new(positional_delete_schema.clone(), vec![ + file_path_col, + pos_col, + ]) + .unwrap(); + + let del_path = format!("{table_location}/encrypted-pos-del.parquet"); + write_encrypted_parquet(&del_path, &batch, encryption_key, Some(aad_prefix)); + + let key_metadata = StandardKeyMetadata::new(encryption_key) + .with_aad_prefix(aad_prefix) + .encode() + .unwrap(); + + let schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + crate::spec::NestedField::required( + 2147483546, + "file_path", + crate::spec::Type::Primitive(crate::spec::PrimitiveType::String), + ) + .into(), + crate::spec::NestedField::required( + 2147483545, + "pos", + crate::spec::Type::Primitive(crate::spec::PrimitiveType::Long), + ) + .into(), + ]) + .build() + .unwrap(), + ); + + let task = FileScanTaskDeleteFile { + file_path: del_path.clone(), + file_size_in_bytes: std::fs::metadata(&del_path).unwrap().len(), + file_type: DataContentType::PositionDeletes, + partition_spec_id: 0, + equality_ids: None, + key_metadata: Some(Box::from(key_metadata.as_ref())), + }; + + let scan_metrics = ScanMetrics::new(); + let delete_file_loader = BasicDeleteFileLoader::new(file_io, scan_metrics); + + let result = delete_file_loader + .read_delete_file(&task, schema) + .await + .unwrap(); + + let batches: Vec<_> = result.try_collect().await.unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 4); + } + + #[tokio::test] + async fn test_read_encrypted_equality_delete_file() { + use std::collections::HashMap; + use std::sync::Arc; + + use arrow_array::{Int64Array, RecordBatch}; + use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + + use crate::encryption::StandardKeyMetadata; + use crate::scan::FileScanTaskDeleteFile; + use crate::spec::DataContentType; + + let encryption_key = b"0123456789abcdef"; + let aad_prefix = b"my-table-uuid!!"; + + let tmp_dir = TempDir::new().unwrap(); + let table_location = tmp_dir.path().to_str().unwrap(); + let file_io = FileIO::new_with_fs(); + + let arrow_schema = Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("id", arrow_schema::DataType::Int64, false).with_metadata( + HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]), + ), + ])); + + let id_col = Arc::new(Int64Array::from(vec![100i64, 200, 300])); + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_col]).unwrap(); + + let del_path = format!("{table_location}/encrypted-eq-del.parquet"); + write_encrypted_parquet(&del_path, &batch, encryption_key, Some(aad_prefix)); + + let key_metadata = StandardKeyMetadata::new(encryption_key) + .with_aad_prefix(aad_prefix) + .encode() + .unwrap(); + + let schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + crate::spec::NestedField::required( + 1, + "id", + crate::spec::Type::Primitive(crate::spec::PrimitiveType::Long), + ) + .into(), + ]) + .build() + .unwrap(), + ); + + let task = FileScanTaskDeleteFile { + file_path: del_path.clone(), + file_size_in_bytes: std::fs::metadata(&del_path).unwrap().len(), + file_type: DataContentType::EqualityDeletes, + partition_spec_id: 0, + equality_ids: Some(vec![1]), + key_metadata: Some(Box::from(key_metadata.as_ref())), + }; + + let scan_metrics = ScanMetrics::new(); + let delete_file_loader = BasicDeleteFileLoader::new(file_io, scan_metrics); + + let result = delete_file_loader + .read_delete_file(&task, schema) + .await + .unwrap(); + + let batches: Vec<_> = result.try_collect().await.unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 3); + } } diff --git a/crates/iceberg/src/arrow/mod.rs b/crates/iceberg/src/arrow/mod.rs index bf53633cfc..4e0286cbde 100644 --- a/crates/iceberg/src/arrow/mod.rs +++ b/crates/iceberg/src/arrow/mod.rs @@ -44,3 +44,6 @@ pub use partition_value_calculator::*; /// Record batch partition splitter for partitioned tables pub mod record_batch_partition_splitter; pub use record_batch_partition_splitter::*; + +#[cfg(test)] +pub(crate) mod test_utils; diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index a4072854e4..f05c3a8b45 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -26,6 +26,7 @@ use std::sync::atomic::AtomicU64; use futures::{StreamExt, TryStreamExt}; use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; use parquet::arrow::{PARQUET_FIELD_ID_META_KEY, ParquetRecordBatchStreamBuilder}; +use parquet::encryption::decrypt::FileDecryptionProperties; use super::{ ArrowFileReader, ArrowReader, ParquetReadOptions, add_fallback_field_ids_to_arrow_schema, @@ -35,6 +36,7 @@ use crate::arrow::caching_delete_file_loader::CachingDeleteFileLoader; use crate::arrow::int96::coerce_int96_timestamps; use crate::arrow::record_batch_transformer::RecordBatchTransformerBuilder; use crate::arrow::scan_metrics::{CountingFileRead, ScanMetrics, ScanResult}; +use crate::encryption::StandardKeyMetadata; use crate::error::Result; use crate::io::{FileIO, FileMetadata, FileRead}; use crate::metadata_columns::{ @@ -122,6 +124,7 @@ impl FileScanTaskReader { task.file_size_in_bytes, parquet_read_options, self.scan_metrics.bytes_read_counter(), + task.key_metadata.as_deref(), ) .await?; @@ -431,6 +434,7 @@ impl ArrowReader { file_size_in_bytes: u64, parquet_read_options: ParquetReadOptions, bytes_read: &Arc, + key_metadata: Option<&[u8]>, ) -> Result<(ArrowFileReader, ArrowReaderMetadata)> { let parquet_file = file_io.new_input(data_file_path)?; let counting_reader = @@ -439,6 +443,7 @@ impl ArrowReader { Box::new(counting_reader), file_size_in_bytes, parquet_read_options, + key_metadata, ) .await } @@ -447,6 +452,7 @@ impl ArrowReader { parquet_reader: Box, file_size_in_bytes: u64, parquet_read_options: ParquetReadOptions, + key_metadata: Option<&[u8]>, ) -> Result<(ArrowFileReader, ArrowReaderMetadata)> { let mut reader = ArrowFileReader::new( FileMetadata { @@ -456,7 +462,9 @@ impl ArrowReader { ) .with_parquet_read_options(parquet_read_options); - let arrow_metadata = ArrowReaderMetadata::load_async(&mut reader, Default::default()) + let arrow_reader_options = Self::build_arrow_reader_options(key_metadata)?; + + let arrow_metadata = ArrowReaderMetadata::load_async(&mut reader, arrow_reader_options) .await .map_err(|e| { Error::new(ErrorKind::Unexpected, "Failed to load Parquet metadata").with_source(e) @@ -464,6 +472,34 @@ impl ArrowReader { Ok((reader, arrow_metadata)) } + + /// Builds `ArrowReaderOptions`, adding `FileDecryptionProperties` when + /// key metadata is present for Parquet Modular Encryption. + fn build_arrow_reader_options(key_metadata: Option<&[u8]>) -> Result { + match key_metadata { + Some(km) => { + let standard_key_metadata = StandardKeyMetadata::decode(km)?; + let mut builder = FileDecryptionProperties::builder( + standard_key_metadata.encryption_key().as_bytes().to_vec(), + ); + if let Some(aad) = standard_key_metadata.aad_prefix() { + builder = builder.with_aad_prefix(aad.to_vec()); + } + let decryption_properties = builder.build().map_err(|e| { + Error::new( + ErrorKind::Unexpected, + "Failed to build Parquet file decryption properties", + ) + .with_source(e) + })?; + Ok( + ArrowReaderOptions::new() + .with_file_decryption_properties(decryption_properties), + ) + } + None => Ok(ArrowReaderOptions::default()), + } + } } #[cfg(test)] @@ -473,7 +509,7 @@ mod tests { use std::sync::Arc; use arrow_array::cast::AsArray; - use arrow_array::{Array, ArrayRef, RecordBatch}; + use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use futures::TryStreamExt; use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY}; @@ -483,6 +519,7 @@ mod tests { use crate::Runtime; use crate::arrow::ArrowReaderBuilder; + use crate::arrow::test_utils::write_encrypted_parquet; use crate::io::FileIO; use crate::scan::{FileScanTask, FileScanTaskStream}; use crate::spec::{DataFileFormat, NestedField, PrimitiveType, Schema, SchemaRef, Type}; @@ -663,6 +700,199 @@ mod tests { } } + #[tokio::test] + async fn test_read_encrypted_parquet() { + let encryption_key = b"0123456789abcdef"; + let aad_prefix = b"aad_prefix"; + + let schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build() + .unwrap(), + ); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + + let tmp_dir = TempDir::new().unwrap(); + let table_location = tmp_dir.path().to_str().unwrap().to_string(); + let file_io = FileIO::new_with_fs(); + + let id_data = Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef; + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap(); + + let file_path = format!("{table_location}/encrypted.parquet"); + write_encrypted_parquet(&file_path, &batch, encryption_key, Some(aad_prefix)); + + let key_metadata = crate::encryption::StandardKeyMetadata::new(encryption_key) + .with_aad_prefix(aad_prefix) + .encode() + .unwrap(); + + let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build(); + + let task = FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len()) + .with_start(0) + .with_length(0) + .with_data_file_path(file_path) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema) + .with_project_field_ids(vec![1]) + .with_case_sensitive(false) + .with_key_metadata(Some(key_metadata)) + .build(); + + let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; + let batches: Vec = reader + .read(tasks) + .unwrap() + .stream() + .try_collect() + .await + .unwrap(); + + assert_eq!(batches.len(), 1); + let ids = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.values(), &[10, 20, 30]); + } + + #[tokio::test] + async fn test_read_encrypted_parquet_without_key_metadata_fails() { + let encryption_key = b"0123456789abcdef"; + + let schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build() + .unwrap(), + ); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + + let tmp_dir = TempDir::new().unwrap(); + let table_location = tmp_dir.path().to_str().unwrap().to_string(); + let file_io = FileIO::new_with_fs(); + + let id_data = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap(); + + let file_path = format!("{table_location}/encrypted_no_key.parquet"); + write_encrypted_parquet(&file_path, &batch, encryption_key, None); + + let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build(); + + let task = FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len()) + .with_start(0) + .with_length(0) + .with_data_file_path(file_path) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema) + .with_project_field_ids(vec![1]) + .with_case_sensitive(false) + .build(); + + let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; + let result: Result, _> = + reader.read(tasks).unwrap().stream().try_collect().await; + + let err = result.unwrap_err(); + assert_eq!(err.kind(), crate::ErrorKind::Unexpected); + let err_str = format!("{err}"); + assert!( + err_str.contains("encrypted footer"), + "Expected error about encrypted footer, got: {err_str}" + ); + assert!( + err_str.contains("decryption properties were not provided"), + "Expected error about missing decryption properties, got: {err_str}" + ); + } + + #[tokio::test] + async fn test_read_encrypted_parquet_with_wrong_key_fails() { + let encryption_key = b"0123456789abcdef"; + let wrong_key = b"fedcba9876543210"; + + let schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build() + .unwrap(), + ); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + "1".to_string(), + )])), + ])); + + let tmp_dir = TempDir::new().unwrap(); + let table_location = tmp_dir.path().to_str().unwrap().to_string(); + let file_io = FileIO::new_with_fs(); + + let id_data = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap(); + + let file_path = format!("{table_location}/encrypted_wrong_key.parquet"); + write_encrypted_parquet(&file_path, &batch, encryption_key, None); + + let wrong_key_metadata = crate::encryption::StandardKeyMetadata::new(wrong_key) + .encode() + .unwrap(); + + let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build(); + + let task = FileScanTask::builder() + .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len()) + .with_start(0) + .with_length(0) + .with_data_file_path(file_path) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema) + .with_project_field_ids(vec![1]) + .with_case_sensitive(false) + .with_key_metadata(Some(wrong_key_metadata)) + .build(); + + let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; + let result: Result, _> = + reader.read(tasks).unwrap().stream().try_collect().await; + + let err = result.unwrap_err(); + assert_eq!(err.kind(), crate::ErrorKind::Unexpected); + let err_str = format!("{err}"); + assert!( + err_str.contains("unable to decrypt parquet footer"), + "Expected error about decryption failure, got: {err_str}" + ); + } + /// Test that concurrency=1 reads all files correctly and in deterministic order. /// This verifies the fast-path optimization for single concurrency. #[tokio::test] diff --git a/crates/iceberg/src/arrow/reader/row_filter.rs b/crates/iceberg/src/arrow/reader/row_filter.rs index a0ad5b627c..623ef71f25 100644 --- a/crates/iceberg/src/arrow/reader/row_filter.rs +++ b/crates/iceberg/src/arrow/reader/row_filter.rs @@ -1148,6 +1148,7 @@ mod tests { partition_spec: None, name_mapping: None, case_sensitive: false, + key_metadata: None, }; let stream = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; @@ -1240,11 +1241,13 @@ mod tests { partition_spec_id: 0, equality_ids: None, file_size_in_bytes: std::fs::metadata(&pos_del_path).unwrap().len(), + key_metadata: None, }], partition: None, partition_spec: None, name_mapping: None, case_sensitive: false, + key_metadata: None, }; let stream_sub2 = diff --git a/crates/iceberg/src/arrow/test_utils.rs b/crates/iceberg/src/arrow/test_utils.rs new file mode 100644 index 0000000000..15ec0dc8da --- /dev/null +++ b/crates/iceberg/src/arrow/test_utils.rs @@ -0,0 +1,48 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fs::File; + +use arrow_array::RecordBatch; +use parquet::arrow::ArrowWriter; +use parquet::basic::Compression; +use parquet::encryption::encrypt::FileEncryptionProperties; +use parquet::file::properties::WriterProperties; + +/// Writes `batch` to `path` as a Parquet file encrypted with `key` and `aad_prefix`. +pub(crate) fn write_encrypted_parquet( + path: &str, + batch: &RecordBatch, + key: &[u8], + aad_prefix: Option<&[u8]>, +) { + let mut builder = FileEncryptionProperties::builder(key.to_vec()); + if let Some(aad) = aad_prefix { + builder = builder.with_aad_prefix(aad.to_vec()); + } + let encryption_properties = builder.build().unwrap(); + + let props = WriterProperties::builder() + .set_compression(Compression::SNAPPY) + .with_file_encryption_properties(encryption_properties) + .build(); + + let file = File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, batch.schema(), Some(props)).unwrap(); + writer.write(batch).expect("Writing batch"); + writer.close().unwrap(); +} diff --git a/crates/iceberg/src/encryption/key_metadata.rs b/crates/iceberg/src/encryption/key_metadata.rs index 4ef66ce394..9e536e4d83 100644 --- a/crates/iceberg/src/encryption/key_metadata.rs +++ b/crates/iceberg/src/encryption/key_metadata.rs @@ -20,7 +20,7 @@ use std::fmt; -use super::SensitiveBytes; +use super::{SecureKey, SensitiveBytes}; use crate::{Error, ErrorKind, Result}; /// Standard key metadata for Iceberg table encryption. @@ -97,7 +97,17 @@ impl StandardKeyMetadata { /// Decodes from Java-compatible format. pub fn decode(bytes: &[u8]) -> Result { - _serde::StandardKeyMetadataV1::decode(bytes).map(Self::from) + let metadata = _serde::StandardKeyMetadataV1::decode(bytes).map(Self::from)?; + // Validate the DEK is a usable AES key (16/24/32 bytes) up front, so a + // malformed key surfaces a clear error. + SecureKey::new(metadata.encryption_key.as_bytes()).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + "Invalid encryption key in key metadata", + ) + .with_source(e) + })?; + Ok(metadata) } } @@ -276,11 +286,33 @@ mod tests { #[test] fn test_roundtrip_without_aad() { - let metadata = StandardKeyMetadata::new(&[1, 2, 3, 4]); + let key = b"0123456789012345"; + let metadata = StandardKeyMetadata::new(key); let serialized = metadata.encode().unwrap(); let parsed = StandardKeyMetadata::decode(&serialized).unwrap(); - assert_eq!(parsed.encryption_key().as_bytes(), &[1, 2, 3, 4]); + assert_eq!(parsed.encryption_key().as_bytes(), key); assert_eq!(parsed.aad_prefix(), None); } + + #[test] + fn test_decode_rejects_invalid_key_length() { + // 24-byte (AES-192) and 32-byte (AES-256) keys are accepted. + for len in [16usize, 24, 32] { + let metadata = StandardKeyMetadata::new(&vec![0u8; len]); + let serialized = metadata.encode().unwrap(); + assert!(StandardKeyMetadata::decode(&serialized).is_ok()); + } + + for len in [0usize, 4, 15, 20, 33] { + let metadata = StandardKeyMetadata::new(&vec![0u8; len]); + let serialized = metadata.encode().unwrap(); + let err = StandardKeyMetadata::decode(&serialized).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!( + err.to_string() + .contains("Invalid encryption key in key metadata") + ); + } + } } diff --git a/crates/iceberg/src/scan/context.rs b/crates/iceberg/src/scan/context.rs index 6a6b2ebd65..8462eb89f4 100644 --- a/crates/iceberg/src/scan/context.rs +++ b/crates/iceberg/src/scan/context.rs @@ -142,6 +142,7 @@ impl ManifestEntryContext { .with_partition_spec(self.partition_spec.clone()) .with_name_mapping(self.name_mapping) .with_case_sensitive(self.case_sensitive) + .with_key_metadata(self.manifest_entry.data_file.key_metadata().map(Box::from)) .build()) } } diff --git a/crates/iceberg/src/scan/task.rs b/crates/iceberg/src/scan/task.rs index f3b556bcbf..faeac51be9 100644 --- a/crates/iceberg/src/scan/task.rs +++ b/crates/iceberg/src/scan/task.rs @@ -118,6 +118,19 @@ pub struct FileScanTask { /// Whether this scan task should treat column names as case-sensitive when binding predicates. pub case_sensitive: bool, + + /// Key metadata for encrypted data files (Parquet Modular Encryption). + /// When present, the reader uses this to build `FileDecryptionProperties`. + /// + /// Note on the trust boundary: for the standard encryption scheme this + /// carries `StandardKeyMetadata`, whose payload is the *plaintext* DEK. + /// Because `FileScanTask` derives `Serialize`, that plaintext DEK is part + /// of the serialized scan plan should these tasks ever be serialized and sent + /// over the network. + #[serde(default)] + #[serde(skip_serializing_if = "Option::is_none")] + #[builder(default)] + pub key_metadata: Option>, } impl FileScanTask { @@ -161,6 +174,13 @@ impl From<&DeleteFileContext> for FileScanTaskDeleteFile { .with_file_type(ctx.manifest_entry.content_type()) .with_partition_spec_id(ctx.partition_spec_id) .with_equality_ids(ctx.manifest_entry.data_file.equality_ids.clone()) + .with_key_metadata( + ctx.manifest_entry + .data_file + .key_metadata + .as_deref() + .map(Box::from), + ) .build() } } @@ -184,4 +204,15 @@ pub struct FileScanTaskDeleteFile { /// equality ids for equality deletes (null for anything other than equality-deletes) #[builder(default)] pub equality_ids: Option>, + + /// Key metadata for encrypted delete files (Parquet Modular Encryption). + /// When present, the reader uses this to build `FileDecryptionProperties`. + /// + /// Same plaintext-DEK trust boundary as [`FileScanTask::key_metadata`]: + /// this is serialized into the scan plan and crosses the planner -> worker + /// channel in the clear for the standard encryption scheme. + #[serde(default)] + #[serde(skip_serializing_if = "Option::is_none")] + #[builder(default)] + pub key_metadata: Option>, } From 788cdea0df61e93a70e3a8b2842851cc1ecee385 Mon Sep 17 00:00:00 2001 From: Christian Date: Thu, 16 Jul 2026 11:20:32 +0200 Subject: [PATCH 06/11] feat: Variant Support (#2188) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? Variant Support. Arrow value support is currently missing as I am unsure how we want to extend `Literal` ## What changes are included in this PR? ### Core: Variant Type - `crates/iceberg/src/spec/datatypes.rs` — new `Variant` type - `crates/iceberg/src/spec/values/literal.rs` — `Variant` literal value - `crates/iceberg/src/spec/schema/` — visitor, index, pruning, mod, id reassigner all handle `Variant` - `crates/iceberg/src/spec/table_metadata.rs` — metadata support ### Avro - `crates/iceberg/src/avro/schema.rs` — read/write `Variant` in Avro ### Arrow - `crates/iceberg/src/arrow/schema.rs` — map `Variant` to Arrow type - `crates/iceberg/src/arrow/reader.rs` — read `Variant` from Arrow - `crates/iceberg/src/arrow/value.rs` — Arrow value conversion - Minor fixes in `caching_delete_file_loader.rs` and `nan_val_cnt_visitor.rs` ### Parquet - `crates/iceberg/src/writer/file_writer/parquet_writer.rs` — write `Variant` columns ### Tests & Dev - `crates/integration_tests/tests/read_variant.rs` — new integration test for reading Variant data - `dev/spark/provision.py` — Spark provisioning to generate Variant test data ## Are these changes tested? Sure! Even integration tested :) --------- Co-authored-by: Shawn Chang --- crates/catalog/glue/src/schema.rs | 43 ++- crates/catalog/hms/src/schema.rs | 40 ++- crates/iceberg/public-api.txt | 28 ++ .../src/arrow/caching_delete_file_loader.rs | 6 +- .../iceberg/src/arrow/nan_val_cnt_visitor.rs | 6 +- crates/iceberg/src/arrow/reader/projection.rs | 97 ++++++- crates/iceberg/src/arrow/schema.rs | 215 ++++++++++++-- crates/iceberg/src/arrow/value.rs | 46 ++- crates/iceberg/src/avro/schema.rs | 27 +- crates/iceberg/src/spec/datatypes.rs | 80 +++++- crates/iceberg/src/spec/partition.rs | 27 ++ .../iceberg/src/spec/schema/id_reassigner.rs | 32 +++ crates/iceberg/src/spec/schema/index.rs | 49 ++++ crates/iceberg/src/spec/schema/mod.rs | 262 ++++++++++++++++++ .../iceberg/src/spec/schema/prune_columns.rs | 34 +++ crates/iceberg/src/spec/schema/visitor.rs | 8 + crates/iceberg/src/spec/sort.rs | 29 ++ crates/iceberg/src/spec/table_metadata.rs | 8 + crates/iceberg/src/spec/values/literal.rs | 4 + .../iceberg/src/transaction/update_schema.rs | 19 +- .../src/writer/file_writer/parquet_writer.rs | 24 +- crates/integrations/datafusion/src/schema.rs | 7 + 22 files changed, 1058 insertions(+), 33 deletions(-) diff --git a/crates/catalog/glue/src/schema.rs b/crates/catalog/glue/src/schema.rs index 864320dae4..aa617f405b 100644 --- a/crates/catalog/glue/src/schema.rs +++ b/crates/catalog/glue/src/schema.rs @@ -25,7 +25,7 @@ pub(crate) const ICEBERG_FIELD_CURRENT: &str = "iceberg.field.current"; use std::collections::HashMap; use aws_sdk_glue::types::Column; -use iceberg::spec::{PrimitiveType, SchemaVisitor, TableMetadata, visit_schema}; +use iceberg::spec::{PrimitiveType, SchemaVisitor, TableMetadata, VariantType, visit_schema}; use iceberg::{Error, ErrorKind, Result}; use crate::error::from_aws_build_error; @@ -182,20 +182,32 @@ impl SchemaVisitor for GlueSchemaBuilder { Ok(glue_type) } + + fn variant(&mut self, _v: &VariantType) -> Result { + Ok("variant".to_string()) + } } #[cfg(test)] mod tests { use iceberg::TableCreation; - use iceberg::spec::{Schema, TableMetadataBuilder}; + use iceberg::spec::{FormatVersion, Schema, TableMetadataBuilder}; use super::*; fn create_metadata(schema: Schema) -> Result { + create_metadata_with_format_version(schema, FormatVersion::V2) + } + + fn create_metadata_with_format_version( + schema: Schema, + format_version: FormatVersion, + ) -> Result { let table_creation = TableCreation::builder() .name("my_table".to_string()) .location("my_location".to_string()) .schema(schema) + .format_version(format_version) .build(); let metadata = TableMetadataBuilder::from_table_creation(table_creation)? .build()? @@ -523,4 +535,31 @@ mod tests { assert_eq!(result, expected); Ok(()) } + + #[test] + fn test_schema_with_variant() -> Result<()> { + let record = r#"{ + "type": "struct", + "schema-id": 1, + "fields": [ + { + "id": 1, + "name": "v", + "required": true, + "type": "variant" + } + ] + }"#; + + let schema = serde_json::from_str::(record)?; + // Variant requires format version 3. + let metadata = create_metadata_with_format_version(schema, FormatVersion::V3)?; + + let result = GlueSchemaBuilder::from_iceberg(&metadata)?.build(); + + let expected = vec![create_column("v", "variant", "1", false)?]; + + assert_eq!(result, expected); + Ok(()) + } } diff --git a/crates/catalog/hms/src/schema.rs b/crates/catalog/hms/src/schema.rs index c23b48719d..32e05fefd1 100644 --- a/crates/catalog/hms/src/schema.rs +++ b/crates/catalog/hms/src/schema.rs @@ -16,7 +16,7 @@ // under the License. use hive_metastore::FieldSchema; -use iceberg::spec::{PrimitiveType, Schema, SchemaVisitor, visit_schema}; +use iceberg::spec::{PrimitiveType, Schema, SchemaVisitor, VariantType, visit_schema}; use iceberg::{Error, ErrorKind, Result}; type HiveSchema = Vec; @@ -139,6 +139,12 @@ impl SchemaVisitor for HiveSchemaBuilder { Ok(hive_type) } + + fn variant(&mut self, _v: &VariantType) -> Result { + // Match iceberg-java's HiveSchemaUtil, which maps VARIANT to "unknown" + // (apache/iceberg#15964). + Ok("unknown".to_string()) + } } #[cfg(test)] @@ -456,4 +462,36 @@ mod tests { Ok(()) } + + #[test] + fn test_schema_with_variant() -> Result<()> { + // VARIANT maps to Hive "unknown", matching iceberg-java's HiveSchemaUtil + // (apache/iceberg#15964). + let record = r#"{ + "type": "struct", + "schema-id": 1, + "fields": [ + { + "id": 1, + "name": "v", + "required": true, + "type": "variant" + } + ] + }"#; + + let schema = serde_json::from_str::(record)?; + + let result = HiveSchemaBuilder::from_iceberg(&schema)?.build(); + + let expected = vec![FieldSchema { + name: Some("v".into()), + r#type: Some("unknown".into()), + comment: None, + }]; + + assert_eq!(result, expected); + + Ok(()) + } } diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 2de70bf265..fcf3cb99e6 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1704,6 +1704,7 @@ pub iceberg::spec::Type::List(iceberg::spec::ListType) pub iceberg::spec::Type::Map(iceberg::spec::MapType) pub iceberg::spec::Type::Primitive(iceberg::spec::PrimitiveType) pub iceberg::spec::Type::Struct(iceberg::spec::StructType) +pub iceberg::spec::Type::Variant(iceberg::spec::VariantType) impl iceberg::spec::Type pub fn iceberg::spec::Type::as_primitive_type(&self) -> core::option::Option<&iceberg::spec::PrimitiveType> pub fn iceberg::spec::Type::decimal(precision: u32, scale: u32) -> iceberg::Result @@ -1713,6 +1714,7 @@ pub fn iceberg::spec::Type::is_floating_type(&self) -> bool pub fn iceberg::spec::Type::is_nested(&self) -> bool pub fn iceberg::spec::Type::is_primitive(&self) -> bool pub fn iceberg::spec::Type::is_struct(&self) -> bool +pub fn iceberg::spec::Type::is_variant(&self) -> bool pub fn iceberg::spec::Type::to_struct_type(self) -> core::option::Option impl core::clone::Clone for iceberg::spec::Type pub fn iceberg::spec::Type::clone(&self) -> iceberg::spec::Type @@ -1727,6 +1729,8 @@ impl core::convert::From for iceberg::spec::Type pub fn iceberg::spec::Type::from(value: iceberg::spec::PrimitiveType) -> Self impl core::convert::From for iceberg::spec::Type pub fn iceberg::spec::Type::from(value: iceberg::spec::StructType) -> Self +impl core::convert::From for iceberg::spec::Type +pub fn iceberg::spec::Type::from(_: iceberg::spec::VariantType) -> Self impl core::fmt::Debug for iceberg::spec::Type pub fn iceberg::spec::Type::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl core::fmt::Display for iceberg::spec::Type @@ -2324,6 +2328,8 @@ impl iceberg::spec::Schema pub fn iceberg::spec::Schema::accessor_by_field_id(&self, field_id: i32) -> core::option::Option> pub fn iceberg::spec::Schema::as_struct(&self) -> &iceberg::spec::StructType pub fn iceberg::spec::Schema::builder() -> iceberg::spec::SchemaBuilder +pub fn iceberg::spec::Schema::calc_min_compatible_format(&self) -> iceberg::spec::FormatVersion +pub fn iceberg::spec::Schema::check_format_compatibility(&self, format_version: iceberg::spec::FormatVersion) -> iceberg::Result<()> pub fn iceberg::spec::Schema::field_by_alias(&self, alias: &str) -> core::option::Option<&iceberg::spec::NestedFieldRef> pub fn iceberg::spec::Schema::field_by_id(&self, field_id: i32) -> core::option::Option<&iceberg::spec::NestedFieldRef> pub fn iceberg::spec::Schema::field_by_name(&self, field_name: &str) -> core::option::Option<&iceberg::spec::NestedFieldRef> @@ -2844,6 +2850,26 @@ impl core::default::Default for iceberg::spec::UnboundPartitionSpecBuilder pub fn iceberg::spec::UnboundPartitionSpecBuilder::default() -> iceberg::spec::UnboundPartitionSpecBuilder impl core::fmt::Debug for iceberg::spec::UnboundPartitionSpecBuilder pub fn iceberg::spec::UnboundPartitionSpecBuilder::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +pub struct iceberg::spec::VariantType +impl core::clone::Clone for iceberg::spec::VariantType +pub fn iceberg::spec::VariantType::clone(&self) -> iceberg::spec::VariantType +impl core::cmp::Eq for iceberg::spec::VariantType +impl core::cmp::PartialEq for iceberg::spec::VariantType +pub fn iceberg::spec::VariantType::eq(&self, other: &iceberg::spec::VariantType) -> bool +impl core::convert::From for iceberg::spec::Type +pub fn iceberg::spec::Type::from(_: iceberg::spec::VariantType) -> Self +impl core::fmt::Debug for iceberg::spec::VariantType +pub fn iceberg::spec::VariantType::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::fmt::Display for iceberg::spec::VariantType +pub fn iceberg::spec::VariantType::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::hash::Hash for iceberg::spec::VariantType +pub fn iceberg::spec::VariantType::hash<__H: core::hash::Hasher>(&self, state: &mut __H) +impl core::marker::Copy for iceberg::spec::VariantType +impl core::marker::StructuralPartialEq for iceberg::spec::VariantType +impl serde_core::ser::Serialize for iceberg::spec::VariantType +pub fn iceberg::spec::VariantType::serialize(&self, serializer: S) -> core::result::Result<::Ok, ::Error> where S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg::spec::VariantType +pub fn iceberg::spec::VariantType::deserialize(deserializer: D) -> core::result::Result::Error> where D: serde_core::de::Deserializer<'de> pub struct iceberg::spec::ViewMetadata impl iceberg::spec::ViewMetadata pub fn iceberg::spec::ViewMetadata::current_schema(&self) -> &iceberg::spec::SchemaRef @@ -2998,6 +3024,7 @@ pub fn iceberg::spec::SchemaVisitor::map(&mut self, map: &iceberg::spec::MapType pub fn iceberg::spec::SchemaVisitor::primitive(&mut self, p: &iceberg::spec::PrimitiveType) -> iceberg::Result pub fn iceberg::spec::SchemaVisitor::schema(&mut self, schema: &iceberg::spec::Schema, value: Self::T) -> iceberg::Result pub fn iceberg::spec::SchemaVisitor::struct(&mut self, struct: &iceberg::spec::StructType, results: alloc::vec::Vec) -> iceberg::Result +pub fn iceberg::spec::SchemaVisitor::variant(&mut self, v: &iceberg::spec::VariantType) -> iceberg::Result pub trait iceberg::spec::SchemaWithPartnerVisitor

pub type iceberg::spec::SchemaWithPartnerVisitor::T pub fn iceberg::spec::SchemaWithPartnerVisitor::after_list_element(&mut self, _field: &iceberg::spec::NestedFieldRef, _partner: &P) -> iceberg::Result<()> @@ -3014,6 +3041,7 @@ pub fn iceberg::spec::SchemaWithPartnerVisitor::map(&mut self, map: &iceberg::sp pub fn iceberg::spec::SchemaWithPartnerVisitor::primitive(&mut self, p: &iceberg::spec::PrimitiveType, partner: &P) -> iceberg::Result pub fn iceberg::spec::SchemaWithPartnerVisitor::schema(&mut self, schema: &iceberg::spec::Schema, partner: &P, value: Self::T) -> iceberg::Result pub fn iceberg::spec::SchemaWithPartnerVisitor::struct(&mut self, struct: &iceberg::spec::StructType, partner: &P, results: alloc::vec::Vec) -> iceberg::Result +pub fn iceberg::spec::SchemaWithPartnerVisitor::variant(&mut self, v: &iceberg::spec::VariantType, partner: &P) -> iceberg::Result pub fn iceberg::spec::deserialize_data_file_from_json(json: &str, partition_spec_id: i32, partition_type: &iceberg::spec::StructType, schema: &iceberg::spec::Schema) -> iceberg::Result pub fn iceberg::spec::prune_columns(schema: &iceberg::spec::Schema, selected: impl core::iter::traits::collect::IntoIterator, select_full_types: bool) -> iceberg::Result pub fn iceberg::spec::read_data_files_from_avro(reader: &mut R, schema: &iceberg::spec::Schema, partition_spec_id: i32, partition_type: &iceberg::spec::StructType, version: iceberg::spec::FormatVersion) -> iceberg::Result> diff --git a/crates/iceberg/src/arrow/caching_delete_file_loader.rs b/crates/iceberg/src/arrow/caching_delete_file_loader.rs index 25704e5bc4..13cf26a23f 100644 --- a/crates/iceberg/src/arrow/caching_delete_file_loader.rs +++ b/crates/iceberg/src/arrow/caching_delete_file_loader.rs @@ -34,7 +34,7 @@ use crate::runtime::Runtime; use crate::scan::{ArrowRecordBatchStream, FileScanTaskDeleteFile}; use crate::spec::{ DataContentType, Datum, ListType, MapType, NestedField, NestedFieldRef, PartnerAccessor, - PrimitiveType, Schema, SchemaRef, SchemaWithPartnerVisitor, StructType, Type, + PrimitiveType, Schema, SchemaRef, SchemaWithPartnerVisitor, StructType, Type, VariantType, visit_schema_with_partner, }; use crate::{Error, ErrorKind, Result}; @@ -564,6 +564,10 @@ impl SchemaWithPartnerVisitor for EqDelColumnProcessor<'_> { fn primitive(&mut self, _primitive: &PrimitiveType, _partner: &ArrayRef) -> Result<()> { Ok(()) } + + fn variant(&mut self, _v: &VariantType, _partner: &ArrayRef) -> Result<()> { + Ok(()) + } } struct EqDelRecordBatchPartnerAccessor; diff --git a/crates/iceberg/src/arrow/nan_val_cnt_visitor.rs b/crates/iceberg/src/arrow/nan_val_cnt_visitor.rs index e514457887..d01f3e9e56 100644 --- a/crates/iceberg/src/arrow/nan_val_cnt_visitor.rs +++ b/crates/iceberg/src/arrow/nan_val_cnt_visitor.rs @@ -28,7 +28,7 @@ use crate::Result; use crate::arrow::{ArrowArrayAccessor, FieldMatchMode}; use crate::spec::{ ListType, MapType, NestedFieldRef, PrimitiveType, Schema, SchemaRef, SchemaWithPartnerVisitor, - StructType, visit_struct_with_partner, + StructType, VariantType, visit_struct_with_partner, }; macro_rules! cast_and_update_cnt_map { @@ -122,6 +122,10 @@ impl SchemaWithPartnerVisitor for NanValueCountVisitor { Ok(()) } + fn variant(&mut self, _v: &VariantType, _col: &ArrayRef) -> Result { + Ok(()) + } + fn after_struct_field(&mut self, field: &NestedFieldRef, partner: &ArrayRef) -> Result<()> { let field_id = field.id; count_float_nans!(partner, self, field_id); diff --git a/crates/iceberg/src/arrow/reader/projection.rs b/crates/iceberg/src/arrow/reader/projection.rs index aad3123f0a..a960814441 100644 --- a/crates/iceberg/src/arrow/reader/projection.rs +++ b/crates/iceberg/src/arrow/reader/projection.rs @@ -81,6 +81,11 @@ impl ArrowReader { Self::include_leaf_field_id(&map_type.key_field, field_ids); Self::include_leaf_field_id(&map_type.value_field, field_ids); } + // Variant projection is rejected earlier (in `get_arrow_projection_mask`); this + // arm only keeps the match exhaustive. Treat it as a leaf, like a primitive. + Type::Variant(_) => { + field_ids.push(field.id); + } } } @@ -119,6 +124,19 @@ impl ArrowReader { return Ok(ProjectionMask::all()); } + // Reading variant columns is not supported yet (see #2188 follow-ups): reject any + // projection that touches a variant, rather than returning a partial/incorrect batch. + for field_id in field_ids { + if let Some(field) = iceberg_schema_of_task.field_by_id(*field_id) + && type_contains_variant(&field.field_type) + { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "Reading variant columns is not supported yet", + )); + } + } + if use_fallback { // Position-based projection necessary because file lacks embedded field IDs Self::get_arrow_projection_mask_fallback(field_ids, parquet_schema) @@ -155,6 +173,8 @@ impl ArrowReader { ) -> Result { let mut column_map = HashMap::new(); let fields = arrow_schema.fields(); + // HashSet for O(1) membership checks instead of O(n) slice scans. + let leaf_field_id_set: HashSet = leaf_field_ids.iter().copied().collect(); // Pre-project only the fields that have been selected, possibly avoiding converting // some Arrow types that are not yet supported. @@ -166,7 +186,7 @@ impl ArrowReader { .and_then(|field_id| i32::from_str(field_id).ok()) .is_some_and(|field_id| { projected_fields.insert((*f).clone(), field_id); - leaf_field_ids.contains(&field_id) + leaf_field_id_set.contains(&field_id) }) }), arrow_schema.metadata().clone(), @@ -245,6 +265,23 @@ impl ArrowReader { } } +/// Whether `field_type` is, or transitively contains, a variant type. +fn type_contains_variant(field_type: &Type) -> bool { + match field_type { + Type::Variant(_) => true, + Type::Struct(s) => s + .fields() + .iter() + .any(|f| type_contains_variant(&f.field_type)), + Type::List(l) => type_contains_variant(&l.element_field.field_type), + Type::Map(m) => { + type_contains_variant(&m.key_field.field_type) + || type_contains_variant(&m.value_field.field_type) + } + Type::Primitive(_) => false, + } +} + /// Build the map of parquet field id to Parquet column index in the schema. /// Returns None if the Parquet file doesn't have field IDs embedded (e.g., migrated tables). pub(super) fn build_field_id_map( @@ -465,7 +502,8 @@ mod tests { use crate::io::FileIO; use crate::scan::{FileScanTask, FileScanTaskStream}; use crate::spec::{ - DataFileFormat, Datum, MappedField, NameMapping, NestedField, PrimitiveType, Schema, Type, + DataFileFormat, Datum, MappedField, NameMapping, NestedField, PrimitiveType, Schema, + StructType, Type, VariantType, }; use crate::{ErrorKind, Runtime}; @@ -561,6 +599,61 @@ message schema { assert_eq!(mask, ProjectionMask::leaves(&parquet_schema, vec![0])); } + #[test] + fn test_arrow_projection_mask_variant_is_unsupported() { + // Reading variant columns is not supported yet: projecting one (top-level or + // nested) must fail loudly rather than return a partial/incorrect batch. + let schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::optional(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(2, "v", Type::Variant(VariantType)).into(), + NestedField::required( + 3, + "s", + Type::Struct(StructType::new(vec![ + NestedField::optional(4, "vv", Type::Variant(VariantType)).into(), + ])), + ) + .into(), + NestedField::required( + 5, + "m", + Type::Map(crate::spec::MapType::required( + 6, + Type::Primitive(PrimitiveType::String), + 7, + Type::Variant(VariantType), + )), + ) + .into(), + ]) + .build() + .unwrap(), + ); + // The parquet/arrow schemas are irrelevant: the variant is rejected before they + // are consulted, so an empty descriptor is enough to drive the code path. + let parquet_schema = SchemaDescriptor::new(Arc::new( + parse_message_type("message schema { optional int32 id = 1; }").unwrap(), + )); + let arrow_schema = Arc::new(ArrowSchema::empty()); + + // 2 = top-level variant, 3 = struct containing a variant, 4 = the nested variant, + // 5 = map, plus a mix with a non-variant sibling. + for projected in [vec![2], vec![3], vec![4], vec![5], vec![1, 2]] { + let err = ArrowReader::get_arrow_projection_mask( + &projected, + &schema, + &parquet_schema, + &arrow_schema, + false, + ) + .expect_err("variant projection must be rejected"); + assert_eq!(err.kind(), ErrorKind::FeatureUnsupported, "{err}"); + } + } + /// Test schema evolution: reading old Parquet file (with only column 'a') /// using a newer table schema (with columns 'a' and 'b'). /// This tests that: diff --git a/crates/iceberg/src/arrow/schema.rs b/crates/iceberg/src/arrow/schema.rs index e6649d926d..f47176bc95 100644 --- a/crates/iceberg/src/arrow/schema.rs +++ b/crates/iceberg/src/arrow/schema.rs @@ -26,7 +26,10 @@ use arrow_array::{ FixedSizeBinaryArray, Float32Array, Float64Array, Int32Array, Int64Array, Scalar, StringArray, TimestampMicrosecondArray, TimestampNanosecondArray, }; -use arrow_schema::{DataType, Field, FieldRef, Fields, Schema as ArrowSchema, TimeUnit}; +use arrow_schema::extension::ExtensionType; +use arrow_schema::{ + ArrowError, DataType, Field, FieldRef, Fields, Schema as ArrowSchema, TimeUnit, +}; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use parquet::file::statistics::Statistics; use uuid::Uuid; @@ -35,7 +38,7 @@ use crate::error::Result; use crate::spec::decimal_utils::i128_from_be_bytes; use crate::spec::{ Datum, FIRST_FIELD_ID, ListType, MapType, NestedField, NestedFieldRef, PrimitiveLiteral, - PrimitiveType, Schema, SchemaVisitor, StructType, Type, + PrimitiveType, Schema, SchemaVisitor, StructType, Type, VariantType, }; use crate::{Error, ErrorKind}; @@ -44,6 +47,59 @@ pub const DEFAULT_MAP_FIELD_NAME: &str = "key_value"; /// UTC time zone for Arrow timestamp type. pub const UTC_TIME_ZONE: &str = "+00:00"; +/// The canonical Arrow [`arrow.parquet.variant`] extension type. +/// +/// Iceberg stores a Variant as a `Struct { metadata: Binary, value: Binary }`. Attaching this +/// extension type to the enclosing field marks that struct as a single logical Variant value, +/// so Arrow consumers treat it as a Variant rather than an anonymous struct. It carries no +/// metadata. +/// +/// [`arrow.parquet.variant`]: https://arrow.apache.org/docs/format/CanonicalExtensions.html#parquet-variant +#[derive(Debug, Clone, Copy, Default)] +pub(crate) struct VariantExtensionType; + +impl ExtensionType for VariantExtensionType { + const NAME: &'static str = "arrow.parquet.variant"; + + type Metadata = (); + + fn metadata(&self) -> &Self::Metadata { + &() + } + + fn serialize_metadata(&self) -> Option { + None + } + + fn deserialize_metadata( + metadata: Option<&str>, + ) -> std::result::Result { + match metadata { + None | Some("") => Ok(()), + Some(other) => Err(ArrowError::InvalidArgumentError(format!( + "arrow.parquet.variant extension type takes no metadata, got {other:?}" + ))), + } + } + + fn supports_data_type(&self, data_type: &DataType) -> std::result::Result<(), ArrowError> { + match data_type { + DataType::Struct(_) => Ok(()), + other => Err(ArrowError::InvalidArgumentError(format!( + "arrow.parquet.variant extension type requires a Struct storage type, got {other}" + ))), + } + } + + fn try_new( + data_type: &DataType, + _metadata: Self::Metadata, + ) -> std::result::Result { + Self.supports_data_type(data_type)?; + Ok(Self) + } +} + /// A post order arrow schema visitor. /// /// For order of methods called, please refer to [`visit_schema`]. @@ -532,9 +588,16 @@ impl SchemaVisitor for ToArrowSchemaConverter { } else { HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), field.id.to_string())]) }; - Ok(ArrowSchemaOrFieldOrType::Field( - Field::new(field.name.clone(), ty, !field.required).with_metadata(metadata), - )) + let arrow_field = + Field::new(field.name.clone(), ty, !field.required).with_metadata(metadata); + // A variant column's storage is a struct; tag the field with the canonical + // `arrow.parquet.variant` extension type so consumers read it as a Variant, not a struct. + let arrow_field = if field.field_type.is_variant() { + arrow_field.with_extension_type(VariantExtensionType) + } else { + arrow_field + }; + Ok(ArrowSchemaOrFieldOrType::Field(arrow_field)) } fn r#struct( @@ -557,25 +620,13 @@ impl SchemaVisitor for ToArrowSchemaConverter { list: &crate::spec::ListType, value: ArrowSchemaOrFieldOrType, ) -> crate::Result { + // `field` already carries the element's field id, doc, and — for a variant element — + // the arrow.parquet.variant extension type. Don't overwrite its metadata here (doing so + // would drop the extension type for `list`). let field = match self.field(&list.element_field, value)? { ArrowSchemaOrFieldOrType::Field(field) => field, _ => unreachable!(), }; - let meta = if let Some(doc) = &list.element_field.doc { - HashMap::from([ - ( - PARQUET_FIELD_ID_META_KEY.to_string(), - list.element_field.id.to_string(), - ), - (ARROW_FIELD_DOC_KEY.to_string(), doc.clone()), - ]) - } else { - HashMap::from([( - PARQUET_FIELD_ID_META_KEY.to_string(), - list.element_field.id.to_string(), - )]) - }; - let field = field.with_metadata(meta); Ok(ArrowSchemaOrFieldOrType::Type(DataType::List(Arc::new( field, )))) @@ -692,6 +743,19 @@ impl SchemaVisitor for ToArrowSchemaConverter { } } } + + fn variant(&mut self, _v: &VariantType) -> crate::Result { + // Variant is stored as a struct of two binary sub-fields (no field IDs on sub-fields). + // Uses Binary (not LargeBinary) matching the Parquet BINARY primitive directly. + // `metadata` is always present; `value` is nullable, since in a shredded variant the + // value may be absent. The enclosing field carries the `arrow.parquet.variant` extension type + // (attached in `field`). + let metadata_field = Field::new("metadata", DataType::Binary, false); + let value_field = Field::new("value", DataType::Binary, true); + Ok(ArrowSchemaOrFieldOrType::Type(DataType::Struct( + vec![metadata_field, value_field].into(), + ))) + } } /// Convert iceberg schema to an arrow schema. @@ -1705,6 +1769,21 @@ mod tests { simple_field("map", map, false, "16"), simple_field("struct", r#struct, false, "17"), simple_field("uuid", DataType::FixedSizeBinary(16), false, "30"), + Field::new( + "v", + DataType::Struct(Fields::from(vec![ + Field::new("metadata", DataType::Binary, false), + Field::new("value", DataType::Binary, true), + ])), + true, + ) + .with_metadata(HashMap::from([ + (PARQUET_FIELD_ID_META_KEY.to_string(), "31".to_string()), + ( + arrow_schema::extension::EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.parquet.variant".to_string(), + ), + ])), ]) } @@ -1888,6 +1967,12 @@ mod tests { "name":"uuid", "required":true, "type":"uuid" + }, + { + "id":31, + "name":"v", + "required":false, + "type":"variant" } ], "identifier-field-ids":[] @@ -1905,6 +1990,96 @@ mod tests { assert_eq!(converted_arrow_schema, arrow_schema); } + #[test] + fn test_variant_type_to_arrow_type() { + // Variant maps to a struct with a required `metadata` and a nullable `value` binary + // field, with no field ids on the sub-fields, matching the Parquet BINARY layout. + let arrow_type = type_to_arrow_type(&Type::Variant(VariantType)).unwrap(); + assert_eq!( + arrow_type, + DataType::Struct(Fields::from(vec![ + Field::new("metadata", DataType::Binary, false), + Field::new("value", DataType::Binary, true), + ])) + ); + } + + #[test] + fn test_variant_field_carries_arrow_extension_type() { + // Converting a schema with a variant column tags the column's field with the + // canonical `arrow.parquet.variant` extension type (the struct storage stays as-is). + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "v", Type::Variant(VariantType)).into(), + ]) + .build() + .unwrap(); + + let arrow_schema = schema_to_arrow_schema(&schema).unwrap(); + let field = arrow_schema.field_with_name("v").unwrap(); + + assert_eq!(field.extension_type_name(), Some("arrow.parquet.variant")); + // Attaching the extension type must not clobber the Iceberg field id. + assert_eq!( + field.metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&"1".to_string()) + ); + assert_eq!( + field.data_type(), + &DataType::Struct(Fields::from(vec![ + Field::new("metadata", DataType::Binary, false), + Field::new("value", DataType::Binary, true), + ])) + ); + } + + #[test] + fn test_variant_nested_in_list_and_map_carries_arrow_extension_type() { + // A variant nested in a list element or map value keeps the arrow.parquet.variant + // extension type. Regression guard: the list converter must not overwrite the + // element field's metadata (which would drop the extension type). + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional( + 1, + "l", + Type::List(ListType::new( + NestedField::optional(2, "element", Type::Variant(VariantType)).into(), + )), + ) + .into(), + NestedField::optional( + 3, + "m", + Type::Map(MapType::new( + NestedField::map_key_element(4, Type::Primitive(PrimitiveType::String)) + .into(), + NestedField::map_value_element(5, Type::Variant(VariantType), false).into(), + )), + ) + .into(), + ]) + .build() + .unwrap(); + + let arrow_schema = schema_to_arrow_schema(&schema).unwrap(); + + let DataType::List(element) = arrow_schema.field_with_name("l").unwrap().data_type() else { + panic!("expected a list"); + }; + assert_eq!(element.extension_type_name(), Some("arrow.parquet.variant")); + + let DataType::Map(entries, _) = arrow_schema.field_with_name("m").unwrap().data_type() + else { + panic!("expected a map"); + }; + let DataType::Struct(kv) = entries.data_type() else { + panic!("expected a key_value struct"); + }; + let value = kv.iter().find(|f| f.name() == "value").unwrap(); + assert_eq!(value.extension_type_name(), Some("arrow.parquet.variant")); + } + #[test] fn test_type_conversion() { // test primitive type diff --git a/crates/iceberg/src/arrow/value.rs b/crates/iceberg/src/arrow/value.rs index c2c0d650d8..2115248e18 100644 --- a/crates/iceberg/src/arrow/value.rs +++ b/crates/iceberg/src/arrow/value.rs @@ -30,7 +30,7 @@ use uuid::Uuid; use super::get_field_id_from_metadata; use crate::spec::{ ListType, Literal, Map, MapType, NestedField, PartnerAccessor, PrimitiveLiteral, PrimitiveType, - SchemaWithPartnerVisitor, Struct, StructType, Type, visit_struct_with_partner, + SchemaWithPartnerVisitor, Struct, StructType, Type, VariantType, visit_struct_with_partner, visit_type_with_partner, }; use crate::{Error, ErrorKind, Result}; @@ -426,6 +426,13 @@ impl SchemaWithPartnerVisitor for ArrowArrayToIcebergStructConverter { } } } + + fn variant(&mut self, _v: &VariantType, _partner: &ArrayRef) -> Result>> { + Err(Error::new( + ErrorKind::FeatureUnsupported, + "Converting variant Arrow array to Iceberg literal is not supported yet", + )) + } } /// Defines how Arrow fields are matched with Iceberg fields when converting data. @@ -1324,6 +1331,43 @@ mod test { assert_eq!(result, vec![None; 0]); } + #[test] + fn test_arrow_variant_to_literal_is_unsupported() { + // Converting a variant Arrow array back to an Iceberg literal is not implemented; + // the visitor must reject it rather than silently decode it incorrectly. + let variant_child = Arc::new(StructArray::from(vec![ + ( + Arc::new(Field::new("metadata", DataType::Binary, false)), + Arc::new(BinaryArray::from(vec![Some(b"m".as_ref())])) as ArrayRef, + ), + ( + Arc::new(Field::new("value", DataType::Binary, false)), + Arc::new(BinaryArray::from(vec![Some(b"v".as_ref())])) as ArrayRef, + ), + ])) as ArrayRef; + + let struct_array = Arc::new(StructArray::from(vec![( + Arc::new( + Field::new("v", variant_child.data_type().clone(), false).with_metadata( + HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]), + ), + ), + variant_child, + )])) as ArrayRef; + + let ty = StructType::new(vec![ + NestedField::required(1, "v", Type::Variant(VariantType)).into(), + ]); + + let err = arrow_struct_to_literal(&struct_array, &ty).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::FeatureUnsupported); + assert!( + err.to_string() + .contains("Converting variant Arrow array to Iceberg literal is not supported yet"), + "{err}" + ); + } + #[test] fn test_find_field_by_id() { // Create Arrow arrays for the nested structure diff --git a/crates/iceberg/src/avro/schema.rs b/crates/iceberg/src/avro/schema.rs index fdbc680977..2ef1e74722 100644 --- a/crates/iceberg/src/avro/schema.rs +++ b/crates/iceberg/src/avro/schema.rs @@ -28,7 +28,7 @@ use serde_json::{Number, Value}; use crate::spec::{ ListType, MapType, NestedField, NestedFieldRef, PrimitiveType, Schema, SchemaVisitor, - StructType, Type, visit_schema, + StructType, Type, VariantType, visit_schema, }; use crate::{Error, ErrorKind, Result, ensure_data_valid}; @@ -243,6 +243,13 @@ impl SchemaVisitor for SchemaToAvroSchema { }; Ok(Either::Left(avro_schema)) } + + fn variant(&mut self, _v: &VariantType) -> Result { + Err(Error::new( + ErrorKind::FeatureUnsupported, + "Converting a variant type to an Avro schema is not supported yet", + )) + } } /// Converting iceberg schema to avro schema. @@ -614,7 +621,9 @@ mod tests { use super::*; use crate::avro::schema::AvroSchemaToSchema; - use crate::spec::{ListType, MapType, NestedField, PrimitiveType, Schema, StructType, Type}; + use crate::spec::{ + ListType, MapType, NestedField, PrimitiveType, Schema, StructType, Type, VariantType, + }; fn read_test_data_file_to_avro_schema(filename: &str) -> AvroSchema { let input = read_to_string(format!( @@ -1212,4 +1221,18 @@ mod tests { converter.primitive(&avro_schema).unwrap().unwrap() ); } + + #[test] + fn test_variant_to_avro_schema_is_unsupported() { + // Converting a variant to an Avro schema is not supported yet; it must error rather + // than emit an incorrect schema. + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "v", Type::Variant(VariantType)).into(), + ]) + .build() + .unwrap(); + let err = schema_to_avro_schema("t", &schema).unwrap_err(); + assert_eq!(err.kind(), crate::ErrorKind::FeatureUnsupported, "{err}"); + } } diff --git a/crates/iceberg/src/spec/datatypes.rs b/crates/iceberg/src/spec/datatypes.rs index 3a112e9f5c..1b45217d8a 100644 --- a/crates/iceberg/src/spec/datatypes.rs +++ b/crates/iceberg/src/spec/datatypes.rs @@ -32,8 +32,8 @@ use serde_json::Value as JsonValue; use super::values::Literal; use crate::ensure_data_valid; use crate::error::Result; -use crate::spec::PrimitiveLiteral; use crate::spec::datatypes::_decimal::{MAX_PRECISION, REQUIRED_LENGTH}; +use crate::spec::{FormatVersion, PrimitiveLiteral}; /// Field name for list type. pub const LIST_FIELD_NAME: &str = "element"; @@ -90,6 +90,8 @@ pub enum Type { List(ListType), /// Map type Map(MapType), + /// Variant Type + Variant(VariantType), } impl fmt::Display for Type { @@ -99,6 +101,7 @@ impl fmt::Display for Type { Type::Struct(s) => write!(f, "{s}"), Type::List(_) => write!(f, "list"), Type::Map(_) => write!(f, "map"), + Type::Variant(_) => write!(f, "variant"), } } } @@ -122,6 +125,30 @@ impl Type { matches!(self, Type::Struct(_) | Type::List(_) | Type::Map(_)) } + /// Whether the type is variant type. + #[inline(always)] + pub fn is_variant(&self) -> bool { + matches!(self, Type::Variant(_)) + } + + /// Minimum [`FormatVersion`] required to support this type, **without** taking + /// nested field types into account. + /// + /// `TimestampNs` / `TimestamptzNs` / `Variant` require [`FormatVersion::V3`]; every + /// other type is valid from [`FormatVersion::V1`]. Mirrors Java's + /// `Schema.MIN_FORMAT_VERSIONS` (a shallow lookup keyed by type id), so it + /// intentionally does not recurse: callers needing the floor for a whole schema + /// iterate its flattened fields (see [`Schema::calc_min_compatible_format`]). + /// + /// [`Schema::calc_min_compatible_format`]: crate::spec::Schema::calc_min_compatible_format + pub(crate) fn min_format_version(&self) -> FormatVersion { + match self { + Type::Primitive(PrimitiveType::TimestampNs | PrimitiveType::TimestamptzNs) + | Type::Variant(_) => FormatVersion::V3, + _ => FormatVersion::V1, + } + } + /// Convert Type to reference of PrimitiveType pub fn as_primitive_type(&self) -> Option<&PrimitiveType> { if let Type::Primitive(primitive_type) = self { @@ -710,6 +737,7 @@ pub(super) mod _serde { use crate::spec::datatypes::Type::Map; use crate::spec::datatypes::{ ListType, MapType, NestedField, NestedFieldRef, PrimitiveType, StructType, Type, + VariantType, }; /// List type for serialization and deserialization @@ -737,6 +765,7 @@ pub(super) mod _serde { value: Cow<'a, Type>, }, Primitive(PrimitiveType), + Variant(VariantType), } impl From> for Type { @@ -775,6 +804,7 @@ pub(super) mod _serde { Self::Struct(StructType::new(fields.into_owned())) } SerdeType::Primitive(p) => Self::Primitive(p), + SerdeType::Variant(v) => Self::Variant(v), } } } @@ -801,6 +831,7 @@ pub(super) mod _serde { fields: Cow::Borrowed(&s.fields), }, Type::Primitive(p) => SerdeType::Primitive(p.clone()), + Type::Variant(v) => SerdeType::Variant(*v), } } } @@ -844,6 +875,42 @@ impl MapType { } } +/// Variant type - can hold semi-structured data of any type. +/// This is an Iceberg V3 feature. +#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)] +pub struct VariantType; + +impl fmt::Display for VariantType { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "variant") + } +} + +impl From for Type { + fn from(_: VariantType) -> Self { + Type::Variant(VariantType) + } +} + +impl Serialize for VariantType { + fn serialize(&self, serializer: S) -> std::result::Result + where S: Serializer { + serializer.serialize_str("variant") + } +} + +impl<'de> Deserialize<'de> for VariantType { + fn deserialize(deserializer: D) -> std::result::Result + where D: Deserializer<'de> { + let s = String::deserialize(deserializer)?; + if s == "variant" { + Ok(VariantType) + } else { + Err(D::Error::custom(format!("expected 'variant', got '{s}'"))) + } + } +} + #[cfg(test)] mod tests { use pretty_assertions::assert_eq; @@ -1276,6 +1343,17 @@ mod tests { } } + #[test] + fn variant_type_serde() { + let json = r#"{"id": 1, "name": "v", "required": true, "type": "variant"}"#; + let field: NestedField = serde_json::from_str(json).unwrap(); + assert_eq!(*field.field_type, Type::Variant(VariantType)); + + let serialized = serde_json::to_string(&field).unwrap(); + let roundtrip: NestedField = serde_json::from_str(&serialized).unwrap(); + assert_eq!(field, roundtrip); + } + #[test] fn struct_type_with_type_field() { // Test that StructType properly deserializes JSON with "type":"struct" field diff --git a/crates/iceberg/src/spec/partition.rs b/crates/iceberg/src/spec/partition.rs index 43b56dcdaa..b022b9af1b 100644 --- a/crates/iceberg/src/spec/partition.rs +++ b/crates/iceberg/src/spec/partition.rs @@ -1432,6 +1432,33 @@ mod tests { .unwrap_err(); } + #[test] + fn test_builder_disallows_variant_source() { + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(crate::spec::PrimitiveType::Int)) + .into(), + NestedField::optional(2, "v", Type::Variant(crate::spec::VariantType)).into(), + ]) + .build() + .unwrap(); + + let err = PartitionSpec::builder(schema) + .with_spec_id(1) + .add_unbound_fields(vec![UnboundPartitionField { + source_id: 2, + field_id: None, + name: "v_part".to_string(), + transform: Transform::Identity, + }]) + .expect_err("variant must not be allowed as a partition source"); + + assert_eq!( + err.message(), + "Cannot partition by non-primitive source field: 'variant'." + ); + } + #[test] fn test_builder_disallows_redundant() { let err = UnboundPartitionSpec::builder() diff --git a/crates/iceberg/src/spec/schema/id_reassigner.rs b/crates/iceberg/src/spec/schema/id_reassigner.rs index 5dbb370001..d817ccbd1f 100644 --- a/crates/iceberg/src/spec/schema/id_reassigner.rs +++ b/crates/iceberg/src/spec/schema/id_reassigner.rs @@ -102,6 +102,7 @@ impl ReassignFieldIds { value_field: Arc::new(value_field), })) } + Type::Variant(v) => Ok(Type::Variant(v)), } } @@ -192,6 +193,37 @@ mod tests { assert_eq!(reassigned_schema.highest_field_id(), 2); } + #[test] + fn test_reassign_ids_variant() { + use crate::spec::VariantType; + + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(5, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(3, "data", Type::Variant(VariantType)).into(), + ]) + .build() + .unwrap(); + + let reassigned = schema + .into_builder() + .with_reassigned_field_ids(0) + .build() + .unwrap(); + + // Variant has no sub-fields, so it survives reassignment unchanged; only the + // top-level field ids shift (id → 0, data → 1). + let expected = Schema::builder() + .with_fields(vec![ + NestedField::required(0, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(1, "data", Type::Variant(VariantType)).into(), + ]) + .build() + .unwrap(); + + pretty_assertions::assert_eq!(expected, reassigned); + } + #[test] fn test_reassigned_ids_nested() { let schema = table_schema_nested(); diff --git a/crates/iceberg/src/spec/schema/index.rs b/crates/iceberg/src/spec/schema/index.rs index d4e77ab2aa..e4358e9ef9 100644 --- a/crates/iceberg/src/spec/schema/index.rs +++ b/crates/iceberg/src/spec/schema/index.rs @@ -17,6 +17,7 @@ use super::utils::try_insert_field; use super::*; +use crate::spec::VariantType; /// Creates a field id to field map. pub fn index_by_id(r#struct: &StructType) -> Result> { @@ -53,6 +54,10 @@ pub fn index_by_id(r#struct: &StructType) -> Result fn primitive(&mut self, _: &PrimitiveType) -> Result { Ok(()) } + + fn variant(&mut self, _v: &VariantType) -> Result { + Ok(()) + } } let mut index = IndexById(HashMap::new()); @@ -145,6 +150,10 @@ pub fn index_parents(r#struct: &StructType) -> Result> { fn primitive(&mut self, _p: &PrimitiveType) -> Result { Ok(()) } + + fn variant(&mut self, _v: &VariantType) -> Result { + Ok(()) + } } let mut index = IndexByParent { @@ -293,6 +302,10 @@ impl SchemaVisitor for IndexByName { fn primitive(&mut self, _p: &PrimitiveType) -> Result { Ok(()) } + + fn variant(&mut self, _v: &VariantType) -> Result { + Ok(()) + } } #[cfg(test)] @@ -315,4 +328,40 @@ mod tests { assert_eq!(result.get(&16).unwrap(), &15); assert_eq!(result.get(&17).unwrap(), &15); } + + #[test] + fn test_index_variant_by_id_and_name() { + // A variant is indexed like a leaf: by id and by full (dotted) name, both at the + // top level and nested inside a struct. Mirrors Java's TestTypeUtil index-by-id / + // index-name-by-id coverage for variant. + let s = StructType::new(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(2, "v", Type::Variant(VariantType)).into(), + NestedField::required( + 3, + "nested", + Type::Struct(StructType::new(vec![ + NestedField::optional(4, "inner_v", Type::Variant(VariantType)).into(), + ])), + ) + .into(), + ]); + + let by_id = index_by_id(&s).unwrap(); + // All four fields (id, v, nested, nested.inner_v) are indexed exactly once. + assert_eq!(by_id.len(), 4); + assert!(by_id[&2].field_type.is_variant()); + assert!(by_id[&4].field_type.is_variant()); + + let (name_to_id, id_to_name) = { + let mut index = IndexByName::default(); + visit_struct(&s, &mut index).unwrap(); + index.indexes() + }; + assert_eq!(id_to_name.len(), 4); + assert_eq!(id_to_name[&2], "v"); + assert_eq!(id_to_name[&4], "nested.inner_v"); + assert_eq!(name_to_id["v"], 2); + assert_eq!(name_to_id["nested.inner_v"], 4); + } } diff --git a/crates/iceberg/src/spec/schema/mod.rs b/crates/iceberg/src/spec/schema/mod.rs index 9109990e19..ac66c2bd72 100644 --- a/crates/iceberg/src/spec/schema/mod.rs +++ b/crates/iceberg/src/spec/schema/mod.rs @@ -39,6 +39,7 @@ pub use self::prune_columns::prune_columns; use super::NestedField; use crate::error::Result; use crate::expr::accessor::StructAccessor; +use crate::spec::FormatVersion; use crate::spec::datatypes::{ LIST_FIELD_NAME, ListType, MAP_KEY_FIELD_NAME, MAP_VALUE_FIELD_NAME, MapType, NestedFieldRef, PrimitiveType, StructType, Type, @@ -53,6 +54,9 @@ pub type SchemaRef = Arc; pub const DEFAULT_SCHEMA_ID: SchemaId = 0; /// Delimiter for schema name, which denotes a nested struct. pub const SCHEMA_NAME_DELIMITER: &str = "."; +/// Minimum format version that allows non-null field default values. +/// Mirrors Java's `Schema.DEFAULT_VALUES_MIN_FORMAT_VERSION`. +pub(crate) const DEFAULT_VALUES_MIN_FORMAT_VERSION: FormatVersion = FormatVersion::V3; /// Defines schema in iceberg. #[derive(Debug, Serialize, Deserialize, Clone)] @@ -421,6 +425,90 @@ impl Schema { pub fn field_id_to_fields(&self) -> &HashMap { &self.id_to_field } + + /// Minimum [`FormatVersion`] required to represent all *types* in this schema. + /// + /// Walks the whole schema — every field, nested ones included — and returns the + /// highest per-type floor, so this is O(fields) rather than a cheap field read. + /// Types only; for initial-default version floors see [`Schema::check_format_compatibility`]. + pub fn calc_min_compatible_format(&self) -> FormatVersion { + // `id_to_field` is flattened, so the max over all fields covers nested ones too. + self.id_to_field + .values() + .map(|f| f.field_type.min_format_version()) + .max() + .unwrap_or(FormatVersion::V1) + } + + /// Returns an error listing every field incompatible with `format_version`. + /// Mirrors Java's `Schema.checkCompatibility()`. Two checks per field: + /// + /// - **Type** — minimum format version required to support that type, without + /// taking nested field types into account. + /// - **Initial default** — a non-null `initial_default` backfills pre-existing rows, + /// so it requires `DEFAULT_VALUES_MIN_FORMAT_VERSION`; `write_default` is not + /// checked, as it only affects newly written rows (read identically at any version). + pub fn check_format_compatibility(&self, format_version: FormatVersion) -> Result<()> { + // (field id, message); sorted by id below for a deterministic error. + let mut problems: Vec<(i32, String)> = Vec::new(); + + // `id_to_field` is flattened, so checking each field by its own type keeps the + // blame on the offending leaf, not its container (mirrors Java's `lazyIdToField`). + for field in self.id_to_field.values() { + let min_version = field.field_type.min_format_version(); + if format_version < min_version { + // Every id in `id_to_field` is also indexed in `id_to_name`; a miss means + // the schema's indexes are inconsistent (a bug), so surface it rather than + // guessing an unqualified name. + let name = self.name_by_field_id(field.id).ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!( + "Field id {} is missing from the schema's name index", + field.id + ), + ) + })?; + problems.push((field.id, format!( + "Invalid type for {name}: {} is not supported until {min_version} but format version is {format_version}.", + field.field_type, + ))); + } + + if let Some(default) = &field.initial_default + && format_version < DEFAULT_VALUES_MIN_FORMAT_VERSION + { + let name = self.name_by_field_id(field.id).ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!( + "Field id {} is missing from the schema's name index", + field.id + ), + ) + })?; + problems.push((field.id, format!( + "Invalid initial default for {name}: non-null default ({default:?}) is not supported until {DEFAULT_VALUES_MIN_FORMAT_VERSION} but format version is {format_version}." + ))); + } + } + + if problems.is_empty() { + return Ok(()); + } + + // Stable sort by id: HashMap order is nondeterministic, and stability keeps a + // field's type problem before its default problem (matches Java's TreeMap order). + let message = problems + .into_iter() + .sorted_by_key(|(id, _)| *id) + .map(|(_, msg)| msg) + .join("\n- "); + Err(Error::new( + ErrorKind::DataInvalid, + format!("Invalid schema for {format_version}:\n- {message}"), + )) + } } impl Display for Schema { @@ -447,6 +535,180 @@ mod tests { use crate::spec::values::Map as MapValue; use crate::spec::{Datum, Literal}; + #[test] + fn test_check_format_compatibility() { + use crate::spec::{FormatVersion, PrimitiveLiteral, VariantType}; + + fn schema_with(fields: Vec) -> Schema { + Schema::builder().with_fields(fields).build().unwrap() + } + + // Variant type requires v3. + let variant = schema_with(vec![ + NestedField::optional(1, "v", Type::Variant(VariantType)).into(), + ]); + assert!( + variant + .check_format_compatibility(FormatVersion::V2) + .is_err() + ); + assert!( + variant + .check_format_compatibility(FormatVersion::V3) + .is_ok() + ); + + // A non-null initial default requires v3, even for a v1-compatible type. + let with_default = schema_with(vec![ + NestedField::optional(1, "a", Type::Primitive(PrimitiveType::Int)) + .with_initial_default(Literal::Primitive(PrimitiveLiteral::Int(1))) + .into(), + ]); + let err = with_default + .check_format_compatibility(FormatVersion::V2) + .unwrap_err(); + assert!( + err.message().contains("Invalid initial default for a"), + "{err}" + ); + assert!( + with_default + .check_format_compatibility(FormatVersion::V3) + .is_ok() + ); + + // No default is fine at any version (an absent/null default never trips). + let no_default = schema_with(vec![ + NestedField::optional(1, "a", Type::Primitive(PrimitiveType::Int)).into(), + ]); + assert!( + no_default + .check_format_compatibility(FormatVersion::V1) + .is_ok() + ); + + // Recursion: a default on a field nested inside a struct is caught. + let nested = schema_with(vec![ + NestedField::required( + 1, + "s", + Type::Struct(StructType::new(vec![ + NestedField::optional(2, "inner", Type::Primitive(PrimitiveType::Long)) + .with_initial_default(Literal::Primitive(PrimitiveLiteral::Long(7))) + .into(), + ])), + ) + .into(), + ]); + let err = nested + .check_format_compatibility(FormatVersion::V2) + .unwrap_err(); + assert!(err.message().contains("inner"), "{err}"); + + // A v3-only type nested in a container blames the leaf, never the container. + let nested_variant = schema_with(vec![ + NestedField::required( + 1, + "container", + Type::Struct(StructType::new(vec![ + NestedField::optional(2, "v", Type::Variant(VariantType)).into(), + ])), + ) + .into(), + ]); + let err = nested_variant + .check_format_compatibility(FormatVersion::V2) + .unwrap_err(); + assert!(err.message().contains("container.v"), "{err}"); + assert!( + !err.message().contains("Invalid type for container:"), + "container must not be blamed: {err}" + ); + } + + #[test] + fn test_calc_min_compatible_format() { + use crate::spec::{FormatVersion, VariantType}; + + fn schema_with(fields: Vec) -> Schema { + Schema::builder().with_fields(fields).build().unwrap() + } + + // All v1-compatible types → V1. + let v1 = schema_with(vec![ + NestedField::required(1, "a", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(2, "b", Type::Primitive(PrimitiveType::String)).into(), + ]); + assert_eq!(v1.calc_min_compatible_format(), FormatVersion::V1); + + // A top-level variant → V3. + let variant = schema_with(vec![ + NestedField::optional(1, "v", Type::Variant(VariantType)).into(), + ]); + assert_eq!(variant.calc_min_compatible_format(), FormatVersion::V3); + + // A v3-only type nested inside a list inside a struct → V3 (flattened fields). + let nested = schema_with(vec![ + NestedField::required( + 1, + "s", + Type::Struct(StructType::new(vec![ + NestedField::optional( + 2, + "l", + Type::List(ListType::new( + NestedField::required( + 3, + "element", + Type::Primitive(PrimitiveType::TimestampNs), + ) + .into(), + )), + ) + .into(), + ])), + ) + .into(), + ]); + assert_eq!(nested.calc_min_compatible_format(), FormatVersion::V3); + } + + #[test] + fn test_check_format_compatibility_message_order() { + use crate::spec::{FormatVersion, PrimitiveLiteral, VariantType}; + + // Fields declared out of id order; field 2 fails both checks (ts_ns type + // AND a non-null initial default). The error must list problems by field id, + // with the type problem before the initial-default problem for field 2. + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(3, "c", Type::Variant(VariantType)).into(), + NestedField::optional(2, "b", Type::Primitive(PrimitiveType::TimestampNs)) + .with_initial_default(Literal::Primitive(PrimitiveLiteral::Long(0))) + .into(), + NestedField::required(1, "a", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build() + .unwrap(); + + let message = schema + .check_format_compatibility(FormatVersion::V2) + .unwrap_err() + .message() + .to_string(); + + let lines: Vec<&str> = message.lines().skip(1).collect(); + assert_eq!( + lines, + vec![ + "- Invalid type for b: timestamp_ns is not supported until v3 but format version is v2.", + "- Invalid initial default for b: non-null default (Primitive(Long(0))) is not supported until v3 but format version is v2.", + "- Invalid type for c: variant is not supported until v3 but format version is v2.", + ], + "{message}" + ); + } + #[test] fn test_construct_schema() { let field1: NestedFieldRef = diff --git a/crates/iceberg/src/spec/schema/prune_columns.rs b/crates/iceberg/src/spec/schema/prune_columns.rs index 14f1bfd25f..af9d332a42 100644 --- a/crates/iceberg/src/spec/schema/prune_columns.rs +++ b/crates/iceberg/src/spec/schema/prune_columns.rs @@ -16,6 +16,7 @@ // under the License. use super::*; +use crate::spec::VariantType; struct PruneColumn { selected: HashSet, @@ -238,6 +239,10 @@ impl SchemaVisitor for PruneColumn { fn primitive(&mut self, _p: &PrimitiveType) -> Result> { Ok(None) } + + fn variant(&mut self, _v: &VariantType) -> Result { + Ok(None) + } } #[cfg(test)] @@ -760,4 +765,33 @@ mod tests { assert!(result.is_ok()); assert_eq!(result.unwrap(), Type::Struct(schema.as_struct().clone())); } + + #[test] + fn test_prune_columns_variant() { + // foo (String, id=1) + v (Variant, id=2). + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "foo", Type::Primitive(PrimitiveType::String)).into(), + NestedField::optional(2, "v", Type::Variant(VariantType)).into(), + ]) + .build() + .unwrap(); + + // A variant is a leaf (like a primitive): selecting it keeps it, the same way + // for select_full_types true and false. + let only_variant = Type::Struct(StructType::new(vec![ + NestedField::optional(2, "v", Type::Variant(VariantType)).into(), + ])); + for full in [false, true] { + let result = prune_columns(&schema, HashSet::from([2]), full).unwrap(); + assert_eq!(result, only_variant, "select_full_types={full}"); + } + + // Selecting a sibling prunes the variant out. + let only_foo = Type::Struct(StructType::new(vec![ + NestedField::optional(1, "foo", Type::Primitive(PrimitiveType::String)).into(), + ])); + let result = prune_columns(&schema, HashSet::from([1]), false).unwrap(); + assert_eq!(result, only_foo); + } } diff --git a/crates/iceberg/src/spec/schema/visitor.rs b/crates/iceberg/src/spec/schema/visitor.rs index 50f7c04caa..16b5fcdca9 100644 --- a/crates/iceberg/src/spec/schema/visitor.rs +++ b/crates/iceberg/src/spec/schema/visitor.rs @@ -16,6 +16,7 @@ // under the License. use super::*; +use crate::spec::VariantType; /// A post order schema visitor. /// @@ -69,6 +70,9 @@ pub trait SchemaVisitor { fn map(&mut self, map: &MapType, key_value: Self::T, value: Self::T) -> Result; /// Called when see a primitive type. fn primitive(&mut self, p: &PrimitiveType) -> Result; + + /// Called when see a variant type. + fn variant(&mut self, v: &VariantType) -> Result; } /// Visiting a type in post order. @@ -99,6 +103,7 @@ pub(crate) fn visit_type(r#type: &Type, visitor: &mut V) -> Re visitor.map(map, key_result, value_result) } Type::Struct(s) => visit_struct(s, visitor), + Type::Variant(v) => visitor.variant(v), } } @@ -185,6 +190,8 @@ pub trait SchemaWithPartnerVisitor

{ ) -> Result; /// Called when see a primitive type. fn primitive(&mut self, p: &PrimitiveType, partner: &P) -> Result; + /// Called when see a variant type. + fn variant(&mut self, v: &VariantType, partner: &P) -> Result; } /// Accessor used to get child partner from parent partner. @@ -242,6 +249,7 @@ pub(crate) fn visit_type_with_partner, A: Part visitor.map(map, partner, key_result, value_result) } Type::Struct(s) => visit_struct_with_partner(s, partner, visitor, accessor), + Type::Variant(v) => visitor.variant(v, partner), } } diff --git a/crates/iceberg/src/spec/sort.rs b/crates/iceberg/src/spec/sort.rs index 379d44cc2d..b27d5a36d3 100644 --- a/crates/iceberg/src/spec/sort.rs +++ b/crates/iceberg/src/spec/sort.rs @@ -452,6 +452,35 @@ mod tests { ) } + #[test] + fn test_build_should_return_err_if_source_field_is_variant() { + let schema = Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::optional(1, "v", Type::Variant(crate::spec::VariantType)).into(), + ]) + .build() + .unwrap(); + + let sort_order_builder_result = SortOrder::builder() + .with_sort_field( + SortField::builder() + .source_id(1) + .direction(SortDirection::Ascending) + .null_order(NullOrder::First) + .transform(Transform::Identity) + .build(), + ) + .build(&schema); + + assert_eq!( + sort_order_builder_result + .expect_err("Expected an Err value") + .message(), + "Cannot sort by non-primitive source field: variant" + ) + } + #[test] fn test_build_should_return_err_if_source_field_type_is_not_supported_by_transform() { let schema = Schema::builder() diff --git a/crates/iceberg/src/spec/table_metadata.rs b/crates/iceberg/src/spec/table_metadata.rs index 607fd98350..597ace956f 100644 --- a/crates/iceberg/src/spec/table_metadata.rs +++ b/crates/iceberg/src/spec/table_metadata.rs @@ -535,6 +535,7 @@ impl TableMetadata { // Normalize location (remove trailing slash) self.location = self.location.trim_end_matches('/').to_string(); self.validate_snapshot_sequence_number()?; + self.validate_schema_format_compatibility()?; self.try_normalize_partition_spec()?; self.try_normalize_sort_order()?; Ok(self) @@ -749,6 +750,13 @@ impl TableMetadata { Ok(()) } + + /// Validates that every type used in the current schema is supported by the + /// table's format version. Delegates to [`Schema::check_format_compatibility`]. + fn validate_schema_format_compatibility(&self) -> Result<()> { + self.current_schema() + .check_format_compatibility(self.format_version) + } } pub(super) mod _serde { diff --git a/crates/iceberg/src/spec/values/literal.rs b/crates/iceberg/src/spec/values/literal.rs index 2eef4eba07..0cc079854a 100644 --- a/crates/iceberg/src/spec/values/literal.rs +++ b/crates/iceberg/src/spec/values/literal.rs @@ -601,6 +601,10 @@ impl Literal { )) } } + Type::Variant(_) => Err(Error::new( + ErrorKind::DataInvalid, + "Variant type is not supported for single-value JSON serialization", + )), } } diff --git a/crates/iceberg/src/transaction/update_schema.rs b/crates/iceberg/src/transaction/update_schema.rs index da843d9b9b..953bcd64ab 100644 --- a/crates/iceberg/src/transaction/update_schema.rs +++ b/crates/iceberg/src/transaction/update_schema.rs @@ -181,6 +181,9 @@ fn assign_fresh_ids(field: &NestedField, next_id: &mut i32) -> NestedFieldRef { fn assign_fresh_ids_to_type(field_type: &Type, next_id: &mut i32) -> Type { match field_type { Type::Primitive(_) => field_type.clone(), + // Variant carries no nested fields, so there is nothing to reassign + // (matches id_reassigner.rs). + Type::Variant(v) => Type::Variant(*v), Type::Struct(struct_type) => { let new_fields: Vec = struct_type .fields() @@ -279,7 +282,7 @@ fn rebuild_field( delete_ids: &HashSet, ) -> NestedFieldRef { match field.field_type.as_ref() { - Type::Primitive(_) => field.clone(), + Type::Primitive(_) | Type::Variant(_) => field.clone(), Type::Struct(s) => { let new_fields = rebuild_fields(s.fields(), adds, delete_ids, Some(field.id)); Arc::new(NestedField { @@ -477,6 +480,7 @@ mod tests { use crate::spec::{ DEFAULT_SCHEMA_ID, Literal, NestedField, PrimitiveType, StructType, TableMetadata, Type, + VariantType, }; use crate::table::Table; use crate::transaction::Transaction; @@ -592,6 +596,19 @@ mod tests { // Existing root-level tests // ----------------------------------------------------------------------- + #[test] + fn test_assign_fresh_ids_variant() { + // Variant carries no sub-fields, so fresh-id assignment only renames the field + // itself and leaves the type untouched. + let mut next_id = 10; + let field = NestedField::optional(1, "data", Type::Variant(VariantType)); + let assigned = super::assign_fresh_ids(&field, &mut next_id); + + assert_eq!(assigned.id, 11); + assert_eq!(*assigned.field_type, Type::Variant(VariantType)); + assert_eq!(next_id, 11); + } + #[tokio::test] async fn test_add_column() { let table = make_v2_table(); diff --git a/crates/iceberg/src/writer/file_writer/parquet_writer.rs b/crates/iceberg/src/writer/file_writer/parquet_writer.rs index eb69f01dfe..7745d337d2 100644 --- a/crates/iceberg/src/writer/file_writer/parquet_writer.rs +++ b/crates/iceberg/src/writer/file_writer/parquet_writer.rs @@ -40,7 +40,7 @@ use crate::io::{FileIO, FileWrite, OutputFile}; use crate::spec::{ DataContentType, DataFileBuilder, DataFileFormat, Datum, ListType, Literal, MapType, NestedFieldRef, PartitionSpec, PrimitiveType, Schema, SchemaRef, SchemaVisitor, Struct, - StructType, TableMetadata, TableProperties, Type, visit_schema, + StructType, TableMetadata, TableProperties, Type, VariantType, visit_schema, }; use crate::transform::create_transform_function; use crate::writer::{CurrentFileStatus, DataFile}; @@ -241,6 +241,13 @@ impl SchemaVisitor for IndexByParquetPathName { Ok(()) } + + fn variant(&mut self, _v: &VariantType) -> Result { + Err(Error::new( + ErrorKind::FeatureUnsupported, + "Writing variant columns to Parquet is not supported yet", + )) + } } /// `ParquetWriter`` is used to write arrow data into parquet file on storage. @@ -822,6 +829,21 @@ mod tests { assert_eq!(visitor.name_to_id, expect); } + #[test] + fn test_index_by_parquet_path_variant_is_unsupported() { + // Writing variant columns to Parquet is not supported yet; indexing a schema that + // contains one must error rather than silently miss-map columns. + let schema = Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "v", Type::Variant(VariantType)).into(), + ]) + .build() + .unwrap(); + let mut visitor = IndexByParquetPathName::new(); + let err = visit_schema(&schema, &mut visitor).unwrap_err(); + assert_eq!(err.kind(), crate::ErrorKind::FeatureUnsupported, "{err}"); + } + #[tokio::test] async fn test_parquet_writer() -> Result<()> { let temp_dir = TempDir::new().unwrap(); diff --git a/crates/integrations/datafusion/src/schema.rs b/crates/integrations/datafusion/src/schema.rs index 91526935c2..c515b3b8af 100644 --- a/crates/integrations/datafusion/src/schema.rs +++ b/crates/integrations/datafusion/src/schema.rs @@ -28,6 +28,7 @@ use futures::StreamExt; use futures::future::try_join_all; use iceberg::arrow::arrow_schema_to_schema_auto_assign_ids; use iceberg::inspect::MetadataTableType; +use iceberg::spec::FormatVersion; use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableCreation, TableIdent}; use crate::table::IcebergTableProvider; @@ -174,10 +175,16 @@ impl SchemaProvider for IcebergSchemaProvider { let iceberg_schema = arrow_schema_to_schema_auto_assign_ids(df_schema.as_ref()) .map_err(to_datafusion_error)?; + // Use at least V2, and upgrade to V3 if the schema requires it (e.g. timestamp_ns / variant). + let format_version = iceberg_schema + .calc_min_compatible_format() + .max(FormatVersion::V2); + // Create the table in the Iceberg catalog let table_creation = TableCreation::builder() .name(name.clone()) .schema(iceberg_schema) + .format_version(format_version) .build(); let catalog = self.catalog.clone(); From 435d104231b651b745c68b76826a053adb16c05a Mon Sep 17 00:00:00 2001 From: Christian Date: Thu, 16 Jul 2026 12:17:06 +0200 Subject: [PATCH 07/11] feat(spec): read timestamp_ns/timestamptz_ns default values in the Literal JSON codec (#2832) ## Which issue does this PR close? We don't have an issue for this. Deserializing Timestamp NS default values failed. ## What changes are included in this PR? `try_into_json` serialized TimestampNs/TimestamptzNs but `try_from_json` had no matching arms, so a nanosecond-timestamp single value (e.g. a column default) failed to deserialize (fell to the DataInvalid catch-all). Add the read arms mirroring the microsecond ones, plus `datetime_to_nanoseconds` / `datetimetz_to_nanoseconds` helpers. Also adds json_timestamp_ns / json_timestamptz_ns tests. ## Are these changes tested? Yes --- crates/iceberg/src/spec/values/literal.rs | 27 +++++++++++++ crates/iceberg/src/spec/values/temporal.rs | 12 ++++++ crates/iceberg/src/spec/values/tests.rs | 47 ++++++++++++++++++++++ 3 files changed, 86 insertions(+) diff --git a/crates/iceberg/src/spec/values/literal.rs b/crates/iceberg/src/spec/values/literal.rs index 0cc079854a..14f711da18 100644 --- a/crates/iceberg/src/spec/values/literal.rs +++ b/crates/iceberg/src/spec/values/literal.rs @@ -493,6 +493,33 @@ impl Literal { )), )))) } + (PrimitiveType::TimestampNs, JsonValue::String(s)) => { + let ndt = NaiveDateTime::parse_from_str(&s, "%Y-%m-%dT%H:%M:%S%.f")?; + let nanos = timestamp::datetime_to_nanoseconds(&ndt).ok_or_else(|| { + Error::new( + crate::ErrorKind::DataInvalid, + format!( + "Timestamp is outside the representable nanosecond range: {ndt}" + ), + ) + })?; + Ok(Some(Literal::Primitive(PrimitiveLiteral::Long(nanos)))) + } + (PrimitiveType::TimestamptzNs, JsonValue::String(s)) => { + let dt = Utc.from_utc_datetime(&NaiveDateTime::parse_from_str( + &s, + "%Y-%m-%dT%H:%M:%S%.f+00:00", + )?); + let nanos = timestamptz::datetimetz_to_nanoseconds(&dt).ok_or_else(|| { + Error::new( + crate::ErrorKind::DataInvalid, + format!( + "Timestamptz is outside the representable nanosecond range: {dt}" + ), + ) + })?; + Ok(Some(Literal::Primitive(PrimitiveLiteral::Long(nanos)))) + } (PrimitiveType::String, JsonValue::String(s)) => { Ok(Some(Literal::Primitive(PrimitiveLiteral::String(s)))) } diff --git a/crates/iceberg/src/spec/values/temporal.rs b/crates/iceberg/src/spec/values/temporal.rs index ad797e127e..c9efcf4c4a 100644 --- a/crates/iceberg/src/spec/values/temporal.rs +++ b/crates/iceberg/src/spec/values/temporal.rs @@ -82,6 +82,12 @@ pub(crate) mod timestamp { pub(crate) fn nanoseconds_to_datetime(nanos: i64) -> NaiveDateTime { DateTime::from_timestamp_nanos(nanos).naive_utc() } + + /// Nanoseconds since the Unix epoch, or `None` if outside the representable `i64` range + /// (roughly the years 1678–2262). + pub(crate) fn datetime_to_nanoseconds(time: &NaiveDateTime) -> Option { + time.and_utc().timestamp_nanos_opt() + } } pub(crate) mod timestamptz { @@ -102,4 +108,10 @@ pub(crate) mod timestamptz { DateTime::from_timestamp(secs, rem as u32).unwrap() } + + /// Nanoseconds since the Unix epoch, or `None` if outside the representable `i64` range + /// (roughly the years 1678–2262). + pub(crate) fn datetimetz_to_nanoseconds(time: &DateTime) -> Option { + time.timestamp_nanos_opt() + } } diff --git a/crates/iceberg/src/spec/values/tests.rs b/crates/iceberg/src/spec/values/tests.rs index a8a6716f9d..f097a74a59 100644 --- a/crates/iceberg/src/spec/values/tests.rs +++ b/crates/iceberg/src/spec/values/tests.rs @@ -214,6 +214,53 @@ fn json_timestamptz() { ); } +#[test] +fn json_timestamp_ns() { + let record = r#""2017-11-16T22:31:08.123456789""#; + + check_json_serde( + record, + Literal::Primitive(PrimitiveLiteral::Long(1510871468123456789)), + &Type::Primitive(PrimitiveType::TimestampNs), + ); +} + +#[test] +fn json_timestamptz_ns() { + let record = r#""2017-11-16T22:31:08.123456789+00:00""#; + + check_json_serde( + record, + Literal::Primitive(PrimitiveLiteral::Long(1510871468123456789)), + &Type::Primitive(PrimitiveType::TimestamptzNs), + ); +} + +#[test] +fn json_timestamptz_ns_rejects_non_utc_offset() { + // Per the spec, timestamptz_ns single-value serialization must use offset "+00:00"; Java's + // SingleValueParser enforces the same (DateTimeUtil.isUTCTimestamptz). A non-UTC offset is not a + // valid encoding and must be rejected, not silently re-based to UTC. + let record = serde_json::Value::String("2017-11-16T22:31:08.123456789+05:00".to_string()); + let result = Literal::try_from_json(record, &Type::Primitive(PrimitiveType::TimestamptzNs)); + assert!( + result.is_err(), + "non-UTC offset must be rejected for timestamptz_ns, got {result:?}" + ); +} + +#[test] +fn json_timestamptz_rejects_non_utc_offset() { + // Micros-precision counterpart, mirroring Java's TestSingleValueParser.testInvalidTimestamptz: + // the offset must be "+00:00", so a non-UTC offset is rejected. + let record = serde_json::Value::String("2017-11-16T22:31:08.123456+05:00".to_string()); + let result = Literal::try_from_json(record, &Type::Primitive(PrimitiveType::Timestamptz)); + assert!( + result.is_err(), + "non-UTC offset must be rejected for timestamptz, got {result:?}" + ); +} + #[test] fn json_string() { let record = r#""iceberg""#; From 4c3d6dacafa06a4ec42236b1dd0d88a822ca2f06 Mon Sep 17 00:00:00 2001 From: Noah Date: Thu, 16 Jul 2026 15:59:56 -0400 Subject: [PATCH 08/11] fix test --- crates/catalog/rest/src/catalog.rs | 2 +- crates/storage/opendal/src/lib.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index da05fc7e3b..468cd40c67 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -525,7 +525,7 @@ impl RestCatalog { } /// All requests and expected responses are derived from the REST catalog API spec: -/// https://github.com/apache/iceberg/blob/main/open-api/rest-catalog-open-api.yaml +/// #[async_trait] impl Catalog for RestCatalog { async fn list_namespaces( diff --git a/crates/storage/opendal/src/lib.rs b/crates/storage/opendal/src/lib.rs index 52f4e68ed3..b50e49d6e8 100644 --- a/crates/storage/opendal/src/lib.rs +++ b/crates/storage/opendal/src/lib.rs @@ -18,8 +18,8 @@ //! OpenDAL-based storage implementation for Apache Iceberg. //! //! This crate provides [`OpenDalStorage`] and [`OpenDalStorageFactory`], -//! which implement the [`Storage`](iceberg::io::Storage) and -//! [`StorageFactory`](iceberg::io::StorageFactory) traits from the `iceberg` crate +//! which implement the [`Storage`] and +//! [`StorageFactory`] traits from the `iceberg` crate //! using [OpenDAL](https://opendal.apache.org/) as the backend. mod utils; From 03f5378231636f5937383672b406db87ed2f71c1 Mon Sep 17 00:00:00 2001 From: Noah Date: Thu, 16 Jul 2026 16:30:37 -0400 Subject: [PATCH 09/11] cleanup --- crates/integrations/datafusion/public-api.txt | 6 ------ 1 file changed, 6 deletions(-) diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index b492624bb4..760a923e13 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -29,7 +29,6 @@ pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::fmt(&self, f: &mut impl datafusion_physical_plan::display::DisplayAs for iceberg_datafusion::physical_plan::IcebergCommitExec pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::fmt_as(&self, t: datafusion_physical_plan::display::DisplayFormatType, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_physical_plan::execution_plan::ExecutionPlan for iceberg_datafusion::physical_plan::IcebergCommitExec -pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::as_any(&self) -> &dyn core::any::Any pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::benefits_from_input_partitioning(&self) -> alloc::vec::Vec pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::children(&self) -> alloc::vec::Vec<&alloc::sync::Arc> pub fn iceberg_datafusion::physical_plan::IcebergCommitExec::execute(&self, partition: usize, context: alloc::sync::Arc) -> datafusion_common::error::Result @@ -46,7 +45,6 @@ pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::fmt(&self, f: &mu impl datafusion_physical_plan::display::DisplayAs for iceberg_datafusion::physical_plan::IcebergMetadataScan pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::fmt_as(&self, _t: datafusion_physical_plan::display::DisplayFormatType, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_physical_plan::execution_plan::ExecutionPlan for iceberg_datafusion::physical_plan::IcebergMetadataScan -pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::as_any(&self) -> &dyn core::any::Any pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::children(&self) -> alloc::vec::Vec<&alloc::sync::Arc> pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::execute(&self, _partition: usize, _context: alloc::sync::Arc) -> datafusion_common::error::Result pub fn iceberg_datafusion::physical_plan::IcebergMetadataScan::name(&self) -> &str @@ -84,7 +82,6 @@ pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::fmt(&self, f: &mut c impl datafusion_physical_plan::display::DisplayAs for iceberg_datafusion::physical_plan::IcebergWriteExec pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::fmt_as(&self, t: datafusion_physical_plan::display::DisplayFormatType, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_physical_plan::execution_plan::ExecutionPlan for iceberg_datafusion::physical_plan::IcebergWriteExec -pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::as_any(&self) -> &dyn core::any::Any pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::benefits_from_input_partitioning(&self) -> alloc::vec::Vec pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::children(&self) -> alloc::vec::Vec<&alloc::sync::Arc> pub fn iceberg_datafusion::physical_plan::IcebergWriteExec::execute(&self, partition: usize, context: alloc::sync::Arc) -> datafusion_common::error::Result @@ -109,7 +106,6 @@ pub fn iceberg_datafusion::physical_plan::PartitionExpr::fmt(&self, f: &mut core impl core::hash::Hash for iceberg_datafusion::physical_plan::PartitionExpr pub fn iceberg_datafusion::physical_plan::PartitionExpr::hash(&self, state: &mut H) impl datafusion_physical_expr_common::physical_expr::PhysicalExpr for iceberg_datafusion::physical_plan::PartitionExpr -pub fn iceberg_datafusion::physical_plan::PartitionExpr::as_any(&self) -> &dyn core::any::Any pub fn iceberg_datafusion::physical_plan::PartitionExpr::children(&self) -> alloc::vec::Vec<&alloc::sync::Arc> pub fn iceberg_datafusion::physical_plan::PartitionExpr::data_type(&self, _input_schema: &arrow_schema::schema::Schema) -> datafusion_common::error::Result pub fn iceberg_datafusion::physical_plan::PartitionExpr::evaluate(&self, batch: &arrow_array::record_batch::RecordBatch) -> datafusion_common::error::Result @@ -162,7 +158,6 @@ pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::clone(& impl core::fmt::Debug for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider -pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::as_any(&self) -> &dyn core::any::Any pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, _projection: core::option::Option<&'life2 alloc::vec::Vec>, _filters: &'life3 [datafusion_expr::expr::Expr], _limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType @@ -245,7 +240,6 @@ pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::clone(& impl core::fmt::Debug for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider -pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::as_any(&self) -> &dyn core::any::Any pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, _projection: core::option::Option<&'life2 alloc::vec::Vec>, _filters: &'life3 [datafusion_expr::expr::Expr], _limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType From 4dcfcbabfecc436b8d2bbfb733cf85b32674e676 Mon Sep 17 00:00:00 2001 From: Noah Date: Fri, 17 Jul 2026 00:40:44 -0400 Subject: [PATCH 10/11] clean --- crates/catalog/rest/src/catalog.rs | 2 +- crates/storage/opendal/src/lib.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 468cd40c67..da05fc7e3b 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -525,7 +525,7 @@ impl RestCatalog { } /// All requests and expected responses are derived from the REST catalog API spec: -/// +/// https://github.com/apache/iceberg/blob/main/open-api/rest-catalog-open-api.yaml #[async_trait] impl Catalog for RestCatalog { async fn list_namespaces( diff --git a/crates/storage/opendal/src/lib.rs b/crates/storage/opendal/src/lib.rs index b50e49d6e8..52f4e68ed3 100644 --- a/crates/storage/opendal/src/lib.rs +++ b/crates/storage/opendal/src/lib.rs @@ -18,8 +18,8 @@ //! OpenDAL-based storage implementation for Apache Iceberg. //! //! This crate provides [`OpenDalStorage`] and [`OpenDalStorageFactory`], -//! which implement the [`Storage`] and -//! [`StorageFactory`] traits from the `iceberg` crate +//! which implement the [`Storage`](iceberg::io::Storage) and +//! [`StorageFactory`](iceberg::io::StorageFactory) traits from the `iceberg` crate //! using [OpenDAL](https://opendal.apache.org/) as the backend. mod utils; From 47538f25653472f152dbbc847cf589ec72cdf099 Mon Sep 17 00:00:00 2001 From: Noah Date: Fri, 17 Jul 2026 21:19:50 -0400 Subject: [PATCH 11/11] cleanup --- crates/integrations/datafusion/src/schema.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/integrations/datafusion/src/schema.rs b/crates/integrations/datafusion/src/schema.rs index c515b3b8af..839c984602 100644 --- a/crates/integrations/datafusion/src/schema.rs +++ b/crates/integrations/datafusion/src/schema.rs @@ -385,7 +385,7 @@ mod tests { .unwrap(); // With config: the table provider carries it (and is therefore distributable). - let config = crate::IcebergCatalogConfig::new("memory", "memory", HashMap::new()); + let config = IcebergCatalogConfig::new("memory", "memory", HashMap::new()); let with_config = IcebergSchemaProvider::try_new(catalog.clone(), Some(config), namespace.clone()) .await