From c1b30c73a568f3f750cedf5f2d8f69ef0f8ca8e4 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sun, 6 Sep 2026 21:47:12 +0300 Subject: [PATCH] fix: limit INSERT field casts to extensions --- datafusion/sql/src/statement.rs | 4 +++- datafusion/sql/tests/common/mod.rs | 23 ++++++++++++++++++++++ datafusion/sql/tests/sql_integration.rs | 26 +++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 1 deletion(-) diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 179d1b9de9b9..11c86a350ae5 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -2961,7 +2961,9 @@ impl SqlToRel<'_, S> { let (_, expr_field) = expr.to_field(source.schema())?; // A storage-type cast alone does not apply extension metadata from the // table schema when the source and target storage types are identical. - let expr = if expr_field.metadata() == target_field.metadata() { + let expr = if target_field.extension_type_name().is_none() + || expr_field.metadata() == target_field.metadata() + { expr } else { match expr { diff --git a/datafusion/sql/tests/common/mod.rs b/datafusion/sql/tests/common/mod.rs index 3f51912899ff..6f40e2f1efb8 100644 --- a/datafusion/sql/tests/common/mod.rs +++ b/datafusion/sql/tests/common/mod.rs @@ -187,6 +187,15 @@ impl ContextProvider for MockContextProvider { Field::new("first_name", DataType::Utf8, false), Field::new("last_name", DataType::Utf8, false), ])), + "string_with_extension" => Ok(Schema::new(vec![ + Field::new("value", DataType::Utf8, false).with_metadata( + [( + "ARROW:extension:name".to_string(), + "example.string".to_string(), + )] + .into(), + ), + ])), "orders" => Ok(Schema::new(vec![ Field::new("order_id", DataType::UInt32, false), Field::new("o_orderkey", DataType::UInt32, false), @@ -227,6 +236,20 @@ impl ContextProvider for MockContextProvider { false, ), ])), + "array_with_field_metadata" => Ok(Schema::new(vec![ + Field::new( + "left", + DataType::List(Arc::new( + Field::new_list_field(DataType::Int64, true).with_metadata( + [("PARQUET:field_id".to_string(), "2".to_string())].into(), + ), + )), + false, + ) + .with_metadata( + [("PARQUET:field_id".to_string(), "1".to_string())].into(), + ), + ])), "lineitem" => Ok(Schema::new(vec![ Field::new("l_orderkey", DataType::UInt32, false), Field::new("l_item_id", DataType::UInt32, false), diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index ee4b0de07a1d..275332a73678 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -816,6 +816,32 @@ fn plan_insert_preserves_target_extension_metadata() { ); } +#[test] +fn plan_insert_preserves_target_extension_metadata_on_type_cast() { + let sql = "INSERT INTO string_with_extension SELECT id FROM test_decimal"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Dml: op=[Insert Into] table=[string_with_extension] + Projection: CAST(test_decimal.id AS Utf8<{"ARROW:extension:name": "example.string"}>) AS value + Projection: test_decimal.id + TableScan: test_decimal + "# + ); +} + +#[test] +fn plan_insert_does_not_promise_ordinary_target_field_metadata() { + let sql = "INSERT INTO array_with_field_metadata SELECT left FROM array"; + let plan = logical_plan(sql).unwrap(); + let LogicalPlan::Dml(dml) = &plan else { + panic!("expected DML plan"); + }; + + assert!(dml.input.schema().field(0).metadata().is_empty()); +} + #[rstest] #[case::duplicate_columns( "INSERT INTO test_decimal (id, price, price) VALUES (1, 2, 3), (4, 5, 6)",