Substrait producer emits non-boolean IfThen clauses for CASE with a base expression

[namanjain24-sudo]

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:

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 <base> WHEN <value> THEN ...,
encoding the base expression as a leading IfClause whose then is left unset
(from_case):

// 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):

// 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

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:

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(())
}

$ 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:

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:

$ 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 <base> = <value> 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
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 [Epic] A collection of Substrait conversion issues #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 Substrait producer emits AGGREGATION_PHASE_UNSPECIFIED for every aggregate and window function #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.

[namanjain24-sudo]

take

[alamb]

SELECT CASE a WHEN 1 THEN 'x' WHEN 2 THEN 'y' ELSE 'z' END FROM data
Produce the plan with to_substrait_plan and inspect the IfThen in the
projection directly, rather than round-tripping it through the consumer:

I don't understand the bug here

this ticket basically says that the substrait plan datafusion produces is wrong but doesn't say what is being affeted

Is this a compatibility problem?

What is the problem you are solving here?

[alamb]

so a DataFusion-to-DataFusion round trip is unaffected and no existing test
fails. An engine that reads the plan as Substrait defines it sees something
different.

What engine are you using?

[namanjain24-sudo]

Yes, compatibility, and this one reproduces on a DataFusion plan with nothing worked around. The
engine is substrait-java io.substrait:core:0.103.0.

SELECT CASE a WHEN 1 THEN 'x' WHEN 2 THEN 'y' ELSE 'z' END FROM data, one nullable i64 column,
through to_substrait_plan, produces an IfThen shaped like this:

clause 0:  if = SELECTION (column a)   then = <unset>
clause 1:  if = LITERAL i64 1          then = "x"
clause 2:  if = LITERAL i64 2          then = "y"
else                                          "z"

Handing that protobuf straight to ProtoPlanConverter:

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())). It converts then
unconditionally and clause 0 does not have one. This plan registers no functions at all, so unlike
#25049 there is no extension-URN problem in the way: it is a straight failure on what we emit today.

Two things are wrong here and only one of them throws, which is worth separating:

1. Clause 0 carries the base expression with no then. That is the hard failure above.
2. The conditions are not boolean. They are a, then 1, then 2, all i64. The spec's if
   expression is if <boolean expression> then .... I checked whether substrait-java objects to
   that on its own — drop clause 0 and feed it the remaining two clauses, and it accepts them.
   It does not type-check if. So that half is a silent divergence, not an error, which is also why
   no round trip anywhere would surface it.

With #25191 every clause becomes equal(a, <value>) carrying output_type: bool and a then, and
the same query is accepted, giving Struct{nullable=false, fields=[Str{nullable=false}]}. That plan
does need the #11545 workaround, because it now registers the equal function and we still write
extension_urn_reference: u32::MAX.

On SwitchExpression, since it is the obvious alternative and I should say why I did not use it:
SwitchExpression.IfValue.if is a Literal, so it cannot express CASE a WHEN b + 1, and our own
consumer answers not_impl_err!("Switch expression not supported")
(logical_plan/consumer/substrait_consumer.rs:348). Emitting it would need consumer work first and
would still need this desugaring as the fallback, so I left it as a follow-up rather than folding it
into the same change.