fix: Emit equality conditions for Substrait CASE base expressions

datafusion/substrait/src/logical_plan/producer/expr/if_then.rs

@@ -17,7 +17,7 @@
 
 use crate::logical_plan::producer::SubstraitProducer;
 use datafusion::common::DFSchemaRef;
-use datafusion::logical_expr::Case;
+use datafusion::logical_expr::{Case, Expr};
 use substrait::proto::Expression;
 use substrait::proto::expression::if_then::IfClause;
 use substrait::proto::expression::{IfThen, RexType};
@@ -32,19 +32,24 @@ pub fn from_case(
         when_then_expr,
         else_expr,
     } = case;
-    let mut ifs: Vec<IfClause> = vec![];
-    // Parse base
-    if let Some(e) = expr {
-        // Base expression exists
-        ifs.push(IfClause {
-            r#if: Some(producer.handle_expr(e, schema)?),
-            then: None,
-        });
-    }
-    // Parse `when`s
-    for (r#if, then) in when_then_expr {
+
+    // Substrait's `IfThen` has no notion of a base expression: every `IfClause`
+    // is a standalone boolean condition. A `CASE <base> WHEN <value> THEN ...`
+    // is therefore emitted as `IfClause`s over `<base> = <value>`, the same
+    // desugaring `from_between` applies to `BETWEEN`. DataFusion matches a base
+    // expression with `=` semantics, so this preserves the plan's meaning,
+    // including a `NULL` `<value>` never matching.
+    let mut ifs: Vec<IfClause> = Vec::with_capacity(when_then_expr.len());
+    for (when, then) in when_then_expr {
+        let condition = match expr {
+            Some(base) => {
+                let eq = Expr::eq(*base.clone(), *when.clone());
+                producer.handle_expr(&eq, schema)?
+            }
+            None => producer.handle_expr(when, schema)?,
+        };
         ifs.push(IfClause {
-            r#if: Some(producer.handle_expr(r#if, schema)?),
+            r#if: Some(condition),
             then: Some(producer.handle_expr(then, schema)?),
         });
     }

datafusion/substrait/tests/cases/roundtrip_logical_plan.rs

@@ -656,12 +656,18 @@ async fn case_without_base_expression() -> Result<()> {
 
 #[tokio::test]
 async fn case_with_base_expression() -> Result<()> {
-    roundtrip(
+    // Substrait has no base expression in `IfThen`, so a base `CASE` is emitted
+    // as conditions over `<base> = <value>` and comes back in that form. The
+    // projection keeps its original name, so the schema is unchanged.
+    assert_expected_plan(
         "SELECT (CASE a
                             WHEN 0 THEN 'zero'
                             WHEN 1 THEN 'one'
                             ELSE 'other'
                            END) FROM data",
+        "Projection: CASE WHEN data.a = Int64(0) THEN Utf8(\"zero\") WHEN data.a = Int64(1) THEN Utf8(\"one\") ELSE Utf8(\"other\") END AS CASE data.a WHEN Int64(0) THEN Utf8(\"zero\") WHEN Int64(1) THEN Utf8(\"one\") ELSE Utf8(\"other\") END\
+        \n  TableScan: data projection=[a]",
+        true,
     )
     .await
 }

datafusion/substrait/tests/cases/serialize.rs

@@ -30,7 +30,8 @@ mod tests {
     use std::{fs, sync::Arc};
     use substrait::proto::expression::field_reference::{ReferenceType, RootType};
     use substrait::proto::expression::reference_segment;
-    use substrait::proto::expression::{ReferenceSegment, RexType};
+    use substrait::proto::expression::{IfThen, ReferenceSegment, RexType};
+    use substrait::proto::extensions::simple_extension_declaration::MappingType;
     use substrait::proto::function_argument::ArgType;
     use substrait::proto::plan_rel::RelType;
     use substrait::proto::rel_common::{Emit, EmitKind};
@@ -321,6 +322,75 @@ mod tests {
         Ok(())
     }
 
+    /// Substrait's `IfThen` has no base expression: every `IfClause` is a
+    /// standalone boolean condition and `then` is the value that clause yields.
+    /// A `CASE <base> WHEN <value> ...` must therefore be emitted as conditions
+    /// over `<base> = <value>`. A round trip cannot catch a regression here,
+    /// because the consumer reads back whatever the producer writes.
+    #[tokio::test]
+    async fn case_with_base_expression_emits_equality_conditions() -> Result<()> {
+        let ctx = create_context().await?;
+        let sql = "SELECT CASE a WHEN 1 THEN 'x' WHEN 2 THEN 'y' ELSE 'z' END FROM data";
+
+        let plan = ctx.sql(sql).await?.into_optimized_plan()?;
+        let proto = to_substrait_plan(&plan, &ctx.state())?;
+
+        let equal_anchors: Vec<u32> = proto
+            .extensions
+            .iter()
+            .filter_map(|e| match e.mapping_type.as_ref().unwrap() {
+                MappingType::ExtensionFunction(f) if f.name == "equal" => {
+                    Some(f.function_anchor)
+                }
+                _ => None,
+            })
+            .collect();
+        assert!(!equal_anchors.is_empty(), "no `equal` function registered");
+
+        let root = match proto.relations.first().unwrap().rel_type.as_ref() {
+            Some(RelType::Root(root)) => root.input.as_ref().unwrap(),
+            _ => panic!("expected Root"),
+        };
+        let Some(rel::RelType::Project(project)) = root.rel_type.as_ref() else {
+            panic!("expected Project")
+        };
+
+        let if_thens: Vec<&IfThen> = project
+            .expressions
+            .iter()
+            .filter_map(|expr| match expr.rex_type.as_ref() {
+                Some(RexType::IfThen(if_then)) => Some(if_then.as_ref()),
+                _ => None,
+            })
+            .collect();
+        assert_eq!(if_thens.len(), 1, "expected one IfThen for `{sql}`");
+        let if_then = if_thens[0];
+
+        // One clause per WHEN, with no extra clause carrying the base expression.
+        assert_eq!(if_then.ifs.len(), 2);
+        assert!(if_then.r#else.is_some());
+
+        for (i, clause) in if_then.ifs.iter().enumerate() {
+            let condition = clause
+                .r#if
+                .as_ref()
+                .unwrap_or_else(|| panic!("clause {i} has no condition"));
+            assert!(clause.then.is_some(), "clause {i} has no `then`");
+
+            match condition.rex_type.as_ref().unwrap() {
+                RexType::ScalarFunction(f) => assert!(
+                    equal_anchors.contains(&f.function_reference),
+                    "clause {i} condition is not an `equal` call"
+                ),
+                other => {
+                    panic!("clause {i} condition is not a scalar function: {other:?}")
+                }
+            }
+        }
+
+        Ok(())
+    }
+
     fn assert_emit(rel_common: Option<&RelCommon>, output_mapping: Vec<i32>) {
         assert_eq!(
             rel_common.unwrap().emit_kind.clone(),