diff --git a/crates/datafusion/src/catalog/schema.rs b/crates/datafusion/src/catalog/schema.rs index 900e874..1668e3d 100644 --- a/crates/datafusion/src/catalog/schema.rs +++ b/crates/datafusion/src/catalog/schema.rs @@ -33,6 +33,7 @@ use iceberg::spec::FormatVersion; use iceberg::{Catalog, NamespaceIdent, TableCreation, TableIdent}; use crate::table_provider::IcebergCatalogTableProvider; +use crate::table_provider::column_defaults::with_write_defaults; use crate::to_datafusion_error; /// Represents a [`SchemaProvider`] for the Iceberg [`Catalog`], managing @@ -159,6 +160,7 @@ impl SchemaProvider for IcebergSchemaProvider { let df_schema = table.schema(); let iceberg_schema = arrow_schema_to_schema_auto_assign_ids(df_schema.as_ref()) .map_err(to_datafusion_error)?; + let iceberg_schema = with_write_defaults(iceberg_schema, table.as_ref())?; // Use at least V2, and upgrade to V3 if the schema requires it (e.g. timestamp_ns / variant). let format_version = iceberg_schema @@ -276,7 +278,9 @@ mod tests { use datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::MemTable; + use datafusion::prelude::lit; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; + use iceberg::spec::Literal; use iceberg::{Catalog, CatalogBuilder, NamespaceIdent}; use tempfile::TempDir; @@ -368,6 +372,43 @@ mod tests { assert!(schema_provider.table_exist("empty_table")); } + #[tokio::test] + async fn test_register_table_stores_column_defaults() { + let (schema_provider, _temp_dir) = create_test_schema_provider().await; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("count", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + ])); + let empty_batch = RecordBatch::new_empty(arrow_schema.clone()); + // The Int32 default for the Int64 column is cast to the column's type. + let mem_table = MemTable::try_new(arrow_schema, vec![vec![empty_batch]]) + .unwrap() + .with_column_defaults(HashMap::from([ + ("count".to_string(), lit(7i32)), + ("name".to_string(), lit("none")), + ])); + + schema_provider + .register_table("defaults_table".to_string(), Arc::new(mem_table)) + .unwrap(); + + let table_ident = + TableIdent::new(schema_provider.namespace.clone(), "defaults_table".into()); + let table = schema_provider + .catalog + .load_table(&table_ident) + .await + .unwrap(); + let schema = table.metadata().current_schema(); + let write_default = + |name: &str| schema.field_by_name(name).unwrap().write_default.clone(); + assert_eq!(write_default("id"), None); + assert_eq!(write_default("count"), Some(Literal::long(7))); + assert_eq!(write_default("name"), Some(Literal::string("none"))); + } + #[tokio::test] async fn test_register_duplicate_table_fails() { let (schema_provider, _temp_dir) = create_test_schema_provider().await; diff --git a/crates/datafusion/src/table_provider/catalog_table_provider.rs b/crates/datafusion/src/table_provider/catalog_table_provider.rs index 6dbb79d..8f9b07e 100644 --- a/crates/datafusion/src/table_provider/catalog_table_provider.rs +++ b/crates/datafusion/src/table_provider/catalog_table_provider.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::num::NonZeroUsize; use std::sync::Arc; @@ -37,6 +38,7 @@ use iceberg::{Catalog, NamespaceIdent, TableIdent}; use crate::error::to_datafusion_error; use crate::metadata::table::IcebergMetadataTableProvider; use crate::read::IcebergDataSource; +use crate::table_provider::column_defaults::column_defaults; use crate::write::commit::IcebergCommitExec; use crate::write::exec::IcebergWriteExec; use crate::write::project::project_with_partition; @@ -59,6 +61,8 @@ pub struct IcebergCatalogTableProvider { table_ident: TableIdent, /// A reference-counted arrow `Schema` (cached at construction) schema: ArrowSchemaRef, + /// The columns' `write-default` values, keyed by column name (cached at construction) + column_defaults: HashMap, } impl IcebergCatalogTableProvider { @@ -78,15 +82,16 @@ impl IcebergCatalogTableProvider { .load_table(&table_ident) .await .map_err(to_datafusion_error)?; + let current_schema = table.metadata().current_schema(); let schema = Arc::new( - schema_to_arrow_schema(table.metadata().current_schema()) - .map_err(to_datafusion_error)?, + schema_to_arrow_schema(current_schema).map_err(to_datafusion_error)?, ); Ok(IcebergCatalogTableProvider { catalog, table_ident, schema, + column_defaults: column_defaults(current_schema), }) } @@ -155,6 +160,10 @@ impl TableProvider for IcebergCatalogTableProvider { 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, diff --git a/crates/datafusion/src/table_provider/column_defaults.rs b/crates/datafusion/src/table_provider/column_defaults.rs new file mode 100644 index 0000000..a930a39 --- /dev/null +++ b/crates/datafusion/src/table_provider/column_defaults.rs @@ -0,0 +1,199 @@ +// 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. + +//! Conversions between DataFusion column defaults and Iceberg `write-default` values. +//! +//! A SQL `DEFAULT` is stored as the field's `write-default` when the table is created, +//! and each `write-default` is handed back to DataFusion, which fills it into columns +//! an INSERT omits. + +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::arrow::array::{RecordBatch, RecordBatchOptions}; +use datafusion::arrow::datatypes::Schema as ArrowSchema; +use datafusion::common::tree_node::TreeNode; +use datafusion::common::{DFSchema, ScalarValue, not_impl_err, plan_err}; +use datafusion::datasource::TableProvider; +use datafusion::error::Result; +use datafusion::logical_expr::execution_props::ExecutionProps; +use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; +use datafusion::logical_expr::{Expr, ExprSchemable, Volatility, lit}; +use datafusion::physical_expr::create_physical_expr; +use iceberg::arrow::{UTC_TIME_ZONE, arrow_primitive_to_literal, type_to_arrow_type}; +use iceberg::spec::{Literal, PrimitiveLiteral, PrimitiveType, Schema, Type}; + +use crate::to_datafusion_error; + +/// The `write-default` of each top-level field of `schema` that has one, as a literal +/// keyed by field name. +/// +/// Defaults that have no DataFusion literal, such as those of non-primitive fields, are +/// left out, so DataFusion fills NULL into those columns instead. +pub(crate) fn column_defaults(schema: &Schema) -> HashMap { + schema + .as_struct() + .fields() + .iter() + .filter_map( + |field| match (field.field_type.as_ref(), &field.write_default) { + (Type::Primitive(ty), Some(Literal::Primitive(value))) => { + Some((field.name.clone(), lit(primitive_to_scalar(ty, value)?))) + } + _ => None, + }, + ) + .collect() +} + +/// `schema` with each top-level field's `write-default` set to the `DEFAULT` that +/// `table` declares for the column of the same name. +/// +/// A `DEFAULT` must be a constant, because the table stores its value rather than the +/// expression. `DEFAULT NULL` is the same as no default. +pub(crate) fn with_write_defaults( + schema: Schema, + table: &dyn TableProvider, +) -> Result { + let fields = schema + .as_struct() + .fields() + .iter() + .map(|field| { + let Some(default) = table.get_column_default(&field.name) else { + return Ok(Arc::clone(field)); + }; + Ok( + match default_to_literal(&field.name, &field.field_type, default)? { + Some(value) => { + Arc::new(field.as_ref().clone().with_write_default(value)) + } + None => Arc::clone(field), + }, + ) + }) + .collect::>>()?; + + Schema::builder() + .with_schema_id(schema.schema_id()) + .with_identifier_field_ids(schema.identifier_field_ids()) + .with_fields(fields) + .build() + .map_err(to_datafusion_error) +} + +/// Evaluates the `DEFAULT` of column `name` to a value of type `ty`, or `None` for NULL. +fn default_to_literal(name: &str, ty: &Type, default: &Expr) -> Result> { + if !ty.is_primitive() { + return not_impl_err!("DEFAULT is not supported for column {name} of type {ty}"); + } + + // Functions that aren't immutable, such as `now()` or `random()`, would be evaluated + // once here instead of on every INSERT. + let is_constant = !default.exists(|expr| { + Ok(match expr { + Expr::Column(_) => true, + Expr::ScalarFunction(function) => { + function.func.signature().volatility != Volatility::Immutable + } + _ => false, + }) + })?; + if !is_constant { + return plan_err!("DEFAULT for column {name} must be a constant, got {default}"); + } + + let data_type = type_to_arrow_type(ty).map_err(to_datafusion_error)?; + let default = default.clone().cast_to(&data_type, &DFSchema::empty())?; + let one_row = RecordBatch::try_new_with_options( + Arc::new(ArrowSchema::empty()), + vec![], + &RecordBatchOptions::new().with_row_count(Some(1)), + )?; + let value = create_physical_expr( + &default, + &DFSchema::empty(), + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )? + .evaluate(&one_row)? + .into_array(1)?; + + Ok(arrow_primitive_to_literal(&value, ty) + .map_err(to_datafusion_error)? + .pop() + .flatten()) +} + +/// `value` as a scalar of the Arrow type that `schema_to_arrow_schema` maps `ty` to, or +/// `None` if `value` isn't a value of `ty`. +fn primitive_to_scalar( + ty: &PrimitiveType, + value: &PrimitiveLiteral, +) -> Option { + let utc = || Some(UTC_TIME_ZONE.into()); + let scalar = match (ty, value) { + (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.into_inner())) + } + (PrimitiveType::Double, PrimitiveLiteral::Double(v)) => { + ScalarValue::Float64(Some(v.into_inner())) + } + (PrimitiveType::Decimal { precision, scale }, PrimitiveLiteral::Int128(v)) => { + ScalarValue::Decimal128( + Some(*v), + u8::try_from(*precision).ok()?, + i8::try_from(*scale).ok()?, + ) + } + (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), utc()) + } + (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => { + ScalarValue::TimestampNanosecond(Some(*v), None) + } + (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => { + ScalarValue::TimestampNanosecond(Some(*v), utc()) + } + (PrimitiveType::String, PrimitiveLiteral::String(v)) => { + ScalarValue::Utf8(Some(v.clone())) + } + (PrimitiveType::Uuid, PrimitiveLiteral::UInt128(v)) => { + ScalarValue::FixedSizeBinary(16, Some(v.to_be_bytes().to_vec())) + } + (PrimitiveType::Fixed(len), PrimitiveLiteral::Binary(v)) => { + ScalarValue::FixedSizeBinary(i32::try_from(*len).ok()?, Some(v.clone())) + } + (PrimitiveType::Binary, PrimitiveLiteral::Binary(v)) => { + ScalarValue::LargeBinary(Some(v.clone())) + } + _ => return None, + }; + Some(scalar) +} diff --git a/crates/datafusion/src/table_provider/mod.rs b/crates/datafusion/src/table_provider/mod.rs index d3ee9f7..b602b38 100644 --- a/crates/datafusion/src/table_provider/mod.rs +++ b/crates/datafusion/src/table_provider/mod.rs @@ -21,6 +21,7 @@ //! The static provider reads a fixed table snapshot. mod catalog_table_provider; +pub(crate) mod column_defaults; mod factory; mod static_table_provider; diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 8d300a3..d6d745f 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -39,7 +39,8 @@ use expect_test::expect; use iceberg::io::LocalFsStorageFactory; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{ - NestedField, PrimitiveType, Schema, StructType, Transform, Type, UnboundPartitionSpec, + Literal, NestedField, PrimitiveType, Schema, StructType, Transform, Type, + UnboundPartitionSpec, }; use iceberg::test_utils::check_record_batches; use iceberg::transaction::{AddColumn, ApplyTransactionAction, Transaction}; @@ -557,6 +558,63 @@ async fn test_insert_into() -> Result<(), Box> { Ok(()) } +#[tokio::test] +async fn test_insert_into_fills_write_defaults() -> Result<(), Box> { + let iceberg_catalog = get_iceberg_catalog().await; + let namespace = NamespaceIdent::new("test_insert_write_defaults".to_string()); + set_test_namespace(&iceberg_catalog, &namespace).await?; + + // Defaults set through the Iceberg API, as another engine would, rather than by SQL. + let schema = Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "quantity", Type::Primitive(PrimitiveType::Int)) + .with_write_default(Literal::int(42)) + .into(), + NestedField::optional(3, "note", Type::Primitive(PrimitiveType::String)) + .with_write_default(Literal::string("none")) + .into(), + NestedField::optional(4, "comment", Type::Primitive(PrimitiveType::String)) + .into(), + ]) + .build()?; + let creation = get_table_creation(temp_path(), "my_table", Some(schema))?; + iceberg_catalog.create_table(&namespace, creation).await?; + + let catalog = + Arc::new(IcebergCatalogProvider::try_new(Arc::new(iceberg_catalog)).await?); + let ctx = SessionContext::new(); + ctx.register_catalog("catalog", catalog); + + ctx.sql("INSERT INTO catalog.test_insert_write_defaults.my_table (id) VALUES (1)") + .await? + .collect() + .await?; + ctx.sql( + "INSERT INTO catalog.test_insert_write_defaults.my_table (id, quantity, note) \ + VALUES (2, 3, 'given')", + ) + .await? + .collect() + .await?; + + let rows = ctx + .sql("SELECT * FROM catalog.test_insert_write_defaults.my_table ORDER BY id") + .await? + .collect() + .await?; + expect![[r#" + +----+----------+-------+---------+ + | id | quantity | note | comment | + +----+----------+-------+---------+ + | 1 | 42 | none | | + | 2 | 3 | given | | + +----+----------+-------+---------+"#]] + .assert_eq(&pretty_format_batches(&rows)?.to_string()); + Ok(()) +} + fn get_nested_struct_type() -> StructType { // Create a nested struct type with: // - address: STRUCT diff --git a/crates/sqllogictest/testdata/slts/df_test/create_table.slt b/crates/sqllogictest/testdata/slts/df_test/create_table.slt index 2eab1b6..b354b49 100644 --- a/crates/sqllogictest/testdata/slts/df_test/create_table.slt +++ b/crates/sqllogictest/testdata/slts/df_test/create_table.slt @@ -88,3 +88,36 @@ SELECT * FROM default.default.nullable_table ---- 1 Value 2 NULL + +# Test CREATE TABLE with column defaults +statement ok +CREATE TABLE default.default.defaults_table (id INT NOT NULL, quantity INT NOT NULL DEFAULT 7, active BOOLEAN NOT NULL DEFAULT true, note STRING DEFAULT 'none', comment STRING) + +# Omitted columns are filled with their defaults, or NULL if they have none +query I +INSERT INTO default.default.defaults_table (id) VALUES (1) +---- +1 + +# Explicit values take precedence over defaults +query I +INSERT INTO default.default.defaults_table (id, quantity, note, comment) VALUES (2, 3, 'given', 'c') +---- +1 + +query IIBTT rowsort +SELECT * FROM default.default.defaults_table +---- +1 7 true none NULL +2 3 true given c + +# Defaults are stored as values, so they must be constant +statement error (?s)DEFAULT for column created must be a constant, got now\(\) +CREATE TABLE default.default.volatile_default_table (id INT NOT NULL, created TIMESTAMP DEFAULT now()) + +statement error (?s)DEFAULT for column x must be a constant, got random\(\) +CREATE TABLE default.default.random_default_table (id INT NOT NULL, x DOUBLE DEFAULT random()) + +# A default that can't be cast to the column's type is rejected when the table is created +statement error (?s)Cannot cast string 'abc' to value of Int32 type +CREATE TABLE default.default.bad_default_table (id INT NOT NULL, x INT DEFAULT 'abc')