From 995bb5299e7e0ea66c34db50c0374651e74c45b9 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 00:21:43 +0300 Subject: [PATCH 01/12] Preserve common CASE result field metadata without recursive planning --- datafusion/expr/src/expr_schema.rs | 237 +++++++++++++++++- .../optimizer/src/analyzer/type_coercion.rs | 28 +++ 2 files changed, 262 insertions(+), 3 deletions(-) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index bb5b7943c0268..0c47d28c41c49 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -18,7 +18,7 @@ use super::{Between, Expr, Like, predicate_bounds}; use crate::ValueOrLambda; use crate::expr::{ - AggregateFunction, AggregateFunctionParams, Alias, BinaryExpr, Cast, InList, + AggregateFunction, AggregateFunctionParams, Alias, BinaryExpr, Case, Cast, InList, InSubquery, Lambda, Placeholder, ScalarFunction, TryCast, Unnest, WindowFunction, WindowFunctionParams, }; @@ -35,7 +35,7 @@ use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_K use datafusion_common::datatype::FieldExt; use datafusion_common::{ Column, DataFusionError, ExprSchema, Result, ScalarValue, Spans, TableReference, - not_impl_err, plan_datafusion_err, plan_err, + internal_err, not_impl_err, plan_datafusion_err, plan_err, }; use datafusion_expr_common::type_coercion::binary::BinaryTypeCoercer; use datafusion_functions_window_common::field::WindowUDFFieldArgs; @@ -116,6 +116,144 @@ fn cast_output_field( Arc::new(f) } +// Resolve metadata bottom-up so nested CASE and CAST branches are visited once. +fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result { + enum Work<'a> { + Visit(&'a Expr), + FinishCase(&'a Case), + FinishCast(&'a FieldRef, bool), + FinishAlias(Option<&'a FieldMetadata>), + } + + struct BranchField { + field: FieldRef, + certainly_null: bool, + } + + fn schedule_case<'a>(case: &'a Case, work: &mut Vec>) { + work.push(Work::FinishCase(case)); + work.extend(case.else_expr.iter().map(|expr| Work::Visit(expr))); + work.extend( + case.when_then_expr + .iter() + .rev() + .map(|(_, then_expr)| Work::Visit(then_expr)), + ); + } + + let mut work = Vec::new(); + let mut fields: Vec = Vec::new(); + schedule_case(case, &mut work); + while let Some(item) = work.pop() { + match item { + Work::Visit(expr) => match expr { + Expr::Case(nested) => schedule_case(nested, &mut work), + Expr::Cast(cast) => { + work.push(Work::FinishCast(&cast.field, false)); + work.push(Work::Visit(&cast.expr)); + } + Expr::TryCast(cast) => { + work.push(Work::FinishCast(&cast.field, true)); + work.push(Work::Visit(&cast.expr)); + } + Expr::Alias(alias) => { + work.push(Work::FinishAlias(alias.metadata.as_ref())); + work.push(Work::Visit(&alias.expr)); + } + Expr::Negative(inner) => work.push(Work::Visit(inner)), + _ => fields.push(BranchField { + field: expr.to_field(schema)?.1, + certainly_null: matches!( + unwrap_certainly_null_expr(expr), + Expr::Literal(value, _) if value.is_null() + ), + }), + }, + Work::FinishCast(target, force_nullable) => { + let Some(source) = fields.pop() else { + return internal_err!("Missing CASE cast input field"); + }; + fields.push(BranchField { + field: cast_output_field(&source.field, target, force_nullable), + certainly_null: source.certainly_null, + }); + } + Work::FinishAlias(metadata) => { + let Some(source) = fields.pop() else { + return internal_err!("Missing CASE alias input field"); + }; + let mut combined = source.field.metadata().clone(); + if let Some(metadata) = metadata { + combined.extend(metadata.to_hashmap()); + } + fields.push(BranchField { + field: Arc::new( + source.field.as_ref().clone().with_metadata(combined), + ), + certainly_null: source.certainly_null, + }); + } + Work::FinishCase(case) => { + let count = + case.when_then_expr.len() + usize::from(case.else_expr.is_some()); + if fields.len() < count { + return internal_err!("Missing CASE result fields"); + } + let start = fields.len() - count; + let mut then_type = DataType::Null; + let mut else_type = DataType::Null; + let mut branch_type = None; + let mut metadata = None; + let mut conflict = false; + let mut certainly_null = true; + for (index, branch) in fields.drain(start..).enumerate() { + let data_type = branch.field.data_type(); + if index < case.when_then_expr.len() { + if then_type.is_null() && !data_type.is_null() { + then_type = data_type.clone(); + } + } else { + else_type = data_type.clone(); + } + certainly_null &= branch.certainly_null; + if branch.certainly_null || data_type.is_null() { + continue; + } + if branch_type.as_ref().is_some_and(|other| other != data_type) + || branch.field.metadata().is_empty() + || metadata + .as_ref() + .is_some_and(|other| other != branch.field.metadata()) + { + conflict = true; + } + branch_type.get_or_insert_with(|| data_type.clone()); + metadata.get_or_insert_with(|| branch.field.metadata().clone()); + } + let data_type = if then_type.is_null() { + else_type + } else { + then_type + }; + if conflict || branch_type.as_ref() != Some(&data_type) { + metadata = None; + } + fields.push(BranchField { + field: Arc::new( + Field::new("", data_type, true) + .with_metadata(metadata.unwrap_or_default()), + ), + certainly_null, + }); + } + } + } + let Some(result) = fields.pop() else { + return internal_err!("Missing CASE output field"); + }; + Ok(result.field) +} + impl ExprSchemable for Expr { /// Returns the [arrow::datatypes::DataType] of the expression /// based on [ExprSchema] @@ -645,11 +783,40 @@ impl ExprSchemable for Expr { Expr::LambdaVariable(LambdaVariable { field: Some(field), .. }) => Ok(Arc::clone(field).renamed(&schema_name)), + Expr::Case(case) => { + let data_type = self.get_type(schema)?; + let nullable = self.nullable(schema)?; + if let Some((_, first_result)) = case.when_then_expr.first() + && matches!( + first_result.as_ref(), + Expr::Column(_) | Expr::Literal(_, _) + ) + { + let field = first_result.to_field(schema)?.1; + if !field.data_type().is_null() + && !matches!(first_result.as_ref(), Expr::Literal(value, _) if value.is_null()) + && field.metadata().is_empty() + { + return Ok(( + relation, + Arc::new(Field::new(&schema_name, data_type, nullable)), + )); + } + } + let branch_field = case_field_metadata(case, schema)?; + let metadata = if branch_field.data_type() == &data_type { + branch_field.metadata().clone() + } else { + Default::default() + }; + Ok(Arc::new( + Field::new(&schema_name, data_type, nullable).with_metadata(metadata), + )) + } Expr::Like(_) | Expr::SimilarTo(_) | Expr::Not(_) | Expr::Between(_) - | Expr::Case(_) | Expr::InList(_) | Expr::InSubquery(_) | Expr::SetComparison(_) @@ -919,6 +1086,13 @@ mod tests { assert_not_nullable(&e, &nullable_schema); assert_not_nullable(&e, ¬_nullable_schema); + let varchar_schema = MockExprSchema::new() + .with_data_type(DataType::Utf8) + .with_nullable(true); + let try_cast = Expr::TryCast(TryCast::new(Box::new(col("x")), DataType::Int32)); + let e = when(col("x").is_not_null(), try_cast).otherwise(lit(0))?; + assert_nullable(&e, &varchar_schema); + // CASE WHEN NOT x IS NULL THEN x ELSE 0 let e = when(not(col("x").is_null()), col("x")).otherwise(lit(0))?; assert_not_nullable(&e, &nullable_schema); @@ -1127,6 +1301,63 @@ mod tests { assert_eq!(meta, outer_ref.metadata(&schema).unwrap()); } + #[test] + fn test_case_field_metadata() -> Result<()> { + let shared = HashMap::from([("type".to_string(), "structured".to_string())]); + let different = HashMap::from([("type".to_string(), "other".to_string())]); + let schema = DFSchema::from_unqualified_fields( + vec![ + Field::new("a", DataType::Int32, false).with_metadata(shared.clone()), + Field::new("b", DataType::Int32, false).with_metadata(shared.clone()), + Field::new("c", DataType::Int32, false).with_metadata(different), + Field::new("d", DataType::Int32, false), + ] + .into(), + HashMap::new(), + )?; + + let same = when(lit(true), col("a")).otherwise(col("b"))?; + assert_eq!(same.to_field(&schema)?.1.metadata(), &shared); + + let null_else = when(lit(true), col("a")).otherwise(lit(ScalarValue::Null))?; + assert_eq!(null_else.to_field(&schema)?.1.metadata(), &shared); + + let coerced_null_else = when(lit(true), col("a")) + .otherwise(lit(ScalarValue::Null).cast_to(&DataType::Int32, &schema)?)?; + assert_eq!(coerced_null_else.to_field(&schema)?.1.metadata(), &shared); + + let try_cast_null_else = when(lit(true), col("a")).otherwise(Expr::TryCast( + TryCast::new(Box::new(lit(ScalarValue::Null)), DataType::Int32), + ))?; + assert_eq!(try_cast_null_else.to_field(&schema)?.1.metadata(), &shared); + + let mut nested = col("a"); + for _ in 0..128 { + nested = when(lit(true), col("a")).otherwise(nested)?; + } + let nested = when(lit(true), nested).otherwise(col("b"))?; + assert_eq!(nested.to_field(&schema)?.1.metadata(), &shared); + + let mut cast_nested = col("a"); + for _ in 0..128 { + let inner = when(lit(true), col("a")).otherwise(cast_nested)?; + cast_nested = Expr::Cast(Cast::new(Box::new(inner), DataType::Int32)); + } + let cast_nested = when(lit(true), cast_nested).otherwise(col("b"))?; + assert_eq!(cast_nested.to_field(&schema)?.1.metadata(), &shared); + + let mismatched = when(lit(true), col("a")).otherwise(col("c"))?; + assert!(mismatched.to_field(&schema)?.1.metadata().is_empty()); + + let unmarked = when(lit(true), col("a")).otherwise(col("d"))?; + assert!(unmarked.to_field(&schema)?.1.metadata().is_empty()); + + let first_unmarked = when(lit(true), col("d")).otherwise(col("a"))?; + assert!(first_unmarked.to_field(&schema)?.1.metadata().is_empty()); + + Ok(()) + } + #[test] fn test_alias_metadata_is_preserved_in_field_metadata() { let schema = MockExprSchema::new().with_data_type(DataType::Int32); diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 6186cc5350bb2..8a73892e7f535 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -2733,6 +2733,34 @@ mod test { } } + #[test] + fn test_case_coercion_preserves_matching_field_metadata() -> Result<()> { + let metadata = std::collections::HashMap::from([( + "structured_type".to_string(), + "ARRAY(NUMBER)".to_string(), + )]); + let schema = DFSchema::from_unqualified_fields( + vec![ + Field::new("structured", DataType::Int32, false) + .with_metadata(metadata.clone()), + ] + .into(), + std::collections::HashMap::new(), + )?; + let case = Case { + expr: None, + when_then_expr: vec![(Box::new(lit(true)), Box::new(col("structured")))], + else_expr: Some(Box::new(lit(ScalarValue::Null))), + }; + let coerced = coerce_case_expression(case, &schema)?; + assert!(matches!(coerced.else_expr.as_deref(), Some(Expr::Cast(_)))); + assert_eq!( + Expr::Case(coerced).to_field(&schema)?.1.metadata(), + &metadata + ); + Ok(()) + } + #[test] fn test_case_expression_coercion() -> Result<()> { let schema = Arc::new(DFSchema::from_unqualified_fields( From c62bde355524d4704c5845f5241eff49192104eb Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 01:08:54 +0300 Subject: [PATCH 02/12] Retain metadata for explicitly typed NULL CASE branches --- datafusion/expr/src/expr_schema.rs | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 0c47d28c41c49..d3643d642df68 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -216,7 +216,9 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result else_type = data_type.clone(); } certainly_null &= branch.certainly_null; - if branch.certainly_null || data_type.is_null() { + if data_type.is_null() + || (branch.certainly_null && branch.field.metadata().is_empty()) + { continue; } if branch_type.as_ref().is_some_and(|other| other != data_type) @@ -1331,6 +1333,13 @@ mod tests { ))?; assert_eq!(try_cast_null_else.to_field(&schema)?.1.metadata(), &shared); + let typed_null = Expr::Cast(Cast::new_from_field( + Box::new(lit(ScalarValue::Null)), + Arc::new(Field::new("", DataType::Int32, true).with_metadata(shared.clone())), + )); + let all_typed_null = when(lit(true), typed_null.clone()).otherwise(typed_null)?; + assert_eq!(all_typed_null.to_field(&schema)?.1.metadata(), &shared); + let mut nested = col("a"); for _ in 0..128 { nested = when(lit(true), col("a")).otherwise(nested)?; From e54eafec2e4d3527e8e6c1ec466e972819f17cf3 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 01:52:39 +0300 Subject: [PATCH 03/12] Avoid repeated CASE metadata planning through binary expressions --- datafusion/expr/src/expr_schema.rs | 35 ++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index d3643d642df68..7312334bdba27 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -121,6 +121,7 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result enum Work<'a> { Visit(&'a Expr), FinishCase(&'a Case), + FinishBinary(&'a BinaryExpr), FinishCast(&'a FieldRef, bool), FinishAlias(Option<&'a FieldMetadata>), } @@ -148,6 +149,11 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result match item { Work::Visit(expr) => match expr { Expr::Case(nested) => schedule_case(nested, &mut work), + Expr::BinaryExpr(binary) => { + work.push(Work::FinishBinary(binary)); + work.push(Work::Visit(&binary.right)); + work.push(Work::Visit(&binary.left)); + } Expr::Cast(cast) => { work.push(Work::FinishCast(&cast.field, false)); work.push(Work::Visit(&cast.expr)); @@ -169,6 +175,26 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result ), }), }, + Work::FinishBinary(binary) => { + let Some(right) = fields.pop() else { + return internal_err!("Missing CASE binary right field"); + }; + let Some(left) = fields.pop() else { + return internal_err!("Missing CASE binary left field"); + }; + let mut coercer = BinaryTypeCoercer::new( + left.field.data_type(), + &binary.op, + right.field.data_type(), + ); + coercer.set_lhs_spans(binary.left.spans().cloned().unwrap_or_default()); + coercer.set_rhs_spans(binary.right.spans().cloned().unwrap_or_default()); + let nullable = left.field.is_nullable() || right.field.is_nullable(); + fields.push(BranchField { + field: Arc::new(Field::new("", coercer.get_result_type()?, nullable)), + certainly_null: false, + }); + } Work::FinishCast(target, force_nullable) => { let Some(source) = fields.pop() else { return internal_err!("Missing CASE cast input field"); @@ -1355,6 +1381,15 @@ mod tests { let cast_nested = when(lit(true), cast_nested).otherwise(col("b"))?; assert_eq!(cast_nested.to_field(&schema)?.1.metadata(), &shared); + let mut binary_nested = col("a"); + for _ in 0..128 { + binary_nested = when(lit(true), binary_nested + lit(1)).otherwise(lit(0))?; + } + let binary_nested = when(lit(true), binary_nested).otherwise(col("b"))?; + let binary_field = binary_nested.to_field(&schema)?.1; + assert_eq!(binary_field.data_type(), &DataType::Int32); + assert!(binary_field.metadata().is_empty()); + let mismatched = when(lit(true), col("a")).otherwise(col("c"))?; assert!(mismatched.to_field(&schema)?.1.metadata().is_empty()); From 9ceb9b3be04beb20b470c67da606216d300805c3 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 01:58:13 +0300 Subject: [PATCH 04/12] Resolve scalar UDF CASE metadata iteratively --- datafusion/expr/src/expr_schema.rs | 65 ++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 7312334bdba27..42478c2807203 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -122,6 +122,7 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result Visit(&'a Expr), FinishCase(&'a Case), FinishBinary(&'a BinaryExpr), + FinishScalar(&'a ScalarFunction), FinishCast(&'a FieldRef, bool), FinishAlias(Option<&'a FieldMetadata>), } @@ -154,6 +155,10 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result work.push(Work::Visit(&binary.right)); work.push(Work::Visit(&binary.left)); } + Expr::ScalarFunction(function) => { + work.push(Work::FinishScalar(function)); + work.extend(function.args.iter().rev().map(Work::Visit)); + } Expr::Cast(cast) => { work.push(Work::FinishCast(&cast.field, false)); work.push(Work::Visit(&cast.expr)); @@ -195,6 +200,34 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result certainly_null: false, }); } + Work::FinishScalar(function) => { + if fields.len() < function.args.len() { + return internal_err!("Missing CASE scalar function arguments"); + } + let start = fields.len() - function.args.len(); + let input_fields = fields + .drain(start..) + .map(|arg| arg.field) + .collect::>(); + let coerced_fields = + verify_function_arguments(function.func.as_ref(), &input_fields)?; + let scalar_arguments = function + .args + .iter() + .map(|arg| match arg { + Expr::Literal(value, _) => Some(value), + _ => None, + }) + .collect::>(); + let field = function.func.return_field_from_args(ReturnFieldArgs { + arg_fields: &coerced_fields, + scalar_arguments: &scalar_arguments, + })?; + fields.push(BranchField { + field, + certainly_null: false, + }); + } Work::FinishCast(target, force_nullable) => { let Some(source) = fields.pop() else { return internal_err!("Missing CASE cast input field"); @@ -1390,6 +1423,38 @@ mod tests { assert_eq!(binary_field.data_type(), &DataType::Int32); assert!(binary_field.metadata().is_empty()); + let identity = Arc::new(crate::expr_fn::create_udf( + "identity", + vec![DataType::Int32], + DataType::Int32, + crate::Volatility::Immutable, + Arc::new(|_| { + Ok( + datafusion_expr_common::columnar_value::ColumnarValue::Scalar( + ScalarValue::Int32(Some(0)), + ), + ) + }), + )); + let mut scalar_nested = col("a"); + for _ in 0..128 { + let nested_case = when(lit(true), scalar_nested).otherwise(col("b"))?; + scalar_nested = Expr::ScalarFunction(ScalarFunction::new_udf( + Arc::clone(&identity), + vec![nested_case], + )); + } + let Expr::Case(scalar_case) = + when(lit(true), scalar_nested).otherwise(col("b"))? + else { + unreachable!(); + }; + assert!( + case_field_metadata(&scalar_case, &schema)? + .metadata() + .is_empty() + ); + let mismatched = when(lit(true), col("a")).otherwise(col("c"))?; assert!(mismatched.to_field(&schema)?.1.metadata().is_empty()); From d582e231303441a688a63a7211b5783262ae0e6f Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 02:06:01 +0300 Subject: [PATCH 05/12] Limit CASE metadata propagation to supported branch forms --- datafusion/expr/src/expr_schema.rs | 91 ++++-------------------------- 1 file changed, 12 insertions(+), 79 deletions(-) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 42478c2807203..ebdff6c1cf7b4 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -116,13 +116,11 @@ fn cast_output_field( Arc::new(f) } -// Resolve metadata bottom-up so nested CASE and CAST branches are visited once. +// Resolve metadata bottom-up for branches whose metadata can be preserved. fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result { enum Work<'a> { Visit(&'a Expr), FinishCase(&'a Case), - FinishBinary(&'a BinaryExpr), - FinishScalar(&'a ScalarFunction), FinishCast(&'a FieldRef, bool), FinishAlias(Option<&'a FieldMetadata>), } @@ -150,15 +148,6 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result match item { Work::Visit(expr) => match expr { Expr::Case(nested) => schedule_case(nested, &mut work), - Expr::BinaryExpr(binary) => { - work.push(Work::FinishBinary(binary)); - work.push(Work::Visit(&binary.right)); - work.push(Work::Visit(&binary.left)); - } - Expr::ScalarFunction(function) => { - work.push(Work::FinishScalar(function)); - work.extend(function.args.iter().rev().map(Work::Visit)); - } Expr::Cast(cast) => { work.push(Work::FinishCast(&cast.field, false)); work.push(Work::Visit(&cast.expr)); @@ -172,62 +161,20 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result work.push(Work::Visit(&alias.expr)); } Expr::Negative(inner) => work.push(Work::Visit(inner)), - _ => fields.push(BranchField { + Expr::Column(_) + | Expr::Literal(_, _) + | Expr::OuterReferenceColumn(_, _) + | Expr::ScalarVariable(_, _) + | Expr::Placeholder(_) + | Expr::LambdaVariable(_) => fields.push(BranchField { field: expr.to_field(schema)?.1, certainly_null: matches!( unwrap_certainly_null_expr(expr), Expr::Literal(value, _) if value.is_null() ), }), + _ => return Ok(Arc::new(Field::new("", DataType::Null, true))), }, - Work::FinishBinary(binary) => { - let Some(right) = fields.pop() else { - return internal_err!("Missing CASE binary right field"); - }; - let Some(left) = fields.pop() else { - return internal_err!("Missing CASE binary left field"); - }; - let mut coercer = BinaryTypeCoercer::new( - left.field.data_type(), - &binary.op, - right.field.data_type(), - ); - coercer.set_lhs_spans(binary.left.spans().cloned().unwrap_or_default()); - coercer.set_rhs_spans(binary.right.spans().cloned().unwrap_or_default()); - let nullable = left.field.is_nullable() || right.field.is_nullable(); - fields.push(BranchField { - field: Arc::new(Field::new("", coercer.get_result_type()?, nullable)), - certainly_null: false, - }); - } - Work::FinishScalar(function) => { - if fields.len() < function.args.len() { - return internal_err!("Missing CASE scalar function arguments"); - } - let start = fields.len() - function.args.len(); - let input_fields = fields - .drain(start..) - .map(|arg| arg.field) - .collect::>(); - let coerced_fields = - verify_function_arguments(function.func.as_ref(), &input_fields)?; - let scalar_arguments = function - .args - .iter() - .map(|arg| match arg { - Expr::Literal(value, _) => Some(value), - _ => None, - }) - .collect::>(); - let field = function.func.return_field_from_args(ReturnFieldArgs { - arg_fields: &coerced_fields, - scalar_arguments: &scalar_arguments, - })?; - fields.push(BranchField { - field, - certainly_null: false, - }); - } Work::FinishCast(target, force_nullable) => { let Some(source) = fields.pop() else { return internal_err!("Missing CASE cast input field"); @@ -1436,24 +1383,10 @@ mod tests { ) }), )); - let mut scalar_nested = col("a"); - for _ in 0..128 { - let nested_case = when(lit(true), scalar_nested).otherwise(col("b"))?; - scalar_nested = Expr::ScalarFunction(ScalarFunction::new_udf( - Arc::clone(&identity), - vec![nested_case], - )); - } - let Expr::Case(scalar_case) = - when(lit(true), scalar_nested).otherwise(col("b"))? - else { - unreachable!(); - }; - assert!( - case_field_metadata(&scalar_case, &schema)? - .metadata() - .is_empty() - ); + let scalar_result = + Expr::ScalarFunction(ScalarFunction::new_udf(identity, vec![col("a")])); + let scalar_case = when(lit(true), scalar_result).otherwise(col("b"))?; + assert!(scalar_case.to_field(&schema)?.1.metadata().is_empty()); let mismatched = when(lit(true), col("a")).otherwise(col("c"))?; assert!(mismatched.to_field(&schema)?.1.metadata().is_empty()); From b1c95777512901b97e51afea4fffc5c92591e5ca Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 11:12:50 +0300 Subject: [PATCH 06/12] Cover CASE metadata through bounded function branches --- datafusion/expr/src/expr_schema.rs | 156 +++++++++++++++++++++++++++++ 1 file changed, 156 insertions(+) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index ebdff6c1cf7b4..7b0aeb6f8165b 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -33,6 +33,7 @@ use arrow::datatypes::FieldRef; use arrow::datatypes::{DataType, Field}; use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::datatype::FieldExt; +use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::{ Column, DataFusionError, ExprSchema, Result, ScalarValue, Spans, TableReference, internal_err, not_impl_err, plan_datafusion_err, plan_err, @@ -118,6 +119,23 @@ fn cast_output_field( // Resolve metadata bottom-up for branches whose metadata can be preserved. fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result { + // `to_field` can revisit nested aliases; delegate only shallow branches. + const MAX_EXACT_BRANCH_DEPTH: usize = 8; + + fn can_infer_exact_branch(expr: &Expr) -> Result { + let mut work = vec![(expr, 1)]; + while let Some((node, depth)) = work.pop() { + if matches!(node, Expr::Case(_)) || depth > MAX_EXACT_BRANCH_DEPTH { + return Ok(false); + } + node.apply_children(|child| { + work.push((child, depth + 1)); + Ok(TreeNodeRecursion::Continue) + })?; + } + Ok(true) + } + enum Work<'a> { Visit(&'a Expr), FinishCase(&'a Case), @@ -148,19 +166,53 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result match item { Work::Visit(expr) => match expr { Expr::Case(nested) => schedule_case(nested, &mut work), + Expr::Cast(cast) + if !cast.field.metadata().is_empty() + && can_infer_exact_branch(expr)? => + { + fields.push(BranchField { + field: expr.to_field(schema)?.1, + certainly_null: false, + }); + } Expr::Cast(cast) => { work.push(Work::FinishCast(&cast.field, false)); work.push(Work::Visit(&cast.expr)); } + Expr::TryCast(cast) + if !cast.field.metadata().is_empty() + && can_infer_exact_branch(expr)? => + { + fields.push(BranchField { + field: expr.to_field(schema)?.1, + certainly_null: false, + }); + } Expr::TryCast(cast) => { work.push(Work::FinishCast(&cast.field, true)); work.push(Work::Visit(&cast.expr)); } + Expr::Alias(alias) + if alias.metadata.is_some() && can_infer_exact_branch(expr)? => + { + fields.push(BranchField { + field: expr.to_field(schema)?.1, + certainly_null: false, + }); + } Expr::Alias(alias) => { work.push(Work::FinishAlias(alias.metadata.as_ref())); work.push(Work::Visit(&alias.expr)); } Expr::Negative(inner) => work.push(Work::Visit(inner)), + Expr::ScalarFunction(_) | Expr::HigherOrderFunction(_) + if can_infer_exact_branch(expr)? => + { + fields.push(BranchField { + field: expr.to_field(schema)?.1, + certainly_null: false, + }); + } Expr::Column(_) | Expr::Literal(_, _) | Expr::OuterReferenceColumn(_, _) @@ -1311,6 +1363,40 @@ mod tests { #[test] fn test_case_field_metadata() -> Result<()> { + #[derive(Debug, PartialEq, Eq, Hash)] + struct MarkedIdentity { + signature: crate::Signature, + } + + impl crate::ScalarUDFImpl for MarkedIdentity { + fn name(&self) -> &str { + "marked_identity" + } + + fn signature(&self) -> &crate::Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int32) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let input = &args.arg_fields[0]; + Ok(Arc::new( + Field::new("marked_identity", DataType::Int32, input.is_nullable()) + .with_metadata(input.metadata().clone()), + )) + } + + fn invoke_with_args( + &self, + _args: crate::ScalarFunctionArgs, + ) -> Result { + Ok(crate::ColumnarValue::Scalar(ScalarValue::Int32(Some(0)))) + } + } + let shared = HashMap::from([("type".to_string(), "structured".to_string())]); let different = HashMap::from([("type".to_string(), "other".to_string())]); let schema = DFSchema::from_unqualified_fields( @@ -1327,6 +1413,19 @@ mod tests { let same = when(lit(true), col("a")).otherwise(col("b"))?; assert_eq!(same.to_field(&schema)?.1.metadata(), &shared); + let no_else = when(lit(true), col("a")).end()?; + assert_eq!(no_else.to_field(&schema)?.1.metadata(), &shared); + + let multiple_when = when(lit(false), col("a")) + .when(lit(true), col("b")) + .otherwise(lit(ScalarValue::Null))?; + assert_eq!(multiple_when.to_field(&schema)?.1.metadata(), &shared); + + let conflicting_when = when(lit(false), col("a")) + .when(lit(true), col("c")) + .otherwise(col("b"))?; + assert!(conflicting_when.to_field(&schema)?.1.metadata().is_empty()); + let null_else = when(lit(true), col("a")).otherwise(lit(ScalarValue::Null))?; assert_eq!(null_else.to_field(&schema)?.1.metadata(), &shared); @@ -1388,6 +1487,63 @@ mod tests { let scalar_case = when(lit(true), scalar_result).otherwise(col("b"))?; assert!(scalar_case.to_field(&schema)?.1.metadata().is_empty()); + let marked = Arc::new(crate::ScalarUDF::from(MarkedIdentity { + signature: crate::Signature::uniform( + 1, + vec![DataType::Int32], + crate::Volatility::Immutable, + ), + })); + let marked_call = |arg| { + Expr::ScalarFunction(ScalarFunction::new_udf(Arc::clone(&marked), vec![arg])) + }; + let marked_case = + when(lit(true), marked_call(col("a"))).otherwise(marked_call(col("b")))?; + assert_eq!(marked_case.to_field(&schema)?.1.metadata(), &shared); + + let nested_arg = when(lit(true), col("a")).otherwise(col("b"))?; + let nested_call = when(lit(true), marked_call(nested_arg)).otherwise(col("b"))?; + assert!(nested_call.to_field(&schema)?.1.metadata().is_empty()); + + let binary = col("d") + lit(1); + let Expr::Alias(alias) = binary.clone().alias("marked") else { + unreachable!(); + }; + let marked_alias = + Expr::Alias(alias.with_metadata(Some(FieldMetadata::from(shared.clone())))); + let aliased_case = when(lit(true), marked_alias).otherwise(col("b"))?; + assert_eq!(aliased_case.to_field(&schema)?.1.metadata(), &shared); + + let marked_cast = Expr::Cast(Cast::new_from_field( + Box::new(binary.clone()), + Arc::new(Field::new("", DataType::Int32, true).with_metadata(shared.clone())), + )); + let cast_case = when(lit(true), marked_cast).otherwise(col("b"))?; + assert_eq!(cast_case.to_field(&schema)?.1.metadata(), &shared); + + let marked_try_cast = Expr::TryCast(TryCast::new_from_field( + Box::new(binary), + Arc::new(Field::new("", DataType::Int32, true).with_metadata(shared.clone())), + )); + let try_cast_case = when(lit(true), marked_try_cast).otherwise(col("b"))?; + assert_eq!(try_cast_case.to_field(&schema)?.1.metadata(), &shared); + + let mut deep_alias = col("a"); + for _ in 0..10 { + deep_alias = deep_alias.alias("nested"); + } + let deep_udf_case = + when(lit(true), marked_call(deep_alias.clone())).otherwise(col("b"))?; + assert!(deep_udf_case.to_field(&schema)?.1.metadata().is_empty()); + + let Expr::Alias(alias) = deep_alias.alias("marked") else { + unreachable!(); + }; + let deep_marked_alias = + Expr::Alias(alias.with_metadata(Some(FieldMetadata::from(shared.clone())))); + let deep_alias_case = when(lit(true), deep_marked_alias).otherwise(col("b"))?; + assert_eq!(deep_alias_case.to_field(&schema)?.1.metadata(), &shared); + let mismatched = when(lit(true), col("a")).otherwise(col("c"))?; assert!(mismatched.to_field(&schema)?.1.metadata().is_empty()); From b2b24b2a70e6419a4f0fad4ae4be4d9d3daeca4a Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 11:45:52 +0300 Subject: [PATCH 07/12] Preserve shallow nested CASE metadata through functions --- datafusion/expr/src/expr_schema.rs | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 7b0aeb6f8165b..95dd88c69ba4c 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -122,10 +122,12 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result // `to_field` can revisit nested aliases; delegate only shallow branches. const MAX_EXACT_BRANCH_DEPTH: usize = 8; - fn can_infer_exact_branch(expr: &Expr) -> Result { + fn can_infer_exact_branch(expr: &Expr, allow_nested_case: bool) -> Result { let mut work = vec![(expr, 1)]; while let Some((node, depth)) = work.pop() { - if matches!(node, Expr::Case(_)) || depth > MAX_EXACT_BRANCH_DEPTH { + if (!allow_nested_case && matches!(node, Expr::Case(_))) + || depth > MAX_EXACT_BRANCH_DEPTH + { return Ok(false); } node.apply_children(|child| { @@ -168,7 +170,7 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result Expr::Case(nested) => schedule_case(nested, &mut work), Expr::Cast(cast) if !cast.field.metadata().is_empty() - && can_infer_exact_branch(expr)? => + && can_infer_exact_branch(expr, false)? => { fields.push(BranchField { field: expr.to_field(schema)?.1, @@ -181,7 +183,7 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result } Expr::TryCast(cast) if !cast.field.metadata().is_empty() - && can_infer_exact_branch(expr)? => + && can_infer_exact_branch(expr, false)? => { fields.push(BranchField { field: expr.to_field(schema)?.1, @@ -193,7 +195,8 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result work.push(Work::Visit(&cast.expr)); } Expr::Alias(alias) - if alias.metadata.is_some() && can_infer_exact_branch(expr)? => + if alias.metadata.is_some() + && can_infer_exact_branch(expr, false)? => { fields.push(BranchField { field: expr.to_field(schema)?.1, @@ -206,7 +209,7 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result } Expr::Negative(inner) => work.push(Work::Visit(inner)), Expr::ScalarFunction(_) | Expr::HigherOrderFunction(_) - if can_infer_exact_branch(expr)? => + if can_infer_exact_branch(expr, true)? => { fields.push(BranchField { field: expr.to_field(schema)?.1, @@ -1503,7 +1506,15 @@ mod tests { let nested_arg = when(lit(true), col("a")).otherwise(col("b"))?; let nested_call = when(lit(true), marked_call(nested_arg)).otherwise(col("b"))?; - assert!(nested_call.to_field(&schema)?.1.metadata().is_empty()); + assert_eq!(nested_call.to_field(&schema)?.1.metadata(), &shared); + + let mut deep_nested_arg = when(lit(true), col("a")).otherwise(col("b"))?; + for _ in 0..8 { + deep_nested_arg = deep_nested_arg.alias("nested"); + } + let deep_nested_call = + when(lit(true), marked_call(deep_nested_arg)).otherwise(col("b"))?; + assert!(deep_nested_call.to_field(&schema)?.1.metadata().is_empty()); let binary = col("d") + lit(1); let Expr::Alias(alias) = binary.clone().alias("marked") else { From cf72dd4d704fb9040bb338144291f97cba56949f Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 12:03:17 +0300 Subject: [PATCH 08/12] Test additional CASE metadata branch forms [skip ci] --- datafusion/expr/src/expr_schema.rs | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 95dd88c69ba4c..f068e0f6ab8cf 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -1408,6 +1408,8 @@ mod tests { Field::new("b", DataType::Int32, false).with_metadata(shared.clone()), Field::new("c", DataType::Int32, false).with_metadata(different), Field::new("d", DataType::Int32, false), + Field::new("e", DataType::Boolean, false).with_metadata(shared.clone()), + Field::new("f", DataType::Boolean, false).with_metadata(shared.clone()), ] .into(), HashMap::new(), @@ -1429,6 +1431,24 @@ mod tests { .otherwise(col("b"))?; assert!(conflicting_when.to_field(&schema)?.1.metadata().is_empty()); + let marked_literal = Expr::Literal( + ScalarValue::Int32(Some(1)), + Some(FieldMetadata::from(shared.clone())), + ); + let literal_case = when(lit(true), marked_literal).otherwise(col("b"))?; + assert_eq!(literal_case.to_field(&schema)?.1.metadata(), &shared); + + let unmarked_literal = when(lit(true), col("a")).otherwise(lit(1_i32))?; + assert!(unmarked_literal.to_field(&schema)?.1.metadata().is_empty()); + + let negative_case = when(lit(true), Expr::Negative(Box::new(col("a")))) + .otherwise(Expr::Negative(Box::new(col("b"))))?; + assert_eq!(negative_case.to_field(&schema)?.1.metadata(), &shared); + + let unsupported_case = + when(lit(true), Expr::Not(Box::new(col("e")))).otherwise(col("f"))?; + assert!(unsupported_case.to_field(&schema)?.1.metadata().is_empty()); + let null_else = when(lit(true), col("a")).otherwise(lit(ScalarValue::Null))?; assert_eq!(null_else.to_field(&schema)?.1.metadata(), &shared); From ede22a646112d0428f5fe7dd24cc73b9be79d417 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 13:49:35 +0300 Subject: [PATCH 09/12] Test CASE metadata across supported leaf expressions --- datafusion/expr/src/expr_schema.rs | 41 ++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index f068e0f6ab8cf..76f2677e23e46 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -1587,6 +1587,47 @@ mod tests { Ok(()) } + #[test] + fn test_case_field_metadata_leaf_branches() -> Result<()> { + let metadata = HashMap::from([("type".to_string(), "structured".to_string())]); + let field = Arc::new( + Field::new("value", DataType::Int32, false).with_metadata(metadata.clone()), + ); + let schema = DFSchema::from_unqualified_fields( + vec![ + Field::new("value", DataType::Int32, false) + .with_metadata(metadata.clone()), + ] + .into(), + HashMap::new(), + )?; + + let branches = [ + Expr::OuterReferenceColumn(Arc::clone(&field), Column::from_name("outer")), + Expr::ScalarVariable(Arc::clone(&field), vec!["value".to_string()]), + Expr::Placeholder(Placeholder::new_with_field( + "$1".to_string(), + Some(Arc::clone(&field)), + )), + Expr::LambdaVariable(LambdaVariable::new("arg".into(), Some(field))), + ]; + for branch in branches { + let case = when(lit(true), col("value")).otherwise(branch)?; + assert_eq!(case.to_field(&schema)?.1.metadata(), &metadata); + } + + let null_first = + when(lit(true), lit(ScalarValue::Null)).otherwise(col("value"))?; + assert_eq!(null_first.to_field(&schema)?.1.metadata(), &metadata); + + let wrong_type = when(lit(true), col("value")).otherwise(Expr::Literal( + ScalarValue::Boolean(Some(false)), + Some(FieldMetadata::from(metadata)), + ))?; + assert!(wrong_type.to_field(&schema)?.1.metadata().is_empty()); + Ok(()) + } + #[test] fn test_alias_metadata_is_preserved_in_field_metadata() { let schema = MockExprSchema::new().with_data_type(DataType::Int32); From 2423c95a393f541a19dd0bfb50ad7db332ca5e99 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 14:00:52 +0300 Subject: [PATCH 10/12] Test CASE field metadata through SQL planning --- datafusion/core/tests/sql/select.rs | 73 +++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs index 96b911e8db130..43cb5350b81ec 100644 --- a/datafusion/core/tests/sql/select.rs +++ b/datafusion/core/tests/sql/select.rs @@ -360,6 +360,79 @@ async fn test_query_parameters_with_metadata() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_sql_case_field_metadata() -> Result<()> { + let ctx = SessionContext::new(); + let metadata = + HashMap::from([("structured_type".to_string(), "ARRAY(NUMBER)".to_string())]); + let other_metadata = + HashMap::from([("structured_type".to_string(), "OBJECT".to_string())]); + let schema = Arc::new(Schema::new(vec![ + Field::new("flag", DataType::Boolean, false), + Field::new("a", DataType::Int32, false).with_metadata(metadata.clone()), + Field::new("b", DataType::Int32, false).with_metadata(metadata.clone()), + Field::new("c", DataType::Int32, false).with_metadata(other_metadata), + ])); + ctx.register_batch( + "t", + RecordBatch::try_new( + schema, + vec![ + Arc::new(BooleanArray::from(vec![true, false])), + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![3, 4])), + Arc::new(Int32Array::from(vec![5, 6])), + ], + )?, + )?; + + let empty_metadata = HashMap::new(); + for (sql, preserve_metadata) in [ + ( + "SELECT CASE WHEN flag THEN a ELSE b END AS value FROM t", + true, + ), + ( + "SELECT CASE WHEN flag THEN a ELSE NULL END AS value FROM t", + true, + ), + ( + "SELECT CASE WHEN flag THEN NULL ELSE a END AS value FROM t", + true, + ), + ( + "SELECT CASE WHEN flag THEN a ELSE c END AS value FROM t", + false, + ), + ] { + let df = ctx.sql(sql).await?; + let optimized = ctx.state().optimize(df.logical_plan())?; + let expected_metadata = if preserve_metadata { + &metadata + } else { + &empty_metadata + }; + for field in [df.schema().field(0), optimized.schema().field(0)] { + assert_eq!(field.metadata(), expected_metadata); + } + if sql == "SELECT CASE WHEN flag THEN a ELSE b END AS value FROM t" { + let batches = df.collect().await?; + datafusion::assert_batches_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| 1 |", + "| 4 |", + "+-------+" + ], + &batches + ); + } + } + Ok(()) +} + #[tokio::test] async fn test_version_function() { let expected_version = format!( From 05bfd056b9ef125a7fdb511fc25a5fc019d3b007 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 16:55:32 +0300 Subject: [PATCH 11/12] Test logical CASE field metadata for marked NULL casts (cherry picked from commit 816119e839317cca1b31ef02ad0e8e302e0c60d0) --- datafusion/expr/src/expr_schema.rs | 33 ++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 76f2677e23e46..fbef7c3f94912 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -1468,6 +1468,39 @@ mod tests { let all_typed_null = when(lit(true), typed_null.clone()).otherwise(typed_null)?; assert_eq!(all_typed_null.to_field(&schema)?.1.metadata(), &shared); + let binary_metadata = HashMap::from([( + EXTENSION_TYPE_NAME_KEY.to_string(), + "geoarrow.wkb".to_string(), + )]); + let binary_schema = DFSchema::from_unqualified_fields( + vec![ + Field::new("binary", DataType::LargeBinary, false) + .with_metadata(binary_metadata.clone()), + ] + .into(), + HashMap::new(), + )?; + let target_field: FieldRef = Arc::new( + Field::new("", DataType::Binary, true).with_metadata(binary_metadata), + ); + for marked_null in [ + Expr::Cast(Cast::new_from_field( + Box::new(lit(ScalarValue::Null)), + Arc::clone(&target_field), + )), + Expr::TryCast(TryCast::new_from_field( + Box::new(lit(ScalarValue::Null)), + Arc::clone(&target_field), + )), + ] { + let nested_all_null = + when(lit(true), marked_null).otherwise(lit(ScalarValue::Null))?; + let type_only_cast = + Expr::Cast(Cast::new(Box::new(nested_all_null), DataType::LargeBinary)); + let outer = when(lit(true), type_only_cast).otherwise(col("binary"))?; + assert!(outer.to_field(&binary_schema)?.1.metadata().is_empty()); + } + let mut nested = col("a"); for _ in 0..128 { nested = when(lit(true), col("a")).otherwise(nested)?; From 52881bf3d2ce11f3f3435754cf419227b2aaca26 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 9 Oct 2026 17:03:29 +0300 Subject: [PATCH 12/12] Test logical CASE metadata depth for marked scalar functions (cherry picked from commit 185ae25f63f622ece719cc6a95e3d8db3d6135ed) --- datafusion/expr/src/expr_schema.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index fbef7c3f94912..863d4347d618b 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -1557,6 +1557,19 @@ mod tests { when(lit(true), marked_call(col("a"))).otherwise(marked_call(col("b")))?; assert_eq!(marked_case.to_field(&schema)?.1.metadata(), &shared); + let mut deep_function = col("a"); + for _ in 0..9 { + deep_function = marked_call(deep_function); + } + let deep_function_case = when(lit(true), deep_function).otherwise(col("b"))?; + assert!( + deep_function_case + .to_field(&schema)? + .1 + .metadata() + .is_empty() + ); + let nested_arg = when(lit(true), col("a")).otherwise(col("b"))?; let nested_call = when(lit(true), marked_call(nested_arg)).otherwise(col("b"))?; assert_eq!(nested_call.to_field(&schema)?.1.metadata(), &shared);