feat: partial project fallback (JVM expression detour)

.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

.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
 

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)

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

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

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

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