Add experimental Project AST JIT for integral add and multiply [fast-ut] [databricks]

integration_tests/src/main/python/ast_test.py

@@ -68,6 +68,12 @@
     not (is_spark_359() or is_spark_403_or_404() or is_spark_412_or_later()))]
 
 _project_ast_enabled_conf = {"spark.rapids.sql.projectAstEnabled": "true"}
+_project_ast_jit_enabled_conf = {"spark.rapids.sql.projectAstJitEnabled": "true"}
+_project_ast_jit_and_legacy_enabled_conf = {
+    "spark.rapids.sql.projectAstEnabled": "true",
+    "spark.rapids.sql.projectAstJitEnabled": "true"
+}
+
 
 def assert_gpu_ast(is_supported, func, conf={}):
     ast_expression = "GpuProjectAstExpression"
@@ -386,6 +392,76 @@ def test_multiplication(data_descr):
             f.lit(-12).cast(data_type) * f.col('b'),
             f.col('a') * f.col('b')))
 
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@disable_ansi_mode
+def test_jit_add_multiply(data_gen):
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: binary_op_df(spark, data_gen).select(
+            f.col('a') + f.col('b'),
+            f.col('a') * f.col('b')),
+        exist_classes=r"GpuProject.*AST_JIT",
+        non_exist_classes="GpuProjectAst",
+        conf=_project_ast_jit_enabled_conf)
+
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@disable_ansi_mode
+def test_jit_does_not_split_unique_unsupported_expression(data_gen):
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: binary_op_df(spark, data_gen).select(
+            (f.col('a') + f.col('b')) - (f.col('a') * f.col('b'))),
+        exist_classes="GpuProject",
+        non_exist_classes="AST_JIT,GpuProjectAst",
+        conf=_project_ast_jit_enabled_conf)
+
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@disable_ansi_mode
+def test_jit_cse_shared_subexpression(data_gen):
+    def project_shared_expression(spark):
+        df = binary_op_df(spark, data_gen)
+        shared = f.col('a') + f.col('b')
+        return df.select(
+            shared.alias('shared'),
+            (shared - f.col('a')).alias('left'),
+            (shared - f.col('b')).alias('right'))
+
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        project_shared_expression,
+        exist_classes=r"GpuProject.*AST_JIT.*AS shared.*AS left.*AS right",
+        non_exist_classes="GpuProjectAst",
+        conf=_project_ast_jit_enabled_conf)
+
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@pytest.mark.parametrize(
+    'tiered_project_enabled', ['true', 'false'], ids=['tiered', 'single_tier'])
+@disable_ansi_mode
+def test_jit_mixed_project_expressions(data_gen, tiered_project_enabled):
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: binary_op_df(spark, data_gen).select(
+            (f.col('a') + f.col('b')).alias('jit'),
+            (f.col('a') - f.col('b')).alias('gpu'),
+            ((f.col('a') * f.col('b')) - f.col('a')).alias('mixed')),
+        exist_classes=r"GpuProject.*AST_JIT.*AS jit.*AS gpu.*AS mixed",
+        non_exist_classes=r"GpuProjectAst,AS gpu.*AST_JIT",
+        conf=copy_and_update(_project_ast_jit_enabled_conf, {
+            'spark.rapids.sql.tiered.project.enabled': tiered_project_enabled
+        }))
+
+@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn)
+@disable_ansi_mode
+def test_jit_and_legacy_ast_mixed_project_expressions(data_gen):
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: binary_op_df(spark, data_gen).select(
+            (f.col('a') + f.col('b')).alias('jit'),
+            (f.col('a') - f.col('b')).alias('legacy'),
+            ((f.col('a') + f.col('b')) -
+                (f.col('a') * f.col('b'))).alias('mixed')),
+        exist_classes=(
+            r"GpuProject.*AST_JIT.*AS jit.*AST\(.*AS legacy.*"
+            r"AST\(.*AS mixed,GpuProjectAstExpression"),
+        non_exist_classes=r"AS legacy.*AST_JIT",
+        conf=_project_ast_jit_and_legacy_enabled_conf)
+
+
 # Each descriptor contains a list of data generators and a corresponding boolean
 # indicating whether that data type is supported by the AST
 # all the below desc are not supported by the AST because ANSI mode is on

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala

