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

integration_tests/src/main/python/ast_test.py

@@ -66,6 +66,7 @@
 ast_acosh_descr = [(double_gen, not (is_spark_403() 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"}
 
 def assert_gpu_ast(is_supported, func, conf={}):
     exist = "GpuProjectAstExec"
@@ -371,6 +372,37 @@ 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)
+
+@disable_ansi_mode
+def test_jit_mixed_nested_subexpressions():
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: binary_op_df(spark, int_gen).select(
+            (f.col('a') + f.col('b')) - (f.col('a') * f.col('b'))),
+        exist_classes=r"GpuProject.*AST_JIT.*- AST_JIT",
+        non_exist_classes="GpuProjectAst",
+        conf=_project_ast_jit_enabled_conf)
+
+@disable_ansi_mode
+def test_jit_mixed_project_expressions():
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: binary_op_df(spark, int_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.*AST_JIT.*AS mixed",
+        non_exist_classes="GpuProjectAst",
+        conf=_project_ast_jit_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,115 @@
+/*
+ * 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.{Scalar, 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.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 wrapMaximalSubtrees(expression: Expression): Expression = expression match {
+    case gpuExpression: GpuExpression
+        if gpuExpression.supportsAstJit && gpuExpression.containsAstJitOperator =>
+      GpuAstJitExpression(gpuExpression)
+    case gpuExpression: GpuExpression =>
+      gpuExpression.mapChildren {
+        case child: GpuExpression => wrapMaximalSubtrees(child)
+        case child => child
+      }
+    case other => other
+  }
+
+  private[rapids] def wrapProjectExpressions(
+      expressions: List[NamedExpression]): List[NamedExpression] = {
+    expressions.map(wrapMaximalSubtrees(_).asInstanceOf[NamedExpression])
+  }
+}
+
+case class GpuAstJitExpression(child: Expression)
+    extends ShimUnaryExpression with GpuExpression with Retryable with AutoCloseable {
+  require(child.isInstanceOf[GpuExpression], "AST JIT child must be a GPU expression")
+
+  @transient private[this] var compiledExpression: CompiledExpression = _
+  @transient private[this] var completionRegistered = false
+
+  override def dataType: DataType = child.dataType
+
+  override def nullable: Boolean = child.nullable
+
+  override def disableTieredProjectCombine: Boolean = true
+
+  override def toString: String = s"AST_JIT($child)"
+
+  override def checkpoint(): Unit = {
+    getCompiledExpression
+  }
+
+  override def restore(): Unit = closeCompiledExpression()
+
+  override def close(): Unit = closeCompiledExpression()
+
+  override def columnarEval(batch: ColumnarBatch): GpuColumnVector = {
+    withResource(tableFromBatch(batch)) { table =>
+      closeOnExcept(getCompiledExpression.computeColumnJit(table)) { result =>
+        GpuColumnVector.from(result, dataType)
+      }
+    }
+  }
+
+  private def getCompiledExpression: CompiledExpression = synchronized {
+    if (compiledExpression == null) {
+      compiledExpression = child.asInstanceOf[GpuExpression]
+        .convertToAst(Int.MaxValue)
+        .compile()
+    }
+    if (!completionRegistered) {
+      Option(TaskContext.get()).foreach { taskContext =>
+        onTaskCompletion(taskContext) {
+          close()
+        }
+        completionRegistered = true
+      }
+    }
+    compiledExpression
+  }
+
+  private def closeCompiledExpression(): Unit = synchronized {
+    Option(compiledExpression).foreach(_.safeClose())
+    compiledExpression = null
+  }
+
+  private def tableFromBatch(batch: ColumnarBatch): Table = {
+    if (batch.numCols() != 0) {
+      GpuColumnVector.from(batch)
+    } else {
+      withResource(Scalar.fromBool(false)) { falseScalar =>
+        withResource(ai.rapids.cudf.ColumnVector.fromScalar(falseScalar, batch.numRows())) {
+          falseColumn => new Table(falseColumn)
+        }
+      }
+    }
+  }
+}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2025, NVIDIA CORPORATION.
+ * Copyright (c) 2019-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.
@@ -268,6 +268,8 @@ 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,21 @@ trait GpuExpression extends Expression {
   def convertToAst(numFirstTableColumns: Int): ast.AstExpression =
     throw new IllegalStateException(s"Cannot convert ${this.getClass.getSimpleName} to AST")
 
+  def selfSupportsAstJit: Boolean = false
+
+  def selfIsAstJitOperator: Boolean = false
+
+  final def supportsAstJit: Boolean = selfSupportsAstJit && children.forall {
+    case child: GpuExpression => child.supportsAstJit
+    case _: AttributeReference => true
+    case _ => false
+  }
+
+  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 {

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

@@ -1247,6 +1247,12 @@ val GPU_COREDUMP_PIPE_PATTERN = conf("spark.rapids.gpu.coreDump.pipePattern")
       .booleanConf
       .createWithDefault(false)
 
+  val ENABLE_PROJECT_AST_JIT = conf("spark.rapids.sql.projectAstJitEnabled")
+      .doc("Enable the experimental cuDF JIT backend for supported project AST expressions.")
+      .internal()
+      .booleanConf
+      .createWithDefault(false)
+
   val ENABLE_TIERED_PROJECT = conf("spark.rapids.sql.tiered.project.enabled")
       .doc("Enable tiered projections.")
       .internal()
@@ -3639,6 +3645,8 @@ class RapidsConf(conf: Map[String, String]) extends Logging {
 
   lazy val isProjectAstEnabled: Boolean = get(ENABLE_PROJECT_AST)
 
+  lazy val isProjectAstJitEnabled: Boolean = get(ENABLE_PROJECT_AST_JIT)
+
   lazy val isTieredProjectEnabled: Boolean = get(ENABLE_TIERED_PROJECT)
 
   lazy val isCombinedExpressionsEnabled: Boolean = get(ENABLE_COMBINED_EXPRESSIONS)

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

@@ -59,6 +59,12 @@ class GpuProjectExecMeta(
     // Force list to avoid recursive Java serialization of lazy list Seq implementation
     val gpuExprs = childExprs.map(_.convertToGpu().asInstanceOf[NamedExpression]).toList
     val gpuChild = childPlans.head.convertIfNeeded()
+    if (conf.isProjectAstJitEnabled) {
+      val jitProjectList = GpuAstJitExpression.wrapProjectExpressions(gpuExprs)
+      if (jitProjectList.exists(_.find(_.isInstanceOf[GpuAstJitExpression]).isDefined)) {
+        return GpuProjectExec(jitProjectList, gpuChild)
+      }
+    }
     if (conf.isProjectAstEnabled) {
       // cuDF requires return column is fixed width
       val allReturnTypesFixedWidth = gpuExprs.forall(e => GpuBatchUtils.isFixedWidth(e.dataType))

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2025, NVIDIA CORPORATION.
+ * Copyright (c) 2019-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.
@@ -645,6 +645,8 @@ object GpuLiteral {
  */
 case class GpuLiteral (value: Any, dataType: DataType) extends GpuLeafExpression {
 
+  override def selfSupportsAstJit: Boolean = true
+
   // Assume this came from Spark Literal and no need to call Literal.validateLiteralValue here.
 
   override def foldable: Boolean = true

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/arithmetic.scala

@@ -314,6 +314,11 @@ abstract class GpuAddBase extends CudfBinaryArithmetic with Serializable {
   override def binaryOp: BinaryOp = BinaryOp.ADD
   override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.ADD)
 
+  override def selfSupportsAstJit: Boolean =
+    !failOnError && (dataType == IntegerType || dataType == LongType)
+
+  override def selfIsAstJitOperator: Boolean = selfSupportsAstJit
+
   override def hasSideEffects: Boolean =
     (failOnError && GpuAnsi.needBasicOpOverflowCheck(dataType)) || super.hasSideEffects
 
@@ -768,6 +773,11 @@ case class GpuMultiply(
   override def binaryOp: BinaryOp = BinaryOp.MUL
   override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.MUL)
 
+  override def selfSupportsAstJit: Boolean =
+    !failOnError && (dataType == IntegerType || dataType == LongType)
+
+  override def selfIsAstJitOperator: Boolean = selfSupportsAstJit
+
   private def multiplyOverflowError(msg: String): ArithmeticException = {
     RapidsErrorUtils.arithmeticOverflowError(msg, origin)
   }

tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala

@@ -0,0 +1,80 @@
+/*
+ * 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 org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import org.apache.spark.sql.rapids.{GpuAdd, GpuMultiply, GpuSubtract}
+import org.apache.spark.sql.types.{FloatType, IntegerType, LongType}
+
+class GpuProjectAstJitSuite extends AnyFunSuite {
+  private def reference(ordinal: Int, dataType: org.apache.spark.sql.types.DataType) =
+    AttributeReference(s"c$ordinal", dataType, nullable = true)()
+
+  private def alias(expression: GpuExpression, name: String) = GpuAlias(expression, name)()
+
+  test("project AST JIT is disabled by default") {
+    assert(!new RapidsConf(Map.empty[String, String]).isProjectAstJitEnabled)
+  }
+
+  test("project AST JIT supports non-ANSI integral add and multiply") {
+    val left = reference(0, LongType)
+    val right = reference(1, LongType)
+    val expression = alias(
+      GpuMultiply(
+        GpuAdd(left, right, failOnError = false)(),
+        right,
+        failOnError = false)(),
+      "result")
+
+    val wrapped = GpuAstJitExpression.wrapProjectExpressions(List(expression))
+    val jit = wrapped.head.asInstanceOf[GpuAlias].child.asInstanceOf[GpuAstJitExpression]
+    assert(jit.child.isInstanceOf[GpuMultiply])
+    assert(jit.child.find(_.isInstanceOf[GpuAstJitExpression]).isEmpty)
+  }
+
+  test("project AST JIT wraps maximal nested subtrees independently") {
+    val left = reference(0, IntegerType)
+    val right = reference(1, IntegerType)
+    val expression = alias(
+      GpuSubtract(
+        GpuAdd(left, right, failOnError = false)(),
+        GpuMultiply(left, right, failOnError = false)(),
+        failOnError = false)(),
+      "result")
+
+    val wrapped = GpuAstJitExpression.wrapProjectExpressions(List(expression))
+    val subtract = wrapped.head.asInstanceOf[GpuAlias].child.asInstanceOf[GpuSubtract]
+    assert(subtract.left.asInstanceOf[GpuAstJitExpression].child.isInstanceOf[GpuAdd])
+    assert(subtract.right.asInstanceOf[GpuAstJitExpression].child.isInstanceOf[GpuMultiply])
+  }
+
+  test("project AST JIT excludes ANSI and floating point arithmetic") {
+    val intLeft = reference(0, IntegerType)
+    val intRight = reference(1, IntegerType)
+    val floatLeft = reference(0, FloatType)
+    val floatRight = reference(1, FloatType)
+
+    val ansiAdd = alias(GpuAdd(intLeft, intRight, failOnError = true)(), "ansi_sum")
+    val floatMultiply = alias(
+      GpuMultiply(floatLeft, floatRight, failOnError = false)(), "float_product")
+
+    val wrapped = GpuAstJitExpression.wrapProjectExpressions(List(ansiAdd, floatMultiply))
+    assert(wrapped.forall(_.find(_.isInstanceOf[GpuAstJitExpression]).isEmpty))
+  }
+}