andygrove opened a new issue, #6676:
URL: https://github.com/apache/datafusion-comet/issues/6676
### Describe the bug
When the input to `CometTakeOrderedAndProjectExec` has more than one
partition and the sort order contains a scalar subquery, the query fails with
`CometRuntimeException: Subquery N not found for plan M`. Spark returns the
right rows.
The per-partition top-K plan in
`CometTakeOrderedAndProjectExec.doExecuteColumnar` is built from
`getTopKNativePlan(child.output, sortOrder, child, limit)` and run through
`CometExec.getCometIterator`, but nothing registers the operator's subqueries
against that iterator's id. The final single-partition plan below it does call
`setSubqueries(it.id, this)` and `cleanSubqueries` (that was the fix for #749),
so a subquery in the project list works and one in the sort keys does not. With
a single input partition the per-partition plan is skipped, which is why the
existing test `subquery execution under CometTakeOrderedAndProjectExec should
not fail` doesn't catch it.
### Steps to reproduce
```scala
sql("CREATE TABLE u USING parquet AS SELECT cast(id AS int) AS a FROM
range(50)")
spark.range(2000).selectExpr("id", "cast(id % 100 AS int) AS a")
.repartition(4).write.saveAsTable("m")
sql("SELECT id FROM m ORDER BY a + (SELECT max(a) FROM u), id LIMIT
7").collect()
```
```
org.apache.comet.CometRuntimeException: Subquery 48 not found for plan 11.
```
Reproduced on `main` at ba08acd81 with the default Spark 4.1 profile.
### Expected behavior
The same seven rows Spark returns.
### Additional context
Registering the subqueries on the per-partition iterator the same way the
final plan does, with `setSubqueries(it.id, this)` plus a task-completion
`cleanSubqueries(it.id, this)`, fixes it locally. The existing
`TakeOrderedAndProjectExec` tests in `CometExecSuite` still pass with that
change.
#5889 makes this easier to hit, because struct-typed scalar subqueries stop
falling back. `ORDER BY a + (SELECT max(a) FROM u), b + (SELECT min(b) FROM u),
id LIMIT 7` is merged into one struct subquery by `MergeScalarSubqueries`. It
falls back on `main` and fails with that PR.
--
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]
