namanjain24-sudo commented on code in PR #25191:
URL: https://github.com/apache/datafusion/pull/25191#discussion_r4022210965
##########
datafusion/substrait/src/logical_plan/producer/expr/if_then.rs:
##########
@@ -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());
Review Comment:
You're right, thanks. I checked `CaseExpr` and it evaluates the base once
into `base_values`, then compares every WHEN against that one value, so the
desugaring changes the meaning of a volatile base. Rejecting it in 6b49064:
```rust
if let Some(base) = expr.as_ref().filter(|base| base.is_volatile()) {
return not_impl_err!(
"Substrait does not support a volatile CASE base expression: {base}"
);
}
```
I left `SwitchExpression` alone for the reason you give: the consumer's
`consume_switch` answers `not_impl_err!("Switch expression not supported")`, so
it would need consumer support and round-trip coverage first, and it could only
carry literal WHEN values anyway.
The regression now uses a counter UDF instead of `random()`, so the
difference is deterministic. `call_counter()` returns 1 on its first call, 2 on
its second, over a single row:
- `CASE call_counter() WHEN 2 THEN 20 WHEN 1 THEN 10 ELSE 99 END` gives `10`
after 1 call: the base is evaluated once, returns 1, and matches the second
WHEN.
- The same CASE after this desugaring gives `99` after 2 calls: 1 does not
equal 2, then 2 does not equal 1, so the row falls through to ELSE.
The test asserts both, and that the producer now rejects the first plan.
With the guard removed it fails, and the plan it then produces contains two
separate `call_counter` calls, one per condition.
One note on the comment above this code, which cites `from_between`: that
desugaring duplicates its input too, but it is not the same bug. DataFusion
plans `BETWEEN` into two comparisons over `Arc::clone(&value_expr)`
(`datafusion/physical-expr/src/planner.rs:446`), so the operand is evaluated
twice there as well and the Substrait form matches. `CaseExpr` is the one that
keeps a single evaluated base.
##########
datafusion/substrait/tests/cases/serialize.rs:
##########
@@ -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() {
Review Comment:
Good point, the assertion was weak. In 6b49064 each condition is now checked
as `equal(<base>, <literal>)` with the operands in that order:
- the call is an `equal` whose anchor is one of the registered `equal`
functions, and it has exactly 2 arguments,
- the left operand is a field reference rooted at the input, pointing at
struct field 0, which is `data.a`,
- the right operand is a literal, asserted per clause as
`LiteralType::I64(1)` and then `LiteralType::I64(2)`, so a swapped operand
order, a wrong literal, or a reference to another column all fail now.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
