fix: reject scalar subqueries inside a codegen-dispatch kernel

[xhumanoid]

A subquery inside a dispatched subtree fails at runtime. SELECT addOne((SELECT max(v) FROM t2)) FROM t dies with "Subquery ... has not finished" from ScalarSubquery.doGenCode on the executor, after the plan has already committed to the kernel and can no longer fall back.

canHandle claimed to support ExecSubqueryExpression on the grounds that the surrounding operator's waitForSubqueries populates the subquery's result before the dispatcher serializes the expression. That ordering never held: the dispatcher closure-serializes the tree during planning, while result is populated at execution time, so the deserialized copy on the executor always has an unset result. The existing reused-subquery test passes only because its subqueries sit under native parents (addOne(x) + (SELECT ...)), where they are evaluated natively and never reach a kernel.

canHandle now refuses them, which turns the runtime failure into a plan-time Spark fallback, and the comment is corrected to say so.

Which issue does this PR close?

Closes #.

How are these changes tested?

add regression test
failed without fix

[mbutrovich]

Thanks for looking at this @xhumanoid. Could we get a deeper investigation into this? I didn't hallucinate that documentation that scalar subqueries work in a codegen'd expression. We have a test for it, but it's specifically for a reused subquery. I wonder if that changes the execution semantics somehow.

datafusion-comet/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala

Line 1146 in 481aefe

test("ScalaUDF composed with reused scalar subquery across projection and filter") {

I'm not sure why the non-reuse scenario would fail differently.

[xhumanoid]

@mbutrovich ok, will check
thanks for pointing one more place

found this bug during investigation of HiveSimpleUDF support
will create pr separately, started from Spark 3.5 they have doCodeGen method
https://issues.apache.org/jira/browse/SPARK-42051
https://issues.apache.org/jira/browse/SPARK-42052

I used ScalaUDF as a base harness and catch it on some of my tests

main different with existing unit test:

existing do evaluation of subquery separately and ScalaUDF independent, it's not a part of UDF kernel:

SELECT addOne(x) + (SELECT max(v) FROM t2) AS r
FROM t 
WHERE addOne(x) < (SELECT max(v) FROM t2) * 2

in my case subquery is a part of UDF parameter

SELECT addOne(   (SELECT max(v) FROM t2)  ) AS r 
FROM t

so during code generation we can't use the value returned from subquery

without patch exception from driver

  Caused by: java.lang.IllegalArgumentException: requirement failed: Subquery subquery#2254, [id=#3678] has not finished
  	at scala.Predef$.require(Predef.scala:337)
  	at org.apache.spark.sql.execution.ScalarSubquery.toLiteral(subquery.scala:108)
  	at org.apache.spark.sql.execution.ScalarSubquery.doGenCode(subquery.scala:104)
  	at org.apache.spark.sql.catalyst.expressions.Expression.$anonfun$genCode$3(Expression.scala:229)
  	at scala.Option.getOrElse(Option.scala:201)
  	at org.apache.spark.sql.catalyst.expressions.Expression.genCode(Expression.scala:224)
  	at org.apache.spark.sql.catalyst.expressions.KnownNotNull.doGenCode(constraintExpressions.scala:48)
  	at org.apache.spark.sql.catalyst.expressions.Expression.$anonfun$genCode$3(Expression.scala:229)
  	at scala.Option.getOrElse(Option.scala:201)
  	at org.apache.spark.sql.catalyst.expressions.Expression.genCode(Expression.scala:224)
  	at org.apache.spark.sql.catalyst.expressions.ScalaUDF.$anonfun$doGenCode$2(ScalaUDF.scala:1118)
  	at scala.collection.immutable.ArraySeq.map(ArraySeq.scala:75)
  	at scala.collection.immutable.ArraySeq.map(ArraySeq.scala:35)
  	at org.apache.spark.sql.catalyst.expressions.ScalaUDF.doGenCode(ScalaUDF.scala:1118)
  	at org.apache.spark.sql.catalyst.expressions.Expression.$anonfun$genCode$3(Expression.scala:229)
  	at scala.Option.getOrElse(Option.scala:201)

query itself still reused as expected

        +- CometProject [r#2255], [if (isnull(Subquery subquery#2254, [id=#3677])) null else addOne(knownnotnull(ReusedSubquery Subquery subquery#2254, [id=#3677])) AS r#2255]

it mention explicitly in the patch

Note this only concerns a subquery *inside* the dispatched subtree: one that is merely a sibling under a
native parent (`udf(x) + (SELECT ...)`) is evaluated natively and never reaches a kernel.

but udf ( select ... ) have to fallback to spark for now

[xhumanoid]

SELECT addOne(   (SELECT 123)  ) AS r 
FROM t

works as expected without patch, so i will keep investigation

ok, it folded into constant before any codegen, codegen itself can happens without any issue

[sunchao]

Correctness

This fixes the dispatch boundary for expressions such as addOne((SELECT max(v) FROM t2)). The bound expression was serialized during planning, before Spark populated the subquery result. The new ExecSubqueryExpression check returns a fallback reason before those bytes are created, so the enclosing operator can stay in Spark. I found no new P1/P2 issue in the change at a10442390b90 against 481aefea9c60.

The maintained Spark 3.5 and 4.0 sources both require the scalar subquery result to be updated before evaluation or code generation. Spark supplies null for zero rows and raises an error for more than one row. Falling back preserves that lifecycle and the result's original type, null handling, and surrounding ANSI behavior. The check is independent of the result type or value.

find visits the dispatched expression and its children, so nesting and a reused subquery plan do not bypass the guard. The native sibling case, addOne(x) + (SELECT ...), still lowers the UDF and scalar subquery separately. Correlated scalar subqueries rewritten into joins expose ordinary attributes to the dispatcher and are not rejected merely because the SQL contained a subquery.

Validation

The new regression uses a table-backed aggregate, compares separately built Spark and Comet plans, and checks that the dispatcher did not compile or reuse a kernel. Using the Scala suite is appropriate because it registers a Scala UDF. The existing sibling-subquery case remains in the suite. I reviewed these paths and the full three-file diff, but did not execute the regression locally.

At the review cutoff, CI, CodeQL, and the Delta gate require approval and have no jobs. The label check passed. There is no native/JNI test result to credit. Canonical source coverage is limited to Spark 3.5 and 4.0, with 3.4 and 4.1 unverified.

Performance

The extra class match uses the existing expression-tree walk. It adds no per-row or per-batch work and avoids serialization and kernel compilation for the rejected shape. The enclosing operator falls back to Spark, so this is a correctness repair with no measured speedup claim. I found no material performance issue requiring a benchmark for this guard.

Design

The shared capability check is the right place to reject the unsupported lifecycle. It covers all callers of the dispatcher and propagates through the existing fallback path. Supporting these subqueries inside kernels would need a separate design for passing resolved values across the dispatch boundary. That is not necessary to make this repair correct.

Abstraction & complexity

The production change adds one case to an existing predicate and introduces no new state or abstraction. The small assertNoCodegen helper mirrors the suite's existing positive assertion and uses the same counters. The change stays focused on the failure boundary.

[xhumanoid]

@mbutrovich ping