diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 1a9072212f2f3..7f010c8872c1e 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, @@ -2953,6 +2953,25 @@ 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 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::>>()?; 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 e7c819bbf64a6..66447afd730b5 100644 --- a/datafusion/sql/tests/common/mod.rs +++ b/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), diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 9f57aaafb0686..98e8f320f4a68 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/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)",