andygrove opened a new issue, #6674:
URL: https://github.com/apache/datafusion-comet/issues/6674
### Describe the bug
When a scalar subquery sits inside an expression that Comet routes to the
JVM codegen dispatcher, such as a Scala UDF, or `to_json` on Spark 4, the task
fails with `requirement failed: Subquery subquery#N, [id=#M] has not finished`.
Spark runs the same query fine.
`CometScalaUDF.emitJvmCodegenDispatch` closure-serializes the bound
expression tree into the `JvmScalarUdf` proto while `CometExecRule` converts
the plan. That happens before the operator's `waitForSubqueries` runs, so the
serialized `ScalarSubquery` still has `updated = false`. On the executor,
`CometBatchKernelCodegen.generateSource` calls `ScalarSubquery.doGenCode`,
which calls `toLiteral` and fails its `require(updated)` check.
The comment in `CometBatchKernelCodegen.canHandle` says
`ExecSubqueryExpression` is accepted because the surrounding operator populates
the result before the bridge serializes the expression, but the serialization
happens at planning time. The existing test `ScalaUDF composed with reused
scalar subquery across projection and filter` keeps the subquery outside the
UDF, so no subquery ever ends up inside a kernel there.
The same plan-time serialization is on `branch-1.0` and `branch-1.1`, where
the dispatcher is also on by default, so the released versions are likely
affected too.
### Steps to reproduce
```scala
spark.udf.register("plus1", (x: Int) => x + 1)
sql("CREATE TABLE t USING parquet AS SELECT cast(id AS int) AS a FROM
range(10)")
sql("SELECT plus1((SELECT max(a) FROM t)) FROM t").collect()
```
```
java.lang.IllegalArgumentException: requirement failed: Subquery subquery#2,
[id=#79] has not finished
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.comet.codegen.CometBatchKernelCodegen$.generateSource(CometBatchKernelCodegen.scala:292)
at
org.apache.comet.codegen.CometBatchKernelCodegen$.compile(CometBatchKernelCodegen.scala:225)
at
org.apache.comet.udf.codegen.CometScalaUDFCodegen.lookupOrCompile(CometScalaUDFCodegen.scala:169)
```
On Spark 4.x, `SELECT to_json(named_struct('m', (SELECT max(a) FROM t)))
FROM t` fails the same way, since `to_json` lowers to an `Invoke` that goes
through the dispatcher.
Reproduced on `main` at 8a7e66441 with the default Spark 4.1 profile. The
same failure shows up on the 3.5 and 4.2 profiles.
### Expected behavior
The query returns Spark's result, either by keeping the subquery out of the
dispatched kernel or by falling back.
### Additional context
Adding `case _: ExecSubqueryExpression => true` to the rejection list in
`CometBatchKernelCodegen.canHandle` makes these queries fall back and match
Spark, and `CometCodegenSuite` still passes with that change. A longer-term
option is to pass the subquery value into the kernel as a data argument that
native evaluates, rather than serializing the subquery into the kernel.
#5889 makes this easier to hit, because struct-typed scalar subqueries,
including the ones `MergeScalarSubqueries` creates, stop falling back. For
example, `SELECT plus1((SELECT max(a) FROM t)), plus1((SELECT min(a) FROM t))
FROM t` falls back on `main` because of the struct-typed merged subquery, and
fails with #5889.
Related: #5572
--
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]
