diff --git a/datafusion/substrait/src/logical_plan/producer/expr/if_then.rs b/datafusion/substrait/src/logical_plan/producer/expr/if_then.rs index 2c10b26436f50..2ee7510a0931a 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/if_then.rs +++ b/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 = 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 WHEN THEN ...` + // is therefore emitted as `IfClause`s over ` = `, 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` `` never matching. + let mut ifs: Vec = 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)?), }); } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 813d0ed6c3489..2c3c989dd2d4d 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/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 ` = ` 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 } diff --git a/datafusion/substrait/tests/cases/serialize.rs b/datafusion/substrait/tests/cases/serialize.rs index 4a8413718edb9..d0aa2718964e9 100644 --- a/datafusion/substrait/tests/cases/serialize.rs +++ b/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 WHEN ...` must therefore be emitted as conditions + /// over ` = `. 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 = 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) { assert_eq!( rel_common.unwrap().emit_kind.clone(),