From b87b85757656b7cd73e6ab1f79b5ff85f622dba3 Mon Sep 17 00:00:00 2001 From: naman Date: Fri, 11 Sep 2026 18:42:10 +0530 Subject: [PATCH] fix: Emit equality conditions for Substrait CASE base expressions Substrait's `IfThen` has no base expression. Every `IfClause` is a standalone boolean condition, and `then` is the value that clause yields. The producer instead encoded `CASE WHEN THEN ...` by pushing a leading `IfClause` that carries the base expression in `if` and leaves `then` unset, followed by one clause per WHEN whose `if` is the raw WHEN operand. For `CASE a WHEN 1 THEN 'x' WHEN 2 THEN 'y' ELSE 'z' END` that emits three clauses whose conditions are `a`, `1` and `2`, none of which is boolean, and a first clause with no result. The convention is private to DataFusion: the consumer reads a `then`-less first clause back as the base expression, so a DataFusion-to-DataFusion round trip is unaffected. Any other engine sees clauses it cannot evaluate. Emit ` = ` as each clause condition instead, the same desugaring `from_between` already applies to `BETWEEN`. DataFusion matches a base expression with `=` semantics, so the plan keeps its meaning, including a NULL WHEN operand never matching. A base `CASE` now round trips as the equivalent searched `CASE`, keeping its original projection name and schema. --- .../src/logical_plan/producer/expr/if_then.rs | 31 ++++---- .../tests/cases/roundtrip_logical_plan.rs | 8 ++- datafusion/substrait/tests/cases/serialize.rs | 72 ++++++++++++++++++- 3 files changed, 96 insertions(+), 15 deletions(-) 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(),