diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index 5fcb4e85d3..e7fcc4e3e7 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -76,6 +76,7 @@ pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafus impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_session::table::TableProvider for iceberg_datafusion::IcebergTableProvider +pub fn iceberg_datafusion::IcebergTableProvider::get_column_default(&self, column: &str) -> core::option::Option<&datafusion_expr::expr::Expr> pub fn iceberg_datafusion::IcebergTableProvider::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, input: alloc::sync::Arc, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait pub fn iceberg_datafusion::IcebergTableProvider::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::IcebergTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef @@ -120,6 +121,7 @@ pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafus impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_session::table::TableProvider for iceberg_datafusion::IcebergTableProvider +pub fn iceberg_datafusion::IcebergTableProvider::get_column_default(&self, column: &str) -> core::option::Option<&datafusion_expr::expr::Expr> pub fn iceberg_datafusion::IcebergTableProvider::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, input: alloc::sync::Arc, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait pub fn iceberg_datafusion::IcebergTableProvider::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::IcebergTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 9de7bcb9c2..603fe22c28 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -28,6 +28,7 @@ pub mod metadata_table; pub mod table_provider_factory; +use std::collections::HashMap; use std::num::NonZeroUsize; use std::sync::Arc; @@ -41,9 +42,12 @@ use datafusion::logical_expr::dml::InsertOp; use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; -use iceberg::arrow::schema_to_arrow_schema; +use datafusion::scalar::ScalarValue; +use iceberg::arrow::{UTC_TIME_ZONE, schema_to_arrow_schema}; use iceberg::inspect::MetadataTableType; -use iceberg::spec::TableProperties; +use iceberg::spec::{ + Literal, PrimitiveLiteral, PrimitiveType, Schema as IcebergSchema, TableProperties, Type, +}; use iceberg::table::Table; use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableIdent}; use metadata_table::IcebergMetadataTableProvider; @@ -72,6 +76,10 @@ pub struct IcebergTableProvider { table_ident: TableIdent, /// A reference-counted arrow `Schema` (cached at construction) schema: ArrowSchemaRef, + /// Column default expressions derived from the schema's `write-default` values + /// (cached at construction). Consulted by DataFusion's insert planner for + /// columns omitted from an `INSERT`. + column_defaults: HashMap, } impl IcebergTableProvider { @@ -88,12 +96,15 @@ impl IcebergTableProvider { // Load table once to get initial schema let table = catalog.load_table(&table_ident).await?; - let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?); + let iceberg_schema = table.metadata().current_schema(); + let schema = Arc::new(schema_to_arrow_schema(iceberg_schema)?); + let column_defaults = column_defaults_from_schema(iceberg_schema); Ok(IcebergTableProvider { catalog, table_ident, schema, + column_defaults, }) } @@ -150,6 +161,10 @@ impl TableProvider for IcebergTableProvider { Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) } + fn get_column_default(&self, column: &str) -> Option<&Expr> { + self.column_defaults.get(column) + } + async fn insert_into( &self, state: &dyn Session, @@ -232,6 +247,89 @@ impl TableProvider for IcebergTableProvider { } } +/// Collects the `write-default` values of a schema's top-level columns as DataFusion +/// expressions, keyed by column name. +/// +/// Per the spec, writers must use `write-default` for columns that are not supplied; +/// DataFusion's insert planner consults these defaults for columns omitted from an +/// `INSERT` and falls back to `NULL` otherwise. Defaults of types that cannot be +/// expressed as a DataFusion scalar are skipped. +fn column_defaults_from_schema(schema: &IcebergSchema) -> HashMap { + schema + .as_struct() + .fields() + .iter() + .filter_map(|field| { + let literal = field.write_default.as_ref()?; + let scalar = literal_to_scalar_value(&field.field_type, literal)?; + Some((field.name.clone(), Expr::Literal(scalar, None))) + }) + .collect() +} + +/// Converts an Iceberg literal of the given type into a DataFusion [`ScalarValue`]. +/// +/// Returns `None` for combinations that have no scalar representation. +/// +/// The scalars this produces carry the same arrow types that +/// [`iceberg::arrow::type_to_arrow_type`] assigns to the columns, so the insert planner +/// has nothing left to reconcile. Building them by going through +/// `create_primitive_array_single_element` in iceberg-core instead would currently drop +/// the `Time`, `Uuid`, `Fixed` and `Binary` defaults, because that function has no arm +/// for `Time64`, `FixedSizeBinary` or `LargeBinary`; keeping the mapping here is a +/// deliberate choice until those arms exist in core. +fn literal_to_scalar_value(field_type: &Type, literal: &Literal) -> Option { + let Type::Primitive(primitive_type) = field_type else { + return None; + }; + let Literal::Primitive(primitive) = literal else { + return None; + }; + Some(match (primitive_type, primitive) { + (PrimitiveType::Boolean, PrimitiveLiteral::Boolean(v)) => ScalarValue::Boolean(Some(*v)), + (PrimitiveType::Int, PrimitiveLiteral::Int(v)) => ScalarValue::Int32(Some(*v)), + (PrimitiveType::Long, PrimitiveLiteral::Long(v)) => ScalarValue::Int64(Some(*v)), + (PrimitiveType::Float, PrimitiveLiteral::Float(v)) => ScalarValue::Float32(Some(v.0)), + (PrimitiveType::Double, PrimitiveLiteral::Double(v)) => ScalarValue::Float64(Some(v.0)), + (PrimitiveType::String, PrimitiveLiteral::String(v)) => ScalarValue::Utf8(Some(v.clone())), + (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => ScalarValue::Date32(Some(*v)), + (PrimitiveType::Time, PrimitiveLiteral::Long(v)) => { + ScalarValue::Time64Microsecond(Some(*v)) + } + (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => { + ScalarValue::TimestampMicrosecond(Some(*v), None) + } + (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => { + ScalarValue::TimestampMicrosecond(Some(*v), Some(UTC_TIME_ZONE.into())) + } + (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => { + ScalarValue::TimestampNanosecond(Some(*v), None) + } + (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => { + ScalarValue::TimestampNanosecond(Some(*v), Some(UTC_TIME_ZONE.into())) + } + (PrimitiveType::Decimal { precision, scale }, PrimitiveLiteral::Int128(v)) => { + ScalarValue::Decimal128(Some(*v), *precision as u8, *scale as i8) + } + (PrimitiveType::Binary, PrimitiveLiteral::Binary(v)) => { + ScalarValue::LargeBinary(Some(v.clone())) + } + (PrimitiveType::Fixed(len), PrimitiveLiteral::Binary(v)) => { + // Size the scalar from the column's declared width rather than the value's + // length, and skip a default whose length contradicts the declaration + let width = i32::try_from(*len).ok()?; + if v.len() != usize::try_from(*len).ok()? { + return None; + } + ScalarValue::FixedSizeBinary(width, Some(v.clone())) + } + (PrimitiveType::Uuid, PrimitiveLiteral::UInt128(v)) => { + ScalarValue::FixedSizeBinary(16, Some(v.to_be_bytes().to_vec())) + } + _ => return None, + }) +} + /// Static table provider for read-only snapshot access. /// /// This provider holds a cached table instance and does not refresh metadata or support @@ -349,6 +447,7 @@ mod tests { use datafusion::common::Column; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionContext; + use iceberg::arrow::type_to_arrow_type; use iceberg::io::FileIO; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{NestedField, PrimitiveType, Schema, Type}; @@ -584,6 +683,195 @@ mod tests { assert!(execution_result.is_ok()); } + async fn get_test_catalog_and_table_with_write_defaults() + -> (Arc, NamespaceIdent, String, TempDir) { + let temp_dir = TempDir::new().unwrap(); + let warehouse_path = temp_dir.path().to_str().unwrap().to_string(); + + let catalog = 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(), + NestedField::optional(2, "category", Type::Primitive(PrimitiveType::String)) + .with_write_default(Literal::string("general")) + .into(), + NestedField::optional(3, "score", Type::Primitive(PrimitiveType::Long)) + .with_write_default(Literal::long(100)) + .into(), + ]) + .build() + .unwrap(); + + let table_creation = TableCreation::builder() + .name("test_table".to_string()) + .location(format!("{warehouse_path}/test_table")) + .schema(schema) + .properties(HashMap::new()) + .build(); + + catalog + .create_table(&namespace, table_creation) + .await + .unwrap(); + + ( + Arc::new(catalog), + namespace, + "test_table".to_string(), + temp_dir, + ) + } + + #[test] + fn test_literal_to_scalar_value_matches_column_arrow_type() { + let uuid = "f79c3e09-677c-4bbd-a479-3f349cb785e7"; + let micros = 1_700_000_000i64; + let decimal = PrimitiveType::Decimal { + precision: 9, + scale: 2, + }; + let cases: Vec<(PrimitiveType, Literal)> = vec![ + (PrimitiveType::Boolean, Literal::bool(true)), + (PrimitiveType::Int, Literal::int(7)), + (PrimitiveType::Long, Literal::long(7i64)), + (PrimitiveType::Float, Literal::float(1.5f32)), + (PrimitiveType::Double, Literal::double(1.5f64)), + (PrimitiveType::String, Literal::string("general")), + (PrimitiveType::Date, Literal::date(19000)), + (PrimitiveType::Time, Literal::time(3_600_000_000i64)), + (PrimitiveType::Timestamp, Literal::timestamp(micros)), + (PrimitiveType::Timestamptz, Literal::timestamptz(micros)), + (PrimitiveType::TimestampNs, Literal::long(micros)), + (PrimitiveType::TimestamptzNs, Literal::long(micros)), + (decimal, Literal::decimal(12345)), + (PrimitiveType::Binary, Literal::binary(vec![1u8, 2, 3])), + (PrimitiveType::Fixed(4), Literal::fixed(vec![1u8, 2, 3, 4])), + (PrimitiveType::Uuid, Literal::uuid_from_str(uuid).unwrap()), + ]; + + for (primitive, literal) in cases { + let field_type = Type::Primitive(primitive); + let scalar = literal_to_scalar_value(&field_type, &literal) + .unwrap_or_else(|| panic!("no scalar produced for {field_type:?}")); + let expected = type_to_arrow_type(&field_type).unwrap(); + assert_eq!( + scalar.data_type(), + expected, + "the scalar must carry the column's arrow type for {field_type:?}" + ); + } + } + + #[test] + fn test_literal_to_scalar_value_skips_fixed_default_of_wrong_width() { + // a default whose length contradicts the declared width has no faithful scalar + let field_type = Type::Primitive(PrimitiveType::Fixed(4)); + assert!(literal_to_scalar_value(&field_type, &Literal::fixed(vec![1u8, 2, 3])).is_none()); + } + + #[tokio::test] + async fn test_insert_fills_write_default_for_omitted_columns() { + let (catalog, namespace, table_name, _temp_dir) = + get_test_catalog_and_table_with_write_defaults().await; + + let provider = + IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), table_name.clone()) + .await + .unwrap(); + + let ctx = SessionContext::new(); + ctx.register_table("test_table", Arc::new(provider)) + .unwrap(); + + // Omitted columns must be filled with their write-default, not NULL + ctx.sql("INSERT INTO test_table (id) VALUES (1)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + // Explicitly provided values must win over the defaults + ctx.sql("INSERT INTO test_table (id, category, score) VALUES (2, 'custom', 7)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let batches = ctx + .sql("SELECT id, category, score FROM test_table ORDER BY id") + .await + .unwrap() + .collect() + .await + .unwrap(); + + datafusion::assert_batches_eq!( + [ + "+----+----------+-------+", + "| id | category | score |", + "+----+----------+-------+", + "| 1 | general | 100 |", + "| 2 | custom | 7 |", + "+----+----------+-------+", + ], + &batches + ); + } + + #[test] + fn test_literal_to_scalar_value() { + assert_eq!( + literal_to_scalar_value( + &Type::Primitive(PrimitiveType::String), + &Literal::string("abc") + ), + Some(ScalarValue::Utf8(Some("abc".to_string()))) + ); + assert_eq!( + literal_to_scalar_value(&Type::Primitive(PrimitiveType::Int), &Literal::int(42)), + Some(ScalarValue::Int32(Some(42))) + ); + assert_eq!( + literal_to_scalar_value(&Type::Primitive(PrimitiveType::Long), &Literal::long(42)), + Some(ScalarValue::Int64(Some(42))) + ); + assert_eq!( + literal_to_scalar_value( + &Type::Primitive(PrimitiveType::Boolean), + &Literal::bool(true) + ), + Some(ScalarValue::Boolean(Some(true))) + ); + assert_eq!( + literal_to_scalar_value( + &Type::Primitive(PrimitiveType::Double), + &Literal::double(1.5) + ), + Some(ScalarValue::Float64(Some(1.5))) + ); + // a type/literal mismatch has no scalar representation + assert_eq!( + literal_to_scalar_value(&Type::Primitive(PrimitiveType::Int), &Literal::string("x")), + None + ); + } + #[tokio::test] async fn test_physical_input_schema_consistent_with_logical_input_schema() { let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await;