diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index c25756a3fcff9..179d1b9de9b91 100644 --- a/datafusion/sql/src/statement.rs +++ b/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, @@ -2958,6 +2958,23 @@ impl 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 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::>>()?; diff --git a/datafusion/sql/tests/cases/params.rs b/datafusion/sql/tests/cases/params.rs index 68c560ead68cd..8b025d3945db1 100644 --- a/datafusion/sql/tests/cases/params.rs +++ b/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) "# ); diff --git a/datafusion/sql/tests/common/mod.rs b/datafusion/sql/tests/common/mod.rs index f8b88498ff51b..3f51912899ff3 100644 --- a/datafusion/sql/tests/common/mod.rs +++ b/datafusion/sql/tests/common/mod.rs @@ -182,6 +182,11 @@ 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), + ])), "orders" => Ok(Schema::new(vec![ Field::new("order_id", DataType::UInt32, false), Field::new("o_orderkey", DataType::UInt32, false), diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index bbda33dac95fd..d387f6decaa8b 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -797,6 +797,22 @@ fn plan_insert_select_expression_from_values() { assert_contains!(plan.display_indent().to_string(), "make_array(column1)"); } +#[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 + "# + ); +} + #[rstest] #[case::duplicate_columns( "INSERT INTO test_decimal (id, price, price) VALUES (1, 2, 3), (4, 5, 6)",