diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 197636aa3e92c..be57aabd0757c 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -2037,10 +2037,27 @@ impl LogicalPlan { self, param_values: &ParamValues, ) -> Result { - self.transform_up_with_subqueries(|plan| { + self.replace_params_with_values_impl(param_values) + .map(|res| res.data) + } + + #[cfg_attr(feature = "recursive_protection", recursive::recursive)] + fn replace_params_with_values_impl( + self, + param_values: &ParamValues, + ) -> Result> { + let plan = self + .map_subqueries(|plan| plan.replace_params_with_values_impl(param_values))? + .transform_sibling(|plan| { + plan.map_children(|input| { + input.replace_params_with_values_impl(param_values) + }) + })?; + let inputs_changed = plan.transformed; + plan.transform_parent(|plan| { let schema = Arc::clone(plan.schema()); let name_preserver = NamePreserver::new(&plan); - plan.map_expressions(|e| { + let expressions = plan.map_expressions(|e| { let (e, has_placeholder) = e.infer_placeholder_types(&schema)?; if !has_placeholder { // Performance optimization: @@ -2094,9 +2111,16 @@ impl LogicalPlan { // Preserve name to avoid breaking column references to this expression Ok(transformed_expr.update_data(|expr| original_name.restore(expr))) } - })? - .map_data(|plan| { - let plan = plan.update_schema_data_type()?; + })?; + let values_changed = inputs_changed || expressions.transformed; + expressions.map_data(|plan| { + // A no-op bind must retain supplied VALUES metadata, rather than + // re-inferring it from type-only storage casts in the rows. + let plan = if !values_changed && matches!(&plan, LogicalPlan::Values(_)) { + plan + } else { + plan.update_schema_data_type()? + }; if let LogicalPlan::Dml(dml) = &plan { validate_bound_object_insert_source(dml)?; } else if let LogicalPlan::Union(union) = &plan { @@ -2105,7 +2129,6 @@ impl LogicalPlan { Ok(plan) }) }) - .map(|res| res.data) } /// Recompute schema fields' data type after replacing params, ensuring fields data type can be @@ -6429,6 +6452,236 @@ mod tests { .unwrap() } + fn timestamp_values_with_child_metadata(tagged: bool) -> Result { + use crate::expr::Cast; + use arrow::array::{Int16Array, StructArray, TimestampMicrosecondArray}; + use arrow::datatypes::{Fields, TimeUnit}; + use arrow_schema::extension::{EXTENSION_TYPE_NAME_KEY, TimestampWithOffset}; + + let children = Fields::from(vec![ + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + false, + ), + Field::new("offset_minutes", DataType::Int16, false), + ]); + let mut source = + Field::new("source_timestamp", DataType::Struct(children.clone()), true); + source.try_with_extension_type(TimestampWithOffset)?; + let mut metadata = source.metadata().clone(); + metadata.insert("embucket.datetime_precision".into(), "6".into()); + source.set_metadata(metadata.clone()); + let target_children = children + .iter() + .enumerate() + .map(|(index, child)| { + let mut metadata = + HashMap::from([("PARQUET:field_id".into(), (index + 2).to_string())]); + if index == 0 { + metadata.insert("embucket.datetime_precision".into(), "6".into()); + } + child.as_ref().clone().with_metadata(metadata) + }) + .collect::>(); + metadata.insert("embucket.datetime_precision".into(), "9".into()); + metadata.insert("PARQUET:field_id".into(), "1".into()); + if !tagged { + metadata.remove(EXTENSION_TYPE_NAME_KEY); + } + let target = Field::new( + "declared_timestamp", + DataType::Struct(target_children.into()), + true, + ) + .with_metadata(metadata); + let schema = Arc::new(DFSchema::from_unqualified_fields( + vec![target].into(), + HashMap::new(), + )?); + let value = ScalarValue::Struct(Arc::new(StructArray::new( + children, + vec![ + Arc::new( + TimestampMicrosecondArray::from(vec![1_577_963_045_123_456]) + .with_timezone("UTC"), + ), + Arc::new(Int16Array::from(vec![-480])), + ], + None, + ))); + let source = Arc::new(source); + LogicalPlanBuilder::values_with_schema( + [value, ScalarValue::Null] + .into_iter() + .map(|value| { + vec![Expr::Cast(Cast::new_from_field( + Box::new(lit(value)), + Arc::clone(&source), + ))] + }) + .collect(), + &schema, + )? + .build() + } + + #[test] + fn test_replace_params_preserves_unchanged_values_schema() -> Result<()> { + use arrow_schema::extension::{EXTENSION_TYPE_NAME_KEY, TimestampWithOffset}; + + for tagged in [true, false] { + let plan = timestamp_values_with_child_metadata(tagged)?; + let schema = Arc::clone(plan.schema()); + let LogicalPlan::Values(values) = &plan else { + unreachable!(); + }; + // The storage casts have already stripped the row expression's extension. + for row in &values.values { + let metadata = row[0].metadata(&DFSchema::empty())?; + assert!(!metadata.inner().contains_key(EXTENSION_TYPE_NAME_KEY)); + assert_eq!(metadata.inner()["embucket.datetime_precision"], "6"); + } + for prepared in [false, true] { + for params in [ + ParamValues::List(vec![]), + ParamValues::Map(HashMap::new()), + ParamValues::from(HashMap::from([( + "timezone", + ScalarValue::from("America/Los_Angeles"), + )])), + ] { + let input = if prepared { + LogicalPlanBuilder::from(plan.clone()) + .prepare("unchanged_values".into(), vec![])? + .build()? + } else { + plan.clone() + }; + let bound = input.with_param_values(params)?; + assert_eq!(bound, plan); + assert!(Arc::ptr_eq(bound.schema(), &schema)); + let field = bound.schema().field(0); + assert_eq!(field.name(), "column1"); + assert!(field.is_nullable()); + assert_eq!(field.metadata()["embucket.datetime_precision"], "9"); + assert_eq!( + field.try_extension_type::().is_ok(), + tagged + ); + } + } + } + Ok(()) + } + + #[test] + fn test_replace_params_preserves_unchanged_values_below_projection() -> Result<()> { + let input = timestamp_values_with_child_metadata(true)?; + let schema = Arc::clone(input.schema()); + let plan = LogicalPlanBuilder::from(input) + .project(vec![col("column1"), placeholder("$1").alias("bound")])? + .build()? + .with_param_values(vec![ScalarValue::Int32(Some(42))])?; + let LogicalPlan::Projection(projection) = plan else { + unreachable!(); + }; + assert!(Arc::ptr_eq(projection.input.schema(), &schema)); + assert_eq!(projection.schema.field(0), schema.field(0)); + assert_eq!(projection.schema.field(1).data_type(), &DataType::Int32); + assert!(!projection.schema.field(1).is_nullable()); + Ok(()) + } + + #[test] + fn test_replace_params_reinfers_changed_values_schema() -> Result<()> { + let declared = Arc::new( + Field::new("declared", DataType::Int32, true) + .with_metadata(HashMap::from([("origin".into(), "declared".into())])), + ); + let schema = Arc::new(DFSchema::from_unqualified_fields( + vec![Arc::clone(&declared)].into(), + HashMap::new(), + )?); + let metadata = HashMap::from([("origin".into(), "bound".into())]); + let plan = LogicalPlanBuilder::values_with_schema( + vec![vec![Expr::Placeholder(Placeholder { + id: "$1".into(), + field: Some(declared), + })]], + &schema, + )? + .project(vec![col("column1").alias("answer")])? + .build()? + .with_param_values(ParamValues::List(vec![ScalarAndMetadata::new( + ScalarValue::Int64(Some(42)), + Some(metadata.clone().into()), + )]))?; + assert!(plan.get_parameter_names()?.is_empty()); + let LogicalPlan::Projection(projection) = plan else { + unreachable!(); + }; + for (schema, name) in [ + (&projection.schema, "answer"), + (projection.input.schema(), "column1"), + ] { + let field = schema.field(0); + assert_eq!(field.name(), name); + assert_eq!(field.data_type(), &DataType::Int64); + assert!(!field.is_nullable()); + assert_eq!(field.metadata(), &metadata); + } + Ok(()) + } + + #[test] + fn test_replace_params_reinfers_values_after_subquery_binding() -> Result<()> { + for typed in [false, true] { + let field = + typed.then(|| { + Arc::new(Field::new("declared", DataType::Int32, true).with_metadata( + HashMap::from([("origin".into(), "declared".into())]), + )) + }); + let subquery = LogicalPlanBuilder::empty(true) + .project(vec![Expr::Placeholder(Placeholder { + id: "$1".into(), + field, + })])? + .build()?; + let metadata = HashMap::from([("origin".into(), "bound".into())]); + let value = if typed { + ScalarValue::Int32(Some(7)) + } else { + ScalarValue::Int64(Some(7)) + }; + let data_type = value.data_type(); + let plan = LogicalPlanBuilder::values(vec![vec![scalar_subquery( + Arc::new(subquery), + )]])? + .project(vec![col("column1").alias("answer")])? + .build()? + .with_param_values(ParamValues::List(vec![ + ScalarAndMetadata::new(value, Some(metadata.clone().into())), + ]))?; + assert!(plan.get_parameter_names()?.is_empty()); + let LogicalPlan::Projection(projection) = plan else { + unreachable!(); + }; + for (schema, name) in [ + (&projection.schema, "answer"), + (projection.input.schema(), "column1"), + ] { + let field = schema.field(0); + assert_eq!(field.name(), name); + assert_eq!(field.data_type(), &data_type); + assert!(field.is_nullable()); + assert_eq!(field.metadata(), &metadata); + } + } + Ok(()) + } + #[test] fn test_replace_invalid_placeholder() { // test empty placeholder