[FEA] Enable per-expression legacy AST projection [databricks] [fast-ut]

integration_tests/src/main/python/ast_test.py

@@ -68,11 +68,12 @@
 _project_ast_enabled_conf = {"spark.rapids.sql.projectAstEnabled": "true"}
 
 def assert_gpu_ast(is_supported, func, conf={}):
-    exist = "GpuProjectAstExec"
-    non_exist = "GpuProjectExec"
+    ast_expression = "GpuProjectAstExpression"
+    exist = ast_expression
+    non_exist = ''
     if not is_supported:
         exist = "GpuProjectExec"
-        non_exist = "GpuProjectAstExec"
+        non_exist = ast_expression
     ast_conf = copy_and_update(conf, _project_ast_enabled_conf)
     assert_cpu_and_gpu_are_equal_collect_with_capture(
         func,
@@ -428,14 +429,29 @@ def test_multi_tier_ast():
         func=lambda spark: spark.range(10).withColumn("x", f.col("id")).repartition(1)\
             .selectExpr("x", "(id < x) == (id < (id + x))"))
 
-
-# MUST NOT use GPU AST when project refers to string type(non-fixed-width),
-# or cudf::compute_column will throw error: Invalid, non-fixed-width type
-# ANSI mode is disabled here due to an overflow issue with integer multiplication on Spark 4.0.0.
 @disable_ansi_mode
