diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index 6a3a00a42d2c0..bf79d12c4df3c 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -1577,6 +1577,9 @@ impl SessionContext { let context = SimplifyContext::builder() .with_schema(Arc::clone(prepared.plan.schema())) .with_config_options(Arc::clone(state.config_options())) + .with_case_equality_planner( + state.execution_props().case_equality_planner.clone(), + ) .with_query_execution_start_time( state.execution_props().query_execution_start_time, ) diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index eaf61c64bf9fa..2daebd5022879 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -822,6 +822,9 @@ impl SessionState { let simplify_context = SimplifyContext::builder() .with_schema(Arc::new(df_schema.clone())) .with_config_options(Arc::clone(config_options)) + .with_case_equality_planner( + self.execution_props().case_equality_planner.clone(), + ) .with_query_execution_start_time( self.execution_props().query_execution_start_time, ) @@ -2023,6 +2026,9 @@ impl ContextProvider for SessionContextProvider<'_> { .ok_or_else(|| plan_datafusion_err!("table function '{name}' not found"))?; let simplify_context = SimplifyContext::builder() .with_config_options(Arc::clone(self.state.config_options())) + .with_case_equality_planner( + self.state.execution_props().case_equality_planner.clone(), + ) .with_query_execution_start_time( self.state.execution_props().query_execution_start_time, ) @@ -2327,6 +2333,12 @@ impl OptimizerConfig for SessionState { Arc::clone(self.config.options()) } + fn case_equality_planner( + &self, + ) -> Option<&Arc> { + self.execution_props.case_equality_planner.as_ref() + } + fn function_registry(&self) -> Option<&dyn FunctionRegistry> { Some(self) } @@ -2408,8 +2420,9 @@ mod tests { use datafusion_common::config::Dialect; use datafusion_common::{DFSchema, ResolvedTableReference, Result, TableReference}; use datafusion_execution::config::SessionConfig; - use datafusion_expr::Expr; use datafusion_expr::HigherOrderUDF; + use datafusion_expr::execution_props::CaseEqualityPlanner; + use datafusion_expr::{Case, Expr, col, lit}; use datafusion_optimizer::Optimizer; use datafusion_optimizer::optimizer::OptimizerRule; use datafusion_physical_plan::display::DisplayableExecutionPlan; @@ -2418,6 +2431,43 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + #[derive(Debug)] + struct NotEqualCasePlanner; + + impl CaseEqualityPlanner for NotEqualCasePlanner { + fn plan_equality( + &self, + _case: &Case, + _schema: &DFSchema, + ) -> Result> { + Ok(Some(col("case_base").not_eq(col("case_when")))) + } + } + + #[test] + fn session_state_simplification_uses_custom_case_equality() -> Result<()> { + let mut state = SessionStateBuilder::new().with_default_features().build(); + state.execution_props_mut().case_equality_planner = + Some(Arc::new(NotEqualCasePlanner)); + let expr = Expr::Case(Case::new( + Some(Box::new(lit(1_i32))), + vec![(Box::new(lit(2_i32)), Box::new(lit("match")))], + Some(Box::new(lit("miss"))), + )); + let schema = Arc::new(Schema::new(vec![Field::new( + "dummy", + DataType::Int32, + false, + )])); + let df_schema = DFSchema::try_from(Arc::clone(&schema))?; + let physical = state.create_physical_expr(expr, &df_schema)?; + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![0]))])?; + let value = physical.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!(value.as_ref(), &StringArray::from(vec!["match"])); + Ok(()) + } + #[test] #[cfg(feature = "sql")] fn test_configured_dialect_names_are_accepted_by_sqlparser() { diff --git a/datafusion/expr/src/execution_props.rs b/datafusion/expr/src/execution_props.rs index 7c5369d1144dd..b396494c2d553 100644 --- a/datafusion/expr/src/execution_props.rs +++ b/datafusion/expr/src/execution_props.rs @@ -17,11 +17,23 @@ use crate::var_provider::{VarProvider, VarType}; use chrono::{DateTime, Utc}; -use datafusion_common::HashMap; use datafusion_common::alias::AliasGenerator; use datafusion_common::config::ConfigOptions; +use datafusion_common::{DFSchema, HashMap, Result}; use std::sync::Arc; +/// Selects an optional equality expression for a simple CASE. The returned +/// expression is planned against synthetic `case_base` and `case_when` fields +/// and must return a boolean result. Like SQL equality, it must never match a +/// null base or null WHEN value. Returning `None` retains Arrow equality. +pub trait CaseEqualityPlanner: std::fmt::Debug + Send + Sync { + fn plan_equality( + &self, + case: &crate::expr::Case, + input_schema: &DFSchema, + ) -> Result>; +} + /// Holds properties and scratch state used while optimizing a [`LogicalPlan`] /// and translating it into an executable physical plan, such as the statement /// start time used during simplification. @@ -59,6 +71,8 @@ pub struct ExecutionProps { pub config_options: Option>, /// Providers for scalar variables pub var_providers: Option>>, + /// Optional simple CASE comparator, selected only during physical planning. + pub case_equality_planner: Option>, } impl Default for ExecutionProps { @@ -75,6 +89,7 @@ impl ExecutionProps { alias_generator: Arc::new(AliasGenerator::new()), config_options: None, var_providers: None, + case_equality_planner: None, } } @@ -143,7 +158,7 @@ mod test { fn debug() { let props = ExecutionProps::new(); assert_eq!( - "ExecutionProps { query_execution_start_time: None, alias_generator: AliasGenerator { next_id: 1 }, config_options: None, var_providers: None }", + "ExecutionProps { query_execution_start_time: None, alias_generator: AliasGenerator { next_id: 1 }, config_options: None, var_providers: None, case_equality_planner: None }", format!("{props:?}") ); } diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index 75041c701454a..ae197feb31689 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -141,7 +141,8 @@ pub use udaf::{ udaf_default_window_function_schema_name, }; pub use udf::{ - ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, StructFieldMapping, + CaseLiteralMatcher, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, + StructFieldMapping, }; pub use udwf::{LimitEffect, ReversedUDWF, WindowUDF, WindowUDFImpl}; pub use window_frame::{WindowFrame, WindowFrameBound, WindowFrameUnits}; diff --git a/datafusion/expr/src/simplify.rs b/datafusion/expr/src/simplify.rs index 522cf122a273c..304090510994b 100644 --- a/datafusion/expr/src/simplify.rs +++ b/datafusion/expr/src/simplify.rs @@ -24,6 +24,7 @@ use chrono::{DateTime, Utc}; use datafusion_common::config::ConfigOptions; use datafusion_common::{DFSchema, DFSchemaRef, Result}; +use crate::execution_props::CaseEqualityPlanner; use crate::{Expr, ExprSchemable}; /// Provides simplification information based on schema, query execution time, @@ -38,6 +39,7 @@ pub struct SimplifyContext { schema: DFSchemaRef, query_execution_start_time: Option>, config_options: Arc, + case_equality_planner: Option>, } /// Builder for [`SimplifyContext`]. @@ -46,6 +48,7 @@ pub struct SimplifyContextBuilder { schema: Option, query_execution_start_time: Option>, config_options: Option>, + case_equality_planner: Option>, } impl Default for SimplifyContext { @@ -54,6 +57,7 @@ impl Default for SimplifyContext { schema: Arc::new(DFSchema::empty()), query_execution_start_time: None, config_options: Arc::new(ConfigOptions::default()), + case_equality_planner: None, } } } @@ -137,6 +141,10 @@ impl SimplifyContext { pub fn config_options(&self) -> &Arc { &self.config_options } + + pub fn case_equality_planner(&self) -> Option<&Arc> { + self.case_equality_planner.as_ref() + } } impl SimplifyContextBuilder { @@ -146,6 +154,15 @@ impl SimplifyContextBuilder { self } + /// Set the simple CASE comparator used during constant folding. + pub fn with_case_equality_planner( + mut self, + planner: Option>, + ) -> Self { + self.case_equality_planner = planner; + self + } + /// Set the schema for this context. pub fn with_schema(mut self, schema: DFSchemaRef) -> Self { self.schema = Some(schema); @@ -175,6 +192,7 @@ impl SimplifyContextBuilder { config_options: self .config_options .unwrap_or_else(|| Arc::new(ConfigOptions::default())), + case_equality_planner: self.case_equality_planner, } } } diff --git a/datafusion/expr/src/udf.rs b/datafusion/expr/src/udf.rs index a56c401fa3b48..75ec5020d77c4 100644 --- a/datafusion/expr/src/udf.rs +++ b/datafusion/expr/src/udf.rs @@ -24,6 +24,7 @@ use crate::simplify::{ExprSimplifyResult, SimplifyContext}; use crate::sort_properties::{ExprProperties, SortProperties}; use crate::udf_eq::UdfEq; use crate::{ColumnarValue, Documentation, Expr, Signature}; +use arrow::array::ArrayRef; use arrow::datatypes::{DataType, Field, FieldRef}; #[cfg(debug_assertions)] use datafusion_common::assert_or_internal_err; @@ -57,6 +58,14 @@ pub struct StructFieldMapping { pub fields: Vec<(Vec, usize)>, } +/// Index for a simple CASE with literal WHEN and result expressions. +/// Implementations must return the first matching WHEN ordinal for each base +/// value, or `else_index` when no WHEN matches, using the UDF's equality +/// semantics. A SQL NULL base must never match. +pub trait CaseLiteralMatcher: Debug + Send + Sync { + fn map_to_when_indices(&self, base: &ArrayRef, else_index: u32) -> Result>; +} + /// Logical representation of a Scalar User Defined Function. /// /// A scalar function produces a single row output for each row of input. This @@ -538,6 +547,17 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { /// Returns this function's name fn name(&self) -> &str; + /// Optionally index literal right-hand values for a simple CASE whose + /// equality comparison is implemented by this UDF. The returned matcher + /// is built once per physical expression; returning `None` keeps the + /// ordinary lazy branch-by-branch evaluation. + fn create_case_literal_matcher( + &self, + _when_values: &[ScalarValue], + ) -> Result>> { + Ok(None) + } + /// Returns any aliases (alternate names) for this function. /// /// Aliases can be used to invoke the same function using different names. diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs index 0abb09bea8768..bde5eb52e1a11 100644 --- a/datafusion/optimizer/src/optimizer.rs +++ b/datafusion/optimizer/src/optimizer.rs @@ -21,6 +21,7 @@ use std::fmt::Debug; use std::sync::Arc; use chrono::{DateTime, Utc}; +use datafusion_expr::execution_props::CaseEqualityPlanner; use datafusion_expr::registry::FunctionRegistry; use datafusion_expr::{InvariantLevel, assert_expected_schema}; use log::{debug, warn}; @@ -153,6 +154,12 @@ pub trait OptimizerConfig { fn options(&self) -> Arc; + /// Optional comparator used to fold simple CASE expressions consistently + /// with physical planning. + fn case_equality_planner(&self) -> Option<&Arc> { + None + } + fn function_registry(&self) -> Option<&dyn FunctionRegistry> { None } diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index 4cb9830aa3185..2f6f0fe302289 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -37,7 +37,9 @@ use datafusion_common::{ HashMap, cast::{as_large_list_array, as_list_array}, metadata::FieldMetadata, - tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRewriter}, + tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter, + }, }; use datafusion_expr::expr::HigherOrderFunction; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; @@ -213,6 +215,8 @@ impl ExprSimplifier { let mut simplifier = Simplifier::new(&self.info); let config_options = Some(Arc::clone(self.info.config_options())); let mut const_evaluator = ConstEvaluator::try_new(config_options)?; + const_evaluator.execution_props.case_equality_planner = + self.info.case_equality_planner().cloned(); let mut shorten_in_list_simplifier = ShortenInListSimplifier::new(); let guarantees_map: HashMap<&Expr, &NullableInterval> = self.guarantees.iter().map(|(k, v)| (k, v)).collect(); @@ -717,6 +721,23 @@ impl ConstEvaluator { Ok(e) => e, Err(err) => return ConstSimplifyResult::SimplifyRuntimeError(err, expr), }; + let mut volatile = false; + if let Err(err) = phys_expr.apply(|node| { + volatile |= node.is_volatile_node(); + Ok(if volatile { + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }) + }) { + return ConstSimplifyResult::SimplifyRuntimeError(err, expr); + } + if volatile { + return ConstSimplifyResult::SimplifyRuntimeError( + exec_datafusion_err!("Cannot fold a volatile expression"), + expr, + ); + } let metadata = phys_expr .return_field(DUMMY_BATCH.schema_ref()) .ok() @@ -2379,6 +2400,7 @@ mod tests { datatypes::{FieldRef, Fields}, }; use datafusion_common::{DFSchemaRef, ToDFSchema, assert_contains}; + use datafusion_expr::execution_props::CaseEqualityPlanner; use datafusion_expr::{ expr::WindowFunction, function::{ @@ -2415,6 +2437,71 @@ mod tests { assert_eq!(expected, simplifier.simplify(expr).unwrap()); } + #[derive(Debug)] + struct NotEqualCasePlanner; + + impl CaseEqualityPlanner for NotEqualCasePlanner { + fn plan_equality( + &self, + _case: &Case, + _schema: &DFSchema, + ) -> Result> { + Ok(Some(col("case_base").not_eq(col("case_when")))) + } + } + + #[test] + fn constant_case_uses_custom_equality_when_configured() -> Result<()> { + let expression = Expr::Case(Case::new( + Some(Box::new(lit(1_i32))), + vec![(Box::new(lit(2_i32)), Box::new(lit("match")))], + Some(Box::new(lit("miss"))), + )); + let ordinary = ExprSimplifier::new(SimplifyContext::default()); + assert_eq!(ordinary.simplify(expression.clone())?, lit("miss")); + + let custom = ExprSimplifier::new( + SimplifyContext::builder() + .with_case_equality_planner(Some(Arc::new(NotEqualCasePlanner))) + .build(), + ); + assert_eq!(custom.simplify(expression)?, lit("match")); + Ok(()) + } + + #[derive(Debug)] + struct VolatileCasePlanner; + + impl CaseEqualityPlanner for VolatileCasePlanner { + fn plan_equality( + &self, + _case: &Case, + _schema: &DFSchema, + ) -> Result> { + let udf = Arc::new(ScalarUDF::new_from_impl(VolatileUdf::new())); + let random = Expr::ScalarFunction(ScalarFunction::new_udf(udf, vec![])); + Ok(Some( + col("case_base").eq(col("case_when")).and(random.eq(lit(0))), + )) + } + } + + #[test] + fn constant_case_with_volatile_comparator_is_not_folded() -> Result<()> { + let expression = Expr::Case(Case::new( + Some(Box::new(lit(1_i32))), + vec![(Box::new(lit(1_i32)), Box::new(lit("match")))], + Some(Box::new(lit("miss"))), + )); + let simplifier = ExprSimplifier::new( + SimplifyContext::builder() + .with_case_equality_planner(Some(Arc::new(VolatileCasePlanner))) + .build(), + ); + assert_eq!(simplifier.simplify(expression.clone())?, expression); + Ok(()) + } + #[test] fn basic_coercion() { let schema = test_schema(); diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs index 0e72a17abc9f7..78d2e7a849729 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs @@ -114,6 +114,7 @@ impl SimplifyExpressions { let info = SimplifyContext::builder() .with_schema(schema) .with_config_options(config.options()) + .with_case_equality_planner(config.case_equality_planner().cloned()) .with_query_execution_start_time(config.query_execution_start_time()) .build(); diff --git a/datafusion/physical-expr-common/src/physical_expr.rs b/datafusion/physical-expr-common/src/physical_expr.rs index 59393e75786bd..d3eb5993aca9c 100644 --- a/datafusion/physical-expr-common/src/physical_expr.rs +++ b/datafusion/physical-expr-common/src/physical_expr.rs @@ -697,6 +697,16 @@ pub mod proto_decode { self.decoder.decode(node, self.schema) } + /// Decode an expression whose columns belong to a nested synthetic + /// scope rather than the enclosing input schema. + pub fn decode_with_schema( + &self, + node: &PhysicalExprNode, + schema: &Schema, + ) -> Result> { + self.decoder.decode(node, schema) + } + /// Decode a required child node, erroring if it is absent. /// /// Proto child expressions are encoded as `Option>`; diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index 29f1151a12cb8..c0e44712a113c 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -48,12 +48,27 @@ use crate::expressions::case::literal_lookup_table::LiteralLookupTable; use arrow::compute::kernels::merge::{MergeIndex, merge, merge_n}; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_physical_expr_common::datum::compare_with_eq; +use datafusion_physical_expr_common::physical_expr::is_volatile; use datafusion_physical_expr_common::utils::scatter; use itertools::Itertools; use std::fmt::{Debug, Formatter}; pub(super) type WhenThen = (Arc, Arc); +fn direct_case_equality_udf( + equality: &Arc, +) -> Option<&ScalarFunctionExpr> { + equality.downcast_ref::().filter(|udf| { + udf.args().len() == 2 + && udf.args()[0] + .downcast_ref::() + .is_some_and(|column| column.index() == 0 && column.name() == "case_base") + && udf.args()[1] + .downcast_ref::() + .is_some_and(|column| column.index() == 1 && column.name() == "case_when") + }) +} + #[derive(Debug, Hash, PartialEq, Eq)] enum EvalMethod { /// CASE WHEN condition THEN result @@ -279,6 +294,8 @@ pub struct CaseExpr { body: CaseBody, /// Evaluation method to use eval_method: EvalMethod, + /// Optional equality expression over two synthetic columns: base and WHEN. + equality: Option>, /// Logical result type and metadata inferred before physical lowering. logical_result_field: Option<(DataType, FieldMetadata)>, } @@ -289,13 +306,16 @@ pub struct CaseExpr { impl Hash for CaseExpr { fn hash(&self, state: &mut H) { self.body.hash(state); + self.equality.hash(state); self.logical_result_field.hash(state); } } impl PartialEq for CaseExpr { fn eq(&self, other: &Self) -> bool { - self.body == other.body && self.logical_result_field == other.logical_result_field + self.body == other.body + && self.equality == other.equality + && self.logical_result_field == other.logical_result_field } } @@ -660,6 +680,30 @@ impl CaseExpr { when_then_expr: Vec, else_expr: Option>, ) -> Result { + Self::try_new_internal(expr, when_then_expr, else_expr, None) + } + + /// Create a simple CASE using a physical boolean comparator over synthetic + /// `case_base` and `case_when` columns. Null operands never match. Ordinary + /// CASEs retain native equality. + pub fn try_new_with_equality( + expr: Arc, + when_then_expr: Vec, + else_expr: Option>, + equality: Arc, + ) -> Result { + Self::try_new_internal(Some(expr), when_then_expr, else_expr, Some(equality)) + } + + fn try_new_internal( + expr: Option>, + when_then_expr: Vec, + else_expr: Option>, + equality: Option>, + ) -> Result { + if equality.is_some() && expr.is_none() { + return internal_err!("Custom CASE equality requires a base expression"); + } // normalize null literals to None in the else_expr (this already happens // during SQL planning, but not necessarily for other use cases) let else_expr = match &else_expr { @@ -680,11 +724,12 @@ impl CaseExpr { else_expr, }; - let eval_method = Self::find_best_eval_method(&body)?; + let eval_method = Self::find_best_eval_method(&body, equality.as_ref())?; Ok(Self { body, eval_method, + equality, logical_result_field: None, }) } @@ -698,9 +743,12 @@ impl CaseExpr { self } - fn find_best_eval_method(body: &CaseBody) -> Result { + fn find_best_eval_method( + body: &CaseBody, + equality: Option<&Arc>, + ) -> Result { if body.expr.is_some() { - if let Some(mapping) = LiteralLookupTable::maybe_new(body) { + if let Some(mapping) = LiteralLookupTable::maybe_new(body, equality) { return Ok(EvalMethod::WithExprScalarLookupTable(mapping)); } @@ -769,6 +817,7 @@ impl CaseBody { &self, batch: &RecordBatch, return_type: &DataType, + equality: Option<&Arc>, ) -> Result { let mut result_builder = ResultBuilder::new(return_type, batch.num_rows()); @@ -834,19 +883,74 @@ impl CaseBody { // The types of case and when expressions will be coerced to match. // We only need to check if the base_value is nested. let base_value_is_nested = base_values.data_type().is_nested(); + let base_field = if equality.is_some() { + let Some(base_expr) = self.expr.as_ref() else { + return internal_err!("Custom CASE equality requires a base"); + }; + Some(base_expr.return_field(batch.schema_ref())?) + } else { + None + }; + let direct_udf = equality.and_then(direct_case_equality_udf); for i in 0..self.when_then_expr.len() { // Evaluate the 'when' predicate for the remainder batch // This results in a boolean array with the same length as the remaining number of rows let when_expr = &self.when_then_expr[i].0; - let when_value = match when_expr.evaluate(&remainder_batch)? { - ColumnarValue::Array(a) => { - compare_with_eq(&a, &base_values, base_value_is_nested) + let evaluated_when = when_expr.evaluate(&remainder_batch)?; + let when_value = if let Some(equality) = equality { + let all_null = match &evaluated_when { + ColumnarValue::Array(array) => { + array.logical_null_count() == array.len() + } + ColumnarValue::Scalar(value) => value.is_null(), + }; + if all_null { + continue; } - ColumnarValue::Scalar(s) => { - compare_with_eq(&s.to_scalar()?, &base_values, base_value_is_nested) + let Some(base_field) = base_field.as_ref() else { + return internal_err!("Custom CASE equality requires a base field"); + }; + let when_field = when_expr.return_field(remainder_batch.schema_ref())?; + let when_array = evaluated_when.into_array(remainder_batch.num_rows())?; + let base_field = Arc::new( + Field::new("case_base", base_values.data_type().clone(), true) + .with_metadata(base_field.metadata().clone()), + ); + let when_field = Arc::new( + Field::new("case_when", when_array.data_type().clone(), true) + .with_metadata(when_field.metadata().clone()), + ); + let compared = if let Some(udf) = direct_udf { + udf.evaluate_with_args( + vec![ + ColumnarValue::Array(Arc::clone(&base_values)), + ColumnarValue::Array(when_array), + ], + vec![base_field, when_field], + remainder_batch.num_rows(), + )? + } else { + let comparison_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![base_field, when_field])), + vec![Arc::clone(&base_values), when_array], + )?; + equality.evaluate(&comparison_batch)? } - }?; + .into_array(remainder_batch.num_rows())?; + as_boolean_array(&compared)?.clone() + } else { + match evaluated_when { + ColumnarValue::Array(a) => { + compare_with_eq(&a, &base_values, base_value_is_nested) + } + ColumnarValue::Scalar(s) => compare_with_eq( + &s.to_scalar()?, + &base_values, + base_value_is_nested, + ), + }? + }; // If the 'when' predicate did not match any rows, continue to the next branch immediately. // Only counts valid slots that are true (masked-null predicate slots are ignored), @@ -1077,11 +1181,14 @@ impl CaseExpr { .collect::>(); if projection.len() < batch.num_columns() { let projected_batch = batch.project(&projection)?; - projected - .body - .case_when_with_expr(&projected_batch, &return_type) + projected.body.case_when_with_expr( + &projected_batch, + &return_type, + self.equality.as_ref(), + ) } else { - self.body.case_when_with_expr(batch, &return_type) + self.body + .case_when_with_expr(batch, &return_type, self.equality.as_ref()) } } @@ -1596,6 +1703,12 @@ impl PhysicalExpr for CaseExpr { } } + fn is_volatile_node(&self) -> bool { + // The comparator is evaluated against synthetic columns, not the outer + // input schema, so it cannot be exposed through `children()`. + self.equality.as_ref().is_some_and(is_volatile) + } + fn children(&self) -> Vec<&Arc> { let mut children = vec![]; if let Some(expr) = &self.body.expr { @@ -1637,10 +1750,11 @@ impl PhysicalExpr for CaseExpr { ), (false, false) => (None, &children[0..children.len()], None), }; - let mut rebuilt = CaseExpr::try_new( + let mut rebuilt = CaseExpr::try_new_internal( expr.cloned(), when_then_expr.iter().cloned().tuples().collect(), else_expr.cloned(), + self.equality.clone(), )?; rebuilt .logical_result_field @@ -1711,6 +1825,11 @@ impl PhysicalExpr for CaseExpr { .as_ref() .map(|(_, metadata)| metadata.to_hashmap()) .unwrap_or_default(), + equality: self + .equality + .as_ref() + .map(|expr| ctx.encode_child(expr).map(Box::new)) + .transpose()?, }, ))), })) @@ -1733,32 +1852,59 @@ impl CaseExpr { "CaseExpr", ); - let mut expr = CaseExpr::try_new( - case.expr - .as_deref() - .map(|expr| ctx.decode(expr)) - .transpose()?, - case.when_then_expr - .iter() - .map(|when_then| { - Ok(( - ctx.decode_required_expression( - when_then.when_expr.as_ref(), - "CaseExpr", - "when_expr", - )?, - ctx.decode_required_expression( - when_then.then_expr.as_ref(), - "CaseExpr", - "then_expr", - )?, - )) - }) - .collect::>>()?, + let base = case + .expr + .as_deref() + .map(|expr| ctx.decode(expr)) + .transpose()?; + let when_then_expr = case + .when_then_expr + .iter() + .map(|when_then| { + Ok(( + ctx.decode_required_expression( + when_then.when_expr.as_ref(), + "CaseExpr", + "when_expr", + )?, + ctx.decode_required_expression( + when_then.then_expr.as_ref(), + "CaseExpr", + "then_expr", + )?, + )) + }) + .collect::>>()?; + let equality = if let Some(equality) = case.equality.as_deref() { + let base_field = base + .as_ref() + .ok_or_else(|| internal_datafusion_err!("Custom CASE requires a base"))? + .return_field(ctx.schema())?; + let when_field = when_then_expr + .first() + .ok_or_else(|| { + internal_datafusion_err!("Custom CASE requires a WHEN branch") + })? + .0 + .return_field(ctx.schema())?; + let schema = Schema::new(vec![ + Field::new("case_base", base_field.data_type().clone(), true) + .with_metadata(base_field.metadata().clone()), + Field::new("case_when", when_field.data_type().clone(), true) + .with_metadata(when_field.metadata().clone()), + ]); + Some(ctx.decode_with_schema(equality, &schema)?) + } else { + None + }; + let mut expr = CaseExpr::try_new_internal( + base, + when_then_expr, case.else_expr .as_deref() .map(|expr| ctx.decode(expr)) .transpose()?, + equality, )?; if let Some(logical_type) = &case.logical_result_type { expr.logical_result_field = Some(( @@ -1901,6 +2047,65 @@ mod tests { } } + #[derive(Debug)] + struct VolatileComparator(Arc); + + impl Hash for VolatileComparator { + fn hash(&self, state: &mut H) { + self.0.hash(state); + } + } + + impl PartialEq for VolatileComparator { + fn eq(&self, other: &Self) -> bool { + self.0.as_ref() == other.0.as_ref() + } + } + + impl Eq for VolatileComparator {} + + impl std::fmt::Display for VolatileComparator { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + std::fmt::Display::fmt(self.0.as_ref(), f) + } + } + + impl PhysicalExpr for VolatileComparator { + fn data_type(&self, schema: &Schema) -> Result { + self.0.data_type(schema) + } + + fn nullable(&self, schema: &Schema) -> Result { + self.0.nullable(schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + self.0.evaluate(batch) + } + + fn is_volatile_node(&self) -> bool { + true + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.0] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!("VolatileComparator requires one child"); + } + Ok(Arc::new(Self(children.remove(0)))) + } + + 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; @@ -1911,6 +2116,9 @@ mod tests { use datafusion_common::plan_err; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_expr::type_coercion::binary::type_union_coercion; + use datafusion_expr::{ + CaseLiteralMatcher, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, + }; use datafusion_expr_common::operator::Operator; use datafusion_physical_expr_common::physical_expr::fmt_sql; use half::f16; @@ -2116,7 +2324,7 @@ mod tests { "marked_identity", vec![DataType::Binary], DataType::Binary, - datafusion_expr::Volatility::Immutable, + Volatility::Immutable, Arc::new(|args| Ok(args[0].clone())), )); let function_field: FieldRef = Arc::new( @@ -2197,6 +2405,254 @@ mod tests { Ok(()) } + #[test] + fn case_with_custom_equality_preserves_first_match_and_nulls() -> Result<()> { + let schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]); + let compare_schema = Schema::new(vec![ + Field::new("case_base", DataType::Int32, true), + Field::new("case_when", DataType::Int32, true), + ]); + let comparison = binary( + col("case_base", &compare_schema)?, + Operator::NotEq, + col("case_when", &compare_schema)?, + &compare_schema, + )?; + let make_case = |equality| { + CaseExpr::try_new_internal( + Some(col("x", &schema)?), + vec![ + (lit(ScalarValue::Null), lit("null")), + (lit(2_i32), lit("first")), + (lit(1_i32), lit("second")), + ], + Some(lit("else")), + equality, + ) + }; + let ordinary = make_case(None)?; + let customized = make_case(Some(comparison))?; + assert_ne!(ordinary, customized); + assert!(matches!( + customized.eval_method, + EvalMethod::WithExpression(_) + )); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(Int32Array::from(vec![Some(1), Some(2), None]))], + )?; + let customized: Arc = Arc::new(customized); + let result = customized.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + as_string_array(&result), + &StringArray::from(vec!["first", "second", "else"]) + ); + + let children = customized.children().into_iter().cloned().collect(); + let rebuilt = customized.with_new_children(children)?; + let result = rebuilt.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + as_string_array(&result), + &StringArray::from(vec!["first", "second", "else"]) + ); + Ok(()) + } + + #[derive(Debug)] + struct NotEqualLiteralMatcher(Vec); + + impl CaseLiteralMatcher for NotEqualLiteralMatcher { + fn map_to_when_indices( + &self, + base: &ArrayRef, + else_index: u32, + ) -> Result> { + let values = as_int32_array(base)?; + Ok((0..values.len()) + .map(|row| { + if values.is_null(row) { + return else_index; + } + self.0 + .iter() + .position(|when| values.value(row) != *when) + .and_then(|index| u32::try_from(index).ok()) + .unwrap_or(else_index) + }) + .collect()) + } + } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct IndexedNotEqualUdf(Signature); + + impl ScalarUDFImpl for IndexedNotEqualUdf { + fn name(&self) -> &str { + "indexed_not_equal" + } + + fn signature(&self) -> &Signature { + &self.0 + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn create_case_literal_matcher( + &self, + when_values: &[ScalarValue], + ) -> Result>> { + let values = when_values + .iter() + .map(|value| match value { + ScalarValue::Int32(Some(value)) => Some(*value), + _ => None, + }) + .collect::>>(); + Ok(values.map(|values| { + Arc::new(NotEqualLiteralMatcher(values)) as Arc + })) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + internal_err!("literal CASE should use the indexed matcher") + } + } + + #[test] + fn custom_case_literal_lookup_preserves_first_match_and_rewrite() -> Result<()> { + let schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]); + let compare_schema = Schema::new(vec![ + Field::new("case_base", DataType::Int32, true), + Field::new("case_when", DataType::Int32, true), + ]); + let udf = Arc::new(datafusion_expr::ScalarUDF::new_from_impl( + IndexedNotEqualUdf(Signature::any(2, Volatility::Immutable)), + )); + let equality: Arc = Arc::new(ScalarFunctionExpr::new( + "indexed_not_equal", + udf, + vec![ + col("case_base", &compare_schema)?, + col("case_when", &compare_schema)?, + ], + Arc::new(Field::new("indexed_not_equal", DataType::Boolean, true)), + Arc::new(datafusion_common::config::ConfigOptions::default()), + )); + let expr: Arc = Arc::new(CaseExpr::try_new_with_equality( + col("x", &schema)?, + vec![(lit(2_i32), lit("first")), (lit(1_i32), lit("second"))], + Some(lit("else")), + equality, + )?); + let case = expr + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("Expected CASE expression"))?; + assert!(matches!( + case.eval_method, + EvalMethod::WithExprScalarLookupTable(_) + )); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(Int32Array::from(vec![ + Some(1), + Some(2), + Some(3), + None, + ]))], + )?; + for candidate in [ + Arc::clone(&expr), + Arc::clone(&expr) + .with_new_children(expr.children().into_iter().cloned().collect())?, + ] { + let result = candidate.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + as_string_array(&result), + &StringArray::from(vec!["first", "second", "first", "else"]) + ); + } + Ok(()) + } + + #[test] + fn case_custom_equality_preserves_udf_array_arguments() -> Result<()> { + let schema = Schema::new(vec![Field::new("x", DataType::Int32, false)]); + let compare_schema = Schema::new(vec![ + Field::new("case_base", DataType::Int32, false), + Field::new("case_when", DataType::Int32, false), + ]); + let udf = Arc::new(datafusion_expr::expr_fn::create_udf( + "scalar_when_eq", + vec![DataType::Int32, DataType::Int32], + DataType::Boolean, + Volatility::Immutable, + Arc::new(|args| { + let [ColumnarValue::Array(base), ColumnarValue::Array(when)] = args + else { + return internal_err!( + "CASE equality must pass two array arguments to UDFs" + ); + }; + let base = as_int32_array(base)?; + let when = as_int32_array(when)?; + let result = BooleanArray::from_iter( + base.values() + .iter() + .zip(when.values()) + .map(|(base, when)| Some(base == when)), + ); + Ok(ColumnarValue::Array(Arc::new(result))) + }), + )); + let comparison: Arc = Arc::new(ScalarFunctionExpr::new( + "scalar_when_eq", + udf, + vec![ + col("case_base", &compare_schema)?, + col("case_when", &compare_schema)?, + ], + Arc::new(Field::new("scalar_when_eq", DataType::Boolean, false)), + Arc::new(datafusion_common::config::ConfigOptions::default()), + )); + let expr = CaseExpr::try_new_with_equality( + col("x", &schema)?, + vec![(lit(2_i32), lit("matched"))], + Some(lit("else")), + comparison, + )?; + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + as_string_array(&result), + &StringArray::from(vec!["else", "matched", "else"]) + ); + Ok(()) + } + + #[test] + fn custom_equality_volatility_is_visible_without_exposing_synthetic_columns() + -> Result<()> { + let case: Arc = Arc::new(CaseExpr::try_new_with_equality( + lit(1_i32), + vec![(lit(1_i32), lit("match"))], + Some(lit("other")), + Arc::new(VolatileComparator(lit(true))), + )?); + assert!(is_volatile(&case)); + assert_eq!(case.children().len(), 4); + let children = case.children().into_iter().cloned().collect(); + let rebuilt = case.with_new_children(children)?; + assert!(is_volatile(&rebuilt)); + assert_eq!(rebuilt.children().len(), 4); + Ok(()) + } + #[test] fn case_with_expr_dictionary() -> Result<()> { let schema = Schema::new(vec![Field::new( @@ -3933,17 +4389,35 @@ mod tests { #[cfg(all(test, feature = "proto"))] mod proto_tests { use super::*; - use crate::expressions::col; + use crate::expressions::{binary, col}; use crate::proto_test_util::{ StubDecoder, StubEncoder, UnreachableDecoder, column_node, }; use arrow::datatypes::Field; + use datafusion_expr::Operator; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecode; use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; use datafusion_proto_models::protobuf; use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalWhenThen}; + use std::cell::RefCell; use std::collections::HashMap; + struct SchemaRecordingDecoder { + field_counts: RefCell>, + } + + impl PhysicalExprDecode for SchemaRecordingDecoder { + fn decode( + &self, + _node: &PhysicalExprNode, + schema: &Schema, + ) -> Result> { + self.field_counts.borrow_mut().push(schema.fields().len()); + Ok(Arc::new(Column::new("decoded", 0))) + } + } + fn proto_case_fixture() -> CaseExpr { let schema = Schema::new(vec![Field::new("a", DataType::Boolean, true)]); CaseExpr::try_new( @@ -3978,6 +4452,7 @@ mod proto_tests { else_expr, logical_result_type: None, logical_result_metadata: Default::default(), + equality: None, }, ))), } @@ -4008,6 +4483,47 @@ mod proto_tests { assert!(case_node.logical_result_metadata.is_empty()); } + #[test] + fn custom_equality_proto_field_roundtrip() -> Result<()> { + let comparison_schema = Schema::new(vec![ + Field::new("case_base", DataType::Boolean, true), + Field::new("case_when", DataType::Boolean, true), + ]); + let comparison = binary( + col("case_base", &comparison_schema)?, + Operator::NotEq, + col("case_when", &comparison_schema)?, + &comparison_schema, + )?; + let schema = Schema::new(vec![Field::new("a", DataType::Boolean, true)]); + let case = CaseExpr::try_new_with_equality( + col("a", &schema)?, + vec![(lit(true), lit(1_i32))], + Some(lit(0_i32)), + comparison, + )?; + let encoder = StubEncoder::ok(); + let node = case + .try_to_proto(&PhysicalExprEncodeCtx::new(&encoder))? + .expect("custom CASE should encode"); + let Some(protobuf::physical_expr_node::ExprType::Case(encoded)) = &node.expr_type + else { + panic!("expected CASE node"); + }; + assert!(encoded.equality.is_some()); + let decoder = SchemaRecordingDecoder { + field_counts: RefCell::new(Vec::new()), + }; + let decoded = CaseExpr::try_from_proto( + &node, + &PhysicalExprDecodeCtx::new(&schema, &decoder), + )?; + let decoded = decoded.downcast_ref::().expect("CaseExpr"); + assert!(decoded.equality.is_some()); + assert_eq!(*decoder.field_counts.borrow(), vec![1, 1, 1, 2, 1]); + Ok(()) + } + #[test] fn proto_preserves_logical_case_field() -> Result<()> { let metadata = HashMap::from([("source".to_string(), "logical".to_string())]); diff --git a/datafusion/physical-expr/src/expressions/case/literal_lookup_table/mod.rs b/datafusion/physical-expr/src/expressions/case/literal_lookup_table/mod.rs index 0d4291ccc934b..dee79524e5568 100644 --- a/datafusion/physical-expr/src/expressions/case/literal_lookup_table/mod.rs +++ b/datafusion/physical-expr/src/expressions/case/literal_lookup_table/mod.rs @@ -20,15 +20,20 @@ mod bytes_like_lookup_table; mod primitive_lookup_table; use crate::expressions::Literal; -use crate::expressions::case::CaseBody; use crate::expressions::case::literal_lookup_table::boolean_lookup_table::BooleanIndexMap; use crate::expressions::case::literal_lookup_table::bytes_like_lookup_table::BytesLikeIndexMap; use crate::expressions::case::literal_lookup_table::primitive_lookup_table::PrimitiveIndexMap; +use crate::expressions::case::{CaseBody, direct_case_equality_udf}; use arrow::array::{Array, ArrayRef, UInt32Array, downcast_primitive}; use arrow::datatypes::DataType; -use datafusion_common::{ScalarValue, arrow_datafusion_err, plan_datafusion_err}; +use datafusion_common::{ + ScalarValue, arrow_datafusion_err, internal_err, plan_datafusion_err, +}; +use datafusion_expr::{CaseLiteralMatcher, Volatility}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use indexmap::IndexMap; use std::fmt::Debug; +use std::sync::Arc; /// Optimization for CASE expressions with literal WHEN and THEN clauses /// @@ -79,7 +84,10 @@ pub(in super::super) struct LiteralLookupTable { } impl LiteralLookupTable { - pub(in super::super) fn maybe_new(body: &CaseBody) -> Option { + pub(in super::super) fn maybe_new( + body: &CaseBody, + equality: Option<&Arc>, + ) -> Option { // We can't use the optimization if we don't have any when then pairs if body.when_then_expr.is_empty() { return None; @@ -128,16 +136,17 @@ impl LiteralLookupTable { // Keep only the first occurrence of each when literal (as the first match is used) // and remove nulls (as they cannot be matched - case NULL WHEN NULL THEN ... ELSE ... END always goes to ELSE) - let (when, then): (Vec, Vec) = { + let (when, then): (Vec, Vec) = if equality.is_some() { + // Custom predicate equality need not be transitive. Preserve every + // literal's original branch ordinal for its matcher. + when_then_exprs_scalars.into_iter().unzip() + } else { let mut map = IndexMap::with_capacity(body.when_then_expr.len()); - - for (when, then) in when_then_exprs_scalars.into_iter() { - // Don't overwrite existing entries as we want to keep the first occurrence + for (when, then) in when_then_exprs_scalars { if !map.contains_key(&when) { map.insert(when, then); } } - map.into_iter().unzip() }; @@ -185,7 +194,20 @@ impl LiteralLookupTable { // The else expression is in the end let else_index = then_and_else_values.len() as u32 - 1; - let lookup = try_creating_lookup_table(when).ok()?; + let lookup: Box = if let Some(equality) = equality { + let udf = direct_case_equality_udf(equality)?; + if udf.fun().signature().volatility != Volatility::Immutable { + return None; + } + Box::new( + udf.fun() + .inner() + .create_case_literal_matcher(&when) + .ok()??, + ) + } else { + try_creating_lookup_table(when).ok()? + }; Some(Self { lookup, @@ -201,6 +223,13 @@ impl LiteralLookupTable { let take_indices = self .lookup .map_to_when_indices(keys_array, self.else_index)?; + if take_indices.len() != keys_array.len() { + return internal_err!( + "CASE literal matcher returned {} rows for {} inputs", + take_indices.len(), + keys_array.len() + ); + } // Zero-copy conversion let take_indices = UInt32Array::from(take_indices); @@ -260,6 +289,16 @@ pub(super) trait WhenLiteralIndexMap: Debug + Send + Sync { ) -> datafusion_common::Result>; } +impl WhenLiteralIndexMap for Arc { + fn map_to_when_indices( + &self, + array: &ArrayRef, + else_index: u32, + ) -> datafusion_common::Result> { + CaseLiteralMatcher::map_to_when_indices(self.as_ref(), array, else_index) + } +} + fn try_creating_lookup_table( unique_non_null_literals: Vec, ) -> datafusion_common::Result> { diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index aeb7545908788..3270e9ce99de0 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -24,7 +24,7 @@ use crate::{ expressions::{self, Column, Literal, binary, like, similar_to}, }; -use arrow::datatypes::Schema; +use arrow::datatypes::{Field, Schema}; use datafusion_common::config::ConfigOptions; use datafusion_common::datatype::FieldExt; use datafusion_common::metadata::FieldMetadata; @@ -338,6 +338,16 @@ pub fn create_physical_expr( similar_to(*negated, *case_insensitive, physical_expr, physical_pattern) } Expr::Case(case) => { + let equality = if case.expr.is_some() { + execution_props + .case_equality_planner + .as_ref() + .map(|planner| planner.plan_equality(case, input_dfschema)) + .transpose()? + .flatten() + } else { + None + }; let expr: Option> = if let Some(e) = &case.expr { Some(create_physical_expr( e.as_ref(), @@ -401,8 +411,40 @@ pub fn create_physical_expr( } else { None }; - let physical_case = - expressions::CaseExpr::try_new(expr, when_then_expr, else_expr)?; + let physical_case = if let Some(equality) = equality { + let base = case.expr.as_ref().ok_or_else(|| { + internal_datafusion_err!("Custom CASE equality requires a base") + })?; + let first_when = case.when_then_expr.first().ok_or_else(|| { + internal_datafusion_err!( + "Custom CASE equality requires a WHEN branch" + ) + })?; + let base_field = base.to_field(input_dfschema)?.1; + let when_field = first_when.0.to_field(input_dfschema)?.1; + let comparison_schema = DFSchema::try_from(Schema::new(vec![ + Field::new("case_base", base_field.data_type().clone(), true) + .with_metadata(base_field.metadata().clone()), + Field::new("case_when", when_field.data_type().clone(), true) + .with_metadata(when_field.metadata().clone()), + ]))?; + let comparison = create_physical_expr( + &equality, + &comparison_schema, + execution_props, + planning_ctx, + )?; + expressions::CaseExpr::try_new_with_equality( + expr.ok_or_else(|| { + internal_datafusion_err!("Custom CASE equality requires a base") + })?, + when_then_expr, + else_expr, + comparison, + )? + } else { + expressions::CaseExpr::try_new(expr, when_then_expr, else_expr)? + }; if common_unmarked_branch { Ok(Arc::new(physical_case)) } else { @@ -759,10 +801,13 @@ pub fn logical2physical(expr: &Expr, schema: &Schema) -> Arc { #[cfg(test)] mod tests { - use arrow::array::{ArrayRef, BooleanArray, RecordBatch, StringArray}; + use arrow::array::{ + ArrayRef, AsArray, BooleanArray, Int32Array, RecordBatch, StringArray, + }; use arrow::datatypes::{DataType, Field}; use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::HashMap; + use datafusion_expr::execution_props::CaseEqualityPlanner; use datafusion_expr::physical_planning_context::{ ScalarSubqueryResults, SubqueryIndex, }; @@ -784,6 +829,63 @@ mod tests { ) } + #[derive(Debug)] + struct NotEqualCasePlanner; + + impl CaseEqualityPlanner for NotEqualCasePlanner { + fn plan_equality( + &self, + _case: &datafusion_expr::expr::Case, + _input_schema: &DFSchema, + ) -> Result> { + Ok(Some(col("case_base").not_eq(col("case_when")))) + } + } + + #[test] + fn custom_simple_case_equality_uses_planner_hook() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let df_schema = DFSchema::try_from(schema.clone())?; + let logical = Expr::Case(datafusion_expr::expr::Case::new( + Some(Box::new(col("a"))), + vec![ + (Box::new(lit(2_i32)), Box::new(lit("first"))), + (Box::new(lit(1_i32)), Box::new(lit("second"))), + ], + Some(Box::new(lit("else"))), + )); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + )?; + let mut props = ExecutionProps::new(); + props.case_equality_planner = Some(Arc::new(NotEqualCasePlanner)); + let custom = create_physical_expr( + &logical, + &df_schema, + &props, + &PhysicalPlanningContext::default(), + )?; + let result = custom.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + result.as_string::(), + &StringArray::from(vec!["first", "second"]) + ); + + let ordinary = create_physical_expr( + &logical, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; + let result = ordinary.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + result.as_string::(), + &StringArray::from(vec!["second", "first"]) + ); + Ok(()) + } + #[test] fn case_metadata_matches_logical_field_after_alias_lowering() -> Result<()> { let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); diff --git a/datafusion/physical-expr/src/scalar_function.rs b/datafusion/physical-expr/src/scalar_function.rs index b46015ac033e3..915384aec6778 100644 --- a/datafusion/physical-expr/src/scalar_function.rs +++ b/datafusion/physical-expr/src/scalar_function.rs @@ -169,6 +169,45 @@ impl ScalarFunctionExpr { &self.config_options } + /// Evaluate precomputed arguments without broadcasting scalar values. + pub(crate) fn evaluate_with_args( + &self, + args: Vec, + arg_fields: Vec, + number_rows: usize, + ) -> Result { + let input_empty = args.is_empty(); + let input_all_scalar = args + .iter() + .all(|arg| matches!(arg, ColumnarValue::Scalar(_))); + + let output = self.eval_fun.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows, + return_field: Arc::clone(&self.return_field), + config_options: Arc::clone(&self.config_options), + })?; + + if let ColumnarValue::Array(array) = &output + && array.len() != number_rows + { + // A one-element array is a scalar only when all inputs were scalar. + let preserve_scalar = array.len() == 1 && !input_empty && input_all_scalar; + return if preserve_scalar { + ScalarValue::try_from_array(array, 0).map(ColumnarValue::Scalar) + } else { + internal_err!( + "UDF {} returned a different number of rows than expected. Expected: {}, Got: {}", + self.name, + number_rows, + array.len() + ) + }; + } + Ok(output) + } + /// Given an arbitrary PhysicalExpr attempt to downcast it to a ScalarFunctionExpr /// and verify that its inner function is of type T. /// If the downcast fails, or the function is not of type T, returns `None`. @@ -261,38 +300,7 @@ impl PhysicalExpr for ScalarFunctionExpr { .map(|e| e.return_field(batch.schema_ref())) .collect::>>()?; - let input_empty = args.is_empty(); - let input_all_scalar = args - .iter() - .all(|arg| matches!(arg, ColumnarValue::Scalar(_))); - - // evaluate the function - let output = self.eval_fun.invoke_with_args(ScalarFunctionArgs { - args, - arg_fields, - number_rows: batch.num_rows(), - return_field: Arc::clone(&self.return_field), - config_options: Arc::clone(&self.config_options), - })?; - - if let ColumnarValue::Array(array) = &output - && array.len() != batch.num_rows() - { - // If the arguments are a non-empty slice of scalar values, we can assume that - // returning a one-element array is equivalent to returning a scalar. - let preserve_scalar = array.len() == 1 && !input_empty && input_all_scalar; - return if preserve_scalar { - ScalarValue::try_from_array(array, 0).map(ColumnarValue::Scalar) - } else { - internal_err!( - "UDF {} returned a different number of rows than expected. Expected: {}, Got: {}", - self.name, - batch.num_rows(), - array.len() - ) - }; - } - Ok(output) + self.evaluate_with_args(args, arg_fields, batch.num_rows()) } fn return_field(&self, _input_schema: &Schema) -> Result { diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 5df5482fb75de..5c39d6ebf4015 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -1417,6 +1417,31 @@ mod tests { use crate::test::exec::StatisticsExec; use arrow::datatypes::{Field, Schema, UnionFields, UnionMode}; + #[test] + fn custom_case_comparator_does_not_leak_synthetic_columns() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)])); + let comparison_schema = Schema::new(vec![ + Field::new("case_base", DataType::Int32, true), + Field::new("case_when", DataType::Int32, true), + ]); + let equality = binary( + col("case_base", &comparison_schema)?, + Operator::Eq, + col("case_when", &comparison_schema)?, + &comparison_schema, + )?; + let predicate: Arc = Arc::new(CaseExpr::try_new_with_equality( + col("x", &schema)?, + vec![(lit(1_i32), lit(true))], + Some(lit(false)), + equality, + )?); + let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let filter = FilterExec::try_new(predicate, input)?; + assert_eq!(filter.schema().fields().len(), 1); + Ok(()) + } + #[tokio::test] async fn collect_columns_predicates() -> Result<()> { let schema = test::aggr_test_schema(); diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index bfca4d42c5cf1..625ad2a39ff94 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1178,6 +1178,7 @@ message PhysicalCaseNode { PhysicalExprNode else_expr = 3; datafusion_common.ArrowType logical_result_type = 4; map logical_result_metadata = 5; + PhysicalExprNode equality = 6; } message PhysicalTryCastNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 82a21febffedf..99af5527ce1cd 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -17835,6 +17835,9 @@ impl serde::Serialize for PhysicalCaseNode { if !self.logical_result_metadata.is_empty() { len += 1; } + if self.equality.is_some() { + 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)?; @@ -17851,6 +17854,9 @@ impl serde::Serialize for PhysicalCaseNode { if !self.logical_result_metadata.is_empty() { struct_ser.serialize_field("logicalResultMetadata", &self.logical_result_metadata)?; } + if let Some(v) = self.equality.as_ref() { + struct_ser.serialize_field("equality", v)?; + } struct_ser.end() } } @@ -17870,6 +17876,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalCaseNode { "logicalResultType", "logical_result_metadata", "logicalResultMetadata", + "equality", ]; #[allow(clippy::enum_variant_names)] @@ -17879,6 +17886,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalCaseNode { ElseExpr, LogicalResultType, LogicalResultMetadata, + Equality, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -17905,6 +17913,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalCaseNode { "elseExpr" | "else_expr" => Ok(GeneratedField::ElseExpr), "logicalResultType" | "logical_result_type" => Ok(GeneratedField::LogicalResultType), "logicalResultMetadata" | "logical_result_metadata" => Ok(GeneratedField::LogicalResultMetadata), + "equality" => Ok(GeneratedField::Equality), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -17929,6 +17938,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalCaseNode { let mut else_expr__ = None; let mut logical_result_type__ = None; let mut logical_result_metadata__ = None; + let mut equality__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Expr => { @@ -17963,6 +17973,12 @@ impl<'de> serde::Deserialize<'de> for PhysicalCaseNode { map_.next_value::>()? ); } + GeneratedField::Equality => { + if equality__.is_some() { + return Err(serde::de::Error::duplicate_field("equality")); + } + equality__ = map_.next_value()?; + } } } Ok(PhysicalCaseNode { @@ -17971,6 +17987,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalCaseNode { else_expr: else_expr__, logical_result_type: logical_result_type__, logical_result_metadata: logical_result_metadata__.unwrap_or_default(), + equality: equality__, }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 47c2d2ff619c0..48f80bedb1bf2 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1833,6 +1833,8 @@ pub struct PhysicalCaseNode { ::prost::alloc::string::String, ::prost::alloc::string::String, >, + #[prost(message, optional, boxed, tag = "6")] + pub equality: ::core::option::Option<::prost::alloc::boxed::Box>, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalTryCastNode { diff --git a/datafusion/proto/tests/cases/plans/misc.rs b/datafusion/proto/tests/cases/plans/misc.rs index 41bb051c28730..9708e1bbb8ece 100644 --- a/datafusion/proto/tests/cases/plans/misc.rs +++ b/datafusion/proto/tests/cases/plans/misc.rs @@ -20,13 +20,18 @@ use super::{roundtrip_test, roundtrip_test_and_return}; use arrow::datatypes::{Fields, TimeUnit}; +use datafusion::arrow::array::{Int32Array, RecordBatch, StringArray}; use datafusion::arrow::compute::kernels::sort::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::logical_expr::Operator; use datafusion::physical_expr::LexOrdering; +use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::analyze::AnalyzeExec; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::explain::ExplainExec; -use datafusion::physical_plan::expressions::{PhysicalSortExpr, col, lit}; +use datafusion::physical_plan::expressions::{ + CaseExpr, PhysicalSortExpr, binary, col, lit, +}; use datafusion::physical_plan::metrics::MetricCategory; use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; @@ -42,7 +47,10 @@ use datafusion::scalar::ScalarValue; use datafusion_common::display::{PlanType, StringifiedPlan}; use datafusion_common::format::ExplainFormat; use datafusion_common::{DataFusionError, Result, UnnestOptions}; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter; +use datafusion_proto::physical_plan::from_proto::parse_physical_expr; +use datafusion_proto::physical_plan::to_proto::{ + serialize_physical_expr, serialize_physical_expr_with_converter, +}; use datafusion_proto::physical_plan::{ AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, PhysicalPlanDecodeContext, @@ -53,6 +61,46 @@ use prost::Message; use std::sync::Arc; use std::vec; +#[test] +fn custom_case_equality_physical_expr_roundtrip() -> Result<()> { + let schema = Schema::new(vec![Field::new("x", DataType::Int32, false)]); + let compare_schema = Schema::new(vec![ + Field::new("case_base", DataType::Int32, true), + Field::new("case_when", DataType::Int32, true), + ]); + let equality = binary( + col("case_base", &compare_schema)?, + Operator::NotEq, + col("case_when", &compare_schema)?, + &compare_schema, + )?; + let expr: Arc = Arc::new(CaseExpr::try_new_with_equality( + col("x", &schema)?, + vec![(lit(2_i32), lit("first")), (lit(1_i32), lit("second"))], + Some(lit("else")), + equality, + )?); + let codec = DefaultPhysicalExtensionCodec {}; + let encoded = serialize_physical_expr(&expr, &codec)?; + let ctx = SessionContext::new(); + let decoded = parse_physical_expr(&encoded, &ctx.task_ctx(), &schema, &codec)?; + assert_eq!(expr.as_ref(), decoded.as_ref()); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + )?; + let result = decoded.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + result + .as_any() + .downcast_ref::() + .expect("strings"), + &StringArray::from(vec!["first", "second"]) + ); + Ok(()) +} + #[test] fn roundtrip_analyze() -> Result<()> { let schema = Arc::new(Schema::new(vec![