fix: preserve target extension metadata in INSERT

datafusion/sql/src/statement.rs

@@ -48,7 +48,7 @@ use datafusion_expr::logical_plan::DdlStatement;
 use datafusion_expr::logical_plan::builder::project;
 use datafusion_expr::utils::expr_to_columns;
 use datafusion_expr::{
-    Analyze, CreateCatalog, CreateCatalogSchema,
+    Analyze, Cast, CreateCatalog, CreateCatalogSchema,
     CreateExternalTable as PlanCreateExternalTable, CreateFunction, CreateFunctionBody,
     CreateIndex as PlanCreateIndex, CreateMemoryTable, CreateView, Deallocate,
     DescribeTable, DmlStatement, DropCatalogSchema, DropFunction, DropTable, DropView,
@@ -2953,6 +2953,25 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                         })
                         .cast_to(target_field.data_type(), &DFSchema::empty())?,
                 };
+                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 target_field.extension_type_name().is_none()
+                    || expr_field.metadata() == target_field.metadata()
+                {
+                    expr
+                } else {
+                    match expr {
+                        Expr::Cast(cast) => Expr::Cast(Cast::new_from_field(
+                            cast.expr,
+                            Arc::clone(target_field),
+                        )),
+                        expr => Expr::Cast(Cast::new_from_field(
+                            Box::new(expr),
+                            Arc::clone(target_field),
+                        )),
+                    }
+                };
                 Ok(expr.alias(target_field.name()))
             })
             .collect::<Result<Vec<Expr>>>()?;

datafusion/sql/tests/cases/params.rs

@@ -837,11 +837,11 @@ fn test_insert_infer_with_metadata() {
         @r#"
     ** Initial Plan:
     Dml: op=[Insert Into] table=[person_with_uuid_extension]
-      Projection: column1 AS id, column2 AS first_name, column3 AS last_name
+      Projection: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name
         Values: ($1, $2, $3)
     ** Final Plan:
     Dml: op=[Insert Into] table=[person_with_uuid_extension]
-      Projection: column1 AS id, column2 AS first_name, column3 AS last_name
+      Projection: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name
         Values: (FixedSizeBinary(16, "1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16") FieldMetadata { inner: {"ARROW:extension:name": "arrow.uuid"} } AS $1, Utf8("Alan") AS $2, Utf8("Turing") AS $3)
     "#
     );
@@ -859,11 +859,11 @@ fn test_insert_infer_with_metadata() {
     ** Initial Plan:
     Prepare: "my_plan" [FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>, Utf8, Utf8]
       Dml: op=[Insert Into] table=[person_with_uuid_extension]
-        Projection: column1 AS id, column2 AS first_name, column3 AS last_name
+        Projection: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name
           Values: ($1, $2, $3)
     ** Final Plan:
     Dml: op=[Insert Into] table=[person_with_uuid_extension]
-      Projection: column1 AS id, column2 AS first_name, column3 AS last_name
+      Projection: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name
         Values: (FixedSizeBinary(16, "1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16") FieldMetadata { inner: {"ARROW:extension:name": "arrow.uuid"} } AS $1, Utf8("Alan") AS $2, Utf8("Turing") AS $3)
     "#
     );

datafusion/sql/tests/common/mod.rs

@@ -174,6 +174,20 @@ impl ContextProvider for MockContextProvider {
                 Field::new("first_name", DataType::Utf8, false),
                 Field::new("last_name", DataType::Utf8, false),
             ])),
+            "person_with_binary_id" => Ok(Schema::new(vec![
+                Field::new("id", DataType::FixedSizeBinary(16), false),
+                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),
@@ -214,6 +228,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),

datafusion/sql/tests/sql_integration.rs

@@ -719,6 +719,48 @@ fn plan_insert_no_target_columns() {
     );
 }
 
+#[test]
+fn plan_insert_preserves_target_extension_metadata() {
+    let sql = "INSERT INTO person_with_uuid_extension \
+               SELECT id, first_name, last_name FROM person_with_binary_id";
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Dml: op=[Insert Into] table=[person_with_uuid_extension]
+      Projection: CAST(person_with_binary_id.id AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, person_with_binary_id.first_name AS first_name, person_with_binary_id.last_name AS last_name
+        Projection: person_with_binary_id.id, person_with_binary_id.first_name, person_with_binary_id.last_name
+          TableScan: person_with_binary_id
+    "#
+    );
+}
+
+#[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)",