From c027b56c1a950334f78057f87fda5b34015195df Mon Sep 17 00:00:00 2001 From: osipovartem Date: Mon, 2 Nov 2026 01:23:26 +0300 Subject: [PATCH] fix: preserve supplied VALUES field metadata [skip ci] --- datafusion/expr/src/logical_plan/builder.rs | 91 ++++++++++++++++++++- datafusion/sql/tests/cases/params.rs | 46 ++++++++++- 2 files changed, 130 insertions(+), 7 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 1d845aa94980d..fbd94ec69bacd 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -275,8 +275,9 @@ impl LogicalPlanBuilder { let n_cols = values[0].len(); let mut fields = ValuesFields::new(); for j in 0..n_cols { - let field_type = schema.field(j).data_type(); - let field_nullable = schema.field(j).is_nullable(); + let field = schema.field(j); + let field_type = field.data_type(); + let field_nullable = field.is_nullable(); for row in values.iter() { let value = &row[j]; let data_type = value.get_type(schema)?; @@ -291,7 +292,15 @@ impl LogicalPlanBuilder { ); } } - fields.push(field_type.to_owned(), field_nullable); + if field.metadata().is_empty() { + fields.push(field_type.to_owned(), field_nullable); + } else { + fields.push_with_metadata( + field_type.to_owned(), + field_nullable, + Some(FieldMetadata::new_from_field(field)), + ); + } } Self::infer_inner(values, fields, schema) @@ -3316,6 +3325,82 @@ mod tests { ); } + #[test] + fn test_values_with_schema_preserves_field_metadata() -> Result<()> { + use arrow::array::{Int16Array, StructArray, TimestampMicrosecondArray}; + use arrow::datatypes::TimeUnit; + use arrow_schema::extension::TimestampWithOffset; + + let storage = DataType::Struct(Fields::from(vec![ + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + false, + ), + Field::new("offset_minutes", DataType::Int16, false), + ])); + let DataType::Struct(children) = &storage else { + unreachable!(); + }; + let value = ScalarValue::Struct(Arc::new(StructArray::new( + children.clone(), + vec![ + Arc::new( + TimestampMicrosecondArray::from(vec![1577934245123456]) + .with_timezone("UTC"), + ), + Arc::new(Int16Array::from(vec![-480])), + ], + None, + ))); + let plain = Field::new("declared_timestamp", storage, true); + let mut tagged = plain.clone(); + tagged.try_with_extension_type(TimestampWithOffset)?; + let numeric = Field::new("declared_number", DataType::Int64, false) + .with_metadata(HashMap::from([("custom".into(), "target".into())])); + for timestamp in [tagged, plain] { + let schema = Arc::new(DFSchema::from_unqualified_fields( + vec![timestamp.clone(), numeric.clone()].into(), + HashMap::new(), + )?); + let plan = LogicalPlanBuilder::values_with_schema( + vec![ + vec![lit(ScalarValue::Null), lit(42_i32)], + vec![lit(value.clone()), lit(7_i32)], + ], + &schema, + )? + .build()?; + for (index, target) in [timestamp, numeric.clone()].iter().enumerate() { + let actual = plan.schema().field(index); + assert_eq!(actual.name(), &format!("column{}", index + 1)); + assert_eq!(actual.data_type(), target.data_type()); + assert_eq!(actual.is_nullable(), target.is_nullable()); + assert_eq!(actual.metadata(), target.metadata()); + } + assert_eq!( + plan.schema() + .field(0) + .try_extension_type::() + .is_ok(), + schema.field(0).extension_type_name().is_some() + ); + let LogicalPlan::Values(values) = plan else { + unreachable!(); + }; + assert_eq!(values.values.len(), 2); + assert_eq!( + values.values[1][0].get_type(&DFSchema::empty())?, + *schema.field(0).data_type() + ); + assert_eq!( + values.values[1][1].get_type(&DFSchema::empty())?, + DataType::Int64 + ); + } + Ok(()) + } + #[test] fn test_values_with_schema_type_mismatch_error_message() { // Date32 field, but the value is a Boolean, which cannot be cast to Date32. diff --git a/datafusion/sql/tests/cases/params.rs b/datafusion/sql/tests/cases/params.rs index 8b025d3945db1..7cf2aeed52a55 100644 --- a/datafusion/sql/tests/cases/params.rs +++ b/datafusion/sql/tests/cases/params.rs @@ -825,6 +825,40 @@ fn test_insert_infer_with_metadata() { ScalarAndMetadata::from(ScalarValue::from("Turing")), ]; + let assert_uuid_metadata = |test: &ParameterTestWithMetadata<'_>| { + let initial = logical_plan(test.sql).unwrap(); + let bound = initial + .clone() + .with_param_values(ParamValues::List(test.param_values.clone())) + .unwrap(); + for plan in [&initial, &bound] { + let plan = match plan { + LogicalPlan::Statement(Statement::Prepare(prepare)) => { + prepare.input.as_ref() + } + plan => plan, + }; + let LogicalPlan::Dml(dml) = plan else { + panic!("expected an INSERT plan"); + }; + let LogicalPlan::Projection(projection) = dml.input.as_ref() else { + panic!("expected an INSERT projection"); + }; + let LogicalPlan::Values(values) = projection.input.as_ref() else { + panic!("expected INSERT VALUES"); + }; + for (schema, name) in + [(&projection.schema, "id"), (&values.schema, "column1")] + { + let field = schema.field(0); + assert_eq!(field.name(), name); + assert_eq!(field.data_type(), uuid_field.data_type()); + assert_eq!(field.is_nullable(), uuid_field.is_nullable()); + assert_eq!(field.metadata(), uuid_field.metadata()); + } + } + }; + // Check a normal insert let test = ParameterTestWithMetadata { sql: "insert into person_with_uuid_extension (id, first_name, last_name) values ($1, $2, $3)", @@ -832,16 +866,18 @@ fn test_insert_infer_with_metadata() { param_values: param_values.clone(), }; + assert_uuid_metadata(&test); + assert_snapshot!( test.run(), @r#" ** Initial Plan: Dml: op=[Insert Into] table=[person_with_uuid_extension] - Projection: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name + Projection: column1 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: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name + Projection: column1 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) "# ); @@ -853,17 +889,19 @@ fn test_insert_infer_with_metadata() { param_values, }; + assert_uuid_metadata(&test); + assert_snapshot!( test.run(), @r#" ** Initial Plan: Prepare: "my_plan" [FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>, Utf8, Utf8] Dml: op=[Insert Into] table=[person_with_uuid_extension] - Projection: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name + Projection: column1 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: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name + Projection: column1 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) "# );