-@ignore_order(local=True)
-def test_refer_to_non_fixed_width_column():
-    gens = [('col_int', int_gen), ('col_string', string_gen)]
-    assert_gpu_and_cpu_are_equal_collect(
-        lambda spark: gen_df(spark, gens).selectExpr("col_int * col_int", "col_string"),
-        conf=_project_ast_enabled_conf)
+@pytest.mark.parametrize(
+    'tiered_project_enabled', ['true', 'false'], ids=['tiered', 'single_tier'])
+def test_project_ast_mixed_expressions(tiered_project_enabled):
+    def project(spark):
+        df = gen_df(spark, [
+            ('a', int_gen),
+            ('b', int_gen),
+            ('c', int_gen),
+            ('d', int_gen),
+            ('col_string', string_gen)
+        ])
+        shared = f.col('a') + f.col('b')
+        return df.select(
+            (shared * f.col('c')).alias('ast_first'),
+            (shared * f.col('d')).alias('ast_second'),
+            f.greatest(shared, f.col('c')).alias('gpu_shared'),
+            f.length(f.col('col_string')).alias('gpu_string'),
+            f.col('col_string').alias('raw_string'))
+
+    assert_gpu_ast(
+        is_supported=True,
+        func=project,
+        conf={
+            'spark.rapids.sql.tiered.project.enabled': tiered_project_enabled
+        })

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.
@@ -135,12 +135,7 @@ object GpuBindReferences extends Logging {
       conf: SQLConf): GpuTieredProject = {
 
     if (RapidsConf.ENABLE_TIERED_PROJECT.get(conf)) {
-      val replaced = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) {
-        GpuEquivalentExpressions.replaceMultiExpressions(expressions, conf)
-      } else {
-        expressions
-      }
-      val exprTiers = GpuEquivalentExpressions.getExprTiers(replaced)
+      val exprTiers = GpuProjectAstExpression.buildExprTiers(expressions, conf)
       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 {

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

@@ -0,0 +1,166 @@
+/*
+ * 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 scala.annotation.tailrec
+
+import ai.rapids.cudf.{Scalar, Table}
+import ai.rapids.cudf.ast.CompiledExpression
+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.internal.SQLConf
+import org.apache.spark.sql.rapids.catalyst.expressions.{
+  GpuEquivalentExpressions, GpuExpressionEquals}
+import org.apache.spark.sql.types.DataType
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+object GpuProjectAstExpression {
+  private[rapids] def wrap(expression: NamedExpression): NamedExpression = {
+    expression match {
+      case alias @ GpuAlias(child: GpuExpression, name) =>
+        GpuAlias(GpuProjectAstExpression(child), name)(
+          alias.exprId, alias.qualifier, alias.explicitMetadata)
+      case other => other
+    }
+  }
+
+  @tailrec
+  private[rapids] def extractTopLevel(expression: Expression): Option[GpuProjectAstExpression] = {
+    expression match {
+      case alias: GpuAlias => extractTopLevel(alias.child)
+      case astExpression: GpuProjectAstExpression => Some(astExpression)
+      case _ => None
+    }
+  }
+
+  private def unwrap(expression: Expression): Expression = expression match {
+    case alias @ GpuAlias(astExpression: GpuProjectAstExpression, name) =>
+      GpuAlias(astExpression.child, name)(
+        alias.exprId, alias.qualifier, alias.explicitMetadata)
+    case astExpression: GpuProjectAstExpression => astExpression.child
+    case other => other
+  }
+
+  private def rewrap(expression: Expression): Expression = expression match {
+    case namedExpression: NamedExpression => wrap(namedExpression)
+    case gpuExpression: GpuExpression => GpuProjectAstExpression(gpuExpression)
+    case other => other
+  }
+
+  private[rapids] def buildExprTiers(
+      expressions: Seq[Expression],
+      conf: SQLConf): Seq[Seq[Expression]] = {
+    val astOutputs = expressions.map(extractTopLevel)
+    val astSubexpressions = astOutputs.flatten.flatMap { astExpression =>
+      astExpression.child.collect {
+        case gpuExpression: GpuExpression => GpuExpressionEquals(gpuExpression)
+      }
+    }.toSet
+    // CSE must see through the marker so AST and non-AST outputs can share the same tiers.
+    val unwrapped = expressions.map(unwrap)
+    val replaced = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) {
+      GpuEquivalentExpressions.replaceMultiExpressions(unwrapped, conf)
+    } else {
+      unwrapped
+    }
+    val tiers = GpuEquivalentExpressions.getExprTiers(
+      replaced,
+      (original, substituted) => (original, substituted) match {
+        case (gpuOriginal: GpuExpression, gpuSubstituted: GpuExpression)
+            if GpuBatchUtils.isFixedWidth(gpuOriginal.dataType) &&
+              astSubexpressions.contains(GpuExpressionEquals(gpuOriginal)) =>
+          GpuProjectAstExpression(gpuSubstituted)
+        case _ => substituted
+      })
+    val finalTier = tiers.last.zip(astOutputs).map {
+      case (expression, Some(_)) => rewrap(expression)
+      case (expression, None) => expression
+    }
+    tiers.dropRight(1) :+ finalTier
+  }
+
+  private[rapids] 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)
+        }
+      }
+    }
+  }
+}
+
+case class GpuProjectAstExpression(child: GpuExpression)
+    extends ShimUnaryExpression with GpuExpression with GpuMetricsInjectable with AutoCloseable {
+  @transient private[this] var compiledExpression: CompiledExpression = _
+  private[this] var opTime: GpuMetric = NoopMetric
+
+  override def dataType: DataType = child.dataType
+
+  override def nullable: Boolean = child.nullable
+
+  override def toString: String = s"AST($child)"
+
+  override def injectMetrics(metrics: Map[String, GpuMetric]): Unit = {
+    opTime = metrics.getOrElse(OP_TIME_LEGACY, NoopMetric)
+  }
+
+  override def close(): Unit = synchronized {
+    Option(compiledExpression).foreach(_.safeClose())
+    compiledExpression = null
+  }
+
+  override def columnarEval(batch: ColumnarBatch): GpuColumnVector = {
+    withResource(GpuProjectAstExpression.tableFromBatch(batch)) { table =>
+      computeColumn(table)
+    }
+  }
+
+  def computeColumn(table: Table): GpuColumnVector = {
+    NvtxIdWithMetrics(NvtxRegistry.PROJECT_AST, opTime) {
+      closeOnExcept(getCompiledExpression.computeColumn(table)) { result =>
+        GpuColumnVector.from(result, dataType)
+      }
+    }
+  }
+
+  private def getCompiledExpression: CompiledExpression = synchronized {
+    if (compiledExpression == null) {
+      val compiled = NvtxIdWithMetrics(NvtxRegistry.COMPILE_ASTS, opTime) {
+        // Project AST has a single input table.
+        child.convertToAst(Int.MaxValue).compile()
+      }
+      closeOnExcept(compiled) { _ =>
+        Option(TaskContext.get()).foreach { taskContext =>
+          onTaskCompletion(taskContext) {
+            close()
+          }
+        }
+        compiledExpression = compiled
+      }
+    }
+    compiledExpression
+  }
+}