Skip to content

feat(datafusion): fill omitted INSERT columns with Iceberg write-default values #2804

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
moomindani wants to merge 1 commit into from
+293 −3
Closed

feat(datafusion): fill omitted INSERT columns with Iceberg write-default values #2804

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions crates/integrations/datafusion/public-api.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<dyn datafusion_physical_plan::execution_plan::ExecutionPlan>, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin<alloc::boxed::Box<(dyn core::future::future::Future<Output = datafusion_common::error::Result<alloc::sync::Arc<dyn datafusion_physical_plan::execution_plan::ExecutionPlan>>> + 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<usize>>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option<usize>) -> core::pin::Pin<alloc::boxed::Box<(dyn core::future::future::Future<Output = datafusion_common::error::Result<alloc::sync::Arc<dyn datafusion_physical_plan::execution_plan::ExecutionPlan>>> + 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
Expand Down Expand Up @@ -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<dyn datafusion_physical_plan::execution_plan::ExecutionPlan>, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin<alloc::boxed::Box<(dyn core::future::future::Future<Output = datafusion_common::error::Result<alloc::sync::Arc<dyn datafusion_physical_plan::execution_plan::ExecutionPlan>>> + 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<usize>>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option<usize>) -> core::pin::Pin<alloc::boxed::Box<(dyn core::future::future::Future<Output = datafusion_common::error::Result<alloc::sync::Arc<dyn datafusion_physical_plan::execution_plan::ExecutionPlan>>> + 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
Expand Down
294 changes: 291 additions & 3 deletions crates/integrations/datafusion/src/table/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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;
Expand Down Expand Up @@ -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<String, Expr>,
}

impl IcebergTableProvider {
Expand All @@ -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,
})
}

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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<String, Expr> {
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<ScalarValue> {

@xanderbailey xanderbailey Aug 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

literal_to_scalar_value re-derives the (PrimitiveType, PrimitiveLiteral) → arrow-type mapping that already exists in iceberg-core in a couple of places: get_arrow_datum (crates/iceberg/src/arrow/schema.rs) and create_primitive_array_single_element (crates/iceberg/src/arrow/value.rs). This is effectively a fourth copy of that knowledge, and it can drift e.g. this match already handles Time, which get_arrow_datum currently doesn't.

One option that reuses the tested mappings and lets DataFusion do the array→scalar step:

let arrow_type = type_to_arrow_type(field_type)?;              // already pub
let array = create_primitive_array_single_element(&arrow_type, &Some(lit))?;
ScalarValue::try_from_array(&array, 0).ok()

Bonus: the resulting scalar's arrow type already matches the column (LargeBinary, FixedSizeBinary(len), UTC tz, …), which removes the reliance on the downstream cast for reconciliation.

Tradeoff: create_primitive_array_single_element is pub(crate), so this needs promoting it to pub in iceberg-core. If you'd rather not expand the core crate's public surface, keeping this match is reasonable, it's self-contained. Flagging mainly so the duplication is a conscious choice. @CTTY or @blackmwk might have stronger opinions.

@moomindani moomindani Aug 21, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for flagging this. I looked into the reuse path before deciding, and the composition would regress today.

type_to_arrow_type maps Time → Time64(Microsecond), Uuid → FixedSizeBinary(16), Fixed(len) → FixedSizeBinary(len) and Binary → LargeBinary. But create_primitive_array_single_element (crates/iceberg/src/arrow/value.rs:627) only has arms for Boolean, Int32, Date32, Int64, Timestamp(us|ns, tz), Float32, Float64, Utf8, Binary, Decimal128 and Struct(None), and falls through to Err("Unsupported constant type combination"). There is no Time64 arm, no FixedSizeBinary arm and no LargeBinary arm. Since column_defaults_from_schema uses filter_map, those four types would silently lose their write-default and fall back to NULL.

So the drift runs in both directions: this match handles Time precisely because core does not, and the same holds for Uuid, Fixed and Binary.

I kept the local match and made the duplication explicit in the doc comment. What I did take from your comment is the arrow-type alignment — see the two replies below: Binary now produces LargeBinary, and Fixed(len) is sized from the declared width, so the scalars carry exactly the types type_to_arrow_type assigns and the downstream cast has nothing left to reconcile. A new unit test asserts scalar.data_type() == type_to_arrow_type(field_type) for all 16 supported combinations, so a future divergence fails the build.

Happy to do the core side as a follow-up if you would like the single source of truth: add the Time64 / FixedSizeBinary / LargeBinary arms to create_primitive_array_single_element, promote it to pub, then reduce this function to your three lines. That expands iceberg-core's public surface, so it seemed better as its own PR than folded in here — but I am happy either way, and would defer to @CTTY / @blackmwk on whether core should export it.

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)) => {

@xanderbailey xanderbailey Aug 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor: Binary maps to ScalarValue::Binary here, whereas iceberg-core's type_to_arrow_type maps Binary → LargeBinary (crates/iceberg/src/arrow/schema.rs). The downstream cast reconciles it, so not a correctness issue but this is another spot where reusing the shared mapping would keep things aligned.

@moomindani moomindani Aug 21, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed — it emits ScalarValue::LargeBinary now, matching type_to_arrow_type. The new test_literal_to_scalar_value_matches_column_arrow_type asserts the scalar's arrow type equals type_to_arrow_type(field_type) for every supported primitive, so this cannot drift back unnoticed.

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
Expand Down Expand Up @@ -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};
Expand Down Expand Up @@ -584,6 +683,195 @@ mod tests {
assert!(execution_result.is_ok());
}

async fn get_test_catalog_and_table_with_write_defaults()
-> (Arc<dyn Catalog>, 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;
Expand Down
Loading