diff --git a/datafusion/common/src/dfschema.rs b/datafusion/common/src/dfschema.rs index 1c072aaa25043..2b30891a67be2 100644 --- a/datafusion/common/src/dfschema.rs +++ b/datafusion/common/src/dfschema.rs @@ -1244,6 +1244,14 @@ pub trait ExprSchema: std::fmt::Debug { // Return the column's field fn field_from_column(&self, col: &Column) -> Result<&FieldRef>; + + /// Return the field and column nullability together. + /// + /// The default preserves custom implementations of [`Self::nullable`]. + fn field_and_nullable(&self, col: &Column) -> Result<(&FieldRef, bool)> { + let field = self.field_from_column(col)?; + Ok((field, self.nullable(col)?)) + } } // Implement `ExprSchema` for `Arc` @@ -1267,6 +1275,10 @@ impl + std::fmt::Debug> ExprSchema for P { fn field_from_column(&self, col: &Column) -> Result<&FieldRef> { self.as_ref().field_from_column(col) } + + fn field_and_nullable(&self, col: &Column) -> Result<(&FieldRef, bool)> { + self.as_ref().field_and_nullable(col) + } } impl ExprSchema for DFSchema { @@ -1276,6 +1288,11 @@ impl ExprSchema for DFSchema { None => self.field_with_unqualified_name(&col.name), } } + + fn field_and_nullable(&self, col: &Column) -> Result<(&FieldRef, bool)> { + let field = self.field_from_column(col)?; + Ok((field, field.is_nullable())) + } } /// DataFusion-specific extensions to [`Schema`]. diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs index fae2fb5bc5e76..3ec35c98b46e2 100644 --- a/datafusion/core/tests/sql/select.rs +++ b/datafusion/core/tests/sql/select.rs @@ -18,6 +18,7 @@ use std::collections::HashMap; use super::*; +use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; use datafusion_common::{ParamValues, ScalarValue, metadata::ScalarAndMetadata}; use insta::assert_snapshot; @@ -387,6 +388,81 @@ 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()), + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "example.case".to_string(), + ), + ]); + let other_metadata = HashMap::from([ + ("structured_type".to_string(), "OBJECT".to_string()), + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "example.other".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])), + ], + )?, + )?; + + for sql in [ + "SELECT CASE WHEN flag THEN a ELSE b END AS value FROM t", + "SELECT CASE flag WHEN TRUE THEN a ELSE b END AS value FROM t", + "SELECT CASE WHEN flag THEN a ELSE NULL END AS value FROM t", + "SELECT CASE WHEN flag THEN NULL ELSE a END AS value FROM t", + ] { + let df = ctx.sql(sql).await?; + let optimized = ctx.state().optimize(df.logical_plan())?; + for field in [df.schema().field(0), optimized.schema().field(0)] { + assert_eq!(field.metadata(), &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 + ); + } + } + + let conflicting = ctx + .sql("SELECT CASE WHEN flag THEN a ELSE c END AS value FROM t") + .await; + let error = match conflicting { + Ok(df) => ctx.state().optimize(df.logical_plan()).unwrap_err(), + Err(error) => error, + }; + assert!(error.to_string().contains("example.case"), "{error}"); + assert!(error.to_string().contains("example.other"), "{error}"); + Ok(()) +} + #[tokio::test] async fn test_limit_offset_parameters_named() -> Result<()> { let ctx = SessionContext::new(); diff --git a/datafusion/expr/Cargo.toml b/datafusion/expr/Cargo.toml index 45b66ee89885e..ccbf454d9a3a8 100644 --- a/datafusion/expr/Cargo.toml +++ b/datafusion/expr/Cargo.toml @@ -78,3 +78,7 @@ insta = { workspace = true } [[bench]] name = "normalize_columns" harness = false + +[[bench]] +name = "case_field_metadata" +harness = false diff --git a/datafusion/expr/benches/case_field_metadata.rs b/datafusion/expr/benches/case_field_metadata.rs new file mode 100644 index 0000000000000..32cbbcfe8c7ba --- /dev/null +++ b/datafusion/expr/benches/case_field_metadata.rs @@ -0,0 +1,81 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::HashMap; +use std::hint::black_box; + +use arrow::datatypes::{DataType, Field}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::DFSchema; +use datafusion_expr::{Expr, ExprSchemable, col, lit, when}; + +fn literal_case(width: usize) -> Expr { + let mut builder = when(lit(true), lit(0_i32)); + for i in 1..width { + builder = builder.when(lit(false), lit(i32::try_from(i).unwrap())); + } + builder.end().unwrap() +} + +fn column_case(width: usize, marked_last: bool) -> (DFSchema, Expr) { + let fields = (0..width) + .map(|i| { + let field = Field::new(format!("c{i}"), DataType::Int32, false); + if marked_last && i == width - 1 { + field.with_metadata(HashMap::from([( + "source".to_owned(), + "marked".to_owned(), + )])) + } else { + field + } + }) + .collect::>(); + let schema = + DFSchema::from_unqualified_fields(fields.into(), HashMap::new()).unwrap(); + let mut builder = when(lit(true), col("c0")); + for i in 1..width { + builder = builder.when(lit(false), col(format!("c{i}"))); + } + (schema, builder.end().unwrap()) +} + +fn case_field_metadata(c: &mut Criterion) { + let mut group = c.benchmark_group("case_field_metadata"); + for width in [10, 100, 1000] { + let literal = literal_case(width); + let empty_schema = DFSchema::empty(); + group.bench_with_input( + BenchmarkId::new("plain_literals", width), + &width, + |b, _| { + b.iter(|| black_box(literal.to_field(black_box(&empty_schema)).unwrap())) + }, + ); + + for (name, marked_last) in [("plain_columns", false), ("marked_last", true)] { + let (schema, expression) = column_case(width, marked_last); + group.bench_with_input(BenchmarkId::new(name, width), &width, |b, _| { + b.iter(|| black_box(expression.to_field(black_box(&schema)).unwrap())) + }); + } + } + group.finish(); +} + +criterion_group!(benches, case_field_metadata); +criterion_main!(benches); diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 2804c736c4b04..1f2d9dd12756f 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, in_subquery_tuple_values, }; @@ -35,12 +35,14 @@ 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, - 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; +use std::collections::HashMap; use std::sync::Arc; /// Trait to allow expr to typable with respect to a schema @@ -138,6 +140,252 @@ fn scalar_argument_for_field(expr: &Expr, arg_field: &FieldRef) -> Option 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, allow_nested_case: bool) -> Result { + let mut work = vec![(expr, 1)]; + while let Some((node, depth)) = work.pop() { + if (!allow_nested_case && 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), + 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) + if !cast.field.metadata().is_empty() + && can_infer_exact_branch(expr, false)? => + { + 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, false)? => + { + 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.as_ref().is_some_and(|meta| !meta.is_empty()) + && can_infer_exact_branch(expr, false)? => + { + 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, true)? => + { + fields.push(BranchField { + field: expr.to_field(schema)?.1, + certainly_null: false, + }); + } + 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::FinishCast(target, force_nullable) => { + let Some(source) = fields.pop() else { + return internal_err!("Missing CASE cast input field"); + }; + let untyped_null = + source.certainly_null && source.field.data_type().is_null(); + let mut field = cast_output_field(&source.field, target, force_nullable); + // A type-only coercion must not make an untyped NULL constrain CASE metadata. + if untyped_null + && target.metadata().is_empty() + && !field.metadata().is_empty() + { + field = + Arc::new(field.as_ref().clone().with_metadata(HashMap::new())); + } + fields.push(BranchField { + field, + 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: Option> = None; + let mut extension: Option<(String, Option)> = None; + let mut type_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 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) { + type_conflict = true; + } + branch_type.get_or_insert_with(|| data_type.clone()); + let branch_metadata = branch.field.metadata(); + if let Some(name) = branch_metadata.get(EXTENSION_TYPE_NAME_KEY) { + match &extension { + Some((first_name, _)) if first_name != name => { + return plan_err!( + "Cannot merge field metadata with differing values of \ + '{EXTENSION_TYPE_NAME_KEY}': '{first_name}' and '{name}'" + ); + } + Some(_) => {} + None => { + extension = Some(( + name.clone(), + branch_metadata + .get(EXTENSION_TYPE_METADATA_KEY) + .cloned(), + )); + } + } + } + for (key, value) in branch_metadata { + if key != EXTENSION_TYPE_NAME_KEY + && key != EXTENSION_TYPE_METADATA_KEY + { + metadata + .get_or_insert_with(HashMap::new) + .entry(key.clone()) + .or_insert_with(|| value.clone()); + } + } + } + let data_type = if then_type.is_null() { + else_type + } else { + then_type + }; + if type_conflict || branch_type.as_ref() != Some(&data_type) { + metadata = None; + extension = None; + } + let mut metadata = metadata.unwrap_or_default(); + if let Some((name, extension_metadata)) = extension { + metadata.insert(EXTENSION_TYPE_NAME_KEY.to_owned(), name); + if let Some(extension_metadata) = extension_metadata { + metadata.insert( + EXTENSION_TYPE_METADATA_KEY.to_owned(), + extension_metadata, + ); + } + } + fields.push(BranchField { + field: Arc::new( + Field::new("", data_type, true).with_metadata(metadata), + ), + 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] @@ -679,6 +927,70 @@ 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 mut all_plain = true; + let mut nullable_then = false; + let mut nullable_else = false; + for (index, result) in case + .when_then_expr + .iter() + .map(|(_, result)| result.as_ref()) + .chain(case.else_expr.iter().map(|result| result.as_ref())) + .enumerate() + { + let (has_metadata, nullable) = match result { + Expr::Column(column) => { + // Defer later nullability checks to Expr::nullable, which can short-circuit. + let (field, nullable) = if nullable_then { + (schema.field_from_column(column)?, false) + } else { + schema.field_and_nullable(column)? + }; + (!field.metadata().is_empty(), nullable) + } + Expr::Literal(value, metadata) => ( + metadata + .as_ref() + .is_some_and(|metadata| !metadata.is_empty()), + value.is_null(), + ), + _ => { + all_plain = false; + break; + } + }; + if has_metadata { + all_plain = false; + break; + } + if index < case.when_then_expr.len() { + nullable_then |= nullable; + } else { + nullable_else = nullable; + } + } + let nullable = if all_plain && !nullable_then { + case.else_expr.is_none() || nullable_else + } else { + self.nullable(schema)? + }; + if all_plain { + 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::Unnest(Unnest { expr, .. }) => { let nullable = self.nullable(schema)?; match expr.get_type(schema)? { @@ -704,7 +1016,6 @@ impl ExprSchemable for Expr { | Expr::SimilarTo(_) | Expr::Not(_) | Expr::Between(_) - | Expr::Case(_) | Expr::InList(_) | Expr::InSubquery(_) | Expr::SetComparison(_) @@ -1086,6 +1397,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); @@ -1278,6 +1596,562 @@ mod tests { assert_eq!(meta, outer_ref.metadata(&schema).unwrap()); } + #[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( + 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), + Field::new("e", DataType::Boolean, false).with_metadata(shared.clone()), + Field::new("f", DataType::Boolean, false).with_metadata(shared.clone()), + ] + .into(), + HashMap::new(), + )?; + + 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_eq!(conflicting_when.to_field(&schema)?.1.metadata(), &shared); + + 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_eq!(unmarked_literal.to_field(&schema)?.1.metadata(), &shared); + + 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); + + 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 Expr::Alias(empty_metadata_alias) = lit(ScalarValue::Null) + .cast_to(&DataType::Int32, &schema)? + .alias("empty_metadata") + else { + unreachable!(); + }; + let empty_metadata_alias = Expr::Alias( + empty_metadata_alias.with_metadata(Some(FieldMetadata::default())), + ); + let aliased_null_else = + when(lit(true), col("a")).otherwise(empty_metadata_alias)?; + assert_eq!(aliased_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 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 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.clone()), + ); + 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_eq!( + outer.to_field(&binary_schema)?.1.metadata(), + &binary_metadata + ); + } + + 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 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 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 scalar_result = + Expr::ScalarFunction(ScalarFunction::new_udf(identity, vec![col("a")])); + let scalar_case = when(lit(true), scalar_result).otherwise(col("b"))?; + assert_eq!(scalar_case.to_field(&schema)?.1.metadata(), &shared); + + 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 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); + + 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 { + 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_eq!(mismatched.to_field(&schema)?.1.metadata(), &shared); + + let unmarked = when(lit(true), col("a")).otherwise(col("d"))?; + assert_eq!(unmarked.to_field(&schema)?.1.metadata(), &shared); + + let first_unmarked = when(lit(true), col("d")).otherwise(col("a"))?; + assert_eq!(first_unmarked.to_field(&schema)?.1.metadata(), &shared); + + Ok(()) + } + + #[test] + fn test_case_extension_metadata_merge_policy() -> Result<()> { + let first = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_owned(), + "geoarrow.wkb".to_owned(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_owned(), + "{\"a\":1,\"b\":2}".to_owned(), + ), + ("shared".to_owned(), "first".to_owned()), + ("only_first".to_owned(), "left".to_owned()), + ]); + let second = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_owned(), + "geoarrow.wkb".to_owned(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_owned(), + "{\"b\":2,\"a\":1}".to_owned(), + ), + ("shared".to_owned(), "second".to_owned()), + ("only_second".to_owned(), "right".to_owned()), + ]); + let schema = DFSchema::from_unqualified_fields( + vec![ + Field::new("first", DataType::Binary, false).with_metadata(first.clone()), + Field::new("second", DataType::Binary, false).with_metadata(second), + Field::new("plain", DataType::Binary, false), + Field::new("bare", DataType::Binary, false).with_metadata(HashMap::from( + [( + EXTENSION_TYPE_NAME_KEY.to_owned(), + "geoarrow.wkb".to_owned(), + )], + )), + Field::new("other", DataType::Binary, false).with_metadata( + HashMap::from([( + EXTENSION_TYPE_NAME_KEY.to_owned(), + "example.other".to_owned(), + )]), + ), + ] + .into(), + HashMap::new(), + )?; + + let combined = when(lit(true), col("first")).otherwise(col("second"))?; + let metadata = combined.to_field(&schema)?.1.metadata().clone(); + assert_eq!( + metadata.get(EXTENSION_TYPE_NAME_KEY), + first.get(EXTENSION_TYPE_NAME_KEY) + ); + assert_eq!( + metadata.get(EXTENSION_TYPE_METADATA_KEY), + first.get(EXTENSION_TYPE_METADATA_KEY) + ); + assert_eq!(metadata.get("shared").map(String::as_str), Some("first")); + assert_eq!(metadata.get("only_first").map(String::as_str), Some("left")); + assert_eq!( + metadata.get("only_second").map(String::as_str), + Some("right") + ); + + let plain = when(lit(true), col("plain")).otherwise(col("first"))?; + assert_eq!(plain.to_field(&schema)?.1.metadata(), &first); + + let bare = when(lit(true), col("bare")).otherwise(col("first"))?; + let metadata = bare.to_field(&schema)?.1.metadata().clone(); + assert_eq!( + metadata.get(EXTENSION_TYPE_NAME_KEY).map(String::as_str), + Some("geoarrow.wkb") + ); + assert!(!metadata.contains_key(EXTENSION_TYPE_METADATA_KEY)); + + let conflicting = when(lit(true), col("first")).otherwise(col("other"))?; + let error = conflicting.to_field(&schema).unwrap_err(); + assert!(error.to_string().contains("geoarrow.wkb"), "{error}"); + assert!(error.to_string().contains("example.other"), "{error}"); + Ok(()) + } + + #[test] + fn test_plain_case_field_nullability() -> Result<()> { + let schema = DFSchema::from_unqualified_fields( + vec![ + Field::new("required", DataType::Int32, false), + Field::new("optional", DataType::Int32, true), + ] + .into(), + HashMap::new(), + )?; + + let no_else = when(lit(true), col("required")).end()?; + assert!(no_else.to_field(&schema)?.1.is_nullable()); + + let required_else = + when(lit(true), col("required")).otherwise(col("required"))?; + assert!(!required_else.to_field(&schema)?.1.is_nullable()); + + let optional_else = + when(lit(true), col("required")).otherwise(col("optional"))?; + assert!(optional_else.to_field(&schema)?.1.is_nullable()); + + let optional_then = + when(lit(true), col("optional")).otherwise(col("required"))?; + assert!(optional_then.to_field(&schema)?.1.is_nullable()); + + let guarded_optional = when(col("optional").is_not_null(), col("optional")) + .otherwise(col("required"))?; + assert!(!guarded_optional.to_field(&schema)?.1.is_nullable()); + Ok(()) + } + + #[test] + fn test_plain_case_preserves_custom_nullable_error() -> Result<()> { + let schema = MockExprSchema::new() + .with_data_type(DataType::Int32) + .with_nullable(false) + .with_error_on_nullable(true); + let expression = when(lit(true), col("x")).otherwise(lit(0_i32))?; + let error = expression.to_field(&schema).unwrap_err(); + assert!(error.to_string().contains("nullable error"), "{error}"); + Ok(()) + } + + #[test] + fn test_plain_case_preserves_nullable_short_circuit() -> Result<()> { + #[derive(Debug)] + struct ShortCircuitSchema { + field: FieldRef, + } + + impl ExprSchema for ShortCircuitSchema { + fn nullable(&self, column: &Column) -> Result { + if column.name == "later" { + return internal_err!("later nullable call"); + } + Ok(true) + } + + fn field_from_column(&self, _column: &Column) -> Result<&FieldRef> { + Ok(&self.field) + } + } + + let schema = ShortCircuitSchema { + field: Arc::new(Field::new("value", DataType::Int32, true)), + }; + let expression = when(lit(true), col("first")) + .when(lit(false), col("later")) + .otherwise(lit(0_i32))?; + assert!(expression.to_field(&schema)?.1.is_nullable()); + 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_case_field_metadata_higher_order_function() -> Result<()> { + #[derive(Debug, PartialEq, Eq, Hash)] + struct MetadataPassthrough { + signature: crate::HigherOrderSignature, + } + + impl crate::HigherOrderUDFImpl for MetadataPassthrough { + fn name(&self) -> &str { + "metadata_passthrough" + } + + fn signature(&self) -> &crate::HigherOrderSignature { + &self.signature + } + + fn lambda_parameters( + &self, + _step: usize, + _fields: &[ValueOrLambda>], + ) -> Result { + Ok(crate::LambdaParametersProgress::Complete(vec![])) + } + + fn return_field_from_args( + &self, + args: HigherOrderReturnFieldArgs, + ) -> Result { + let ValueOrLambda::Value(field) = &args.arg_fields[0] else { + unreachable!(); + }; + Ok(Arc::clone(field)) + } + + fn invoke_with_args( + &self, + _args: crate::HigherOrderFunctionArgs, + ) -> Result { + unreachable!() + } + } + + let metadata = HashMap::from([("kind".to_string(), "marked".to_string())]); + let schema = DFSchema::from_unqualified_fields( + vec![ + Field::new("a", DataType::Int32, false).with_metadata(metadata.clone()), + Field::new("b", DataType::Int32, false).with_metadata(metadata.clone()), + ] + .into(), + HashMap::new(), + )?; + let udf = Arc::new(crate::HigherOrderUDF::new_from_impl(MetadataPassthrough { + signature: crate::HigherOrderSignature::any(1, crate::Volatility::Immutable), + })); + let call = |arg| { + Expr::HigherOrderFunction(crate::expr::HigherOrderFunction::new( + Arc::clone(&udf), + vec![arg], + )) + }; + + let shallow = when(lit(true), call(col("a"))).otherwise(col("b"))?; + assert_eq!(shallow.to_field(&schema)?.1.metadata(), &metadata); + + let mut deep = col("a"); + for _ in 0..9 { + deep = call(deep); + } + let deep_case = when(lit(true), deep).otherwise(col("b"))?; + assert!(deep_case.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 abc756879a20d..402656a8fe78f 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -3166,6 +3166,83 @@ 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, None)?; + 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_coercion_preserves_metadata_with_annotated_null() -> Result<()> { + let metadata = std::collections::HashMap::from([ + ( + "ARROW:extension:name".to_string(), + "example.case".to_string(), + ), + ("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 conflicting = std::collections::HashMap::from([ + ( + "ARROW:extension:name".to_string(), + "example.other".to_string(), + ), + ("structured_type".to_string(), "OBJECT".to_string()), + ]); + for null_metadata in [metadata.clone(), conflicting] { + let annotated_null = Expr::Literal( + ScalarValue::Null, + Some(expr::FieldMetadata::from(null_metadata)), + ); + let case = Case { + expr: None, + when_then_expr: vec![(Box::new(lit(true)), Box::new(col("structured")))], + else_expr: Some(Box::new(annotated_null)), + }; + assert_eq!( + Expr::Case(case.clone()).to_field(&schema)?.1.metadata(), + &metadata + ); + + let coerced = coerce_case_expression(case, &schema, None)?; + 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(