Codegen dispatcher single-ordinal null short-circuit swallows ANSI errors from literal subtrees

[andygrove]

Describe the bug

The JVM codegen dispatcher's null short-circuit can swallow an ANSI error that Spark raises, returning a row where Spark fails the query.

CometBatchKernelCodegen.canShortCircuitNulls allows the short-circuit when the dispatched tree reads exactly one input ordinal, on the reasoning quoted in its own scaladoc:

One ordinal: there is nothing left for Spark to evaluate ahead of that ordinal's own null check.

That is not true. A literal-only subtree between the root and the ordinal can still raise. ConstantFolding normally folds such a subtree away, but it deliberately leaves it in place when the subtree throws and sits inside a conditional branch (it tags the node FAILED_TO_EVALUATE and moves on), so the throwing expression survives into the physical plan.

This is the residual hole in #5219 / #5218. That fix added rootChildrenAreLeaves for the multi-ordinal case but left the single-ordinal branch unguarded.

Found by @sunchao while reviewing #5607.

Steps to reproduce

Spark 4.1, ANSI enabled (the default on 4.x):

CREATE TABLE t (flag BOOLEAN, n INT) USING parquet;
INSERT INTO t VALUES (true, NULL), (false, NULL);

SELECT IF(flag, upper(substring('abc', CAST(1L DIV 0L AS INT), n)), NULL) FROM t;

Spark raises [DIVIDE_BY_ZERO]. Comet returns a row.

Three things have to line up, and all three are ordinary:

• ConstantFolding refuses to fold 1L DIV 0L because it sits under an If branch, so the throwing literal reaches execution.
• TernaryExpression.nullSafeCodeGen emits Substring's pos code before it tests len's null, so Spark evaluates the division even though n is NULL.
• Upper, Substring, Cast and IntegralDivide are all null-intolerant and the dispatched tree reads exactly one ordinal, so canShortCircuitNulls takes the single-ordinal branch and the kernel writes NULL before ev.code runs.

upper is just a convenient witness. Any dispatched null-intolerant root with a throwing foldable subtree between it and its single input reproduces it, which is most of the ~70 expressions that route through the dispatcher. #5607 (routing length / bit_length / octet_length on BinaryType through the dispatcher) exposes three more:

SELECT IF(flag, length(substring(X'00', CAST(1L DIV 0L AS INT), n)), 0) FROM t;
SELECT IF(flag, bit_length(substring(X'00', CAST(1L DIV 0L AS INT), n)), 0) FROM t;
SELECT IF(flag, octet_length(substring(X'00', CAST(1L DIV 0L AS INT), n)), 0) FROM t;

Expected behavior

Comet raises the same [DIVIDE_BY_ZERO] Spark does.

Suggested fix

Tighten canShortCircuitNulls: in addition to the existing conditions, require that no node in the tree other than a Literal is foldable. After ConstantFolding has run, a surviving foldable non-Literal node is precisely one that threw during folding, which is exactly the dangerous case.

This keeps the existing fast paths, none of which contain a foldable non-Literal node: upper(substring(s, 1, 2)), pmod(a, b), a + b, conv(a, b, c), make_timestamp(...).

Regression test

test("single-input short-circuit does not swallow an ANSI error from a literal subtree") {
  // `canShortCircuitNulls` assumes a single input ordinal leaves Spark nothing to evaluate ahead
  // of that ordinal's null check. Not true when the tree carries a foldable subtree that throws:
  // `ConstantFolding` leaves it in place under a conditional branch, and `Substring`'s `pos` is
  // evaluated before `len`'s null is tested, so Spark raises on the (true, NULL) row while the
  // kernel short-circuits and returns NULL.
  withTable("t") {
    sql("CREATE TABLE t (flag BOOLEAN, n INT) USING parquet")
    sql("INSERT INTO t VALUES (true, NULL), (false, NULL)")
    withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
      Seq(
        "upper(substring('abc', CAST(1L DIV 0L AS INT), n))",
        "length(substring(X'00', CAST(1L DIV 0L AS INT), n))",
        "bit_length(substring(X'00', CAST(1L DIV 0L AS INT), n))",
        "octet_length(substring(X'00', CAST(1L DIV 0L AS INT), n))").foreach { e =>
        val (sparkErr, cometErr) =
          checkSparkAnswerMaybeThrows(sql(s"SELECT IF(flag, $e, NULL) FROM t"))
        assert(sparkErr.isDefined, s"$e: Spark no longer raises, the row is not a witness")
        assert(cometErr.isDefined, s"$e: Comet returned a value where Spark raised")
        assert(cometErr.get.getMessage.contains("DIVIDE_BY_ZERO"))
      }
    }
  }
}

The length / bit_length / octet_length cases only apply once #5607 lands.

Additional context

Verified end to end on Spark 4.1 against the #5607 branch. The upper case is unaffected by that PR and reproduces on main.