Honor column defaults on Insert

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;

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<String, Expr>,
 }
 
 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,

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<String, Expr> {
+    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<Schema> {
+    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::<Result<Vec<_>>>()?;
+
+    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<Option<Literal>> {
+    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<ScalarValue> {
+    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)
+}

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;