From c512f20d6d1ee759b516380ef671ba245d32d5cf Mon Sep 17 00:00:00 2001 From: Scott Schenkein Date: Sat, 4 Jul 2026 11:09:55 -0400 Subject: [PATCH 1/3] feat: partial project fallback (JVM expression detour) When a projection or filter contains an expression with no native translation, evaluate just that subtree in the JVM via Comet's Arrow-direct codegen dispatcher instead of falling the whole operator back to Spark. The operator, and the native island around it, stays native. Gated behind spark.comet.exec.jvmDetour.enabled (default off, experimental). - CometConf: new COMET_EXEC_JVM_DETOUR_ENABLED config. - QueryPlanSerde: a withJvmDetour thread-local gate plus a last-resort .orElse hook in exprToProtoInternal that retries any unsupported node via emitJvmCodegenDispatch and clears FALLBACK_REASONS on success. - CometScalaUDF: refuse subquery-bearing trees in emitJvmCodegenDispatch (they are closure-serialized at plan time, before waitForSubqueries) so they fall back cleanly instead of failing in the kernel. - operators: wrap CometProjectExec/CometFilterExec convert bodies in withJvmDetour so the detour is scoped to those two contexts. - tests: CometPartialProjectFallbackSuite (serde, parity, island preservation, gate scoping), CometCodegenFuzzSuite hook fuzz, and CometPartialProjectFallbackBenchmark. No proto or native changes: reuses the existing JvmScalarUdfExpr JVM callback machinery. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01JnvYuZKVMTeY2rewk94ZFe --- .../scala/org/apache/comet/CometConf.scala | 16 + .../apache/comet/serde/CometScalaUDF.scala | 18 +- .../apache/comet/serde/QueryPlanSerde.scala | 66 +++ .../apache/spark/sql/comet/operators.scala | 5 +- .../apache/comet/CometCodegenFuzzSuite.scala | 52 +++ .../CometPartialProjectFallbackSuite.scala | 422 ++++++++++++++++++ ...CometPartialProjectFallbackBenchmark.scala | 195 ++++++++ 7 files changed, 771 insertions(+), 3 deletions(-) create mode 100644 spark/src/test/scala/org/apache/comet/CometPartialProjectFallbackSuite.scala create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometPartialProjectFallbackBenchmark.scala diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 8e47151358..6900fcaf89 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -381,6 +381,22 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_EXEC_JVM_DETOUR_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.jvmDetour.enabled") + .category(CATEGORY_EXEC) + .doc("Experimental. When a projection or filter contains an expression with no native " + + "translation, evaluate just that subtree in the JVM via Comet's Arrow-direct codegen " + + "dispatcher (Spark's own `doGenCode`/`eval` run inside the Comet pipeline) instead of " + + "falling the whole operator back to Spark. Only the argument columns of the detoured " + + "subtree cross to the JVM over Arrow FFI, so the pipeline above and below stays native. " + + "Requires " + COMET_SCALA_UDF_CODEGEN_ENABLED.key + "=true. In a filter predicate the " + + "detoured subexpression is evaluated for all rows, the same as any native Comet " + + "predicate; under ANSI mode this can differ from Spark's per-row AND/OR short-circuit " + + "(e.g. a subexpression that raises, guarded by a preceding conjunct). Disabled by " + + "default while the feature is stabilized.") + .booleanConf + .createWithDefault(false) + val COMET_EXEC_SHUFFLE_WITH_HASH_PARTITIONING_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.native.shuffle.partitioning.hash.enabled") .category(CATEGORY_SHUFFLE) diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index 3df275f057..d4a64d5902 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -20,7 +20,7 @@ package org.apache.comet.serde import org.apache.spark.SparkEnv -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, Literal, RuntimeReplaceable, ScalaUDF} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, Literal, PlanExpression, RuntimeReplaceable, ScalaUDF} import org.apache.spark.sql.types.BinaryType import org.apache.comet.CometConf @@ -88,6 +88,22 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { case other => other } + // A subquery inside the dispatched tree cannot be evaluated in the kernel. The bound tree is + // closure-serialized into arg 0 here at plan time, before the enclosing operator's + // `waitForSubqueries` populates the subquery result at execution, so the kernel would + // deserialize an unpopulated subquery and its `eval` would throw "Subquery ... has not + // finished". Refuse so the operator falls back cleanly. Match the catalyst base + // `PlanExpression` (rather than the narrower `ExecSubqueryExpression`) so every subquery form + // is caught. (A subquery that is a *sibling* of the dispatched node, e.g. + // `dispatched(x) + (SELECT ...)`, is unaffected: it is evaluated natively, never in this tree.) + if (target.exists(_.isInstanceOf[PlanExpression[_]])) { + withFallbackReason( + expr, + "codegen dispatch: a subquery inside the dispatched expression is not supported " + + "(its result is not populated at plan-time serialization)") + return None + } + // Bind against only the AttributeReferences the tree actually reads, so ordinals align with // the data args we ship. val attrs = target.collect { case a: AttributeReference => a }.distinct diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 7146eaec9b..a108413cb2 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -36,6 +36,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.comet.CometConf +import org.apache.comet.CometExplainInfo import org.apache.comet.CometSparkSessionExtensions.{withFallbackReason, withInfo} import org.apache.comet.expressions._ import org.apache.comet.parquet.CometParquetUtils @@ -50,6 +51,44 @@ import org.apache.comet.shims.{CometExprShim, CometTypeShim} */ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { + /** + * True while serializing an expression context (a projection or a filter predicate) where a + * last-resort JVM expression detour is permitted. Set by [[withJvmDetour]] around the bodies of + * `CometProjectExec.convert` / `CometFilterExec.convert`. Scoping the detour this way keeps it + * out of join keys, sort keys, aggregate/group expressions, and partitioning expressions until + * those contexts are benchmarked, without threading a flag through every serde method. + * + * A plain (non-inheritable) `ThreadLocal`, deliberately NOT `scala.util.DynamicVariable`: that + * is backed by an `InheritableThreadLocal`, so any thread spawned while a `withJvmDetour` scope + * is active would copy `true` at construction and never restore it, leaking the detour into + * unrelated serde contexts (join/sort/aggregate/partitioning) on that child thread. Serde runs + * synchronously on the planning thread, so a non-inheritable thread-local with save/restore is + * both correct and leak-free. + */ + private val jvmDetourContext: ThreadLocal[Boolean] = ThreadLocal.withInitial(() => false) + + /** + * Evaluate `body` with the JVM expression detour permitted for any [[exprToProto]] calls made + * within it. Used by the projection and filter operator serdes; see [[jvmDetourContext]]. + */ + def withJvmDetour[T](body: => T): T = { + val previous = jvmDetourContext.get() + jvmDetourContext.set(true) + try body + finally jvmDetourContext.set(previous) + } + + /** + * True when the last-resort JVM expression detour may fire for the node currently being + * serialized: we are inside a detour-eligible context ([[jvmDetourContext]]) and the feature is + * enabled via [[CometConf.COMET_EXEC_JVM_DETOUR_ENABLED]]. The dispatcher's own gate + * ([[CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED]]) and plan-time eligibility + * (`CometBatchKernelCodegen.canHandle`) are enforced inside + * `CometScalaUDF.emitJvmCodegenDispatch`. + */ + private def jvmDetourAllowed: Boolean = + jvmDetourContext.get() && CometConf.COMET_EXEC_JVM_DETOUR_ENABLED.get() + private[comet] val arrayExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map( // ArrayAppend is a concrete expression only on Spark 3.x. Spark 4.0+ marks it // RuntimeReplaceable and rewrites it to ArrayInsert before serde, so this entry is @@ -847,6 +886,33 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { None } }) + .orElse { + // Last-resort JVM expression detour. Any `None` above — no registered handler, an + // `Unsupported`/opt-out `Incompatible` support level, or a handler whose `convert` + // returned `None` — means this node has no native translation. Inside a detour-eligible + // context (projection / filter, see `jvmDetourAllowed`) retry the *outermost* unsupported + // node as a JVM detour: `emitJvmCodegenDispatch` ships its argument columns over Arrow FFI + // and evaluates it with Spark's own codegen inside the Comet pipeline, so the operator — + // and the native island around it — stays native. Supported ancestors keep converting + // natively and reference this detour's output. If the detour cannot fire (feature off, + // dispatcher off, or `CometBatchKernelCodegen.canHandle` rejects the tree), `None` + // propagates unchanged and the operator falls back exactly as before. + if (jvmDetourAllowed) { + CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding).map { proto => + // A successful detour accelerates this node, so it is not a fallback. The conversion + // chain above already recorded a fallback reason (this hook runs last, unlike the + // `CodegenDispatchFallback` path which records a reason only when dispatch fails). + // Clear that reason so the node carries only the `withInfo` tag, preserving the + // `withInfo`-xor-`withFallbackReason` invariant the rest of the serde upholds and + // keeping a detoured sibling out of an operator-level fallback rollup. + expr.unsetTagValue(CometExplainInfo.FALLBACK_REASONS) + withInfo(expr, "Expression evaluated in the JVM via the codegen dispatcher") + proto + } + } else { + None + } + } .map { protoExpr => // Attach QueryContext and expr_id to the expression val builder = protoExpr.toBuilder diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index e4d6b53770..19bf1f7eb5 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -925,7 +925,8 @@ object CometProjectExec extends CometOperatorSerde[ProjectExec] { op: ProjectExec, builder: Operator.Builder, childOp: Operator*): Option[OperatorOuterClass.Operator] = { - val exprs = op.projectList.map(exprToProto(_, op.child.output)) + val exprs = + QueryPlanSerde.withJvmDetour(op.projectList.map(exprToProto(_, op.child.output))) if (exprs.forall(_.isDefined) && childOp.nonEmpty) { val projectBuilder = OperatorOuterClass.Projection @@ -985,7 +986,7 @@ object CometFilterExec extends CometOperatorSerde[FilterExec] { op: FilterExec, builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { - val cond = exprToProto(op.condition, op.child.output) + val cond = QueryPlanSerde.withJvmDetour(exprToProto(op.condition, op.child.output)) if (cond.isDefined && childOp.nonEmpty) { val filterBuilder = OperatorOuterClass.Filter diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenFuzzSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenFuzzSuite.scala index 5c40dccaa3..73c468baf3 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenFuzzSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenFuzzSuite.scala @@ -413,6 +413,58 @@ class CometCodegenFuzzSuite } } + // --- Detour hook fuzz: force serde-unsupported roots by disabling a registered built-in's serde + // so the last-resort `exprToProto` fallthrough (partial project fallback, + // `spark.comet.exec.jvmDetour.enabled`) is the path under test rather than the ScalaUDF entry + // point. Downstream (`emitJvmCodegenDispatch` -> Janino kernel) is shared, so this fuzzes the + // hook's routing and the FFI marshaling of the detoured argument columns. ANSI is disabled so + // `abs(Long.MinValue)`-style overflow wraps identically in both engines rather than throwing. --- + + private def withDetourAbsDisabled(f: => Unit): Unit = + withSQLConf( + CometConf.COMET_EXEC_JVM_DETOUR_ENABLED.key -> "true", + CometConf.getExprEnabledConfigKey("Abs") -> "false", + SQLConf.ANSI_ENABLED.key -> "false")(f) + + test("detour hook fuzz: disabled abs over every numeric column of the random schema") { + val numericFields = spark.table("t1").schema.fields.filter { f => + f.dataType match { + case ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType => true + case _: DecimalType => true + case _ => false + } + } + assert(numericFields.nonEmpty, "expected at least one numeric column in the random schema") + withDetourAbsDisabled { + for (field <- numericFields) { + withClue(s"column ${field.name}: ${field.dataType}: ") { + assertCodegenRan { + checkSparkAnswerAndOperator(s"SELECT abs(${field.name}) FROM t1") + } + } + } + } + } + + test("detour hook fuzz: disabled abs over the decimal precision/scale/null-density sweep") { + withDetourAbsDisabled { + for { + density <- nullDensities + (precision, scale) <- decimalShapes + } { + val seed = (((precision * 31L) + scale) * 31L + density.hashCode) * 17L + 3 + val values = generateDecimals(seed, precision, scale, density) + withDecimalTable(s"DECIMAL($precision, $scale)", values) { + withClue(s"precision=$precision scale=$scale nullDensity=$density: ") { + assertCodegenRan { + checkSparkAnswerAndOperator(sql("SELECT abs(d) FROM t")) + } + } + } + } + } + } + /** * Randomized output-writer coverage (#4539). Generates a random nested output type and a random * catalyst value of that type, wraps it in a `Literal`, and drives it through the kernel output diff --git a/spark/src/test/scala/org/apache/comet/CometPartialProjectFallbackSuite.scala b/spark/src/test/scala/org/apache/comet/CometPartialProjectFallbackSuite.scala new file mode 100644 index 0000000000..c28a145aa1 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometPartialProjectFallbackSuite.scala @@ -0,0 +1,422 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, DataFrame} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Sentences} +import org.apache.spark.sql.comet.{CometFilterExec, CometProjectExec} +import org.apache.spark.sql.execution.{ProjectExec, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.types.{ArrayType, StringType} + +import org.apache.comet.serde.QueryPlanSerde +import org.apache.comet.udf.codegen.CometScalaUDFCodegen + +/** + * Tests for the partial project fallback (JVM expression detour). When a projection or filter + * contains an expression with no native translation, [[CometConf.COMET_EXEC_JVM_DETOUR_ENABLED]] + * routes just that subtree through the Arrow-direct codegen dispatcher instead of falling the + * whole operator back to Spark. + * + * Two ways of producing an "unsupported" expression are used: + * - `sentences(str)`, which Comet has no native serde for (so it exercises the `case _ =>` + * fallthrough the hook backstops), is deterministic, and is a `CodegenFallback` returning + * `array>` (interpreted eval + a nested-complex result across the FFI + * boundary). + * - a registered built-in whose per-expression serde is disabled + * (`spark.comet.expression..enabled=false`), which routes a genuinely-supported + * expression through the same fallthrough. This lets the parity tests below ship a chosen + * input type (decimal, timestamp, array, map, a subquery, a nondeterministic child) across + * the FFI boundary through the hook while staying cheap and stable as native coverage + * evolves. + */ +class CometPartialProjectFallbackSuite + extends CometTestBase + with AdaptiveSparkPlanHelper + with CometCodegenAssertions { + + override protected def sparkConf: SparkConf = + super.sparkConf + .set(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key, "true") + + private def withStringCol(values: String*)(f: => Unit): Unit = { + withTable("t") { + sql("CREATE TABLE t (s STRING) USING parquet") + val rows = values + .map(v => if (v == null) "(NULL)" else s"('${v.replace("'", "''")}')") + .mkString(", ") + if (values.nonEmpty) sql(s"INSERT INTO t VALUES $rows") + f + } + } + + private def withDetour[T](enabled: Boolean)(f: => T): T = + withSQLConf(CometConf.COMET_EXEC_JVM_DETOUR_ENABLED.key -> enabled.toString)(f) + + /** Whether the (already-executed) DataFrame's plan kept a native `CometProjectExec`. */ + private def hasCometProject(df: DataFrame): Boolean = + collect(stripAQEPlan(df.queryExecution.executedPlan)) { case p: CometProjectExec => + p + }.nonEmpty + + test("unsupported projection expression stays native with the detour enabled") { + withStringCol("Hello World. Foo bar.", "One sentence here.", null) { + withDetour(enabled = true) { + assertCodegenRan { + checkSparkAnswerAndOperator( + sql("SELECT sentences(s) FROM t"), + includeClasses = Seq(classOf[CometProjectExec])) + } + } + } + } + + test("same expression falls the operator back to Spark with the detour disabled (default)") { + withStringCol("Hello World. Foo bar.", null) { + withDetour(enabled = false) { + CometScalaUDFCodegen.resetStats() + // Results still match Spark (the projection just runs on the JVM Spark path), but the + // project operator is NOT accelerated and the dispatcher never fires. + val df = sql("SELECT sentences(s) FROM t") + checkSparkAnswer(df) + assert( + !hasCometProject(df), + "expected the projection to fall back to Spark with the detour disabled") + val after = CometScalaUDFCodegen.stats() + assert( + after.compileCount == 0 && after.cacheHitCount == 0, + s"expected no dispatcher activity with the detour disabled, got $after") + } + } + } + + test("detour fires at the outermost unsupported node, keeping supported ancestors native") { + // `size(sentences(s))`: `size` has a native translation, only `sentences` does not. The hook + // must detour `sentences` alone and let the native `size` reference its output. Proof: a + // compiled kernel exists whose signature is `sentences`'s — one `VarCharVector` input and an + // `array>` output. If the detour granularity were the whole top-level + // expression, the only kernel would have an IntegerType output (`size`'s) and this signature + // would be absent. Presence is a robust check against the append-only JVM-wide signature set. + withStringCol("Hello World. Foo bar.", "Solo.", null) { + withDetour(enabled = true) { + assertCodegenRan { + checkSparkAnswerAndOperator( + sql("SELECT size(sentences(s)) FROM t"), + includeClasses = Seq(classOf[CometProjectExec])) + } + // Structural match on the signature (ignoring `containsNull` flags, which vary by + // context): one VarCharVector input, output `array>` — the sentences subtree. + val sigs = CometScalaUDFCodegen.snapshotCompiledSignatures() + val present = sigs.exists { + case (inputs, ArrayType(ArrayType(StringType, _), _)) => + inputs.map(_.getSimpleName) == IndexedSeq("VarCharVector") + case _ => false + } + assert( + present, + "expected a kernel (VarCharVector -> array>) for the sentences subtree " + + s"only; cache had ${sigs.map { case (c, d) => (c.map(_.getSimpleName), d) }}") + } + } + } + + test("unsupported expression in a filter predicate keeps the filter native") { + // The detour gate is set around `CometFilterExec.convert` as well. `size(sentences(s)) > 0` + // has an unsupported `sentences` leaf under an otherwise-native predicate. + withStringCol("Hello World. Foo bar.", "", "Solo.", null) { + withDetour(enabled = true) { + assertCodegenRan { + checkSparkAnswerAndOperator( + sql("SELECT s FROM t WHERE size(sentences(s)) > 0"), + includeClasses = Seq(classOf[CometFilterExec])) + } + } + } + } + + test("detour refused by canHandle still falls back cleanly") { + // Force `CometBatchKernelCodegen.canHandle` to reject the tree at plan time via the + // `spark.sql.codegen.maxFields` gate. `sentences(s)` has 2 nested fields (the string input + + // the array> output collapses to 1), so maxFields=1 refuses it. The detour hook + // must surface that as a clean whole-operator fallback, not a mid-execution Janino failure. + withStringCol("Hello World.", null) { + withDetour(enabled = true) { + withSQLConf("spark.sql.codegen.maxFields" -> "1") { + CometScalaUDFCodegen.resetStats() + val df = sql("SELECT sentences(s) FROM t") + checkSparkAnswer(df) + assert( + !hasCometProject(df), + "expected fallback when canHandle refuses the detoured subtree") + val after = CometScalaUDFCodegen.stats() + assert( + after.compileCount == 0, + s"expected no kernel compile when the detour is refused, got $after") + } + } + } + } + + test("the detour gate is scoped: it does not fire outside a projection or filter context") { + // Serde-level assertion of the `DynamicVariable` scoping. `exprToProto` on an unsupported + // expression returns None on its own, and only returns a `JvmScalarUdf` proto when the call is + // wrapped in `withJvmDetour` (which `CometProjectExec` / `CometFilterExec` do) AND the feature + // is enabled. This is what keeps the detour out of join keys, sort keys, and aggregate + // expressions in v1 without threading a flag through every serde method. + val attr = AttributeReference("s", StringType)() + val expr = Sentences(attr) + val inputs = Seq(attr) + + withDetour(enabled = true) { + // Outside a detour-eligible context: no detour, even though the feature is enabled. + assert( + QueryPlanSerde.exprToProto(expr, inputs).isEmpty, + "expected no detour outside withJvmDetour") + + // Inside a detour-eligible context: the hook fires and emits a JvmScalarUdf proto. + val detoured = QueryPlanSerde.withJvmDetour(QueryPlanSerde.exprToProto(expr, inputs)) + assert(detoured.isDefined, "expected a detour proto inside withJvmDetour") + assert(detoured.get.hasJvmScalarUdf, s"expected a JvmScalarUdf proto, got ${detoured.get}") + } + + withDetour(enabled = false) { + // Feature disabled: even inside a detour-eligible context, no detour. + assert( + QueryPlanSerde.withJvmDetour(QueryPlanSerde.exprToProto(expr, inputs)).isEmpty, + "expected no detour when the feature flag is disabled") + } + } + + // --------------------------------------------------------------------------------------------- + // Parity suites: a supported built-in routed through the hook (by disabling its serde) ships a + // chosen input type across the FFI boundary. Results run through Spark's own `doGenCode` inside + // the kernel, so parity with Spark is the invariant. Each asserts the operator stays native. + // --------------------------------------------------------------------------------------------- + + /** + * Enable the detour + dispatcher and disable the given built-in serdes so they route through + * it. + */ + private def withDetouredSerdes(disabled: String*)(f: => Unit): Unit = { + val confs = Seq( + CometConf.COMET_EXEC_JVM_DETOUR_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") ++ + disabled.map(c => CometConf.getExprEnabledConfigKey(c) -> "false") + withSQLConf(confs: _*)(f) + } + + test("parity: NULL-heavy input through the detour hook") { + withTable("t") { + sql("CREATE TABLE t (x INT) USING parquet") + val rows = (0 until 300) + .map(i => if (i % 3 == 0) "(NULL)" else s"(${i - 150})") + .mkString(", ") + sql(s"INSERT INTO t VALUES $rows") + withDetouredSerdes("Abs") { + assertCodegenRan { + checkSparkAnswerAndOperator( + sql("SELECT abs(x) FROM t"), + includeClasses = Seq(classOf[CometProjectExec])) + } + } + } + } + + test("parity: empty input keeps the detour native") { + withTable("t") { + sql("CREATE TABLE t (x INT) USING parquet") // no rows: exercises the empty-batch path + withDetouredSerdes("Abs") { + // No assertCodegenRan: with zero rows the kernel never compiles. We only assert the plan + // stays native and produces the same (empty) result as Spark. + checkSparkAnswerAndOperator( + sql("SELECT abs(x) FROM t"), + includeClasses = Seq(classOf[CometProjectExec])) + } + } + } + + test("parity: decimals cross the FFI boundary through the detour hook") { + withTable("t") { + sql("CREATE TABLE t (d DECIMAL(18,6), e DECIMAL(38,10)) USING parquet") + sql( + "INSERT INTO t VALUES (1.5, 2.5), (-3.25, -4.75), (NULL, NULL), (0, 0), " + + "(999999999999.999999, 9999999999999999999999999999.0)") + withDetouredSerdes("Abs") { + assertCodegenRan { + checkSparkAnswerAndOperator( + sql("SELECT abs(d), abs(e) FROM t"), + includeClasses = Seq(classOf[CometProjectExec])) + } + } + } + } + + test("parity: timestamps cross the FFI boundary through the detour hook") { + withTable("t") { + sql("CREATE TABLE t (ts TIMESTAMP) USING parquet") + sql( + "INSERT INTO t VALUES (TIMESTAMP'2024-01-01 13:45:00'), " + + "(TIMESTAMP'1970-01-01 00:00:01'), (NULL), (TIMESTAMP'2024-06-15 23:59:59')") + withDetouredSerdes("Hour") { + assertCodegenRan { + checkSparkAnswerAndOperator( + sql("SELECT hour(ts) FROM t"), + includeClasses = Seq(classOf[CometProjectExec])) + } + } + } + } + + test("parity: nested array/map types cross the FFI boundary through the detour hook") { + withTable("t") { + sql("CREATE TABLE t (arr ARRAY, m MAP) USING parquet") + sql( + "INSERT INTO t VALUES (array(1, 2, 3), map('a', 1)), (array(), map()), " + + "(NULL, NULL), (array(5, 6), map('x', 9, 'y', 10))") + withDetouredSerdes("Size") { + assertCodegenRan { + checkSparkAnswerAndOperator( + sql("SELECT size(arr), size(m) FROM t"), + includeClasses = Seq(classOf[CometProjectExec])) + } + } + } + } + + test("a subquery inside the detoured subtree falls back cleanly, not to a kernel failure") { + 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 (10), (20)") + withDetouredSerdes("Abs") { + // The detoured tree is closure-serialized at plan time, before the operator's + // `waitForSubqueries` populates the `ScalarSubquery` result, so evaluating it in the + // kernel would throw "Subquery ... has not finished". The detour refuses a subquery-bearing + // tree, so the projection falls back to Spark and results still match. (A subquery that is + // a *sibling* of the detoured node is unaffected; it stays native.) + val df = sql("SELECT abs(x - (SELECT max(v) FROM t2)) FROM t") + checkSparkAnswer(df) + assert( + !hasCometProject(df), + "expected the projection to fall back when the detoured tree contains a subquery") + } + } + } + + test("a subquery that is a sibling of the detoured node stays native") { + // With `Abs` disabled, `abs(x)` detours (no subquery inside its tree) while the sibling + // `ScalarSubquery` under the native `+` is evaluated natively (Comet populates it via + // `waitForSubqueries`). The whole projection therefore stays native. This pins the + // sibling-subquery behavior the subquery guard's comment relies on. + 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 (10), (20)") + withDetouredSerdes("Abs") { + assertCodegenRan { + checkSparkAnswerAndOperator( + sql("SELECT abs(x) + (SELECT max(v) FROM t2) FROM t"), + includeClasses = Seq(classOf[CometProjectExec])) + } + } + } + } + + test("a mixed projection falls back whole when one sibling cannot detour") { + // `abs(x)` detours, but `make_interval(x)` returns CalendarInterval, which `canHandle` rejects + // (unsupported output type) and which Comet has no serde for, so it cannot be detoured. The + // per-expression detour of `abs(x)` succeeds, yet `exprs.forall(_.isDefined)` is false, so the + // whole projection falls back to a Spark `ProjectExec`. Results must still match Spark (the + // detour never changes results, only whether an operator stays native). + withTable("t") { + sql("CREATE TABLE t (x INT) USING parquet") + sql("INSERT INTO t VALUES (1), (2), (3)") + withDetouredSerdes("Abs") { + val df = sql("SELECT abs(x) AS a, make_interval(x) AS b FROM t") + checkSparkAnswer(df) + assert( + !hasCometProject(df), + "expected the whole projection to fall back when a sibling cannot detour") + } + } + } + + test("parity: nondeterministic detoured tree stays native with per-partition seeding") { + withTempPath { dir => + // Four partitions so per-partition `Rand` seeding (kernel `init(partitionIndex)`) matters: + // matching Spark on a fixed seed proves the seed advances per partition, not from 0 each time. + spark.range(0, 4096, 1, numPartitions = 4).write.parquet(dir.getAbsolutePath) + spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("t") + withDetouredSerdes("Abs") { + assertCodegenRan { + checkSparkAnswerAndOperator( + sql("SELECT id, abs(rand(42)) AS r FROM t"), + includeClasses = Seq(classOf[CometProjectExec])) + } + } + spark.catalog.dropTempView("t") + } + } + + test("island preservation: the projection stays native with no row transition below it") { + // Match every columnar->row transition kind by simple name (Spark's `ColumnarToRowExec` plus + // Comet's `CometColumnarToRowExec` / `CometNativeColumnarToRowExec`) so the check does not + // depend on which one Comet inserts. + def isRowTransition(p: SparkPlan): Boolean = p.getClass.getSimpleName match { + case "ColumnarToRowExec" | "CometColumnarToRowExec" | "CometNativeColumnarToRowExec" => true + case _ => false + } + withTable("t") { + sql("CREATE TABLE t (x INT) USING parquet") + sql("INSERT INTO t VALUES (1), (2), (3)") + + // Detour off: the projection falls back to a Spark `ProjectExec`, so a columnar->row + // transition feeds it and the native island above the scan is broken. + withSQLConf( + CometConf.COMET_EXEC_JVM_DETOUR_ENABLED.key -> "false", + CometConf.getExprEnabledConfigKey("Abs") -> "false") { + val offPlan = stripAQEPlan(sql("SELECT abs(x) FROM t").queryExecution.executedPlan) + assert( + collect(offPlan) { case p: ProjectExec => p }.nonEmpty, + "expected a Spark ProjectExec (fallback) with the detour off") + assert( + collect(offPlan) { case p if isRowTransition(p) => p }.nonEmpty, + "expected a columnar->row transition feeding the fallback project") + } + + // Detour on: the projection stays native and no columnar->row transition appears anywhere in + // its subtree, so the island above the scan is preserved end to end. + withDetouredSerdes("Abs") { + val onPlan = stripAQEPlan(sql("SELECT abs(x) FROM t").queryExecution.executedPlan) + val proj = collect(onPlan) { case p: CometProjectExec => p } + assert(proj.nonEmpty, "expected the projection to stay native with the detour on") + assert( + collect(onPlan) { case p: ProjectExec => p }.isEmpty, + "expected no Spark ProjectExec with the detour on") + assert( + proj.forall(p => collect(p) { case c if isRowTransition(c) => c }.isEmpty), + "expected no columnar->row transition below the native projection (island preserved)") + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPartialProjectFallbackBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPartialProjectFallbackBenchmark.scala new file mode 100644 index 0000000000..8ed48ce25c --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPartialProjectFallbackBenchmark.scala @@ -0,0 +1,195 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.benchmark + +import org.apache.spark.SparkConf +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.SparkSession + +import org.apache.comet.{CometConf, CometSparkSessionExtensions} + +/** + * Benchmark to measure the effect of the partial project fallback / JVM expression detour + * (`spark.comet.exec.jvmDetour.enabled`) on projections and filters that contain a single + * expression with no native translation. + * + * Without the detour, one unsupported expression falls the whole operator back to Spark, which + * breaks the native island above the scan: a `ColumnarToRow` transition materializes every read + * column and the supported sibling expressions run row-wise in Spark WSCG. With the detour, only + * the unsupported subtree crosses to the JVM (its argument columns over Arrow FFI); the operator + * — and the pipeline around it — stays native. + * + * The unsupported expression is simulated by disabling `sqrt`'s native serde + * (`spark.comet.expression.Sqrt.enabled=false`) in both Comet cases, so `sqrt` reaches the + * last-resort detour hook exactly as a genuinely-unsupported expression would. Only the detour + * flag differs between the two Comet cases, so the delta isolates the island-preservation win + * rather than the cost of `sqrt` itself (which runs in the JVM either way). Using a cheap, + * registered expression keeps the measurement stable as Comet's native coverage evolves. + * + * Three shapes: a lone unsupported expression over a wide native scan (pure transition cost), a + * projection mixing many native expressions with one unsupported (the compounding case — native + * siblings stay native), and a filter whose predicate contains an unsupported subexpression. To + * run: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometPartialProjectFallbackBenchmark + * }}} + */ +object CometPartialProjectFallbackBenchmark extends CometBenchmarkBase { + + override def getSparkSession: SparkSession = { + // Mirrors CometTestBase's memory/scan setup so native execution is actually active. Without + // an enabled Comet memory mode (on-heap here) the native runtime is disabled and the whole + // plan — scan included — falls back to Spark, leaving no native island to preserve. The Spark + // baseline case turns Comet off per-run via withSQLConf. + val conf = new SparkConf() + .setAppName("CometPartialProjectFallbackBenchmark") + .set("spark.master", "local[5]") + .setIfMissing("spark.driver.memory", "3g") + .setIfMissing("spark.executor.memory", "3g") + .set( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "2g") + .set(CometConf.COMET_ENABLED.key, "true") + .set(CometConf.COMET_ONHEAP_ENABLED.key, "true") + .set(CometConf.COMET_EXEC_ENABLED.key, "true") + .set(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key, "true") + .set(CometConf.COMET_SPARK_TO_ARROW_ENABLED.key, "true") + .set(CometConf.COMET_NATIVE_SCAN_ENABLED.key, "true") + .set(CometConf.COMET_ONHEAP_MEMORY_OVERHEAD.key, "2g") + + SparkSession.builder + .config(conf) + .withExtensions(new CometSparkSessionExtensions) + .getOrCreate() + } + + // Disabling `sqrt`'s serde makes it reach the detour hook, standing in for any expression with + // no native translation. Present in BOTH Comet cases so only the detour flag varies. + private val cometConfigs: Map[String, String] = Map( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.getExprEnabledConfigKey("Sqrt") -> "false") + + /** + * Print whether the detour-on plan is fully Comet native, so a misconfigured run is obvious. + */ + private def reportDetourPlan(query: String): Unit = { + val configs = cometConfigs + (CometConf.COMET_EXEC_JVM_DETOUR_ENABLED.key -> "true") + withSQLConf(configs.toSeq: _*) { + val df = spark.sql(query) + df.noop() + val plan = stripAQEPlan(df.queryExecution.executedPlan) + // scalastyle:off println + findFirstNonCometOperator(plan) match { + case Some(op) => + println( + s" [detour on] NOTE: plan not fully Comet native (first: ${op.nodeName}) — " + + "island not preserved for this shape") + case None => + println(" [detour on] plan is fully Comet native (island preserved)") + } + // scalastyle:on println + } + } + + private def benchmarkQuery(name: String, cardinality: Long, query: String): Unit = { + reportDetourPlan(query) + val benchmark = new Benchmark(name, cardinality, output = output) + + benchmark.addCase("Spark") { _ => + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.sql(query).noop() + } + } + + benchmark.addCase("Comet (detour off - operator falls back)") { _ => + val configs = cometConfigs + (CometConf.COMET_EXEC_JVM_DETOUR_ENABLED.key -> "false") + withSQLConf(configs.toSeq: _*) { + spark.sql(query).noop() + } + } + + benchmark.addCase("Comet (detour on - island preserved)") { _ => + val configs = cometConfigs + (CometConf.COMET_EXEC_JVM_DETOUR_ENABLED.key -> "true") + withSQLConf(configs.toSeq: _*) { + spark.sql(query).noop() + } + } + + benchmark.run() + } + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + val rows = 16 * 1024 * 1024 + + withTempPath { dir => + withTempTable("t") { + // Wide numeric table: `x` feeds the (disabled) sqrt; a..f are native passthrough / + // arithmetic inputs that stay columnar only when the operator stays native. + spark + .range(rows) + .selectExpr( + "id", + "cast(id as double) AS x", + "cast(id % 1000 as double) AS a", + "cast(id % 997 as double) AS b", + "cast(id % 991 as double) AS c", + "cast(id % 977 as double) AS d", + "cast(id % 100 as double) AS e", + "cast(id % 7 as double) AS f") + .write + .parquet(dir.getAbsolutePath) + spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("t") + + runBenchmark( + "Partial project fallback - lone unsupported expression over a native scan") { + // Off: the whole projection (sqrt + six passthrough columns) falls back to Spark, so + // every read column crosses a ColumnarToRow transition. On: native project, sqrt detours. + benchmarkQuery( + "lone unsupported expr, wide passthrough", + rows, + "SELECT sqrt(x) AS s, a, b, c, d, e, f FROM t") + } + + runBenchmark("Partial project fallback - mixed native + one unsupported expression") { + // Off: all eight expressions run row-wise in Spark. On: the seven native arithmetic + // expressions stay native, only sqrt detours to the JVM kernel. + benchmarkQuery( + "mixed projection (7 native + 1 unsupported)", + rows, + "SELECT sqrt(x) AS s, a + b AS n1, a * c AS n2, b - d AS n3, c * e AS n4, " + + "(a + b) * (c - d) AS n5, e + f AS n6, a - e AS n7 FROM t") + } + + runBenchmark("Partial project fallback - filter with an unsupported subexpression") { + // Off: the filter falls back, breaking the scan->filter island. On: the filter stays + // native with sqrt detoured under an otherwise-native predicate. + benchmarkQuery( + "filter predicate with unsupported subexpr", + rows, + "SELECT id, x FROM t WHERE sqrt(x) > 1.0 AND a > 5.0") + } + } + } + } +} From b9285337db9bff1cf074951708f20cb950813f38 Mon Sep 17 00:00:00 2001 From: Scott Schenkein Date: Sat, 4 Jul 2026 11:17:46 -0400 Subject: [PATCH 2/3] ci: register CometPartialProjectFallbackSuite in PR build workflows dev/ci/check-suites.py requires every test suite to be enumerated in the PR build workflows. Add the new suite to pr_build_linux.yml and pr_build_macos.yml next to the other codegen suites. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01JnvYuZKVMTeY2rewk94ZFe --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 85a2a0fa6a..248413d6f0 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -391,6 +391,7 @@ jobs: org.apache.comet.CometCodegenHOFSuite org.apache.comet.CometFuzzMathSuite org.apache.comet.CometCodegenFuzzSuite + org.apache.comet.CometPartialProjectFallbackSuite org.apache.comet.CometStringDecodeSuite org.apache.comet.CometWidthBucketSuite fail-fast: false diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index abde1554f6..02b12cd9c0 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -207,6 +207,7 @@ jobs: org.apache.comet.CometCodegenHOFSuite org.apache.comet.CometFuzzMathSuite org.apache.comet.CometCodegenFuzzSuite + org.apache.comet.CometPartialProjectFallbackSuite org.apache.comet.CometStringDecodeSuite org.apache.comet.CometWidthBucketSuite From e6873f6d277f7da14b276e249b12eda3be7ede6b Mon Sep 17 00:00:00 2001 From: Scott Schenkein Date: Sat, 4 Jul 2026 19:11:19 -0400 Subject: [PATCH 3/3] fix: make withDetour test helper compile on Spark 3.x `withDetour[T](enabled)(f: => T): T = withSQLConf(...)(f)` failed to compile on Spark 3.4/3.5: those versions type `withSQLConf(...)(f)` as `Unit`, which cannot unify with the declared return `T` (Spark 4.x infers it, so only the 3.x jobs were red). Capture the block's result in a local instead of returning `withSQLConf`'s value, which compiles on all supported Spark versions (verified against spark-3.5 and spark-4.1). Co-Authored-By: Claude Opus 4.8 --- .../comet/CometPartialProjectFallbackSuite.scala | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/CometPartialProjectFallbackSuite.scala b/spark/src/test/scala/org/apache/comet/CometPartialProjectFallbackSuite.scala index c28a145aa1..02ac64a01c 100644 --- a/spark/src/test/scala/org/apache/comet/CometPartialProjectFallbackSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometPartialProjectFallbackSuite.scala @@ -68,8 +68,16 @@ class CometPartialProjectFallbackSuite } } - private def withDetour[T](enabled: Boolean)(f: => T): T = - withSQLConf(CometConf.COMET_EXEC_JVM_DETOUR_ENABLED.key -> enabled.toString)(f) + private def withDetour[T](enabled: Boolean)(f: => T): T = { + // Capture the block's result explicitly rather than returning `withSQLConf`'s value: + // Spark 3.4/3.5 (Scala 2.12) type `withSQLConf(...)(f)` as `Unit`, so returning it directly + // fails to unify with `T`. Spark 4.x infers it, but this form compiles on all versions. + var result: Option[T] = None + withSQLConf(CometConf.COMET_EXEC_JVM_DETOUR_ENABLED.key -> enabled.toString) { + result = Some(f) + } + result.get + } /** Whether the (already-executed) DataFrame's plan kept a native `CometProjectExec`. */ private def hasCometProject(df: DataFrame): Boolean =