Support custom simple CASE equality with ordered literal matching

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,
             )

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<dyn datafusion_expr::execution_props::CaseEqualityPlanner>> {
+        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<Option<Expr>> {
+            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() {

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<Option<crate::Expr>>;
+}
+
 /// 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<Arc<ConfigOptions>>,
     /// Providers for scalar variables
     pub var_providers: Option<HashMap<VarType, Arc<dyn VarProvider + Send + Sync>>>,
+    /// Optional simple CASE comparator, selected only during physical planning.
+    pub case_equality_planner: Option<Arc<dyn CaseEqualityPlanner>>,
 }
 
 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:?}")
         );
     }

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};

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<DateTime<Utc>>,
     config_options: Arc<ConfigOptions>,
+    case_equality_planner: Option<Arc<dyn CaseEqualityPlanner>>,
 }
 
 /// Builder for [`SimplifyContext`].
@@ -46,6 +48,7 @@ pub struct SimplifyContextBuilder {
     schema: Option<DFSchemaRef>,
     query_execution_start_time: Option<DateTime<Utc>>,
     config_options: Option<Arc<ConfigOptions>>,
+    case_equality_planner: Option<Arc<dyn CaseEqualityPlanner>>,
 }
 
 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<ConfigOptions> {
         &self.config_options
     }
+
+    pub fn case_equality_planner(&self) -> Option<&Arc<dyn CaseEqualityPlanner>> {
+        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<Arc<dyn CaseEqualityPlanner>>,
+    ) -> 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,
         }
     }
 }

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<ScalarValue>, 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<Vec<u32>>;
+}
+
 /// 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<Option<Arc<dyn CaseLiteralMatcher>>> {
+        Ok(None)
+    }
+
     /// Returns any aliases (alternate names) for this function.
     ///
     /// Aliases can be used to invoke the same function using different names.

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<ConfigOptions>;
 
+    /// Optional comparator used to fold simple CASE expressions consistently
+    /// with physical planning.
+    fn case_equality_planner(&self) -> Option<&Arc<dyn CaseEqualityPlanner>> {
+        None
+    }
+
     fn function_registry(&self) -> Option<&dyn FunctionRegistry> {
         None
     }