From a10442390b90d414a1d406b9a57bc0cb035ab886 Mon Sep 17 00:00:00 2001 From: Aliaksei Dziomin Date: Mon, 14 Sep 2026 22:18:32 +0700 Subject: [PATCH] fix: reject scalar subqueries inside a codegen-dispatch kernel 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. --- .../codegen/CometBatchKernelCodegen.scala | 16 ++++++++++------ .../apache/comet/CometCodegenAssertions.scala | 13 +++++++++++++ .../org/apache/comet/CometCodegenSuite.scala | 18 ++++++++++++++++++ 3 files changed, 41 insertions(+), 6 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala index 83fbca6b635..3b51b6d3919 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -156,11 +156,14 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come // instance with a single `init(partitionIndex)` call, so `Rand` / `MonotonicallyIncreasingID` // state advances correctly across batches. // - // `ExecSubqueryExpression` (`ScalarSubquery`, `InSubqueryExec`) is accepted: the surrounding - // Comet operator's inherited `SparkPlan.waitForSubqueries` populates the subquery's - // `result` field before evaluation. The closure serializer captures that value into the - // arg-0 bytes, and the dispatcher keys its compile cache on those bytes, so distinct subquery - // results produce distinct cache entries. + // `ExecSubqueryExpression` (`ScalarSubquery`, `InSubqueryExec`) is rejected. Its `result` field + // is populated by the surrounding operator's `SparkPlan.waitForSubqueries` at *execution* time, + // whereas the dispatcher closure-serializes the tree into the arg-0 bytes during planning. The + // deserialized copy on the executor therefore still has an unset `result`, and + // `ScalarSubquery.doGenCode` fails the kernel compile with "Subquery ... has not finished". + // Refusing here turns that runtime failure into a plan-time Spark fallback. 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. // // `Unevaluable`: rejected by default. `isCodegenInertUnevaluable` exempts version-specific // leaves that are `Unevaluable` but never invoked by codegen (e.g. Spark 4.0's @@ -168,6 +171,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come boundExpr.find { case _: org.apache.spark.sql.catalyst.expressions.aggregate.AggregateFunction => true case _: org.apache.spark.sql.catalyst.expressions.Generator => true + case _: org.apache.spark.sql.execution.ExecSubqueryExpression => true case u: Unevaluable if isCodegenInertUnevaluable(u) => false case _: Unevaluable => true case _ => false @@ -175,7 +179,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come case Some(bad) => return Some( s"codegen dispatch: expression ${bad.getClass.getSimpleName} not supported " + - "(aggregate, generator, or unevaluable)") + "(aggregate, generator, subquery, or unevaluable)") case None => } val badRef = boundExpr.collectFirst { diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala b/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala index bce8bfc5986..f28d8f44a26 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenAssertions.scala @@ -42,6 +42,19 @@ trait CometCodegenAssertions { s"expected codegen dispatcher activity, got $after") } + /** + * Asserts the dispatcher did not run during `f`, for the cases where the expression is expected + * to be routed back to Spark rather than into a kernel. + */ + protected def assertNoCodegen(f: => Unit): Unit = { + CometScalaUDFCodegen.resetStats() + f + val after = CometScalaUDFCodegen.stats() + assert( + after.compileCount + after.cacheHitCount == 0, + s"expected no codegen dispatcher activity, got $after") + } + /** * Asserts the composed subtree fused into one kernel signature, not N (one per sub-expression). * Uses the JVM-wide signature set rather than `compileCount` because per-task `boundExpr` diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 5c1f957453a..e978dec1216 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -1143,6 +1143,24 @@ class CometCodegenSuite checkSparkAnswerAndOperator(df) } + test("a scalar subquery inside the dispatched subtree falls back to Spark") { + // The dispatcher serializes the tree at plan time, but a subquery's `result` is populated by + // the surrounding operator's `waitForSubqueries` at execution time. The deserialized copy on + // the executor therefore has no result, and `ScalarSubquery.doGenCode` used to fail the kernel + // compile with "Subquery ... has not finished" -- a runtime failure, after the plan had already + // committed to the kernel. `canHandle` now refuses it, so the operator falls back cleanly. + spark.udf.register("addOne", (i: Int) => i + 1) + withTable("t", "t2") { + sql("CREATE TABLE t (x INT) USING parquet") + sql("INSERT INTO t VALUES (1), (2), (3), (4), (5)") + sql("CREATE TABLE t2 (v INT) USING parquet") + sql("INSERT INTO t2 VALUES (2), (4)") + assertNoCodegen { + checkSparkAnswer(sql("SELECT addOne((SELECT max(v) FROM t2)) AS r FROM t")) + } + } + } + test("ScalaUDF composed with reused scalar subquery across projection and filter") { // The same scalar subquery appears in two sites: the projection (which the dispatcher // compiles into a fused kernel) and the filter (a separate operator). Each site holds its