From 7804ffed9f481c57bc7a565928bfa83430c2a91a Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 21 Sep 2026 17:11:45 +0200 Subject: [PATCH 1/2] refactor: use DataFusion errors in integration --- Cargo.lock | 1 - crates/datafusion/Cargo.toml | 1 - crates/datafusion/src/catalog.rs | 7 +- crates/datafusion/src/error.rs | 102 ++++++++++++++++-- crates/datafusion/src/physical_plan/commit.rs | 44 ++++---- .../datafusion/src/physical_plan/project.rs | 17 +-- .../src/physical_plan/repartition.rs | 11 +- crates/datafusion/src/physical_plan/scan.rs | 14 +-- crates/datafusion/src/physical_plan/sort.rs | 18 ++-- crates/datafusion/src/physical_plan/write.rs | 60 +++++------ crates/datafusion/src/schema.rs | 72 ++++--------- crates/datafusion/src/table/metadata_table.rs | 6 +- crates/datafusion/src/table/mod.rs | 89 +++++++-------- .../src/table/table_provider_factory.rs | 38 ++++--- crates/datafusion/src/task_writer.rs | 34 +++--- .../tests/integration_datafusion_test.rs | 27 ++--- 16 files changed, 304 insertions(+), 237 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 673525b..8ac0c1a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2111,7 +2111,6 @@ dependencies = [ name = "datafusion-iceberg" version = "0.10.1" dependencies = [ - "anyhow", "async-trait", "dashmap", "datafusion", diff --git a/crates/datafusion/Cargo.toml b/crates/datafusion/Cargo.toml index 9647aed..b4c3b92 100644 --- a/crates/datafusion/Cargo.toml +++ b/crates/datafusion/Cargo.toml @@ -29,7 +29,6 @@ license = { workspace = true } repository = { workspace = true } [dependencies] -anyhow = { workspace = true } async-trait = { workspace = true } dashmap = "6" datafusion = { workspace = true } diff --git a/crates/datafusion/src/catalog.rs b/crates/datafusion/src/catalog.rs index 2c6e1ff..a96f731 100644 --- a/crates/datafusion/src/catalog.rs +++ b/crates/datafusion/src/catalog.rs @@ -19,10 +19,12 @@ use std::collections::HashMap; use std::sync::Arc; use datafusion::catalog::{CatalogProvider, SchemaProvider}; +use datafusion::error::Result; use futures::future::try_join_all; -use iceberg::{Catalog, NamespaceIdent, Result}; +use iceberg::{Catalog, NamespaceIdent}; use crate::schema::IcebergSchemaProvider; +use crate::to_datafusion_error; /// Provides an interface to manage and access multiple schemas /// within an Iceberg [`Catalog`]. @@ -51,7 +53,8 @@ impl IcebergCatalogProvider { // As of right now; schemas might become stale. let schema_names: Vec<_> = client .list_namespaces(None) - .await? + .await + .map_err(to_datafusion_error)? .iter() .flat_map(|ns| ns.as_ref().clone()) .collect(); diff --git a/crates/datafusion/src/error.rs b/crates/datafusion/src/error.rs index 5bc4701..89674fe 100644 --- a/crates/datafusion/src/error.rs +++ b/crates/datafusion/src/error.rs @@ -15,18 +15,106 @@ // specific language governing permissions and limitations // under the License. -use anyhow::anyhow; +use datafusion::error::DataFusionError; use iceberg::{Error, ErrorKind}; /// Converts a datafusion error into an iceberg error. -pub fn from_datafusion_error(error: datafusion::error::DataFusionError) -> Error { +pub fn from_datafusion_error(error: DataFusionError) -> Error { + let fallback_message = error.to_string(); + let DataFusionError::Context(context, error) = error else { + return unexpected_datafusion_error(fallback_message); + }; + + let Some(kind) = parse_iceberg_error_kind(&context) else { + return unexpected_datafusion_error(fallback_message); + }; + let DataFusionError::Execution(message) = error.as_ref() else { + return unexpected_datafusion_error(fallback_message); + }; + + Error::new(kind, strip_error_kind(kind, message)) +} + +/// Converts an iceberg error into a datafusion error. +pub fn to_datafusion_error(error: Error) -> DataFusionError { + DataFusionError::Context( + format!("IcebergError({})", error.kind()), + Box::new(DataFusionError::Execution(error.to_string())), + ) +} + +fn unexpected_datafusion_error(message: String) -> Error { Error::new( ErrorKind::Unexpected, - "Operation failed for hitting datafusion error".to_string(), + format!("DataFusion execution failed: {message}"), ) - .with_source(anyhow!("datafusion error: {error:?}")) } -/// Converts an iceberg error into a datafusion error. -pub fn to_datafusion_error(error: Error) -> datafusion::error::DataFusionError { - datafusion::error::DataFusionError::External(error.into()) + +fn parse_iceberg_error_kind(context: &str) -> Option { + let kind = context.strip_prefix("IcebergError(")?.strip_suffix(')')?; + + match kind { + "PreconditionFailed" => Some(ErrorKind::PreconditionFailed), + "Unexpected" => Some(ErrorKind::Unexpected), + "DataInvalid" => Some(ErrorKind::DataInvalid), + "NamespaceAlreadyExists" => Some(ErrorKind::NamespaceAlreadyExists), + "TableAlreadyExists" => Some(ErrorKind::TableAlreadyExists), + "NamespaceNotFound" => Some(ErrorKind::NamespaceNotFound), + "TableNotFound" => Some(ErrorKind::TableNotFound), + "FeatureUnsupported" => Some(ErrorKind::FeatureUnsupported), + "CatalogCommitConflicts" => Some(ErrorKind::CatalogCommitConflicts), + _ => None, + } +} + +fn strip_error_kind(kind: ErrorKind, message: &str) -> String { + let kind = kind.into_static(); + if message == kind { + String::new() + } else { + message + .strip_prefix(&format!("{kind} => ")) + .unwrap_or(message) + .to_string() + } +} + +#[cfg(test)] +mod tests { + use datafusion::error::DataFusionError; + use iceberg::{Error, ErrorKind}; + + use super::{from_datafusion_error, to_datafusion_error}; + + #[test] + fn roundtrips_iceberg_error_kind_and_message() { + let error = Error::new(ErrorKind::DataInvalid, "invalid manifest"); + let roundtripped = from_datafusion_error(to_datafusion_error(error)); + + assert_eq!(roundtripped.kind(), ErrorKind::DataInvalid); + assert_eq!(roundtripped.to_string(), "DataInvalid => invalid manifest"); + } + + #[test] + fn encodes_iceberg_errors_with_native_datafusion_variants() { + let error = to_datafusion_error(Error::new(ErrorKind::DataInvalid, "invalid manifest")); + + assert!(matches!( + error, + DataFusionError::Context(context, inner) + if context == "IcebergError(DataInvalid)" + && matches!(inner.as_ref(), DataFusionError::Execution(message) if message == "DataInvalid => invalid manifest") + )); + } + + #[test] + fn maps_non_iceberg_datafusion_errors_to_unexpected() { + let error = from_datafusion_error(DataFusionError::Execution("worker failed".to_string())); + + assert_eq!(error.kind(), ErrorKind::Unexpected); + assert_eq!( + error.to_string(), + "Unexpected => DataFusion execution failed: Execution error: worker failed" + ); + } } diff --git a/crates/datafusion/src/physical_plan/commit.rs b/crates/datafusion/src/physical_plan/commit.rs index d080f61..dcbee3a 100644 --- a/crates/datafusion/src/physical_plan/commit.rs +++ b/crates/datafusion/src/physical_plan/commit.rs @@ -22,8 +22,10 @@ use datafusion::arrow::array::{Array, ArrayRef, RecordBatch, StringArray, UInt64 use datafusion::arrow::datatypes::{ DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, }; -use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{DataFusionError, Result as DFResult}; +use datafusion::common::{ + internal_datafusion_err, internal_err, tree_node::TreeNodeRecursion, +}; +use datafusion::error::{DataFusionError, Result}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -84,7 +86,7 @@ impl IcebergCommitExec { } // Create a record batch with just the count of rows written - fn make_count_batch(count: u64) -> DFResult { + fn make_count_batch(count: u64) -> Result { let count_array = Arc::new(UInt64Array::from(vec![count])) as ArrayRef; RecordBatch::try_from_iter_with_nullable(vec![("count", count_array, false)]) @@ -142,8 +144,8 @@ impl ExecutionPlan for IcebergCommitExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> DFResult, - ) -> DFResult { + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -163,12 +165,12 @@ impl ExecutionPlan for IcebergCommitExec { fn with_new_children( self: Arc, children: Vec>, - ) -> DFResult> { + ) -> Result> { if children.len() != 1 { - return Err(DataFusionError::Internal(format!( + return internal_err!( "IcebergCommitExec expects exactly one child, but provided {}", children.len() - ))); + ); } Ok(Arc::new(IcebergCommitExec::new( @@ -183,12 +185,12 @@ impl ExecutionPlan for IcebergCommitExec { &self, partition: usize, context: Arc, - ) -> DFResult { + ) -> Result { // IcebergCommitExec only has one partition (partition 0) if partition != 0 { - return Err(DataFusionError::Internal(format!( + return internal_err!( "IcebergCommitExec only has one partition, but got partition {partition}" - ))); + ); } let table = self.table.clone(); @@ -215,15 +217,15 @@ impl ExecutionPlan for IcebergCommitExec { let files_array = batch .column_by_name(DATA_FILES_COL_NAME) .ok_or_else(|| { - DataFusionError::Internal( - "Expected 'data_files' column in input batch".to_string(), + internal_datafusion_err!( + "Expected 'data_files' column in input batch" ) })? .as_any() .downcast_ref::() .ok_or_else(|| { - DataFusionError::Internal( - "Expected 'data_files' column to be StringArray".to_string(), + internal_datafusion_err!( + "Expected 'data_files' column to be StringArray" ) })?; @@ -231,7 +233,7 @@ impl ExecutionPlan for IcebergCommitExec { let batch_files: Vec = files_array .into_iter() .flatten() - .map(|f| -> DFResult { + .map(|f| -> Result { // Parse JSON to DataFileSerde and convert to DataFile deserialize_data_file_from_json( f, @@ -241,7 +243,7 @@ impl ExecutionPlan for IcebergCommitExec { ) .map_err(to_datafusion_error) }) - .collect::>()?; + .collect::>()?; // add record_counts from the current batch to total record count total_record_count += @@ -361,15 +363,15 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> DFResult, - ) -> DFResult { + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { Ok(TreeNodeRecursion::Continue) } fn with_new_children( self: Arc, _children: Vec>, - ) -> datafusion::common::Result> { + ) -> Result> { Ok(self) } @@ -377,7 +379,7 @@ mod tests { &self, _partition: usize, _context: Arc, - ) -> datafusion::common::Result { + ) -> Result { // Create a record batch with the serialized data files let array = Arc::new(StringArray::from(self.data_files_json.clone())) as ArrayRef; diff --git a/crates/datafusion/src/physical_plan/project.rs b/crates/datafusion/src/physical_plan/project.rs index d8ab9b3..f127d8b 100644 --- a/crates/datafusion/src/physical_plan/project.rs +++ b/crates/datafusion/src/physical_plan/project.rs @@ -21,7 +21,8 @@ use std::sync::Arc; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::{DataType, Schema as ArrowSchema}; -use datafusion::common::{DataFusionError, Result as DFResult}; +use datafusion::common::plan_err; +use datafusion::error::Result; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::expressions::Column; use datafusion::physical_plan::projection::ProjectionExec; @@ -51,7 +52,7 @@ use crate::to_datafusion_error; pub fn project_with_partition( input: Arc, table: &Table, -) -> DFResult> { +) -> Result> { let metadata = table.metadata(); let partition_spec = metadata.default_partition_spec(); let table_schema = metadata.current_schema(); @@ -72,11 +73,11 @@ pub fn project_with_partition( .map_err(to_datafusion_error)?; if input_schema_cleaned != expected_schema_cleaned { - return Err(DataFusionError::Plan(format!( + return plan_err!( "Input schema does not match Iceberg table schema.\n\ Expected schema: {expected_schema_cleaned}\n\ Input schema: {input_schema_cleaned}" - ))); + ); } let calculator = @@ -129,15 +130,15 @@ impl PartialEq for PartitionExpr { impl Eq for PartitionExpr {} impl PhysicalExpr for PartitionExpr { - fn data_type(&self, _input_schema: &ArrowSchema) -> DFResult { + fn data_type(&self, _input_schema: &ArrowSchema) -> Result { Ok(self.calculator.partition_arrow_type().clone()) } - fn nullable(&self, _input_schema: &ArrowSchema) -> DFResult { + fn nullable(&self, _input_schema: &ArrowSchema) -> Result { Ok(false) } - fn evaluate(&self, batch: &RecordBatch) -> DFResult { + fn evaluate(&self, batch: &RecordBatch) -> Result { let array = self .calculator .calculate(batch) @@ -152,7 +153,7 @@ impl PhysicalExpr for PartitionExpr { fn with_new_children( self: Arc, _children: Vec>, - ) -> DFResult> { + ) -> Result> { Ok(self) } diff --git a/crates/datafusion/src/physical_plan/repartition.rs b/crates/datafusion/src/physical_plan/repartition.rs index e6cbc13..0dbda0c 100644 --- a/crates/datafusion/src/physical_plan/repartition.rs +++ b/crates/datafusion/src/physical_plan/repartition.rs @@ -18,7 +18,8 @@ use std::num::NonZeroUsize; use std::sync::Arc; -use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::common::plan_err; +use datafusion::error::Result; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::expressions::Column; use datafusion::physical_plan::repartition::RepartitionExec; @@ -90,7 +91,7 @@ pub(crate) fn repartition( input: Arc, table_metadata: TableMetadataRef, target_partitions: NonZeroUsize, -) -> DFResult> { +) -> Result> { let partitioning_strategy = determine_partitioning_strategy(&input, &table_metadata, target_partitions)?; @@ -125,7 +126,7 @@ fn determine_partitioning_strategy( input: &Arc, table_metadata: &TableMetadata, target_partitions: NonZeroUsize, -) -> DFResult { +) -> Result { let partition_spec = table_metadata.default_partition_spec(); let input_schema = input.schema(); let target_partition_count = target_partitions.get(); @@ -158,10 +159,10 @@ fn determine_partitioning_strategy( } // Case 2: Partitioned table missing _partition column (normally this should not happen) - (true, Err(_)) => Err(DataFusionError::Plan(format!( + (true, Err(_)) => plan_err!( "Partitioned table input missing {PROJECTED_PARTITION_VALUE_COLUMN} column. \ Ensure projection happens before repartitioning." - ))), + ), // Case 3: Unpartitioned table, always use RoundRobinBatch (false, _) => Ok(Partitioning::RoundRobinBatch(target_partition_count)), diff --git a/crates/datafusion/src/physical_plan/scan.rs b/crates/datafusion/src/physical_plan/scan.rs index 1796304..b9b15c5 100644 --- a/crates/datafusion/src/physical_plan/scan.rs +++ b/crates/datafusion/src/physical_plan/scan.rs @@ -22,7 +22,7 @@ use std::vec; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::error::Result as DFResult; +use datafusion::error::Result; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -128,15 +128,15 @@ impl ExecutionPlan for IcebergTableScan { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> DFResult, - ) -> DFResult { + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { Ok(TreeNodeRecursion::Continue) } fn with_new_children( self: Arc, _children: Vec>, - ) -> DFResult> { + ) -> Result> { Ok(self) } @@ -148,7 +148,7 @@ impl ExecutionPlan for IcebergTableScan { &self, _partition: usize, _context: Arc, - ) -> DFResult { + ) -> Result { let fut = get_batch_stream( self.table.clone(), self.snapshot_id, @@ -158,7 +158,7 @@ impl ExecutionPlan for IcebergTableScan { let stream = futures::stream::once(fut).try_flatten(); // Apply limit if specified - let limited_stream: Pin> + Send>> = + let limited_stream: Pin> + Send>> = if let Some(limit) = self.limit { let mut remaining = limit; Box::pin(stream.try_filter_map(move |batch| { @@ -217,7 +217,7 @@ async fn get_batch_stream( snapshot_id: Option, column_names: Option>, predicates: Option, -) -> DFResult> + Send>>> { +) -> Result> + Send>>> { let scan_builder = match snapshot_id { Some(snapshot_id) => table.scan().snapshot_id(snapshot_id), None => table.scan(), diff --git a/crates/datafusion/src/physical_plan/sort.rs b/crates/datafusion/src/physical_plan/sort.rs index fce9210..1993e08 100644 --- a/crates/datafusion/src/physical_plan/sort.rs +++ b/crates/datafusion/src/physical_plan/sort.rs @@ -20,8 +20,8 @@ use std::sync::Arc; use datafusion::arrow::compute::SortOptions; -use datafusion::common::Result as DFResult; -use datafusion::error::DataFusionError; +use datafusion::common::plan_datafusion_err; +use datafusion::error::Result; use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::expressions::Column; @@ -44,17 +44,15 @@ use iceberg::arrow::PROJECTED_PARTITION_VALUE_COLUMN; /// * `Err` - If the partition column is not found pub(crate) fn sort_by_partition( input: Arc, -) -> DFResult> { +) -> Result> { let schema = input.schema(); // Find the partition column in the schema let (partition_column_index, _partition_field) = schema .column_with_name(PROJECTED_PARTITION_VALUE_COLUMN) - .ok_or_else(|| { - DataFusionError::Plan(format!( - "Partition column '{PROJECTED_PARTITION_VALUE_COLUMN}' not found in schema. Ensure the plan has been extended with partition values using project_with_partition." - )) - })?; + .ok_or_else(|| plan_datafusion_err!( + "Partition column '{PROJECTED_PARTITION_VALUE_COLUMN}' not found in schema. Ensure the plan has been extended with partition values using project_with_partition." + ))?; // Create a single sort expression for the partition column let column_expr = Arc::new(Column::new( @@ -70,9 +68,7 @@ pub(crate) fn sort_by_partition( // Create a SortExec with preserve_partitioning=true to ensure the output partitioning // is the same as the input partitioning, and the data is sorted within each partition let lex_ordering = LexOrdering::new(vec![sort_expr]).ok_or_else(|| { - DataFusionError::Plan( - "Failed to create LexOrdering from sort expression".to_string(), - ) + plan_datafusion_err!("Failed to create LexOrdering from sort expression") })?; let sort_exec = SortExec::new(lex_ordering, input).with_preserve_partitioning(true); diff --git a/crates/datafusion/src/physical_plan/write.rs b/crates/datafusion/src/physical_plan/write.rs index c5a4206..96c0ac0 100644 --- a/crates/datafusion/src/physical_plan/write.rs +++ b/crates/datafusion/src/physical_plan/write.rs @@ -23,9 +23,9 @@ use datafusion::arrow::array::{ArrayRef, RecordBatch, StringArray}; use datafusion::arrow::datatypes::{ DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, }; -use datafusion::common::Result as DFResult; -use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{internal_err, not_impl_err, tree_node::TreeNodeRecursion}; use datafusion::error::DataFusionError; +use datafusion::error::Result; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -44,7 +44,6 @@ use iceberg::writer::file_writer::location_generator::{ DefaultFileNameGenerator, DefaultLocationGenerator, }; use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder; -use iceberg::{Error, ErrorKind}; use uuid::Uuid; use crate::physical_plan::DATA_FILES_COL_NAME; @@ -94,7 +93,7 @@ impl IcebergWriteExec { } // Create a record batch with serialized data files - fn make_result_batch(data_files: Vec) -> DFResult { + fn make_result_batch(data_files: Vec) -> Result { let files_array = Arc::new(StringArray::from(data_files)) as ArrayRef; RecordBatch::try_new(Self::make_result_schema(), vec![files_array]).map_err(|e| { @@ -161,20 +160,20 @@ impl ExecutionPlan for IcebergWriteExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> DFResult, - ) -> DFResult { + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { Ok(TreeNodeRecursion::Continue) } fn with_new_children( self: Arc, children: Vec>, - ) -> DFResult> { + ) -> Result> { if children.len() != 1 { - return Err(DataFusionError::Internal(format!( + return internal_err!( "IcebergWriteExec expects exactly one child, but provided {}", children.len() - ))); + ); } Ok(Arc::new(Self::new( @@ -208,7 +207,7 @@ impl ExecutionPlan for IcebergWriteExec { &self, partition: usize, context: Arc, - ) -> DFResult { + ) -> Result { let partition_type = self.table.metadata().default_partition_type().clone(); let format_version = self.table.metadata().format_version(); @@ -222,12 +221,9 @@ impl ExecutionPlan for IcebergWriteExec { let file_format = DataFileFormat::from_str(&write_format_default) .map_err(to_datafusion_error)?; if file_format != DataFileFormat::Parquet { - return Err(to_datafusion_error(Error::new( - ErrorKind::FeatureUnsupported, - format!( - "File format {file_format} is not supported for insert_into yet!" - ), - ))); + return not_impl_err!( + "File format {file_format} is not supported for insert_into yet!" + ); } // Build the writer from the already-parsed table properties so it honors @@ -275,8 +271,7 @@ impl ExecutionPlan for IcebergWriteExec { fanout_enabled, schema.clone(), partition_spec, - ) - .map_err(to_datafusion_error)?; + )?; // Get input data let data = execute_input_stream( @@ -293,13 +288,10 @@ impl ExecutionPlan for IcebergWriteExec { while let Some(batch) = input_stream.next().await { let batch = batch?; - task_writer - .write(batch) - .await - .map_err(to_datafusion_error)?; + task_writer.write(batch).await?; } - let data_files = task_writer.close().await.map_err(to_datafusion_error)?; + let data_files = task_writer.close().await?; // Convert builders to data files and then to JSON strings let data_files_strs: Vec = data_files @@ -312,7 +304,7 @@ impl ExecutionPlan for IcebergWriteExec { ) .map_err(to_datafusion_error) }) - .collect::>>()?; + .collect::>>()?; Self::make_result_batch(data_files_strs) }) @@ -335,7 +327,7 @@ mod tests { use datafusion::arrow::datatypes::{ DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef, }; - use datafusion::common::Result as DFResult; + use datafusion::error::Result; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -350,7 +342,8 @@ mod tests { deserialize_data_file_from_json, }; use iceberg::{ - Catalog, CatalogBuilder, MemoryCatalog, NamespaceIdent, Result, TableCreation, + Catalog, CatalogBuilder, Error, ErrorKind, MemoryCatalog, NamespaceIdent, + Result as IcebergResult, TableCreation, }; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use tempfile::TempDir; @@ -414,15 +407,15 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> DFResult, - ) -> DFResult { + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { Ok(TreeNodeRecursion::Continue) } fn with_new_children( self: Arc, _children: Vec>, - ) -> DFResult> { + ) -> Result> { Ok(self) } @@ -430,7 +423,7 @@ mod tests { &self, _partition: usize, _context: Arc, - ) -> DFResult { + ) -> Result { let batches = self.batches.clone(); let stream = stream::iter(batches.into_iter().map(Ok)); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -458,7 +451,7 @@ mod tests { } /// Helper function to create a test table schema - fn get_test_schema() -> Result { + fn get_test_schema() -> IcebergResult { Schema::builder() .with_schema_id(0) .with_fields(vec![ @@ -485,7 +478,7 @@ mod tests { } #[tokio::test] - async fn test_iceberg_write_exec() -> Result<()> { + async fn test_iceberg_write_exec() -> Result<(), Box> { // 1. Set up test environment let iceberg_catalog = get_iceberg_catalog().await; let namespace = NamespaceIdent::new("test_namespace".to_string()); @@ -654,7 +647,8 @@ mod tests { } #[tokio::test] - async fn test_iceberg_write_exec_advertises_result_schema() -> Result<()> { + async fn test_iceberg_write_exec_advertises_result_schema() + -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; let namespace = NamespaceIdent::new("test_namespace".to_string()); iceberg_catalog diff --git a/crates/datafusion/src/schema.rs b/crates/datafusion/src/schema.rs index 04bea0b..6bc4b4d 100644 --- a/crates/datafusion/src/schema.rs +++ b/crates/datafusion/src/schema.rs @@ -20,18 +20,17 @@ use std::sync::Arc; use async_trait::async_trait; use dashmap::DashMap; use datafusion::catalog::SchemaProvider; +use datafusion::common::{exec_datafusion_err, exec_err, plan_datafusion_err}; use datafusion::datasource::TableProvider; -use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::error::Result; use datafusion::execution::TaskContext; use datafusion::prelude::SessionContext; -use futures::StreamExt; +use futures::TryStreamExt; 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 iceberg::{Catalog, NamespaceIdent, TableCreation, TableIdent}; use crate::table::IcebergTableProvider; use crate::to_datafusion_error; @@ -69,7 +68,8 @@ impl IcebergSchemaProvider { // As of right now; tables might become stale. let table_names: Vec<_> = client .list_tables(&namespace) - .await? + .await + .map_err(to_datafusion_error)? .iter() .map(|tbl| tbl.name().to_string()) .collect(); @@ -122,15 +122,12 @@ impl SchemaProvider for IcebergSchemaProvider { } } - async fn table(&self, name: &str) -> DFResult>> { + async fn table(&self, name: &str) -> Result>> { if let Some((table_name, metadata_table_name)) = name.split_once('$') { let metadata_table_type = MetadataTableType::try_from(metadata_table_name) - .map_err(DataFusionError::Plan)?; + .map_err(|e| plan_datafusion_err!("{e}"))?; if let Some(table) = self.tables.get(table_name) { - let metadata_table = table - .metadata_table(metadata_table_type) - .await - .map_err(to_datafusion_error)?; + let metadata_table = table.metadata_table(metadata_table_type).await?; return Ok(Some(Arc::new(metadata_table))); } else { return Ok(None); @@ -147,12 +144,10 @@ impl SchemaProvider for IcebergSchemaProvider { &self, name: String, table: Arc, - ) -> DFResult>> { + ) -> Result>> { // Check if table already exists if self.table_exist(name.as_str()) { - return Err(DataFusionError::Execution(format!( - "Table {name} already exists" - ))); + return exec_err!("Table {name} already exists"); } // Convert DataFusion schema to Iceberg schema @@ -184,9 +179,7 @@ impl SchemaProvider for IcebergSchemaProvider { let rt = tokio::runtime::Handle::current(); rt.block_on(async move { // Verify the input table is empty - CREATE TABLE only accepts schema definition - ensure_table_is_empty(&table) - .await - .map_err(to_datafusion_error)?; + ensure_table_is_empty(&table).await?; catalog .create_table(&namespace, table_creation) @@ -199,8 +192,7 @@ impl SchemaProvider for IcebergSchemaProvider { namespace.clone(), name_clone.clone(), ) - .await - .map_err(to_datafusion_error)?; + .await?; // Store the new table provider tables.insert(name_clone, Arc::new(table_provider)); @@ -211,12 +203,11 @@ impl SchemaProvider for IcebergSchemaProvider { // Block on the spawned task to get the result // This is safe because spawn_blocking moves the blocking to a dedicated thread pool - futures::executor::block_on(result).map_err(|e| { - DataFusionError::Execution(format!("Failed to create Iceberg table: {e}")) - })? + futures::executor::block_on(result) + .map_err(|e| exec_datafusion_err!("Failed to create Iceberg table: {e}"))? } - fn deregister_table(&self, name: &str) -> DFResult>> { + fn deregister_table(&self, name: &str) -> Result>> { // Check if table exists if !self.table_exist(name) { return Ok(None); @@ -248,9 +239,8 @@ impl SchemaProvider for IcebergSchemaProvider { }) }); - futures::executor::block_on(result).map_err(|e| { - DataFusionError::Execution(format!("Failed to drop Iceberg table: {e}")) - })? + futures::executor::block_on(result) + .map_err(|e| exec_datafusion_err!("Failed to drop Iceberg table: {e}"))? } } @@ -258,32 +248,16 @@ impl SchemaProvider for IcebergSchemaProvider { /// Returns an error if the table has any rows. async fn ensure_table_is_empty(table: &Arc) -> Result<()> { let session_ctx = SessionContext::new(); - let exec_plan = table - .scan(&session_ctx.state(), None, &[], Some(1)) - .await - .map_err(|e| { - Error::new(ErrorKind::Unexpected, format!("Failed to scan table: {e}")) - })?; + let exec_plan = table.scan(&session_ctx.state(), None, &[], Some(1)).await?; let task_ctx = Arc::new(TaskContext::default()); - let stream = exec_plan.execute(0, task_ctx).map_err(|e| { - Error::new( - ErrorKind::Unexpected, - format!("Failed to execute scan: {e}"), - ) - })?; + let stream = exec_plan.execute(0, task_ctx)?; - let batches: Vec<_> = stream.collect().await; - let has_data = batches - .into_iter() - .filter_map(|r| r.ok()) - .any(|batch| batch.num_rows() > 0); + let batches: Vec<_> = stream.try_collect().await?; + let has_data = batches.iter().any(|batch| batch.num_rows() > 0); if has_data { - return Err(Error::new( - ErrorKind::Unexpected, - "register_table does not support tables with data.", - )); + return exec_err!("register_table does not support tables with data."); } Ok(()) diff --git a/crates/datafusion/src/table/metadata_table.rs b/crates/datafusion/src/table/metadata_table.rs index 56d0c10..e9ef9fa 100644 --- a/crates/datafusion/src/table/metadata_table.rs +++ b/crates/datafusion/src/table/metadata_table.rs @@ -22,7 +22,7 @@ use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; use datafusion::catalog::Session; use datafusion::datasource::{TableProvider, TableType}; -use datafusion::error::Result as DFResult; +use datafusion::error::Result; use datafusion::logical_expr::Expr; use datafusion::physical_plan::ExecutionPlan; use futures::TryStreamExt; @@ -64,13 +64,13 @@ impl TableProvider for IcebergMetadataTableProvider { _projection: Option<&Vec>, _filters: &[Expr], _limit: Option, - ) -> DFResult> { + ) -> Result> { Ok(Arc::new(IcebergMetadataScan::new(self.clone()))) } } impl IcebergMetadataTableProvider { - pub async fn scan(self) -> DFResult>> { + pub async fn scan(self) -> Result>> { let metadata_table = self.table.inspect(); let stream = match self.r#type { MetadataTableType::Snapshots => metadata_table.snapshots().scan().await, diff --git a/crates/datafusion/src/table/mod.rs b/crates/datafusion/src/table/mod.rs index a1db3c8..69f4b3a 100644 --- a/crates/datafusion/src/table/mod.rs +++ b/crates/datafusion/src/table/mod.rs @@ -34,9 +34,9 @@ use std::sync::Arc; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; use datafusion::catalog::Session; -use datafusion::common::DataFusionError; +use datafusion::common::{config_datafusion_err, exec_datafusion_err, not_impl_err}; use datafusion::datasource::{TableProvider, TableType}; -use datafusion::error::Result as DFResult; +use datafusion::error::Result; use datafusion::logical_expr::dml::InsertOp; use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; use datafusion::physical_plan::ExecutionPlan; @@ -45,7 +45,7 @@ use iceberg::arrow::schema_to_arrow_schema; use iceberg::inspect::MetadataTableType; use iceberg::spec::TableProperties; use iceberg::table::Table; -use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableIdent}; +use iceberg::{Catalog, NamespaceIdent, TableIdent}; use metadata_table::IcebergMetadataTableProvider; use crate::error::to_datafusion_error; @@ -87,8 +87,14 @@ impl IcebergTableProvider { let table_ident = TableIdent::new(namespace, name.into()); // Load table once to get initial schema - let table = catalog.load_table(&table_ident).await?; - let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?); + let table = catalog + .load_table(&table_ident) + .await + .map_err(to_datafusion_error)?; + let schema = Arc::new( + schema_to_arrow_schema(table.metadata().current_schema()) + .map_err(to_datafusion_error)?, + ); Ok(IcebergTableProvider { catalog, @@ -102,7 +108,11 @@ impl IcebergTableProvider { r#type: MetadataTableType, ) -> Result { // Load fresh table metadata for metadata table access - let table = self.catalog.load_table(&self.table_ident).await?; + let table = self + .catalog + .load_table(&self.table_ident) + .await + .map_err(to_datafusion_error)?; Ok(IcebergMetadataTableProvider { table, r#type }) } } @@ -123,7 +133,7 @@ impl TableProvider for IcebergTableProvider { projection: Option<&Vec>, filters: &[Expr], limit: Option, - ) -> DFResult> { + ) -> Result> { // Load fresh table metadata from catalog let table = self .catalog @@ -145,7 +155,7 @@ impl TableProvider for IcebergTableProvider { fn supports_filters_pushdown( &self, filters: &[&Expr], - ) -> DFResult> { + ) -> Result> { // Push down all filters, as a single source of truth, the scanner will drop the filters which couldn't be push down Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) } @@ -155,11 +165,11 @@ impl TableProvider for IcebergTableProvider { state: &dyn Session, input: Arc, _insert_op: InsertOp, - ) -> DFResult> { + ) -> Result> { if _insert_op != InsertOp::Append { - return Err(DataFusionError::NotImplemented(format!( + return not_impl_err!( "IcebergTableProvider supports only append inserts, got {_insert_op}" - ))); + ); } // Load fresh table metadata from catalog @@ -181,9 +191,7 @@ impl TableProvider for IcebergTableProvider { // Step 2: Repartition for parallel processing let target_partitions = NonZeroUsize::new(state.config().target_partitions()) .ok_or_else(|| { - DataFusionError::Configuration( - "target_partitions must be greater than 0".to_string(), - ) + config_datafusion_err!("target_partitions must be greater than 0") })?; let repartitioned_plan = @@ -195,19 +203,12 @@ impl TableProvider for IcebergTableProvider { .properties() .get(TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED) .map(|value| { - value - .parse::() - .map_err(|e| { - Error::new( - ErrorKind::DataInvalid, - format!( - "Invalid value for {}, expected 'true' or 'false'", - TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED - ), - ) - .with_source(e) - }) - .map_err(to_datafusion_error) + value.parse::().map_err(|e| { + config_datafusion_err!( + "Invalid value for {}, expected 'true' or 'false': {e}", + TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED + ) + }) }) .transpose()? .unwrap_or(TableProperties::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED_DEFAULT); @@ -255,7 +256,10 @@ impl IcebergStaticTableProvider { /// /// Uses the table's current snapshot for all queries. Does not support write operations. pub async fn try_new_from_table(table: Table) -> Result { - let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?); + let schema = Arc::new( + schema_to_arrow_schema(table.metadata().current_schema()) + .map_err(to_datafusion_error)?, + ); Ok(IcebergStaticTableProvider { table, snapshot_id: None, @@ -275,16 +279,16 @@ impl IcebergStaticTableProvider { .metadata() .snapshot_by_id(snapshot_id) .ok_or_else(|| { - Error::new( - ErrorKind::Unexpected, - format!( - "snapshot id {snapshot_id} not found in table {}", - table.identifier().name() - ), + exec_datafusion_err!( + "snapshot id {snapshot_id} not found in table {}", + table.identifier().name() ) })?; - let table_schema = snapshot.schema(table.metadata())?; - let schema = Arc::new(schema_to_arrow_schema(&table_schema)?); + let table_schema = snapshot + .schema(table.metadata()) + .map_err(to_datafusion_error)?; + let schema = + Arc::new(schema_to_arrow_schema(&table_schema).map_err(to_datafusion_error)?); Ok(IcebergStaticTableProvider { table, snapshot_id: Some(snapshot_id), @@ -309,7 +313,7 @@ impl TableProvider for IcebergStaticTableProvider { projection: Option<&Vec>, filters: &[Expr], limit: Option, - ) -> DFResult> { + ) -> Result> { // Use cached table (no refresh) Ok(Arc::new(IcebergTableScan::new( self.table.clone(), @@ -324,7 +328,7 @@ impl TableProvider for IcebergStaticTableProvider { fn supports_filters_pushdown( &self, filters: &[&Expr], - ) -> DFResult> { + ) -> Result> { // Push down all filters, as a single source of truth, the scanner will drop the filters which couldn't be push down Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) } @@ -334,13 +338,11 @@ impl TableProvider for IcebergStaticTableProvider { _state: &dyn Session, _input: Arc, _insert_op: InsertOp, - ) -> DFResult> { - Err(to_datafusion_error(Error::new( - ErrorKind::FeatureUnsupported, + ) -> Result> { + not_impl_err!( "Write operations are not supported on IcebergStaticTableProvider. \ Use IcebergTableProvider with a catalog for write support." - .to_string(), - ))) + ) } } @@ -350,6 +352,7 @@ mod tests { use std::sync::Arc; use datafusion::common::Column; + use datafusion::error::DataFusionError; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionContext; use iceberg::io::FileIO; diff --git a/crates/datafusion/src/table/table_provider_factory.rs b/crates/datafusion/src/table/table_provider_factory.rs index 656854d..1e7704b 100644 --- a/crates/datafusion/src/table/table_provider_factory.rs +++ b/crates/datafusion/src/table/table_provider_factory.rs @@ -21,12 +21,12 @@ use std::sync::Arc; use async_trait::async_trait; use datafusion::catalog::{Session, TableProvider, TableProviderFactory}; -use datafusion::common::TableReference; -use datafusion::error::Result as DFResult; +use datafusion::common::{TableReference, not_impl_err}; +use datafusion::error::Result; use datafusion::logical_expr::CreateExternalTable; +use iceberg::TableIdent; use iceberg::io::{FileIOBuilder, LocalFsStorageFactory, StorageFactory}; use iceberg::table::StaticTable; -use iceberg::{Error, ErrorKind, Result, TableIdent}; use super::IcebergStaticTableProvider; use crate::to_datafusion_error; @@ -122,8 +122,8 @@ impl TableProviderFactory for IcebergTableProviderFactory { &self, _state: &dyn Session, cmd: &CreateExternalTable, - ) -> DFResult> { - let metadata_file_path = check_cmd(cmd).map_err(to_datafusion_error)?; + ) -> Result> { + let metadata_file_path = check_cmd(cmd)?; let table_name = &cmd.name; let options = &cmd.options; @@ -141,13 +141,10 @@ impl TableProviderFactory for IcebergTableProviderFactory { options, storage_factory, ) - .await - .map_err(to_datafusion_error)? + .await? .into_table(); - let provider = IcebergStaticTableProvider::try_new_from_table(table) - .await - .map_err(to_datafusion_error)?; + let provider = IcebergStaticTableProvider::try_new_from_table(table).await?; Ok(Arc::new(provider)) } @@ -171,18 +168,16 @@ fn check_cmd(cmd: &CreateExternalTable) -> Result<&str> { || !column_defaults.is_empty(); if is_invalid { - return Err(Error::new( - ErrorKind::FeatureUnsupported, - "Currently we only support reading existing icebergs tables in external table command. To create new table, please use catalog provider.", - )); + return not_impl_err!( + "Currently we only support reading existing icebergs tables in external table command. To create new table, please use catalog provider." + ); } match cmd.locations.as_slice() { [location] => Ok(location), - _ => Err(Error::new( - ErrorKind::FeatureUnsupported, - "Iceberg external tables require exactly one metadata location.", - )), + _ => not_impl_err!( + "Iceberg external tables require exactly one metadata location." + ), } } @@ -213,11 +208,14 @@ async fn create_static_table( props: &HashMap, storage_factory: Arc, ) -> Result { - let table_ident = TableIdent::from_strs(table_name.to_vec())?; + let table_ident = + TableIdent::from_strs(table_name.to_vec()).map_err(to_datafusion_error)?; let file_io = FileIOBuilder::new(storage_factory) .with_props(props) .build(); - StaticTable::from_metadata_file(metadata_file_path, table_ident, file_io).await + StaticTable::from_metadata_file(metadata_file_path, table_ident, file_io) + .await + .map_err(to_datafusion_error) } #[cfg(test)] diff --git a/crates/datafusion/src/task_writer.rs b/crates/datafusion/src/task_writer.rs index 3e7c713..09a26b7 100644 --- a/crates/datafusion/src/task_writer.rs +++ b/crates/datafusion/src/task_writer.rs @@ -21,7 +21,7 @@ //! of RecordBatch data to Iceberg tables. use datafusion::arrow::array::RecordBatch; -use iceberg::Result; +use datafusion::error::Result; use iceberg::arrow::RecordBatchPartitionSplitter; use iceberg::spec::{DataFile, PartitionSpecRef, SchemaRef}; use iceberg::writer::IcebergWriterBuilder; @@ -30,6 +30,8 @@ use iceberg::writer::partitioning::clustered_writer::ClusteredWriter; use iceberg::writer::partitioning::fanout_writer::FanoutWriter; use iceberg::writer::partitioning::unpartitioned_writer::UnpartitionedWriter; +use crate::to_datafusion_error; + /// High-level writer for DataFusion that handles partitioning and routing of RecordBatch data. /// /// TaskWriter coordinates writing data to Iceberg tables by: @@ -138,7 +140,8 @@ impl TaskWriter { RecordBatchPartitionSplitter::try_new_with_precomputed_values( schema.clone(), partition_spec.clone(), - )?, + ) + .map_err(to_datafusion_error)?, ) } else { None @@ -183,7 +186,7 @@ impl TaskWriter { match &mut self.writer { SupportedWriter::Unpartitioned(writer) => { // Unpartitioned: write directly without splitting - writer.write(batch).await + writer.write(batch).await.map_err(to_datafusion_error) } SupportedWriter::Fanout(writer) => { Self::write_partitioned_batches(writer, &self.partition_splitter, &batch) @@ -220,11 +223,14 @@ impl TaskWriter { let splitter = partition_splitter .as_ref() .expect("Partition splitter should be initialized"); - let partitioned_batches = splitter.split(batch)?; + let partitioned_batches = splitter.split(batch).map_err(to_datafusion_error)?; // Write each partition for (partition_key, partition_batch) in partitioned_batches { - writer.write(partition_key, partition_batch).await?; + writer + .write(partition_key, partition_batch) + .await + .map_err(to_datafusion_error)?; } Ok(()) @@ -254,9 +260,11 @@ impl TaskWriter { /// ``` pub async fn close(self) -> Result> { match self.writer { - SupportedWriter::Unpartitioned(writer) => writer.close().await, - SupportedWriter::Fanout(writer) => writer.close().await, - SupportedWriter::Clustered(writer) => writer.close().await, + SupportedWriter::Unpartitioned(writer) => { + writer.close().await.map_err(to_datafusion_error) + } + SupportedWriter::Fanout(writer) => writer.close().await.map_err(to_datafusion_error), + SupportedWriter::Clustered(writer) => writer.close().await.map_err(to_datafusion_error), } } } @@ -287,7 +295,7 @@ mod tests { use super::*; - fn create_test_schema() -> Result> { + fn create_test_schema() -> iceberg::Result> { Ok(Arc::new( iceberg::spec::Schema::builder() .with_schema_id(1) @@ -358,7 +366,7 @@ mod tests { fn create_writer_builder( temp_dir: &TempDir, schema: Arc, - ) -> Result< + ) -> iceberg::Result< DataFileWriterBuilder< ParquetWriterBuilder, DefaultLocationGenerator, @@ -386,7 +394,7 @@ mod tests { } #[tokio::test] - async fn test_task_writer_unpartitioned() -> Result<()> { + async fn test_task_writer_unpartitioned() -> Result<(), Box> { let temp_dir = TempDir::new()?; let schema = create_test_schema()?; let arrow_schema = create_arrow_schema(); @@ -453,7 +461,7 @@ mod tests { } #[tokio::test] - async fn test_task_writer_partitioned_fanout() -> Result<()> { + async fn test_task_writer_partitioned_fanout() -> Result<(), Box> { let temp_dir = TempDir::new()?; let schema = create_test_schema()?; let arrow_schema = create_arrow_schema_with_partition(); @@ -504,7 +512,7 @@ mod tests { } #[tokio::test] - async fn test_task_writer_partitioned_clustered() -> Result<()> { + async fn test_task_writer_partitioned_clustered() -> Result<(), Box> { let temp_dir = TempDir::new()?; let schema = create_test_schema()?; let arrow_schema = create_arrow_schema_with_partition(); diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 41753ab..dcea6a5 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -18,6 +18,7 @@ //! Integration tests for Iceberg Datafusion with Hive Metastore. use std::collections::HashMap; +use std::error::Error; use std::sync::Arc; use std::vec; @@ -34,8 +35,8 @@ use iceberg::spec::{ }; use iceberg::test_utils::check_record_batches; use iceberg::{ - Catalog, CatalogBuilder, MemoryCatalog, NamespaceIdent, Result, TableCreation, - TableIdent, + Catalog, CatalogBuilder, MemoryCatalog, NamespaceIdent, Result as IcebergResult, + TableCreation, TableIdent, }; use tempfile::TempDir; @@ -65,7 +66,7 @@ fn get_struct_type() -> StructType { async fn set_test_namespace( catalog: &MemoryCatalog, namespace: &NamespaceIdent, -) -> Result<()> { +) -> IcebergResult<()> { let properties = HashMap::new(); catalog.create_namespace(namespace, properties).await?; @@ -77,7 +78,7 @@ fn get_table_creation( location: impl ToString, name: impl ToString, schema: Option, -) -> Result { +) -> IcebergResult { let schema = match schema { None => Schema::builder() .with_schema_id(0) @@ -102,7 +103,7 @@ fn get_table_creation( } #[tokio::test] -async fn test_provider_plan_stream_schema() -> Result<()> { +async fn test_provider_plan_stream_schema() -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; let namespace = NamespaceIdent::new("test_provider_get_table_schema".to_string()); set_test_namespace(&iceberg_catalog, &namespace).await?; @@ -155,7 +156,7 @@ async fn test_provider_plan_stream_schema() -> Result<()> { } #[tokio::test] -async fn test_provider_list_table_names() -> Result<()> { +async fn test_provider_list_table_names() -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; let namespace = NamespaceIdent::new("test_provider_list_table_names".to_string()); set_test_namespace(&iceberg_catalog, &namespace).await?; @@ -188,7 +189,7 @@ async fn test_provider_list_table_names() -> Result<()> { } #[tokio::test] -async fn test_provider_list_schema_names() -> Result<()> { +async fn test_provider_list_schema_names() -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; let namespace = NamespaceIdent::new("test_provider_list_schema_names".to_string()); set_test_namespace(&iceberg_catalog, &namespace).await?; @@ -213,7 +214,7 @@ async fn test_provider_list_schema_names() -> Result<()> { } #[tokio::test] -async fn test_table_projection() -> Result<()> { +async fn test_table_projection() -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; let namespace = NamespaceIdent::new("ns".to_string()); set_test_namespace(&iceberg_catalog, &namespace).await?; @@ -282,7 +283,7 @@ async fn test_table_projection() -> Result<()> { } #[tokio::test] -async fn test_table_predict_pushdown() -> Result<()> { +async fn test_table_predict_pushdown() -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; let namespace = NamespaceIdent::new("ns".to_string()); set_test_namespace(&iceberg_catalog, &namespace).await?; @@ -328,7 +329,7 @@ async fn test_table_predict_pushdown() -> Result<()> { } #[tokio::test] -async fn test_metadata_table() -> Result<()> { +async fn test_metadata_table() -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; let namespace = NamespaceIdent::new("ns".to_string()); set_test_namespace(&iceberg_catalog, &namespace).await?; @@ -455,7 +456,7 @@ async fn test_metadata_table() -> Result<()> { } #[tokio::test] -async fn test_insert_into() -> Result<()> { +async fn test_insert_into() -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; let namespace = NamespaceIdent::new("test_insert_into".to_string()); set_test_namespace(&iceberg_catalog, &namespace).await?; @@ -581,7 +582,7 @@ fn get_nested_struct_type() -> StructType { } #[tokio::test] -async fn test_insert_into_nested() -> Result<()> { +async fn test_insert_into_nested() -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; let namespace = NamespaceIdent::new("test_insert_nested".to_string()); set_test_namespace(&iceberg_catalog, &namespace).await?; @@ -839,7 +840,7 @@ async fn test_insert_into_nested() -> Result<()> { } #[tokio::test] -async fn test_insert_into_partitioned() -> Result<()> { +async fn test_insert_into_partitioned() -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; let namespace = NamespaceIdent::new("test_partitioned_write".to_string()); set_test_namespace(&iceberg_catalog, &namespace).await?; From b54edaec1ece35bd1e1aa22b535d4a96c502bcfe Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 22 Sep 2026 16:43:13 +0200 Subject: [PATCH 2/2] test: update DataFusion error expectations --- crates/datafusion/src/error.rs | 7 +++++-- crates/datafusion/src/task_writer.rs | 14 ++++++++++---- .../slts/df_test/timestamp_predicate_pushdown.slt | 4 ++-- 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/crates/datafusion/src/error.rs b/crates/datafusion/src/error.rs index 89674fe..9733521 100644 --- a/crates/datafusion/src/error.rs +++ b/crates/datafusion/src/error.rs @@ -97,7 +97,8 @@ mod tests { #[test] fn encodes_iceberg_errors_with_native_datafusion_variants() { - let error = to_datafusion_error(Error::new(ErrorKind::DataInvalid, "invalid manifest")); + let error = + to_datafusion_error(Error::new(ErrorKind::DataInvalid, "invalid manifest")); assert!(matches!( error, @@ -109,7 +110,9 @@ mod tests { #[test] fn maps_non_iceberg_datafusion_errors_to_unexpected() { - let error = from_datafusion_error(DataFusionError::Execution("worker failed".to_string())); + let error = from_datafusion_error(DataFusionError::Execution( + "worker failed".to_string(), + )); assert_eq!(error.kind(), ErrorKind::Unexpected); assert_eq!( diff --git a/crates/datafusion/src/task_writer.rs b/crates/datafusion/src/task_writer.rs index 09a26b7..45df287 100644 --- a/crates/datafusion/src/task_writer.rs +++ b/crates/datafusion/src/task_writer.rs @@ -263,8 +263,12 @@ impl TaskWriter { SupportedWriter::Unpartitioned(writer) => { writer.close().await.map_err(to_datafusion_error) } - SupportedWriter::Fanout(writer) => writer.close().await.map_err(to_datafusion_error), - SupportedWriter::Clustered(writer) => writer.close().await.map_err(to_datafusion_error), + SupportedWriter::Fanout(writer) => { + writer.close().await.map_err(to_datafusion_error) + } + SupportedWriter::Clustered(writer) => { + writer.close().await.map_err(to_datafusion_error) + } } } } @@ -461,7 +465,8 @@ mod tests { } #[tokio::test] - async fn test_task_writer_partitioned_fanout() -> Result<(), Box> { + async fn test_task_writer_partitioned_fanout() + -> Result<(), Box> { let temp_dir = TempDir::new()?; let schema = create_test_schema()?; let arrow_schema = create_arrow_schema_with_partition(); @@ -512,7 +517,8 @@ mod tests { } #[tokio::test] - async fn test_task_writer_partitioned_clustered() -> Result<(), Box> { + async fn test_task_writer_partitioned_clustered() + -> Result<(), Box> { let temp_dir = TempDir::new()?; let schema = create_test_schema()?; let arrow_schema = create_arrow_schema_with_partition(); diff --git a/crates/sqllogictest/testdata/slts/df_test/timestamp_predicate_pushdown.slt b/crates/sqllogictest/testdata/slts/df_test/timestamp_predicate_pushdown.slt index ffa7417..8b1f833 100644 --- a/crates/sqllogictest/testdata/slts/df_test/timestamp_predicate_pushdown.slt +++ b/crates/sqllogictest/testdata/slts/df_test/timestamp_predicate_pushdown.slt @@ -174,10 +174,10 @@ DROP TABLE default.default.test_timestamp_micros # Test with TIMESTAMP(3) - millisecond precision # This should fail because Iceberg doesn't support millisecond precision -statement error DataFusion error: External error: DataInvalid => Unsupported Arrow data type: Timestamp\(ms\) +statement error (?s)DataFusion error: IcebergError\(DataInvalid\).*Unsupported Arrow data type: Timestamp\(ms\) CREATE TABLE default.default.test_timestamp_millis (id INT NOT NULL, ts TIMESTAMP(3)) # Test with TIMESTAMP(0) - second precision # This should fail because Iceberg doesn't support second precision -statement error DataFusion error: External error: DataInvalid => Unsupported Arrow data type: Timestamp\(s\) +statement error (?s)DataFusion error: IcebergError\(DataInvalid\).*Unsupported Arrow data type: Timestamp\(s\) CREATE TABLE default.default.test_timestamp_seconds (id INT NOT NULL, ts TIMESTAMP(0))