diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index 17288a9737699..29f1151a12cb8 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -20,16 +20,19 @@ mod literal_lookup_table; use super::{Column, Literal}; use crate::PhysicalExpr; use crate::expressions::{ - CastExpr, LambdaVariable, NegativeExpr, NotExpr, lit, try_cast, + CastExpr, LambdaVariable, NegativeExpr, NotExpr, TryCastExpr, lit, try_cast, }; +use crate::{HigherOrderFunctionExpr, ScalarFunctionExpr}; use arrow::array::*; use arrow::compute::kernels::zip::zip; use arrow::compute::{ FilterBuilder, FilterPredicate, is_not_null, not, nullif, prep_null_mask_filter, }; -use arrow::datatypes::{DataType, Schema, UInt32Type, UnionMode}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema, UInt32Type, UnionMode}; use arrow::error::ArrowError; +use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::cast::as_boolean_array; +use datafusion_common::metadata::FieldMetadata; use datafusion_common::{ DataFusionError, Result, ScalarValue, assert_or_internal_err, exec_err, internal_datafusion_err, internal_err, @@ -276,6 +279,8 @@ pub struct CaseExpr { body: CaseBody, /// Evaluation method to use eval_method: EvalMethod, + /// Logical result type and metadata inferred before physical lowering. + logical_result_field: Option<(DataType, FieldMetadata)>, } // eval_method is functionally derived from body, so excluding it from @@ -284,12 +289,13 @@ pub struct CaseExpr { impl Hash for CaseExpr { fn hash(&self, state: &mut H) { self.body.hash(state); + self.logical_result_field.hash(state); } } impl PartialEq for CaseExpr { fn eq(&self, other: &Self) -> bool { - self.body == other.body + self.body == other.body && self.logical_result_field == other.logical_result_field } } @@ -676,7 +682,20 @@ impl CaseExpr { let eval_method = Self::find_best_eval_method(&body)?; - Ok(Self { body, eval_method }) + Ok(Self { + body, + eval_method, + logical_result_field: None, + }) + } + + /// Preserve the logical CASE result metadata across physical lowering. + pub fn with_logical_result_field(mut self, field: &FieldRef) -> Self { + self.logical_result_field = Some(( + field.data_type().clone(), + FieldMetadata::from(field.metadata().clone()), + )); + self } fn find_best_eval_method(body: &CaseBody) -> Result { @@ -1247,11 +1266,243 @@ impl CaseExpr { } } +type Metadata = std::collections::HashMap; + +// Resolve nested CASE fields bottom-up so field names and nullability are not +// recomputed for every suffix of a nested CASE chain. +fn case_result_field_metadata( + case: &CaseExpr, + input_schema: &Schema, +) -> Result<(DataType, Metadata)> { + const MAX_EXACT_BRANCH_DEPTH: usize = 8; + + fn can_infer_exact_function_branch(expr: &dyn PhysicalExpr) -> bool { + let mut work = vec![(expr, 1)]; + while let Some((node, depth)) = work.pop() { + if depth > MAX_EXACT_BRANCH_DEPTH { + return false; + } + for child in node.children() { + work.push((child.as_ref(), depth + 1)); + } + } + true + } + + struct BranchField { + data_type: DataType, + metadata: Metadata, + certainly_null: bool, + } + + enum Work<'a> { + Visit(&'a dyn PhysicalExpr), + FinishCase(&'a CaseExpr), + FinishCast(&'a DataType), + } + + fn schedule_case<'a>(case: &'a CaseExpr, work: &mut Vec>) { + work.push(Work::FinishCase(case)); + work.extend( + case.body + .else_expr + .iter() + .map(|expr| Work::Visit(expr.as_ref())), + ); + work.extend( + case.body + .when_then_expr + .iter() + .rev() + .map(|(_, expr)| Work::Visit(expr.as_ref())), + ); + } + + let mut work = Vec::new(); + let mut fields = Vec::new(); + schedule_case(case, &mut work); + while let Some(item) = work.pop() { + match item { + Work::Visit(expr) => { + if let Some(nested) = expr.downcast_ref::() { + schedule_case(nested, &mut work); + } else if let Some(cast) = expr.downcast_ref::() { + if let Some(metadata) = cast.target_metadata() { + fields.push(BranchField { + data_type: cast.cast_type().clone(), + metadata: metadata.clone(), + certainly_null: false, + }); + } else { + work.push(Work::FinishCast(cast.cast_type())); + work.push(Work::Visit(cast.expr().as_ref())); + } + } else if let Some(cast) = expr.downcast_ref::() { + if let Some(metadata) = cast.target_metadata() { + fields.push(BranchField { + data_type: cast.cast_type().clone(), + metadata: metadata.clone(), + certainly_null: false, + }); + } else { + work.push(Work::FinishCast(cast.cast_type())); + work.push(Work::Visit(cast.expr().as_ref())); + } + } else if let Some(negative) = expr.downcast_ref::() { + work.push(Work::Visit(negative.arg().as_ref())); + } else if expr.is::() { + fields.push(BranchField { + data_type: DataType::Boolean, + metadata: Metadata::new(), + certainly_null: false, + }); + } else if expr.is::() + || expr.is::() + { + let field = expr.return_field(input_schema)?; + let metadata = if field.metadata().is_empty() + || can_infer_exact_function_branch(expr) + { + field.metadata().clone() + } else { + Metadata::new() + }; + fields.push(BranchField { + data_type: field.data_type().clone(), + metadata, + certainly_null: false, + }); + } else { + let field = expr.return_field(input_schema)?; + fields.push(BranchField { + data_type: field.data_type().clone(), + metadata: field.metadata().clone(), + certainly_null: expr + .downcast_ref::() + .is_some_and(|literal| literal.value().is_null()), + }); + } + } + Work::FinishCast(target) => { + let Some(mut source) = fields.pop() else { + return internal_err!("Missing CASE cast input field"); + }; + source.data_type = target.clone(); + source.metadata.remove(EXTENSION_TYPE_NAME_KEY); + source.metadata.remove(EXTENSION_TYPE_METADATA_KEY); + fields.push(source); + } + Work::FinishCase(case) => { + let count = case.body.when_then_expr.len() + + usize::from(case.body.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 marked_type = None; + let mut metadata = None; + let mut conflict = false; + let mut certainly_null = true; + for (index, branch) in fields.drain(start..).enumerate() { + if index < case.body.when_then_expr.len() { + if then_type.is_null() && !branch.data_type.is_null() { + then_type = branch.data_type.clone(); + } + } else { + else_type = branch.data_type.clone(); + } + certainly_null &= branch.certainly_null; + if branch.data_type.is_null() + || (branch.certainly_null && branch.metadata.is_empty()) + { + continue; + } + if marked_type + .as_ref() + .is_some_and(|other| other != &branch.data_type) + || branch.metadata.is_empty() + || metadata + .as_ref() + .is_some_and(|other| other != &branch.metadata) + { + conflict = true; + } + marked_type.get_or_insert_with(|| branch.data_type.clone()); + metadata.get_or_insert(branch.metadata); + } + let data_type = if then_type.is_null() { + else_type + } else { + then_type + }; + if conflict || marked_type.as_ref() != Some(&data_type) { + metadata = None; + } + let mut metadata = metadata.unwrap_or_default(); + if let Some((logical_type, logical_metadata)) = &case.logical_result_field + { + metadata = if logical_type == &data_type { + logical_metadata.to_hashmap() + } else { + Metadata::new() + }; + } + fields.push(BranchField { + data_type, + metadata, + certainly_null, + }); + } + } + } + let Some(result) = fields.pop() else { + return internal_err!("Missing CASE output field"); + }; + Ok((result.data_type, result.metadata)) +} + impl PhysicalExpr for CaseExpr { fn data_type(&self, input_schema: &Schema) -> Result { self.body.data_type(input_schema) } + fn return_field(&self, input_schema: &Schema) -> Result { + let data_type = self.data_type(input_schema)?; + let nullable = self.nullable(input_schema)?; + if let Some((logical_type, logical_metadata)) = &self.logical_result_field { + let field = Field::new(format!("{self}"), data_type, nullable); + return Ok(Arc::new(if logical_type == field.data_type() { + field.with_metadata(logical_metadata.to_hashmap()) + } else { + field + })); + } + if let Some((_, first_result)) = self.body.when_then_expr.first() + && (first_result.is::() || first_result.is::()) + { + let first_field = first_result.return_field(input_schema)?; + if !first_field.data_type().is_null() + && !first_result + .downcast_ref::() + .is_some_and(|literal| literal.value().is_null()) + && first_field.metadata().is_empty() + { + return Ok(Arc::new(Field::new(format!("{self}"), data_type, nullable))); + } + } + let (branch_type, metadata) = case_result_field_metadata(self, input_schema)?; + let metadata = if branch_type == data_type { + metadata + } else { + Metadata::new() + }; + Ok(Arc::new( + Field::new(format!("{self}"), data_type, nullable).with_metadata(metadata), + )) + } + fn nullable(&self, input_schema: &Schema) -> Result { let nullable_then = self .body @@ -1386,11 +1637,15 @@ impl PhysicalExpr for CaseExpr { ), (false, false) => (None, &children[0..children.len()], None), }; - Ok(Arc::new(CaseExpr::try_new( + let mut rebuilt = CaseExpr::try_new( expr.cloned(), when_then_expr.iter().cloned().tuples().collect(), else_expr.cloned(), - )?)) + )?; + rebuilt + .logical_result_field + .clone_from(&self.logical_result_field); + Ok(Arc::new(rebuilt)) } } @@ -1446,6 +1701,16 @@ impl PhysicalExpr for CaseExpr { .else_expr() .map(|expr| ctx.encode_child(expr).map(Box::new)) .transpose()?, + logical_result_type: self + .logical_result_field + .as_ref() + .map(|(data_type, _)| data_type.try_into()) + .transpose()?, + logical_result_metadata: self + .logical_result_field + .as_ref() + .map(|(_, metadata)| metadata.to_hashmap()) + .unwrap_or_default(), }, ))), })) @@ -1468,7 +1733,7 @@ impl CaseExpr { "CaseExpr", ); - Ok(Arc::new(CaseExpr::try_new( + let mut expr = CaseExpr::try_new( case.expr .as_deref() .map(|expr| ctx.decode(expr)) @@ -1494,7 +1759,16 @@ impl CaseExpr { .as_deref() .map(|expr| ctx.decode(expr)) .transpose()?, - )?)) + )?; + if let Some(logical_type) = &case.logical_result_type { + expr.logical_result_field = Some(( + logical_type.try_into()?, + FieldMetadata::from(case.logical_result_metadata.clone()), + )); + } else if !case.logical_result_metadata.is_empty() { + return internal_err!("CaseExpr logical metadata requires a result type"); + } + Ok(Arc::new(expr)) } } @@ -1575,12 +1849,65 @@ pub fn case( mod tests { use super::*; + use std::collections::HashMap; + + std::thread_local! { + static COUNTED_COLUMN_DISPLAYS: std::cell::Cell = const { std::cell::Cell::new(0) }; + static COUNTED_COLUMN_FIELDS: std::cell::Cell = const { std::cell::Cell::new(0) }; + } + + #[derive(Debug, Hash, PartialEq, Eq)] + struct CountedColumn(Column); + + impl std::fmt::Display for CountedColumn { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + COUNTED_COLUMN_DISPLAYS.with(|count| count.set(count.get() + 1)); + std::fmt::Display::fmt(&self.0, f) + } + } + + impl PhysicalExpr for CountedColumn { + fn data_type(&self, input_schema: &Schema) -> Result { + self.0.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result { + self.0.nullable(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + self.0.evaluate(batch) + } + + fn return_field(&self, input_schema: &Schema) -> Result { + COUNTED_COLUMN_FIELDS.with(|count| count.set(count.get() + 1)); + self.0.return_field(input_schema) + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert!(children.is_empty()); + Ok(self) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.0.fmt_sql(f) + } + } + use crate::expressions; use crate::expressions::{BinaryExpr, binary, cast, col, is_not_null}; use arrow::buffer::Buffer; use arrow::datatypes::DataType::Float64; use arrow::datatypes::Field; use datafusion_common::cast::{as_float64_array, as_int32_array}; + use datafusion_common::metadata::FieldMetadata; use datafusion_common::plan_err; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_expr::type_coercion::binary::type_union_coercion; @@ -1588,6 +1915,258 @@ mod tests { use datafusion_physical_expr_common::physical_expr::fmt_sql; use half::f16; + #[test] + fn case_result_field_metadata() -> Result<()> { + let metadata = HashMap::from([( + "ARROW:extension:name".to_string(), + "geoarrow.wkb".to_string(), + )]); + let other = HashMap::from([( + "ARROW:extension:name".to_string(), + "arrow.uuid".to_string(), + )]); + let schema = Schema::new(vec![ + Field::new("a", DataType::Binary, false).with_metadata(metadata.clone()), + Field::new("b", DataType::Binary, false).with_metadata(metadata.clone()), + Field::new("c", DataType::Binary, false).with_metadata(other), + Field::new("d", DataType::Binary, false), + ]); + let empty = HashMap::new(); + + for (else_column, expected_metadata) in [ + (Some("b"), &metadata), + (Some("c"), &empty), + (Some("d"), &empty), + (None, &metadata), + ] { + let else_expr = else_column.map(|name| col(name, &schema)).transpose()?; + let expr = case(None, vec![(lit(true), col("a", &schema)?)], else_expr)?; + assert_eq!(expr.return_field(&schema)?.metadata(), expected_metadata); + } + + let null_else = case( + None, + vec![(lit(true), col("a", &schema)?)], + Some(lit(ScalarValue::Null)), + )?; + assert_eq!(null_else.return_field(&schema)?.metadata(), &metadata); + + let null_first = case( + None, + vec![(lit(true), lit(ScalarValue::Null))], + Some(col("a", &schema)?), + )?; + assert_eq!(null_first.return_field(&schema)?.metadata(), &metadata); + + let typed_null = Arc::new(Literal::new_with_metadata( + ScalarValue::Binary(None), + Some(FieldMetadata::from(metadata.clone())), + )); + let typed_null_case = case( + None, + vec![(lit(true), typed_null)], + Some(col("a", &schema)?), + )?; + assert_eq!(typed_null_case.return_field(&schema)?.metadata(), &metadata); + + let cast_metadata = + HashMap::from([("logical".to_string(), "number".to_string())]); + let cast_schema = Schema::new(vec![ + Field::new("small", DataType::Int32, false) + .with_metadata(cast_metadata.clone()), + Field::new("large", DataType::Int64, false) + .with_metadata(cast_metadata.clone()), + ]); + let nested = case( + None, + vec![(lit(true), col("small", &cast_schema)?)], + Some(col("small", &cast_schema)?), + )?; + let casted = cast(nested, &cast_schema, DataType::Int64)?; + let cast_case = case( + None, + vec![(lit(true), casted)], + Some(col("large", &cast_schema)?), + )?; + assert_eq!( + cast_case.return_field(&cast_schema)?.metadata(), + &cast_metadata + ); + + let try_cast_null = + Arc::new(TryCastExpr::new(lit(ScalarValue::Null), DataType::Int64)); + let try_cast_case = case( + None, + vec![(lit(true), try_cast_null)], + Some(col("large", &cast_schema)?), + )?; + assert_eq!( + try_cast_case.return_field(&cast_schema)?.metadata(), + &cast_metadata + ); + Ok(()) + } + + #[test] + fn case_logical_field_survives_rewriting_children() -> Result<()> { + let metadata = HashMap::from([("source".to_string(), "physical".to_string())]); + let schema = Schema::new(vec![ + Field::new("a", DataType::Binary, false).with_metadata(metadata.clone()), + Field::new("b", DataType::Binary, false).with_metadata(metadata), + ]); + let expr: Arc = Arc::new( + CaseExpr::try_new( + None, + vec![(lit(true), col("a", &schema)?)], + Some(col("b", &schema)?), + )? + .with_logical_result_field(&Arc::new(Field::new( + "logical_case", + DataType::Binary, + true, + ))), + ); + assert!(expr.return_field(&schema)?.metadata().is_empty()); + let children = expr.children().into_iter().cloned().collect(); + let rewritten = expr.with_new_children(children)?; + assert!(rewritten.return_field(&schema)?.metadata().is_empty()); + Ok(()) + } + + #[test] + fn case_not_does_not_inherit_input_metadata() -> Result<()> { + let metadata = HashMap::from([("logical".to_string(), "boolean".to_string())]); + let schema = Schema::new(vec![ + Field::new("a", DataType::Boolean, false).with_metadata(metadata.clone()), + Field::new("b", DataType::Boolean, false).with_metadata(metadata), + ]); + let negated: Arc = Arc::new(NotExpr::new(col("a", &schema)?)); + let expr = case(None, vec![(lit(true), negated)], Some(col("b", &schema)?))?; + assert!(expr.return_field(&schema)?.metadata().is_empty()); + + let column: Arc = Arc::new(CountedColumn(Column::new("a", 0))); + for depth in [16, 64] { + let mut nested = Arc::clone(&column); + for _ in 0..depth { + nested = case( + None, + vec![(lit(true), Arc::new(NotExpr::new(nested)))], + Some(col("b", &schema)?), + )?; + } + COUNTED_COLUMN_FIELDS.with(|count| count.set(0)); + let nested_case = nested.downcast_ref::().ok_or_else(|| { + internal_datafusion_err!("Expected nested CASE expression") + })?; + super::case_result_field_metadata(nested_case, &schema)?; + assert_eq!(COUNTED_COLUMN_FIELDS.with(|count| count.get()), 0); + } + Ok(()) + } + + #[test] + fn case_explicitly_marked_null_cast_matches_logical_field() -> Result<()> { + let metadata = HashMap::from([( + EXTENSION_TYPE_NAME_KEY.to_string(), + "geoarrow.wkb".to_string(), + )]); + let schema = Schema::new(vec![ + Field::new("a", DataType::LargeBinary, false).with_metadata(metadata.clone()), + ]); + let target_field: FieldRef = Arc::new( + Field::new("marked_null", DataType::Binary, true).with_metadata(metadata), + ); + let marked_nulls: [Arc; 2] = [ + Arc::new(CastExpr::new_with_target_field( + lit(ScalarValue::Null), + Arc::clone(&target_field), + None, + )), + Arc::new(TryCastExpr::new_with_target_field( + lit(ScalarValue::Null), + target_field, + )), + ]; + for marked_null in marked_nulls { + let all_null = case( + None, + vec![(lit(true), marked_null)], + Some(lit(ScalarValue::Null)), + )?; + let type_only_cast: Arc = + Arc::new(CastExpr::new(all_null, DataType::LargeBinary, None)); + let outer = case( + None, + vec![(lit(true), type_only_cast)], + Some(col("a", &schema)?), + )?; + assert!(outer.return_field(&schema)?.metadata().is_empty()); + } + Ok(()) + } + + #[test] + fn case_function_metadata_respects_logical_depth_limit() -> Result<()> { + let metadata = HashMap::from([("logical".to_string(), "binary".to_string())]); + let schema = Schema::new(vec![ + Field::new("a", DataType::Binary, false).with_metadata(metadata.clone()), + Field::new("b", DataType::Binary, false).with_metadata(metadata.clone()), + ]); + let udf = Arc::new(datafusion_expr::expr_fn::create_udf( + "marked_identity", + vec![DataType::Binary], + DataType::Binary, + datafusion_expr::Volatility::Immutable, + Arc::new(|args| Ok(args[0].clone())), + )); + let function_field: FieldRef = Arc::new( + Field::new("marked_identity", DataType::Binary, true) + .with_metadata(metadata.clone()), + ); + let options = Arc::new(datafusion_common::config::ConfigOptions::default()); + for (depth, expected_marked) in [(2, true), (9, false)] { + let mut branch = col("a", &schema)?; + for _ in 0..depth { + branch = Arc::new(ScalarFunctionExpr::new( + "marked_identity", + Arc::clone(&udf), + vec![branch], + Arc::clone(&function_field), + Arc::clone(&options), + )); + } + let expr = case(None, vec![(lit(true), branch)], Some(col("b", &schema)?))?; + let field = expr.return_field(&schema)?; + if expected_marked { + assert_eq!(field.metadata(), &metadata); + } else { + assert!(field.metadata().is_empty()); + } + } + Ok(()) + } + + #[test] + fn nested_case_metadata_field_names_scale_linearly() -> Result<()> { + let metadata = HashMap::from([("logical".to_string(), "binary".to_string())]); + let schema = Schema::new(vec![ + Field::new("a", DataType::Binary, false).with_metadata(metadata.clone()), + ]); + let column: Arc = Arc::new(CountedColumn(Column::new("a", 0))); + let mut display_counts = Vec::new(); + for depth in [16, 64] { + let mut expr = Arc::clone(&column); + for _ in 0..depth { + expr = case(None, vec![(lit(true), Arc::clone(&column))], Some(expr))?; + } + COUNTED_COLUMN_DISPLAYS.with(|count| count.set(0)); + assert_eq!(expr.return_field(&schema)?.metadata(), &metadata); + display_counts.push(COUNTED_COLUMN_DISPLAYS.with(|count| count.get())); + } + assert!(display_counts[1] <= display_counts[0] * 5); + Ok(()) + } + #[test] fn case_with_expr() -> Result<()> { let batch = case_test_batch()?; @@ -3363,6 +3942,7 @@ mod proto_tests { use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; use datafusion_proto_models::protobuf; use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalWhenThen}; + use std::collections::HashMap; fn proto_case_fixture() -> CaseExpr { let schema = Schema::new(vec![Field::new("a", DataType::Boolean, true)]); @@ -3396,6 +3976,8 @@ mod proto_tests { expr, when_then_expr, else_expr, + logical_result_type: None, + logical_result_metadata: Default::default(), }, ))), } @@ -3422,6 +4004,37 @@ mod proto_tests { assert!(case_node.when_then_expr[0].when_expr.is_some()); assert!(case_node.when_then_expr[0].then_expr.is_some()); assert!(case_node.else_expr.is_some()); + assert!(case_node.logical_result_type.is_none()); + assert!(case_node.logical_result_metadata.is_empty()); + } + + #[test] + fn proto_preserves_logical_case_field() -> Result<()> { + let metadata = HashMap::from([("source".to_string(), "logical".to_string())]); + let case = proto_case_fixture().with_logical_result_field(&Arc::new( + Field::new("case", DataType::Int32, false).with_metadata(metadata.clone()), + )); + let encoder = StubEncoder::ok(); + let encoded = case + .try_to_proto(&PhysicalExprEncodeCtx::new(&encoder))? + .expect("CaseExpr should encode"); + let decoder = StubDecoder::ok(); + let decoded = CaseExpr::try_from_proto( + &encoded, + &PhysicalExprDecodeCtx::new(&Schema::empty(), &decoder), + )?; + let decoded = decoded.downcast_ref::().expect("CaseExpr"); + assert_eq!(decoded.logical_result_field, case.logical_result_field); + assert_eq!( + decoded + .logical_result_field + .as_ref() + .unwrap() + .1 + .to_hashmap(), + metadata + ); + Ok(()) } #[test] diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index edf10ac476ac6..a0902fffe6dab 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -41,7 +41,7 @@ use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::var_provider::VarType; use datafusion_expr::var_provider::is_system_variables; use datafusion_expr::{ - Between, BinaryExpr, Expr, Like, Operator, TryCast, binary_expr, lit, + Between, BinaryExpr, Expr, ExprSchemable, Like, Operator, TryCast, binary_expr, lit, }; /// [PhysicalExpr] evaluate DataFusion expressions such as `A + 1`, or `CAST(c1 @@ -371,6 +371,25 @@ pub fn create_physical_expr( .zip(then_expr.iter()) .map(|(w, t)| (Arc::clone(w), Arc::clone(t))) .collect(); + // Match the logical CASE fast path without inferring its field twice. + let common_unmarked_branch = + match (case.when_then_expr.first(), when_then_expr.first()) { + (Some((_, logical_result)), Some((_, physical_result))) + if matches!( + logical_result.as_ref(), + Expr::Column(_) | Expr::Literal(_, _) + ) => + { + let field = physical_result.return_field(input_schema)?; + !field.data_type().is_null() + && !matches!( + logical_result.as_ref(), + Expr::Literal(value, _) if value.is_null() + ) + && field.metadata().is_empty() + } + _ => false, + }; let else_expr: Option> = if let Some(e) = &case.else_expr { Some(create_physical_expr( @@ -382,7 +401,16 @@ pub fn create_physical_expr( } else { None }; - Ok(expressions::case(expr, when_then_expr, else_expr)?) + let physical_case = + expressions::CaseExpr::try_new(expr, when_then_expr, else_expr)?; + if common_unmarked_branch { + Ok(Arc::new(physical_case)) + } else { + Ok(Arc::new( + physical_case + .with_logical_result_field(&e.to_field(input_dfschema)?.1), + )) + } } Expr::Cast(Cast { expr, field }) => expressions::cast_with_target_field( create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?, @@ -738,7 +766,7 @@ mod tests { use datafusion_expr::physical_planning_context::{ ScalarSubqueryResults, SubqueryIndex, }; - use datafusion_expr::{LogicalPlanBuilder, col, scalar_subquery}; + use datafusion_expr::{LogicalPlanBuilder, col, scalar_subquery, when}; use super::*; @@ -756,6 +784,206 @@ mod tests { ) } + #[test] + fn case_metadata_matches_logical_field_after_alias_lowering() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let df_schema = DFSchema::try_from(schema.clone())?; + let mut aliased = col("a"); + for _ in 0..10 { + aliased = aliased.alias("nested"); + } + let marked = |arg| { + datafusion_functions::core::expr_fn::with_metadata(vec![ + arg, + lit("source"), + lit("test"), + ]) + }; + let logical = when(lit(true), marked(aliased)).otherwise(marked(col("a")))?; + let logical_field = logical.to_field(&df_schema)?.1; + assert!(logical_field.metadata().is_empty()); + + let physical = create_physical_expr( + &logical, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; + assert_eq!( + physical.return_field(&schema)?.metadata(), + logical_field.metadata() + ); + Ok(()) + } + + #[test] + fn nested_case_metadata_matches_logical_field() -> Result<()> { + let metadata = std::collections::HashMap::from([( + "source".to_string(), + "nested".to_string(), + )]); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false).with_metadata(metadata), + ]); + let df_schema = DFSchema::try_from(schema.clone())?; + let inner = when(lit(true), col("a")).otherwise(col("a"))?; + let outer = when(lit(true), inner.clone()).otherwise(inner)?; + let logical_field = outer.to_field(&df_schema)?.1; + let physical = create_physical_expr( + &outer, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; + assert_eq!( + physical.return_field(&schema)?.metadata(), + logical_field.metadata() + ); + + let mut aliased = col("a"); + for _ in 0..10 { + aliased = aliased.alias("nested"); + } + let marked = |arg| { + datafusion_functions::core::expr_fn::with_metadata(vec![ + arg, + lit("source"), + lit("test"), + ]) + }; + let inner = when(lit(true), marked(aliased)).otherwise(marked(col("a")))?; + let outer = when(lit(true), inner.clone()).otherwise(inner)?; + let logical_field = outer.to_field(&df_schema)?.1; + let physical = create_physical_expr( + &outer, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; + assert_eq!( + physical.return_field(&schema)?.metadata(), + logical_field.metadata() + ); + Ok(()) + } + + #[test] + fn nested_case_marked_typed_null_matches_logical_field() -> Result<()> { + use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; + + let metadata = std::collections::HashMap::from([( + EXTENSION_TYPE_NAME_KEY.to_string(), + "geoarrow.wkb".to_string(), + )]); + let schema = Schema::new(vec![ + Field::new("binary", DataType::LargeBinary, false) + .with_metadata(metadata.clone()), + ]); + let df_schema = DFSchema::try_from(schema.clone())?; + let typed_null = Expr::Cast(Cast::new( + Box::new(lit(ScalarValue::Null)), + DataType::Binary, + )); + let marked_null = typed_null + .alias_with_metadata("marked_null", Some(FieldMetadata::from(metadata))); + let nested = when(lit(true), marked_null).otherwise(lit(ScalarValue::Null))?; + let cast_nested = Expr::Cast(Cast::new(Box::new(nested), DataType::LargeBinary)); + let outer = when(lit(true), cast_nested).otherwise(col("binary"))?; + let logical_field = outer.to_field(&df_schema)?.1; + assert!(logical_field.metadata().is_empty()); + let physical = create_physical_expr( + &outer, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; + assert_eq!( + physical.return_field(&schema)?.metadata(), + logical_field.metadata() + ); + Ok(()) + } + + #[test] + fn case_null_scalar_variable_matches_logical_field() -> Result<()> { + use datafusion_expr::var_provider::{VarProvider, VarType}; + + #[derive(Debug)] + struct NullBinaryProvider; + + impl VarProvider for NullBinaryProvider { + fn get_value(&self, _names: Vec) -> Result { + Ok(ScalarValue::Binary(None)) + } + + fn get_type(&self, _names: &[String]) -> Option { + Some(DataType::Binary) + } + } + + let metadata = std::collections::HashMap::from([( + "source".to_string(), + "marked".to_string(), + )]); + let schema = Schema::new(vec![ + Field::new("binary", DataType::Binary, false).with_metadata(metadata), + ]); + let df_schema = DFSchema::try_from(schema.clone())?; + let variable = Expr::ScalarVariable( + Arc::new(Field::new("variable", DataType::Binary, true)), + vec!["@variable".to_string()], + ); + let logical = when(lit(true), variable).otherwise(col("binary"))?; + let logical_field = logical.to_field(&df_schema)?.1; + assert!(logical_field.metadata().is_empty()); + let mut props = ExecutionProps::new(); + props.add_var_provider(VarType::UserDefined, Arc::new(NullBinaryProvider)); + let physical = create_physical_expr( + &logical, + &df_schema, + &props, + &PhysicalPlanningContext::default(), + )?; + assert_eq!( + physical.return_field(&schema)?.metadata(), + logical_field.metadata() + ); + Ok(()) + } + + #[test] + fn case_marked_not_matches_logical_field() -> Result<()> { + let existing = std::collections::HashMap::from([( + "origin".to_string(), + "input".to_string(), + )]); + let schema = Schema::new(vec![ + Field::new("a", DataType::Boolean, false).with_metadata(existing), + Field::new("b", DataType::Boolean, false), + ]); + let df_schema = DFSchema::try_from(schema.clone())?; + let marked = |arg| { + datafusion_functions::core::expr_fn::with_metadata(vec![ + arg, + lit("source"), + lit("test"), + ]) + }; + let logical = when(lit(true), marked(!col("a"))).otherwise(marked(col("b")))?; + let logical_field = logical.to_field(&df_schema)?.1; + let physical = create_physical_expr( + &logical, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; + assert_eq!( + physical.return_field(&schema)?.metadata(), + logical_field.metadata() + ); + Ok(()) + } + fn as_planner_cast(physical: &Arc) -> &expressions::CastExpr { physical .downcast_ref::() diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 43a90264c2b1f..f608ca9f1af1f 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1176,6 +1176,8 @@ message PhysicalCaseNode { PhysicalExprNode expr = 1; repeated PhysicalWhenThen when_then_expr = 2; PhysicalExprNode else_expr = 3; + datafusion_common.ArrowType logical_result_type = 4; + map logical_result_metadata = 5; } message PhysicalTryCastNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 908f9752b7f18..732474ff4243a 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -17775,6 +17775,12 @@ impl serde::Serialize for PhysicalCaseNode { if self.else_expr.is_some() { len += 1; } + if self.logical_result_type.is_some() { + len += 1; + } + if !self.logical_result_metadata.is_empty() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalCaseNode", len)?; if let Some(v) = self.expr.as_ref() { struct_ser.serialize_field("expr", v)?; @@ -17785,6 +17791,12 @@ impl serde::Serialize for PhysicalCaseNode { if let Some(v) = self.else_expr.as_ref() { struct_ser.serialize_field("elseExpr", v)?; } + if let Some(v) = self.logical_result_type.as_ref() { + struct_ser.serialize_field("logicalResultType", v)?; + } + if !self.logical_result_metadata.is_empty() { + struct_ser.serialize_field("logicalResultMetadata", &self.logical_result_metadata)?; + } struct_ser.end() } } @@ -17800,6 +17812,10 @@ impl<'de> serde::Deserialize<'de> for PhysicalCaseNode { "whenThenExpr", "else_expr", "elseExpr", + "logical_result_type", + "logicalResultType", + "logical_result_metadata", + "logicalResultMetadata", ]; #[allow(clippy::enum_variant_names)] @@ -17807,6 +17823,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalCaseNode { Expr, WhenThenExpr, ElseExpr, + LogicalResultType, + LogicalResultMetadata, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -17831,6 +17849,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalCaseNode { "expr" => Ok(GeneratedField::Expr), "whenThenExpr" | "when_then_expr" => Ok(GeneratedField::WhenThenExpr), "elseExpr" | "else_expr" => Ok(GeneratedField::ElseExpr), + "logicalResultType" | "logical_result_type" => Ok(GeneratedField::LogicalResultType), + "logicalResultMetadata" | "logical_result_metadata" => Ok(GeneratedField::LogicalResultMetadata), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -17853,6 +17873,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalCaseNode { let mut expr__ = None; let mut when_then_expr__ = None; let mut else_expr__ = None; + let mut logical_result_type__ = None; + let mut logical_result_metadata__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Expr => { @@ -17873,12 +17895,28 @@ impl<'de> serde::Deserialize<'de> for PhysicalCaseNode { } else_expr__ = map_.next_value()?; } + GeneratedField::LogicalResultType => { + if logical_result_type__.is_some() { + return Err(serde::de::Error::duplicate_field("logicalResultType")); + } + logical_result_type__ = map_.next_value()?; + } + GeneratedField::LogicalResultMetadata => { + if logical_result_metadata__.is_some() { + return Err(serde::de::Error::duplicate_field("logicalResultMetadata")); + } + logical_result_metadata__ = Some( + map_.next_value::>()? + ); + } } } Ok(PhysicalCaseNode { expr: expr__, when_then_expr: when_then_expr__.unwrap_or_default(), else_expr: else_expr__, + logical_result_type: logical_result_type__, + logical_result_metadata: logical_result_metadata__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index ba00577ab9a1b..51d6c88a5dc1b 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1826,6 +1826,13 @@ pub struct PhysicalCaseNode { pub when_then_expr: ::prost::alloc::vec::Vec, #[prost(message, optional, boxed, tag = "3")] pub else_expr: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, tag = "4")] + pub logical_result_type: ::core::option::Option, + #[prost(map = "string, string", tag = "5")] + pub logical_result_metadata: ::std::collections::HashMap< + ::prost::alloc::string::String, + ::prost::alloc::string::String, + >, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalTryCastNode { diff --git a/datafusion/sqllogictest/test_files/metadata.slt b/datafusion/sqllogictest/test_files/metadata.slt index 0fc74fa6cf602..0fceff32068c6 100644 --- a/datafusion/sqllogictest/test_files/metadata.slt +++ b/datafusion/sqllogictest/test_files/metadata.slt @@ -430,6 +430,26 @@ select arrow_metadata(with_metadata(id, 'unit', 'ms'), 'metadata_key') from tabl ---- the id field +# CASE must expose the same field metadata directly and through a projection. +query T +select arrow_metadata( + case when id = 1 then id + else with_metadata(id, 'metadata_key', 'the id field') end, + 'metadata_key') +from table_with_metadata where id = 1; +---- +the id field + +query T +select arrow_metadata(chosen, 'metadata_key') +from ( + select case when id = 1 then id + else with_metadata(id, 'metadata_key', 'the id field') end as chosen + from table_with_metadata where id = 1 +); +---- +the id field + # with_metadata: new keys overwrite existing on collision query T select arrow_metadata(with_metadata(id, 'metadata_key', 'overridden'), 'metadata_key') from table_with_metadata limit 1;