diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 4b5d279a87f43..17a9f8d2c3d8e 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -36,7 +36,6 @@ 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, Diagnostic, ExprSchema, Result, ScalarValue, Spans, TableReference, internal_err, not_impl_err, plan_datafusion_err, plan_err, @@ -121,32 +120,46 @@ fn cast_output_field( Arc::new(f) } -// 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, 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) - })?; +// Keep the first-branch result so strict CASE fast paths need no second type walk. +fn case_result_type(case: &Case, schema: &dyn ExprSchema) -> Result<(DataType, bool)> { + for (index, (_, then_expr)) in case.when_then_expr.iter().enumerate() { + let data_type = then_expr.get_type(schema)?; + if !data_type.is_null() { + return Ok((data_type, index == 0)); } - Ok(true) } + let data_type = case + .else_expr + .as_ref() + .map_or(Ok(DataType::Null), |expr| expr.get_type(schema))?; + Ok((data_type, false)) +} + +fn binary_output_field( + binary: &BinaryExpr, + left: &Field, + right: &Field, + name: &str, +) -> Result { + let mut coercer = + BinaryTypeCoercer::new(left.data_type(), &binary.op, right.data_type()); + coercer.set_lhs_spans(binary.left.spans().cloned().unwrap_or_default()); + coercer.set_rhs_spans(binary.right.spans().cloned().unwrap_or_default()); + Ok(Arc::new(Field::new( + name, + coercer.get_result_type()?, + left.is_nullable() || right.is_nullable(), + ))) +} +// 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), FinishCast(&'a FieldRef, bool), FinishAlias(Option<&'a FieldMetadata>), + FinishBinary(&'a BinaryExpr), } struct BranchField { @@ -172,67 +185,31 @@ 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, 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.is_some() - && 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::BinaryExpr(binary) => { + work.push(Work::FinishBinary(binary)); + work.push(Work::Visit(&binary.right)); + work.push(Work::Visit(&binary.left)); } - Expr::Column(_) - | Expr::Literal(_, _) - | Expr::OuterReferenceColumn(_, _) - | Expr::ScalarVariable(_, _) - | Expr::Placeholder(_) - | Expr::LambdaVariable(_) => fields.push(BranchField { + _ => fields.push(BranchField { field: expr.to_field(schema)?.1, certainly_null: matches!( - unwrap_certainly_null_expr(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 { @@ -240,7 +217,7 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result }; fields.push(BranchField { field: cast_output_field(&source.field, target, force_nullable), - certainly_null: source.certainly_null, + certainly_null: source.certainly_null && target.metadata().is_empty(), }); } Work::FinishAlias(metadata) => { @@ -255,7 +232,19 @@ fn case_field_metadata(case: &Case, schema: &dyn ExprSchema) -> Result field: Arc::new( source.field.as_ref().clone().with_metadata(combined), ), - certainly_null: source.certainly_null, + certainly_null: source.certainly_null && metadata.is_none(), + }); + } + 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"); + }; + fields.push(BranchField { + field: binary_output_field(binary, &left.field, &right.field, "")?, + certainly_null: false, }); } Work::FinishCase(case) => { @@ -378,15 +367,7 @@ impl ExprSchemable for Expr { Expr::ScalarVariable(field, _) => Ok(field.data_type().clone()), Expr::Literal(l, _) => Ok(l.data_type()), Expr::Case(case) => { - for (_, then_expr) in &case.when_then_expr { - let then_type = then_expr.get_type(schema)?; - if !then_type.is_null() { - return Ok(then_type); - } - } - case.else_expr - .as_ref() - .map_or(Ok(DataType::Null), |e| e.get_type(schema)) + case_result_type(case, schema).map(|(data_type, _)| data_type) } Expr::Cast(Cast { field, .. }) | Expr::TryCast(TryCast { field, .. }) => { Ok(field.data_type().clone()) @@ -760,22 +741,12 @@ impl ExprSchemable for Expr { .clone() .with_nullable(true), )), - Expr::BinaryExpr(BinaryExpr { left, right, op }) => { - let (left_field, right_field) = - (left.to_field(schema)?.1, right.to_field(schema)?.1); - - let (lhs_type, lhs_nullable) = - (left_field.data_type(), left_field.is_nullable()); - let (rhs_type, rhs_nullable) = - (right_field.data_type(), right_field.is_nullable()); - let mut coercer = BinaryTypeCoercer::new(lhs_type, op, rhs_type); - coercer.set_lhs_spans(left.spans().cloned().unwrap_or_default()); - coercer.set_rhs_spans(right.spans().cloned().unwrap_or_default()); - Ok(Arc::new(Field::new( - &schema_name, - coercer.get_result_type()?, - lhs_nullable || rhs_nullable, - ))) + Expr::BinaryExpr(binary) => { + let (left_field, right_field) = ( + binary.left.to_field(schema)?.1, + binary.right.to_field(schema)?.1, + ); + binary_output_field(binary, &left_field, &right_field, &schema_name) } Expr::WindowFunction(window_function) => { let WindowFunction { @@ -864,19 +835,22 @@ impl ExprSchemable for Expr { field: Some(field), .. }) => Ok(Arc::clone(field).renamed(&schema_name)), Expr::Case(case) => { - let data_type = self.get_type(schema)?; + let (data_type, first_result_is_typed) = case_result_type(case, 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() - { + if let Some((_, first_result)) = case.when_then_expr.first() { + let empty_metadata = match first_result.as_ref() { + // Strict merging discards all metadata for a typed, + // unmarked result; do not re-infer nested binary CASE fields. + Expr::BinaryExpr(_) | Expr::Not(_) => first_result_is_typed, + Expr::Column(_) | Expr::Literal(_, _) => { + let field = first_result.to_field(schema)?.1; + !field.data_type().is_null() + && !matches!(first_result.as_ref(), Expr::Literal(value, _) if value.is_null()) + && field.metadata().is_empty() + } + _ => false, + }; + if empty_metadata { return Ok(( relation, Arc::new(Field::new(&schema_name, data_type, nullable)), @@ -1645,6 +1619,20 @@ mod tests { .otherwise(Expr::Negative(Box::new(col("b"))))?; assert_eq!(negative_case.to_field(&schema)?.1.metadata(), &shared); + let null_binary = lit(ScalarValue::Null) + lit(ScalarValue::Null); + assert_eq!(null_binary.get_type(&schema)?, DataType::Int64); + let marked_int64 = Expr::Literal( + ScalarValue::Int64(Some(1)), + Some(FieldMetadata::from(shared.clone())), + ); + for (then_expr, else_expr) in [ + (null_binary.clone(), marked_int64.clone()), + (marked_int64, null_binary), + ] { + let case = when(lit(true), then_expr).otherwise(else_expr)?; + assert!(case.to_field(&schema)?.1.metadata().is_empty()); + } + 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()); @@ -1683,6 +1671,10 @@ mod tests { let target_field: FieldRef = Arc::new( Field::new("", DataType::Binary, true).with_metadata(binary_metadata), ); + let Expr::Alias(null_alias) = lit(ScalarValue::Binary(None)).alias("marked_null") + else { + unreachable!(); + }; for marked_null in [ Expr::Cast(Cast::new_from_field( Box::new(lit(ScalarValue::Null)), @@ -1692,6 +1684,10 @@ mod tests { Box::new(lit(ScalarValue::Null)), Arc::clone(&target_field), )), + Expr::Alias( + null_alias + .with_metadata(Some(FieldMetadata::from(target_field.metadata()))), + ), ] { let nested_all_null = when(lit(true), marked_null).otherwise(lit(ScalarValue::Null))?; @@ -1762,13 +1758,7 @@ mod tests { 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() - ); + assert_eq!(deep_function_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"))?; @@ -1780,7 +1770,7 @@ mod tests { } 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()); + assert_eq!(deep_nested_call.to_field(&schema)?.1.metadata(), &shared); let binary = col("d") + lit(1); let Expr::Alias(alias) = binary.clone().alias("marked") else { @@ -1811,7 +1801,7 @@ mod tests { } 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()); + assert_eq!(deep_udf_case.to_field(&schema)?.1.metadata(), &shared); let Expr::Alias(alias) = deep_alias.alias("marked") else { unreachable!(); @@ -1833,6 +1823,186 @@ mod tests { Ok(()) } + #[test] + fn test_case_field_metadata_not_null_is_a_contributor() -> Result<()> { + let metadata = HashMap::from([("logical".to_owned(), "boolean".to_owned())]); + let schema = DFSchema::from_unqualified_fields( + vec![ + Field::new("value", DataType::Boolean, false) + .with_metadata(metadata.clone()), + ] + .into(), + HashMap::new(), + )?; + for null in [ + lit(ScalarValue::Null), + lit(ScalarValue::Boolean(None)), + Expr::Literal( + ScalarValue::Boolean(None), + Some(FieldMetadata::from(metadata.clone())), + ), + ] { + let plain = when(lit(true), null.clone()).otherwise(col("value"))?; + assert_eq!(plain.to_field(&schema)?.1.metadata(), &metadata); + + let negated = Expr::Not(Box::new(null)); + for (then_expr, else_expr) in + [(negated.clone(), col("value")), (col("value"), negated)] + { + let expr = when(lit(true), then_expr).otherwise(else_expr)?; + let field = expr.to_field(&schema)?.1; + assert_eq!(field.data_type(), &DataType::Boolean); + assert!(field.is_nullable()); + assert!(field.metadata().is_empty()); + } + } + Ok(()) + } + + #[test] + fn test_case_field_metadata_preserves_aliased_null_cast_metadata() -> Result<()> { + let schema = DFSchema::empty(); + let inner = + FieldMetadata::from(HashMap::from([("inner".to_owned(), "x".to_owned())])); + let outer = + FieldMetadata::from(HashMap::from([("outer".to_owned(), "y".to_owned())])); + let expected = HashMap::from([ + ("inner".to_owned(), "x".to_owned()), + ("outer".to_owned(), "y".to_owned()), + ]); + let null = Expr::Literal(ScalarValue::Null, Some(inner)); + for cast in [ + Expr::Cast(Cast::new(Box::new(null.clone()), DataType::Int32)), + Expr::TryCast(TryCast::new(Box::new(null), DataType::Int32)), + ] { + let Expr::Alias(alias) = cast.alias("null_value") else { + unreachable!(); + }; + let aliased = Expr::Alias(alias.with_metadata(Some(outer.clone()))); + assert_eq!(aliased.to_field(&schema)?.1.metadata(), &expected); + let matching = Expr::Literal( + ScalarValue::Int32(Some(0)), + Some(FieldMetadata::from(expected.clone())), + ); + let case = when(lit(true), aliased).otherwise(matching)?; + assert_eq!(case.to_field(&schema)?.1.metadata(), &expected); + } + Ok(()) + } + + #[test] + fn test_case_unmarked_binary_skips_metadata_fields() -> Result<()> { + #[derive(Debug)] + struct MetadataLookupSchema { + field: FieldRef, + lookups: std::cell::Cell, + type_lookups: std::cell::Cell, + } + + impl ExprSchema for MetadataLookupSchema { + fn nullable(&self, _column: &Column) -> Result { + Ok(self.field.is_nullable()) + } + + fn data_type(&self, _column: &Column) -> Result<&DataType> { + self.type_lookups.set(self.type_lookups.get() + 1); + Ok(self.field.data_type()) + } + + fn field_from_column(&self, _column: &Column) -> Result<&FieldRef> { + self.lookups.set(self.lookups.get() + 1); + Ok(&self.field) + } + } + + let schema = MetadataLookupSchema { + field: Arc::new(Field::new("value", DataType::Int32, false).with_metadata( + HashMap::from([("source".to_owned(), "marked".to_owned())]), + )), + lookups: std::cell::Cell::new(0), + type_lookups: std::cell::Cell::new(0), + }; + let expression = + when(lit(true), col("value") + lit(1_i32)).otherwise(lit(0_i32))?; + let field = expression.to_field(&schema)?.1; + assert_eq!(field.data_type(), &DataType::Int32); + assert!(!field.is_nullable()); + assert!(field.metadata().is_empty()); + assert_eq!(schema.lookups.get(), 0); + assert_eq!(schema.type_lookups.get(), 1); + Ok(()) + } + + #[test] + fn test_nested_case_field_inference_is_linear() -> Result<()> { + #[derive(Debug)] + struct CountingSchema { + field: FieldRef, + calls: std::cell::Cell, + } + + impl CountingSchema { + fn visit(&self) -> Result<()> { + let calls = self.calls.get() + 1; + self.calls.set(calls); + assert_or_internal_err!(calls <= 4096, "Repeated CASE field inference"); + Ok(()) + } + } + + impl ExprSchema for CountingSchema { + fn nullable(&self, _column: &Column) -> Result { + self.visit()?; + Ok(false) + } + + fn field_from_column(&self, _column: &Column) -> Result<&FieldRef> { + self.visit()?; + Ok(&self.field) + } + } + + let metadata = HashMap::from([("source".to_owned(), "marked".to_owned())]); + let empty = HashMap::new(); + for binary_branch in [false, true] { + for depth in [8, 32, 128] { + let schema = CountingSchema { + field: Arc::new( + Field::new("value", DataType::Int32, false) + .with_metadata(metadata.clone()), + ), + calls: std::cell::Cell::new(0), + }; + let mut expression = col("value"); + for _ in 0..depth { + expression = if binary_branch { + when(lit(true), expression + col("value")) + .otherwise(lit(0_i32))? + } else { + let branch = Expr::Cast(Cast::new( + Box::new(expression.alias("nested")), + DataType::Int32, + )); + when(lit(true), col("value")).otherwise(branch)? + }; + } + let field = expression.to_field(&schema)?.1; + assert_eq!(field.data_type(), &DataType::Int32); + assert_eq!( + field.metadata(), + if binary_branch { &empty } else { &metadata } + ); + assert!(!field.is_nullable()); + assert!( + schema.calls.get() <= 6 * depth + 8, + "binary {binary_branch}, depth {depth}: {} schema lookups", + schema.calls.get() + ); + } + } + Ok(()) + } + #[test] fn test_alias_metadata_is_preserved_in_field_metadata() { let schema = MockExprSchema::new().with_data_type(DataType::Int32); diff --git a/datafusion/expr/src/udf.rs b/datafusion/expr/src/udf.rs index 75ec5020d77c4..ec9b0e6b52075 100644 --- a/datafusion/expr/src/udf.rs +++ b/datafusion/expr/src/udf.rs @@ -1147,6 +1147,13 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { self.inner.name() } + fn create_case_literal_matcher( + &self, + when_values: &[ScalarValue], + ) -> Result>> { + self.inner.create_case_literal_matcher(when_values) + } + fn display_name(&self, args: &[Expr]) -> Result { #[expect(deprecated)] self.inner.display_name(args) @@ -1273,6 +1280,9 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { #[cfg(test)] mod tests { use super::*; + use arrow::array::Int32Array; + use datafusion_common::cast::as_int32_array; + use datafusion_common::plan_err; use datafusion_expr_common::signature::Volatility; use std::hash::DefaultHasher; @@ -1300,6 +1310,113 @@ mod tests { } } + #[derive(Debug)] + struct TestCaseLiteralMatcher; + + impl CaseLiteralMatcher for TestCaseLiteralMatcher { + fn map_to_when_indices( + &self, + base: &ArrayRef, + else_index: u32, + ) -> Result> { + Ok(as_int32_array(base.as_ref())? + .iter() + .map(|value| if value == Some(7) { 0 } else { else_index }) + .collect()) + } + } + + #[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)] + enum MatcherMode { + Indexed, + Fallback, + Reject, + } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestCaseMatcherUdf { + signature: Signature, + mode: MatcherMode, + } + + impl ScalarUDFImpl for TestCaseMatcherUdf { + fn name(&self) -> &str { + "test_case_matcher" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + unreachable!("this test only exercises matcher construction") + } + + fn create_case_literal_matcher( + &self, + when_values: &[ScalarValue], + ) -> Result>> { + assert_eq!( + when_values, + &[ + ScalarValue::Int32(Some(7)), + ScalarValue::Int32(None), + ScalarValue::Int32(Some(7)) + ] + ); + match self.mode { + MatcherMode::Indexed => Ok(Some(Arc::new(TestCaseLiteralMatcher))), + MatcherMode::Fallback => Ok(None), + MatcherMode::Reject => plan_err!("test matcher rejected literals"), + } + } + } + + #[test] + fn test_case_literal_matcher_through_aliases() -> Result<()> { + let when_values = [ + ScalarValue::Int32(Some(7)), + ScalarValue::Int32(None), + ScalarValue::Int32(Some(7)), + ]; + let base: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), Some(8), None])); + for mode in [ + MatcherMode::Indexed, + MatcherMode::Fallback, + MatcherMode::Reject, + ] { + let udf = ScalarUDF::from(TestCaseMatcherUdf { + signature: Signature::any(2, Volatility::Immutable), + mode, + }); + let alias = udf.clone().with_aliases(["alternate"]); + let nested_alias = alias.clone().with_aliases(["another"]); + for udf in [udf, alias, nested_alias] { + let matcher = udf.inner().create_case_literal_matcher(&when_values); + match mode { + MatcherMode::Indexed => assert_eq!( + matcher? + .expect("aliases must preserve the matcher") + .map_to_when_indices(&base, 3)?, + vec![0, 3, 3] + ), + MatcherMode::Fallback => assert!(matcher?.is_none()), + MatcherMode::Reject => assert!( + matcher + .expect_err("aliases must preserve errors") + .to_string() + .contains("test matcher rejected literals") + ), + } + } + } + Ok(()) + } + // PartialEq and Hash must be consistent, and also PartialEq and PartialOrd // must be consistent, so they are tested together. #[test] diff --git a/datafusion/expr/tests/case_field_metadata_perf.rs b/datafusion/expr/tests/case_field_metadata_perf.rs new file mode 100644 index 0000000000000..19ae43704cdd2 --- /dev/null +++ b/datafusion/expr/tests/case_field_metadata_perf.rs @@ -0,0 +1,212 @@ +// 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. + +//! Planning-only inference on prebuilt expressions, not SQL execution or an +//! Arrow 59 release query-performance benchmark. Copy unchanged to 78e3/current +//! for paired fresh-process runs with identical compiler/profile/features. +//! Only equivalent, correct outputs are compared; depth-9 UDF changes are excluded. + +use std::collections::HashMap; +use std::hint::black_box; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::instant::Instant; +use datafusion_common::{DFSchema, Result}; +use datafusion_expr::{ + Case, Cast, ColumnarValue, Expr, ExprSchemable, ReturnFieldArgs, ScalarFunctionArgs, + ScalarUDF, ScalarUDFImpl, Signature, Volatility, col, lit, when, +}; + +#[derive(Debug, PartialEq, Eq, Hash)] +struct MetadataIdentity { + signature: Signature, +} + +impl ScalarUDFImpl for MetadataIdentity { + fn name(&self) -> &str { + "metadata_identity" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int32) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + Ok(Arc::clone(&args.arg_fields[0])) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + Ok(args.args[0].clone()) + } +} + +fn flat_case( + width: usize, + result: impl Fn(usize) -> Expr, + otherwise: Option, +) -> Expr { + let branches = (0..width) + .map(|i| (Box::new(lit(i == 0)), Box::new(result(i)))) + .collect(); + Expr::Case(Case::new(None, branches, otherwise.map(Box::new))) +} + +fn nested_case(depth: usize, wrap: impl Fn(Expr) -> Expr) -> Expr { + let mut expr = col("c0"); + for _ in 0..depth { + expr = when(lit(true), col("c1")).otherwise(wrap(expr)).unwrap(); + } + expr +} + +fn measure( + name: &str, + expr: &Expr, + schema: &DFSchema, + nullable: bool, + metadata: &HashMap, +) { + let (_, field) = expr.to_field(schema).unwrap(); + assert_eq!(field.data_type(), &DataType::Int32, "{name}: type"); + assert_eq!(field.is_nullable(), nullable, "{name}: nullable"); + assert_eq!(field.metadata(), metadata, "{name}: metadata"); + drop(field); + for _ in 0..20 { + black_box(black_box(expr).to_field(black_box(schema)).unwrap()); + } + let mut samples = Vec::with_capacity(101); + for _ in 0..101 { + let start = Instant::now(); + let result = black_box(black_box(expr).to_field(black_box(schema))); + let elapsed = start.elapsed(); + // Result checking, destruction, sample storage and JSON are not timed. + black_box(result.unwrap()); + samples.push(u64::try_from(elapsed.as_nanos()).expect("duration exceeds u64 ns")); + } + samples.sort_unstable(); + println!( + "CASE_FIELD_PERF {}", + serde_json::json!({ + "name": name, + "median_ns": samples[samples.len() / 2], + "min_ns": samples[0], + "max_ns": samples[samples.len() - 1], + "samples": samples.len(), + }) + ); +} + +#[test] +#[ignore = "planning microbenchmark; paired fresh processes with --ignored --nocapture"] +fn case_field_metadata_perf() { + eprintln!( + "CASE_FIELD_PERF_CONTEXT manifest_dir={} source={}", + env!("CARGO_MANIFEST_DIR"), + file!() + ); + let empty = HashMap::new(); + let shared = HashMap::from([("source".to_owned(), "marked".to_owned())]); + let make_schema = |metadata: &HashMap| { + let fields = (0..64) + .map(|i| { + Field::new(format!("c{i}"), DataType::Int32, false) + .with_metadata(metadata.clone()) + }) + .collect::>(); + DFSchema::from_unqualified_fields(fields.into(), HashMap::new()).unwrap() + }; + let plain = make_schema(&empty); + let marked = make_schema(&shared); + // Sizes are fixed: at most 1000 flat branches or 128 linear nesting levels. + for width in [64, 1000] { + let expr = flat_case(width, |i| lit(i32::try_from(i).unwrap()), None); + measure( + &format!("scalar_literals_{width}"), + &expr, + &plain, + true, + &empty, + ); + } + // Native strict merge: identical branch maps survive; otherwise output is empty. + for (name, schema, metadata) in [ + ("plain_columns_64", &plain, &empty), + ("all_marked_columns_64", &marked, &shared), + ] { + let expr = flat_case(64, |i| col(format!("c{i}")), Some(col("c0"))); + measure(name, &expr, schema, false, metadata); + } + let identity = ScalarUDF::from(MetadataIdentity { + signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable), + }); + // Shallow controls stay below 78e3's depth-8 guard; the UDF is never evaluated. + for depth in [1, 4] { + let wrap = |mut expr| { + for _ in 0..depth { + expr = identity.call(vec![expr]); + } + expr + }; + let expr = when(lit(true), wrap(col("c0"))) + .otherwise(wrap(col("c1"))) + .unwrap(); + measure( + &format!("shallow_case_udf_{depth}"), + &expr, + &marked, + false, + &shared, + ); + } + for depth in [32, 128] { + for (name, expr) in [ + ( + "nested_case_aliases", + nested_case(depth, |e| e.alias("nested")), + ), + ( + "nested_case_casts", + nested_case(depth, |e| { + Expr::Cast(Cast::new(Box::new(e), DataType::Int32)) + }), + ), + ] { + measure(&format!("{name}_{depth}"), &expr, &marked, false, &shared); + } + } + for depth in [8, 32, 128] { + let mut expr = col("c0"); + for _ in 0..depth { + // Binary arithmetic clears source metadata; ELSE is also unmarked. + expr = when(lit(true), expr + col("c1")) + .otherwise(lit(0_i32)) + .unwrap(); + } + measure( + &format!("nested_case_binary_{depth}"), + &expr, + &marked, + false, + &empty, + ); + } +} diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index c0e44712a113c..4c788f086d7e9 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -18,11 +18,10 @@ mod literal_lookup_table; use super::{Column, Literal}; -use crate::PhysicalExpr; use crate::expressions::{ CastExpr, LambdaVariable, NegativeExpr, NotExpr, TryCastExpr, lit, try_cast, }; -use crate::{HigherOrderFunctionExpr, ScalarFunctionExpr}; +use crate::{PhysicalExpr, ScalarFunctionExpr}; use arrow::array::*; use arrow::compute::kernels::zip::zip; use arrow::compute::{ @@ -1381,21 +1380,6 @@ 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, @@ -1463,22 +1447,6 @@ fn case_result_field_metadata( 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 { @@ -2272,6 +2240,43 @@ mod tests { Ok(()) } + #[test] + fn case_not_null_is_a_contributor() -> Result<()> { + let metadata = HashMap::from([("logical".to_owned(), "boolean".to_owned())]); + let schema = Schema::new(vec![ + Field::new("value", DataType::Boolean, false).with_metadata(metadata.clone()), + ]); + let nulls: [Arc; 3] = [ + lit(ScalarValue::Null), + lit(ScalarValue::Boolean(None)), + Arc::new(Literal::new_with_metadata( + ScalarValue::Boolean(None), + Some(FieldMetadata::from(metadata.clone())), + )), + ]; + for null in nulls { + let plain = case( + None, + vec![(lit(true), Arc::clone(&null))], + Some(col("value", &schema)?), + )?; + assert_eq!(plain.return_field(&schema)?.metadata(), &metadata); + + let negated: Arc = Arc::new(NotExpr::new(null)); + for (then_expr, else_expr) in [ + (Arc::clone(&negated), col("value", &schema)?), + (col("value", &schema)?, negated), + ] { + let expr = case(None, vec![(lit(true), then_expr)], Some(else_expr))?; + let field = expr.return_field(&schema)?; + assert_eq!(field.data_type(), &DataType::Boolean); + assert!(field.is_nullable()); + assert!(field.metadata().is_empty()); + } + } + Ok(()) + } + #[test] fn case_explicitly_marked_null_cast_matches_logical_field() -> Result<()> { let metadata = HashMap::from([( @@ -2314,7 +2319,7 @@ mod tests { } #[test] - fn case_function_metadata_respects_logical_depth_limit() -> Result<()> { + fn case_function_metadata_preserves_cached_fields_at_any_depth() -> 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()), @@ -2332,7 +2337,7 @@ mod tests { .with_metadata(metadata.clone()), ); let options = Arc::new(datafusion_common::config::ConfigOptions::default()); - for (depth, expected_marked) in [(2, true), (9, false)] { + for depth in [2, 9, 32, 128] { let mut branch = col("a", &schema)?; for _ in 0..depth { branch = Arc::new(ScalarFunctionExpr::new( @@ -2345,11 +2350,7 @@ mod tests { } 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()); - } + assert_eq!(field.metadata(), &metadata, "depth {depth}"); } Ok(()) } diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index 3270e9ce99de0..281742523c0c3 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -903,7 +903,13 @@ mod tests { }; 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()); + assert_eq!( + logical_field.metadata(), + &std::collections::HashMap::from([( + "source".to_string(), + "test".to_string() + )]) + ); let physical = create_physical_expr( &logical,