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

docs/dev/nvtx_ranges.md

@@ -115,6 +115,7 @@ ProjectExec|Executing projection operation on columnar batch
 Columnar batch serialize row only|Serializing row-only batch (no GPU data)
 disk spill|Spilling data from host memory to disk
 Async Shuffle Buffer|Asynchronous shuffle buffering operation
+Compile AST JIT|Compiling an AST expression for JIT evaluation
 CSV decode|Decoding CSV data
 get batch|Getting join batch
 Spark Task|Spark task execution range for stage and task tracking
@@ -190,6 +191,7 @@ Join gather|Gathering join results
 waitForCPU|Waiting for CPU batch in hybrid execution
 parquet get blocks with filter|Retrieving Parquet blocks after applying filters
 dynamic sort heuristic|Applying dynamic sort heuristic for aggregation
+Project AST JIT|Applying JIT-compiled AST projection to batch
 shuffle concat load batch|Concatenating and loading batch in shuffle operation
 parquet read footer bytes|Reading raw footer bytes from Parquet file
 spill batch|Spilling join batch

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,107 @@
+/*
+ * 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.{ColumnVector, Table}
+import ai.rapids.cudf.ast.CompiledExpression
+import com.nvidia.spark.Retryable
+
+import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression}
+
+object GpuAstJitExpression {
+  private def canUseAstJit(expression: GpuExpression): Boolean =
+    GpuBatchUtils.isFixedWidth(expression.dataType) &&
+      expression.supportsAstJit && expression.containsAstJitOperator
+
+  private def asAstJit(child: GpuExpression): Option[GpuAstJitExpression] = child match {
+    case jitExpression: GpuAstJitExpression => Some(jitExpression)
+    case astExpression: GpuProjectAstExpression => asAstJit(astExpression.child)
+    case other if canUseAstJit(other) => Some(GpuAstJitExpression(other))
+    case _ => None
+  }
+
+  private[rapids] def wrapTierExpression(expression: Expression): Expression = expression match {
+    case alias @ GpuAlias(child: GpuExpression, _) =>
+      asAstJit(child).map(GpuProjectAstExpressionBase.replaceChild(alias, _)).getOrElse(alias)
+    case other => other
+  }
+
+  private[rapids] def wrapProjectExpressions(
+      expressions: List[NamedExpression]): List[NamedExpression] = {
+    expressions.map(wrapTierExpression(_).asInstanceOf[NamedExpression])
+  }
+
+  /** Extracts a Project AST JIT wrapper after unwrapping any top-level aliases. */
+  private[rapids] def extractTopLevel(expression: Expression): Option[GpuAstJitExpression] =
+    GpuProjectAstExpressionBase.extractTopLevel(expression).collect {
+      case jitExpression: GpuAstJitExpression => jitExpression
+    }
+
+  private def hasJitCandidate(expression: Expression): Boolean = expression.find {
+    case gpuExpression: GpuExpression =>
+      gpuExpression.supportsAstJit && gpuExpression.containsAstJitOperator
+    case _ => false
+  }.isDefined
+
+  private def finalBackend(expression: Expression): String = {
+    GpuProjectAstExpressionBase.extractTopLevel(expression) match {
+      case Some(_: GpuAstJitExpression) => "Project AST JIT"
+      case Some(_: GpuProjectAstExpression) => "legacy Project AST"
+      case _ => "the regular GPU projection"
+    }
+  }
+
+  private[rapids] def explainFinalSelections(
+      expressionTiers: Seq[Seq[Expression]],
+      all: Boolean): String = {
+    expressionTiers.zipWithIndex.flatMap { case (expressions, tier) =>
+      val explanations = expressions.iterator.collect {
+        case expression if all ||
+            (extractTopLevel(expression).isEmpty && hasJitCandidate(expression)) =>
+          s"    $expression final backend: ${finalBackend(expression)}\n"
+      }.mkString
+      if (explanations.nonEmpty) {
+        Some(s"  TIER $tier\n$explanations")
+      } else {
+        None
+      }
+    }.mkString
+  }
+}
+
+case class GpuAstJitExpression(child: GpuExpression)
+    extends GpuProjectAstExpressionBase with Retryable {
+
+  override protected def backendName: String = "AST JIT"
+
+  override protected def compileNvtxId: NvtxId = NvtxRegistry.COMPILE_AST_JIT
+
+  override protected def computeNvtxId: NvtxId = NvtxRegistry.PROJECT_AST_JIT
+
+  override protected def evaluate(
+      compiled: CompiledExpression,
+      table: Table): ColumnVector = compiled.computeColumnJit(table)
+
+  override def toString: String = s"AST_JIT($child)"
+
+  override def checkpoint(): Unit = {
+    getCompiledExpression
+  }
+
+  // Compiled ASTs are immutable and remain valid across retry attempts.
+  override def restore(): Unit = ()
+}

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

@@ -42,6 +42,19 @@ trait GpuBind {
 
 object GpuBindReferences extends Logging {
 
+  private def explainFinalProjectAstJitSelection(
+      tieredProject: GpuTieredProject,
+      conf: SQLConf): Unit = {
+    val explain = RapidsConf.EXPLAIN.get(conf)
+    if (RapidsConf.shouldExplain(explain)) {
+      val explanation = GpuAstJitExpression.explainFinalSelections(
+        tieredProject.exprTiers, RapidsConf.shouldExplainAll(explain))
+      if (explanation.nonEmpty) {
+        logWarning(s"FINAL PROJECT AST JIT SELECTION\n$explanation")
+      }
+    }
+  }
+
   /**
    * An alternative to `Expression.transformDown`, but when a result is returned by `rule` it is
    * assumed that it handled processing exp and all of its children, so rule will not be called on
@@ -124,18 +137,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 tieredProject = if (RapidsConf.ENABLE_TIERED_PROJECT.get(conf)) {
+      val exprTiers = GpuProjectAstExpressionBase.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,8 +188,44 @@ 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)))
     }
+    if (enableProjectAstJit) {
+      explainFinalProjectAstJitSelection(tieredProject, conf)
+    }
+    tieredProject
+  }
+
+  /**
+   * 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) ==========
@@ -257,12 +307,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})"