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