@@ -0,0 +1,144 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.
+ *
+ * Licensed 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 com.nvidia.spark.rapids
+
+import ai.rapids.cudf.Table
+import ai.rapids.cudf.ast.CompiledExpression
+import com.nvidia.spark.Retryable
+import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
+import com.nvidia.spark.rapids.GpuMetric.OP_TIME_LEGACY
+import com.nvidia.spark.rapids.RapidsPluginImplicits._
+import com.nvidia.spark.rapids.ScalableTaskCompletion.onTaskCompletion
+import com.nvidia.spark.rapids.shims.ShimUnaryExpression
+
+import org.apache.spark.TaskContext
+import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression}
+import org.apache.spark.sql.types.DataType
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+object GpuAstJitExpression {
+  private def canUseAstJit(expression: GpuExpression): Boolean =
+    GpuBatchUtils.isFixedWidth(expression.dataType) &&
+      expression.supportsAstJit && expression.containsAstJitOperator
+
+  private[rapids] def wrapTierExpression(expression: Expression): Expression = expression match {
+    case alias @ GpuAlias(_: GpuAstJitExpression, _) => alias
+    case alias @ GpuAlias(astExpression: GpuProjectAstExpression, _)
+        if canUseAstJit(astExpression.child) =>
+      GpuProjectAstExpression.replaceChild(alias, GpuAstJitExpression(astExpression.child))
+    case alias @ GpuAlias(child: GpuExpression, _) if canUseAstJit(child) =>
+      GpuProjectAstExpression.replaceChild(alias, GpuAstJitExpression(child))
+    case other => other
+  }
+
+  private[rapids] def wrapProjectExpressions(
+      expressions: List[NamedExpression]): List[NamedExpression] = {
+    expressions.map(wrapTierExpression(_).asInstanceOf[NamedExpression])
+  }
+
+  private[rapids] def contains(expression: Expression): Boolean =
+    GpuProjectAstExpressionBase.extractTopLevel(expression)
+      .exists(_.isInstanceOf[GpuAstJitExpression])
+}
+
+case class GpuAstJitExpression(child: GpuExpression)
+    extends ShimUnaryExpression with GpuProjectAstExpressionBase
+        with GpuMetricsInjectable with Retryable with AutoCloseable {
+
+  @transient private[this] var compiledExpression: CompiledExpression = _
+  @transient private[this] var completionRegistered = false
+  private[this] var opTime: GpuMetric = NoopMetric
+
+  override def dataType: DataType = child.dataType
+
+  override def nullable: Boolean = child.nullable
+
+  override def toString: String = s"AST_JIT($child)"
+
+  override def injectMetrics(metrics: Map[String, GpuMetric]): Unit = {
+    opTime = metrics.getOrElse(OP_TIME_LEGACY, NoopMetric)
+  }
+
+  override def checkpoint(): Unit = {
+    getCompiledExpression
+  }
+
+  // Compiled ASTs are immutable and remain valid across retry attempts.
+  override def restore(): Unit = ()
+
+  override def close(): Unit = closeCompiledExpression()
+
+  override def columnarEval(batch: ColumnarBatch): GpuColumnVector = {
+    withResource(GpuProjectAstExpression.tableFromBatch(batch)) { table =>
+      computeColumn(table)
+    }
+  }
+
+  private[rapids] override def computeColumn(table: Table): GpuColumnVector = {
+    val compiled = getCompiledExpression
+    NvtxIdWithMetrics(NvtxRegistry.PROJECT_AST, opTime) {
+      closeOnExcept(compiled.computeColumnJit(table)) { result =>
+        GpuColumnVector.from(result, dataType)
+      }
+    }
+  }
+
+  private def getCompiledExpression: CompiledExpression = synchronized {
+    if (compiledExpression == null) {
+      compiledExpression = NvtxIdWithMetrics(NvtxRegistry.COMPILE_ASTS, opTime) {
+        // Force every bound reference to the left table; Project AST has one input table.
+        child.convertToAst(Int.MaxValue).compile()
+      }
+    }
+    if (!completionRegistered) {
+      Option(TaskContext.get()).foreach { taskContext =>
+        completionRegistered = true
+        try {
+          onTaskCompletion(taskContext) {
+            clearRegistrationAndClose()
+          }
+          if (!completionRegistered) {
+            throw new IllegalStateException(
+              "Task completed while registering the AST JIT cleanup callback")
+          }
+        } catch {
+          case t: Throwable =>
+            clearRegistrationAndClose(t)
+            throw t
+        }
+      }
+    }
+    compiledExpression
+  }
+
+  private def clearRegistrationAndClose(error: Throwable = null): Unit =
+    closeCompiledExpression(error, clearRegistration = true)
+
+  private def closeCompiledExpression(
+      error: Throwable = null,
+      clearRegistration: Boolean = false): Unit = {
+    val toClose = synchronized {
+      if (clearRegistration) {
+        completionRegistered = false
+      }
+      val current = compiledExpression
+      compiledExpression = null
+      current
+    }
+    Option(toClose).foreach(_.safeClose(error))
+  }
+}

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuBoundAttribute.scala

