Substrait producer emits non-boolean IfThen clauses for CASE with a base expression
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
### Describe the bug
Substrait's `IfThen` has no notion of a base expression. Every `IfClause` is a
standalone boolean condition, and `then` is the value that clause yields:
```proto
message IfThen {
// A list of one or more IfClauses
repeated IfClause ifs = 1;
// The returned Expression if no IfClauses are satisified
Expression else = 2;
message IfClause {
Expression if = 1;
Expression then = 2;
}
}
```
The producer nevertheless uses `IfThen` for `CASE WHEN THEN ...`,
encoding the base expression as a leading `IfClause` whose `then` is left unset
([`from_case`](https://github.com/apache/datafusion/blob/517ac6d53133621c99a9a7dda470549a26348f6a/datafusion/substrait/src/logical_plan/producer/expr/if_then.rs#L36-L50)):
```rust
// 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 {
ifs.push(IfClause {
r#if: Some(producer.handle_expr(r#if, schema)?),
then: Some(producer.handle_expr(then, schema)?),
});
}
```
The following clauses then carry the raw WHEN operand in `if`. Neither the
leading clause nor the WHEN operands are boolean, so the emitted `IfThen` is not
something a spec-conforming consumer can evaluate.
That convention is private to DataFusion. The consumer decodes it back
([`from_if_then`](https://github.com/apache/datafusion/blob/517ac6d53133621c99a9a7dda470549a26348f6a/datafusion/substrait/src/logical_plan/consumer/expr/if_then.rs#L29-L43)):
```rust
// If the first element does not have a `then` part, then we can assume it's a base expression
```
so a DataFusion-to-DataFusion round trip is unaffected and no existing test
fails.
Another implementation reading the same plan does not get that far. substrait-java
(`io.substrait:core:0.103.0`) rejects it outright, because `ProtoExpressionConverter`
converts each clause's `then` unconditionally and the leading clause has none:
```
java.lang.IllegalArgumentException: Unknown type: REXTYPE_NOT_SET
at io.substrait.expression.proto.ProtoExpressionConverter.from(ProtoExpressionConverter.java:364)
at io.substrait.expression.proto.ProtoExpressionConverter.lambda$from$1(ProtoExpressionConverter.java:212)
```
Line 212 is `ExpressionCreator.ifThenClause(from(t.getIf()), from(t.getThen()))`.
No extension URNs are involved: this plan registers no functions at all, so the
failure is on the encoding itself rather than on #11545.
The non-boolean conditions are a second, quieter problem. substrait-java does not
type-check `IfClause.if`, so removing the leading clause and feeding it the two
`i64` conditions is accepted. That half is a silent divergence from the spec
rather than an error, which is why no round trip anywhere surfaces it.
### To Reproduce
```sql
SELECT CASE a WHEN 1 THEN 'x' WHEN 2 THEN 'y' ELSE 'z' END FROM data
```
Save this as `datafusion/substrait/examples/case_base_probe.rs`. It produces the
plan and prints the `IfThen` as it appears in the protobuf, without going through
the consumer:
```rust
use datafusion::common::Result;
use datafusion::prelude::*;
use datafusion_substrait::logical_plan::producer::to_substrait_plan;
use prost::Message;
use substrait::proto::expression::RexType;
use substrait::proto::{plan_rel, rel};
#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
let ctx = SessionContext::new();
ctx.sql("CREATE TABLE data(a BIGINT) AS VALUES (1), (2), (3)")
.await?
.collect()
.await?;
let plan = ctx
.sql("SELECT CASE a WHEN 1 THEN 'x' WHEN 2 THEN 'y' ELSE 'z' END FROM data")
.await?
.into_optimized_plan()?;
let proto = to_substrait_plan(&plan, &ctx.state())?;
std::fs::write("case.pb", proto.encode_to_vec())?;
let Some(plan_rel::RelType::Root(root)) = proto.relations[0].rel_type.as_ref() else {
panic!("expected Root")
};
let Some(rel::RelType::Project(project)) =
root.input.as_ref().unwrap().rel_type.as_ref()
else {
panic!("expected Project")
};
for expr in &project.expressions {
let Some(RexType::IfThen(if_then)) = expr.rex_type.as_ref() else {
continue;
};
println!("IfThen with {} clauses, else present: {}",
if_then.ifs.len(), if_then.r#else.is_some());
for (i, clause) in if_then.ifs.iter().enumerate() {
let kind = match clause.r#if.as_ref().unwrap().rex_type.as_ref().unwrap() {
RexType::Literal(l) => format!("Literal({:?})", l.literal_type.as_ref().unwrap()),
RexType::Selection(_) => "Selection(field reference)".to_string(),
RexType::ScalarFunction(f) => {
format!("ScalarFunction(anchor {}, output_type {:?})",
f.function_reference,
f.output_type.as_ref().map(|t| t.kind.as_ref().unwrap()))
}
other => format!("{other:?}"),
};
println!(" ifs[{i}]: if = {kind}");
println!(" then set: {}", clause.then.is_some());
}
}
Ok(())
}
```
```console
$ cargo run --locked -p datafusion-substrait --example case_base_probe
IfThen with 3 clauses, else present: true
ifs[0]: if = Selection(field reference)
then set: false
ifs[1]: if = Literal(I64(1))
then set: true
ifs[2]: if = Literal(I64(2))
then set: true
```
Three clauses, with conditions `a`, `1` and `2`. All three are `i64`, and the
first has no result at all.
The `case.pb` the probe writes is what substrait-java rejects. To see that
directly, on a JDK with no build tool:
```java
import io.substrait.proto.Plan;
import io.substrait.plan.ProtoPlanConverter;
import java.nio.file.*;
public class CaseProbe {
public static void main(String[] args) throws Exception {
Plan plan = Plan.parseFrom(Files.readAllBytes(Path.of("case.pb")));
try {
new ProtoPlanConverter().from(plan);
System.out.println("accepted");
} catch (Exception e) {
System.out.println(e.getClass().getName() + ": " + e.getMessage());
for (int i = 0; i < 2; i++) System.out.println(" at " + e.getStackTrace()[i]);
}
}
}
```
With `io.substrait:core:0.103.0` and its runtime dependencies on the classpath:
```console
$ javac -cp "lib/*" -d . CaseProbe.java && java -cp "lib/*:." CaseProbe
java.lang.IllegalArgumentException: Unknown type: REXTYPE_NOT_SET
at io.substrait.expression.proto.ProtoExpressionConverter.from(ProtoExpressionConverter.java:364)
at io.substrait.expression.proto.ProtoExpressionConverter.lambda$from$1(ProtoExpressionConverter.java:212)
```
For comparison, the searched form `CASE WHEN a = 1 THEN 'x' ELSE 'z' END` is
emitted correctly, as a single clause whose condition is an `equal` call with
`output_type` boolean.
### Expected behavior
Every emitted `IfClause` should have a boolean condition and a `then` value.
Substrait has no switch-style construct that fits this case in general:
`SwitchExpression` exists, but its `IfValue.if` is a `Literal`, so it cannot
express `CASE a WHEN b + 1 THEN ...`, and DataFusion's own consumer currently
rejects it with `not_impl_err!("Switch expression not supported")`. The general
translation is therefore to desugar the base expression, emitting ` =
` as each clause condition:
```
ifs[0]: if = equal(a, 1), then = 'x'
ifs[1]: if = equal(a, 2), then = 'y'
else : 'z'
```
DataFusion matches a base expression with `=` semantics
([`compare_with_eq`](https://github.com/apache/datafusion/blob/517ac6d53133621c99a9a7dda470549a26348f6a/datafusion/physical-expr-common/src/datum.rs#L138)
uses Arrow's `eq`), so this preserves the plan's meaning, including a NULL WHEN
operand never matching. The producer already applies exactly this kind of
desugaring to `BETWEEN` in `from_between`.
A base `CASE` would then round trip as the equivalent searched `CASE`, keeping
its projection name and schema.
### Additional context
- Distinct from the round-trip failures collected under #16248: this round trip
succeeds. The problem is only visible by reading the emitted protobuf, because
producer and consumer share the private convention.
- Same shape as #25100 (`phase` left at `AGGREGATION_PHASE_UNSPECIFIED`): a
producer field whose value is wrong by the spec but invisible to DataFusion's
own consumer.
- Adding `SwitchExpression` support, on both sides, would preserve the base-CASE
structure for the subset where every WHEN operand is a literal. That is worth
doing separately; it does not cover the general case, and the emitted plan
should be valid regardless.
- Verified against the `algebra.proto` shipped in the pinned `substrait` 0.63.0
crate.
Contributor guide
Assessment
This issue has not been assessed yet.