@@ -124,18 +124,19 @@ object GpuBindReferences extends Logging {
   }
 
   /**
-   * Binding method for tiered expressions without metric injection.
-   * This is for use by GpuBind implementations and should not be called directly
-   * from SparkPlan nodes. Use the public API that requires metrics instead, except
-   * when absolutely needed.
+   * Shared implementation for generic and Project-specific tiered binding.
+   *
+   * @param enableProjectAstJit whether eligible tiers may use Project AST JIT
    */
-  def bindGpuReferencesTieredNoMetrics[A <: Expression](
+  private def bindGpuReferencesTieredNoMetricsInternal[A <: Expression](
       expressions: Seq[A],
       input: AttributeSeq,
-      conf: SQLConf): GpuTieredProject = {
+      conf: SQLConf,
+      enableProjectAstJit: Boolean): GpuTieredProject = {
 
     if (RapidsConf.ENABLE_TIERED_PROJECT.get(conf)) {
-      val exprTiers = GpuProjectAstExpression.buildExprTiers(expressions, conf)
+      val exprTiers = GpuProjectAstExpression.buildExprTiers(
+        expressions, conf, enableProjectAstJit)
       val inputTiers = GpuEquivalentExpressions.getInputTiers(exprTiers, input)
       // Update ExprTiers to include the columns that are pass through and drop unneeded columns
       val newExprTiers = exprTiers.zipWithIndex.map {
@@ -174,10 +175,42 @@ object GpuBindReferences extends Logging {
       }
       GpuTieredProject(tiered)
     } else {
-      GpuTieredProject(Seq(GpuBindReferences.bindGpuReferencesNoMetrics(expressions, input)))
+      val projectExpressions = if (enableProjectAstJit) {
+        expressions.map(GpuAstJitExpression.wrapTierExpression)
+      } else {
+        expressions
+      }
+      GpuTieredProject(Seq(
+        GpuBindReferences.bindGpuReferencesNoMetrics(projectExpressions, input)))
     }
   }
 
+  /**
+   * Binding method for tiered expressions without metric injection.
+   * This is for use by GpuBind implementations and should not be called directly
+   * from SparkPlan nodes. Use the public API that requires metrics instead, except
+   * when absolutely needed.
+   */
+  def bindGpuReferencesTieredNoMetrics[A <: Expression](
+      expressions: Seq[A],
+      input: AttributeSeq,
+      conf: SQLConf): GpuTieredProject = {
+    bindGpuReferencesTieredNoMetricsInternal(
+      expressions, input, conf, enableProjectAstJit = false)
+  }
+
+  /**
+   * Project-specific tiered binding without metric injection. Unlike the generic binder,
+   * this path allows configured Project AST JIT selection.
+   */
+  private[rapids] def bindGpuProjectReferencesTieredNoMetrics[A <: Expression](
+      expressions: Seq[A],
+      input: AttributeSeq,
+      conf: SQLConf): GpuTieredProject = {
+    bindGpuReferencesTieredNoMetricsInternal(
+      expressions, input, conf, RapidsConf.ENABLE_PROJECT_AST_JIT.get(conf))
+  }
+
   // ========== Public "Front Door" APIs (for use by SparkPlan nodes) ==========
   // These methods require metrics and inject them after binding
 
@@ -257,12 +290,32 @@ object GpuBindReferences extends Logging {
     bound.injectMetrics(metrics)
     bound
   }
+
+  /**
+   * Bind Project expressions in a tiered manner and inject metrics. Project AST JIT selection is
+   * confined to this entry point so generic tiered binders do not enable it for other operators.
+   * @param expressions The expressions to bind
+   * @param input The input schema
+   * @param conf SQL configuration
+   * @param metrics Metrics to inject into the bound expressions
+   */
+  def bindGpuProjectReferencesTiered[A <: Expression](
+      expressions: Seq[A],
+      input: AttributeSeq,
+      conf: SQLConf,
+      metrics: Map[String, GpuMetric]): GpuTieredProject = {
+    val bound = bindGpuProjectReferencesTieredNoMetrics(expressions, input, conf)
+    bound.injectMetrics(metrics)
+    bound
+  }
 }
 
 case class GpuBoundReference(ordinal: Int, dataType: DataType, nullable: Boolean)
     (val exprId: ExprId, val name: String)
   extends GpuLeafExpression with ShimExpression {
 
+  override def selfSupportsAstJit: Boolean = true
+
   override def toString: String =
     s"input[$ordinal, ${dataType.simpleString}, $nullable]($name#${exprId.id})"
 

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpressions.scala

@@ -197,6 +197,31 @@ trait GpuExpression extends Expression {
   def convertToAst(numFirstTableColumns: Int): ast.AstExpression =
     throw new IllegalStateException(s"Cannot convert ${this.getClass.getSimpleName} to AST")
 
+  /**
+   * Whether this node supports AST JIT for its current semantics and types, excluding its
+   * children. Operator overrides must validate their execution modes and local input/output types.
+   */
+  def selfSupportsAstJit: Boolean = false
+
+  /**
+   * Whether this node is an operation, rather than an AST-compatible leaf. Literals and references
+   * leave this false so they do not trigger compilation without useful work to JIT.
+   */
+  def selfIsAstJitOperator: Boolean = false
+
+  /** Whether this node and its complete expression subtree support AST JIT. */
+  final def supportsAstJit: Boolean = selfSupportsAstJit && children.forall {
+    case child: GpuExpression => child.supportsAstJit
+    case _: AttributeReference => true
+    case _ => false
+  }
+
+  /** Whether this expression subtree contains an operation that makes AST JIT useful. */
+  final def containsAstJitOperator: Boolean = selfIsAstJitOperator || children.exists {
+    case child: GpuExpression => child.containsAstJitOperator
+    case _ => false
+  }
+
   /** Could evaluating this expression cause side-effects, such as throwing an exception? */
   def hasSideEffects: Boolean =
     children.exists {
@@ -391,6 +416,8 @@ trait CudfBinaryExpression extends GpuBinaryExpression {
   def castOutputAtEnd: Boolean = false
   def astOperator: Option[ast.BinaryOperator] = None
 
+  override def selfIsAstJitOperator: Boolean = selfSupportsAstJit
+
   def outputType(l: BinaryOperable, r: BinaryOperable): DType = {
     val over = outputTypeOverride
     if (over == null) {