diff --git a/docs/dev/nvtx_ranges.md b/docs/dev/nvtx_ranges.md index f3c31cf4324..187622cd3f1 100644 --- a/docs/dev/nvtx_ranges.md +++ b/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 diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index d96b836efda..15b81c2fa2c 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/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,110 @@ 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_partial_project_with_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="GpuProjectAst", + conf=_project_ast_jit_enabled_conf) + +@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn) +@disable_ansi_mode +def test_jit_multi_output_shared_subtree(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 * 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 left.*AST_JIT.*AS right", + 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_regular_jit_waves(data_gen): + def project_waves(spark): + df = binary_op_df(spark, data_gen) + shared = f.col('a') + f.col('b') + regular = f.greatest(shared, f.col('a')) + return df.select( + (shared * f.col('a')).alias('early'), + regular.alias('regular'), + (regular * f.col('b')).alias('late')) + + assert_cpu_and_gpu_are_equal_collect_with_capture( + project_waves, + exist_classes=r"GpuProject.*AST_JIT.*AS early.*AS regular.*AS late", + 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_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 diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala new file mode 100644 index 00000000000..cf6a10a678e --- /dev/null +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala @@ -0,0 +1,124 @@ +/* + * 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.{AstExpression, CompiledExpression} +import com.nvidia.spark.Retryable +import com.nvidia.spark.rapids.Arm.withResource + +import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression} +import org.apache.spark.sql.vectorized.ColumnarBatch + +object GpuAstJitExpression { + private[rapids] def canUseAstJit(expression: Expression): Boolean = expression match { + case gpuExpression: GpuExpression => + GpuBatchUtils.isFixedWidth(expression.dataType) && + gpuExpression.supportsAstJit && gpuExpression.containsAstJitOperator + case _ => false + } + + 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]) + } + + private[rapids] def computeColumns( + expressions: Seq[GpuAstJitExpression], + table: Table): ColumnarBatch = { + require(expressions.size > 1, "Multi-output AST JIT requires at least two expressions") + val compiledExpressions = expressions.map(_.getCompiledExpression).toArray + expressions.head.withComputeMetrics { + withResource(CompiledExpression.computeTableJit(table, compiledExpressions: _*)) { result => + GpuColumnVector.from(result, expressions.map(_.dataType).toArray) + } + } + } + + /** 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 def disableTieredProjectCombine: Boolean = true + + 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 compileAst(ast: AstExpression): CompiledExpression = ast.compileJit() + + 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 = () +} diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitProjectPlanner.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitProjectPlanner.scala new file mode 100644 index 00000000000..79f8a0beefa --- /dev/null +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitProjectPlanner.scala @@ -0,0 +1,270 @@ +/* + * 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.collection.mutable + +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression, Literal} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.rapids.catalyst.expressions.{ + GpuEquivalentExpressions, GpuExpressionEquals} + +/** + * Splits a Project expression forest at AST JIT/regular GPU boundaries. Expressions on the same + * dependency frontier remain in the same physical tier so the JIT executor can evaluate all of + * their roots together. Values referenced across frontiers become tier outputs. + */ +private[rapids] object GpuAstJitProjectPlanner { + private sealed trait Backend + private case object AstJit extends Backend + private case object RegularGpu extends Backend + + private final case class Candidate( + id: Int, + expression: Expression, + backend: Backend, + depth: Int, + alias: GpuAlias) + + private final case class FinalRoot( + expression: Expression, + child: Expression, + backend: Option[Backend], + depth: Int) + + def buildExprTiers(expressions: Seq[Expression], conf: SQLConf): Seq[Seq[Expression]] = { + val combined = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) { + GpuEquivalentExpressions.replaceMultiExpressions(expressions, conf) + } else { + expressions + } + new Planner(combined).build() + } + + private final class Planner(expressions: Seq[Expression]) { + private val candidatesByExpression = + mutable.HashMap.empty[GpuExpressionEquals, Candidate] + private val candidates = mutable.ArrayBuffer.empty[Candidate] + + private def stripAlias(expression: Expression): Expression = expression match { + case alias: GpuAlias => stripAlias(alias.child) + case other => other + } + + private def rootChild(expression: Expression): Expression = expression match { + case alias: GpuAlias => stripAlias(alias.child) + case other => stripAlias(other) + } + + private def replaceRootChild(expression: Expression, child: Expression): Expression = { + expression match { + case alias: GpuAlias => GpuProjectAstExpressionBase.replaceChild(alias, child) + case other => other + } + } + + private def backendOf(expression: Expression): Option[Backend] = stripAlias(expression) match { + case _: AttributeReference | _: GpuBoundReference | _: GpuLiteral | _: Literal => None + case gpuExpression: GpuExpression + if gpuExpression.deterministic && + GpuBatchUtils.isFixedWidth(gpuExpression.dataType) && + gpuExpression.selfSupportsAstJit && gpuExpression.selfIsAstJitOperator => + Some(AstJit) + case _ => Some(RegularGpu) + } + + private def canDescend(expression: Expression): Boolean = stripAlias(expression) match { + case gpuExpression: GpuExpression => !gpuExpression.disableTieredProjectCombine + case _ => true + } + + private def canCrossBoundary(expression: Expression, backend: Backend): Boolean = { + backend != AstJit || (stripAlias(expression) match { + case gpuExpression: GpuExpression => !gpuExpression.hasSideEffects + case _ => false + }) + } + + private def addDependency( + dependencies: mutable.LinkedHashMap[Int, Candidate], + candidate: Candidate): Unit = { + dependencies.getOrElseUpdate(candidate.id, candidate) + } + + private def collectDependencies( + expression: Expression, + ownerBackend: Backend): Seq[Candidate] = { + val dependencies = mutable.LinkedHashMap.empty[Int, Candidate] + + def visit(node: Expression): Unit = { + val child = stripAlias(node) + backendOf(child) match { + case Some(childBackend) + if childBackend != ownerBackend && canCrossBoundary(child, childBackend) => + addDependency(dependencies, candidateFor(child)) + case _ if canDescend(child) => child.children.foreach(visit) + case _ => + } + } + + stripAlias(expression).children.foreach(visit) + dependencies.values.toSeq + } + + private def candidateFor(expression: Expression): Candidate = { + val child = stripAlias(expression) + require(child.deterministic, s"Cannot materialize non-deterministic expression $child") + val key = GpuExpressionEquals(child) + candidatesByExpression.getOrElse(key, { + val backend = backendOf(child).getOrElse(RegularGpu) + val dependencies = collectDependencies(child, backend) + val depth = dependencies.map(_.depth + 1).foldLeft(0)(math.max) + val id = candidates.size + val alias = GpuAlias(child, s"project_wave_$id")() + val candidate = Candidate(id, child, backend, depth, alias) + candidatesByExpression.put(key, candidate) + candidates += candidate + candidate + }) + } + + private def findCandidate(expression: Expression): Option[Candidate] = { + val child = stripAlias(expression) + if (child.deterministic) { + candidatesByExpression.get(GpuExpressionEquals(child)) + } else { + None + } + } + + private def rewrite(expression: Expression, currentDepth: Int): Expression = { + def recurse(node: Expression, isRoot: Boolean): Expression = { + val child = stripAlias(node) + val earlierCandidate = if (isRoot) { + None + } else { + findCandidate(child).filter(_.depth < currentDepth) + } + earlierCandidate.map(_.alias.toAttribute).getOrElse { + if (canDescend(child)) { + child.mapChildren(recurse(_, isRoot = false)) + } else { + child + } + } + } + + recurse(stripAlias(expression), isRoot = true) + } + + private def regularCseTiers(regularExpressions: Seq[Expression]): Seq[Seq[Expression]] = { + if (regularExpressions.isEmpty) { + Seq.empty + } else { + GpuEquivalentExpressions.getExprTiers(regularExpressions) + } + } + + private def buildCandidateWave(waveCandidates: Seq[Candidate]): Seq[Seq[Expression]] = { + val rewritten = waveCandidates.map { candidate => + val child = rewrite(candidate.expression, candidate.depth) + candidate -> GpuProjectAstExpressionBase.replaceChild(candidate.alias, child) + } + val regularAliases = rewritten.collect { + case (candidate, alias) if candidate.backend == RegularGpu => alias + } + val jitAliases = rewritten.collect { + case (candidate, alias) if candidate.backend == AstJit => alias + } + val regularTiers = regularCseTiers(regularAliases) + if (regularTiers.isEmpty) { + Seq(jitAliases) + } else { + regularTiers.dropRight(1) :+ (regularTiers.last ++ jitAliases) + } + } + + private def buildFinalTiers( + finalRoots: Seq[FinalRoot], + finalProducers: Map[Int, Candidate], + finalDepth: Int): Seq[Seq[Expression]] = { + val rewrittenFinals = finalRoots.zipWithIndex.map { case (root, index) => + val child = finalProducers.get(index) + .map(_.alias.toAttribute) + .getOrElse(rewrite(root.child, finalDepth)) + replaceRootChild(root.expression, child) + } + val regularIndexes = rewrittenFinals.indices.filter { index => + !GpuAstJitExpression.canUseAstJit(rootChild(rewrittenFinals(index))) + } + val regularFinals = regularIndexes.map(rewrittenFinals(_)) + val regularTiers = regularCseTiers(regularFinals) + if (regularTiers.isEmpty) { + Seq(rewrittenFinals) + } else { + val rewrittenRegularFinals = regularTiers.last.iterator + val regularIndexSet = regularIndexes.toSet + val finalTier = rewrittenFinals.indices.map { index => + if (regularIndexSet.contains(index)) { + rewrittenRegularFinals.next() + } else { + rewrittenFinals(index) + } + } + regularTiers.dropRight(1) :+ finalTier + } + } + + def build(): Seq[Seq[Expression]] = { + val finalRoots = expressions.map { expression => + val child = rootChild(expression) + val backend = backendOf(child) + val dependencies = backend.map(collectDependencies(child, _)).getOrElse(Seq.empty) + val depth = dependencies.map(_.depth + 1).foldLeft(0)(math.max) + FinalRoot(expression, child, backend, depth) + } + + val hasJitCandidate = finalRoots.exists(_.backend.contains(AstJit)) || + candidates.exists(_.backend == AstJit) + if (!hasJitCandidate) { + return GpuEquivalentExpressions.getExprTiers(expressions) + } + + val finalDepth = finalRoots.map(_.depth).foldLeft(0)(math.max) + val finalProducers = mutable.LinkedHashMap.empty[Int, Candidate] + finalRoots.zipWithIndex.foreach { case (root, index) => + findCandidate(root.child).filter(_.depth < finalDepth).foreach { candidate => + finalProducers.put(index, candidate) + } + if (!finalProducers.contains(index) && root.depth < finalDepth && + root.backend.nonEmpty && root.child.deterministic) { + finalProducers.put(index, candidateFor(root.child)) + } + } + + val prioritized = finalProducers.values.toSeq.distinct + val priorityIds = prioritized.map(_.id).toSet + val orderedCandidates = prioritized ++ candidates.filterNot(c => priorityIds.contains(c.id)) + val candidateTiers = (0 until finalDepth).flatMap { depth => + val wave = orderedCandidates.filter(_.depth == depth) + if (wave.nonEmpty) buildCandidateWave(wave) else Seq.empty + } + (candidateTiers ++ buildFinalTiers(finalRoots, finalProducers.toMap, finalDepth)) + .map(_.toList).toList + } + } +} diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuBoundAttribute.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuBoundAttribute.scala index 5c6058cdafa..1341385b356 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuBoundAttribute.scala +++ b/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})" diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpressions.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpressions.scala index 9227e0b7a76..3a68a00f820 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpressions.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExpressions.scala @@ -197,6 +197,32 @@ 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. Returning true requires `convertToAst` to work for the same execution modes and + * 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 { @@ -321,6 +347,14 @@ object CudfUnaryExpression { trait CudfUnaryExpression extends GpuUnaryExpression { def unaryOp: UnaryOp + override final def selfSupportsAstJit: Boolean = + CudfUnaryExpression.opToAstMap.contains(unaryOp) && astJitCompatible + + /** Whether this operator's execution modes and types work with AST JIT. */ + protected def astJitCompatible: Boolean = false + + override final def selfIsAstJitOperator: Boolean = selfSupportsAstJit + override def doColumnar(input: GpuColumnVector): ColumnVector = input.getBase.unaryOp(unaryOp) override def convertToAst(numFirstTableColumns: Int): ast.AstExpression = { @@ -391,6 +425,13 @@ trait CudfBinaryExpression extends GpuBinaryExpression { def castOutputAtEnd: Boolean = false def astOperator: Option[ast.BinaryOperator] = None + override final def selfSupportsAstJit: Boolean = astOperator.isDefined && astJitCompatible + + /** Whether this operator's execution modes and types work with AST JIT. */ + protected def astJitCompatible: Boolean = false + + override final def selfIsAstJitOperator: Boolean = selfSupportsAstJit + def outputType(l: BinaryOperable, r: BinaryOperable): DType = { val over = outputTypeOverride if (over == null) { diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala index 9583ee85b95..47ec34bb1e9 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala @@ -4927,8 +4927,7 @@ object GpuOverrides extends Logging { } else { wrap.runAfterTagRules() wrap.tagForExplain() - val shouldExplainAll = explain.equalsIgnoreCase("ALL") - wrap.explain(shouldExplainAll) + wrap.explain(RapidsConf.shouldExplainAll(explain)) } } diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuProjectAstExpression.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuProjectAstExpression.scala index 79b5f8ed13a..162e804e910 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuProjectAstExpression.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuProjectAstExpression.scala @@ -19,7 +19,7 @@ package com.nvidia.spark.rapids import scala.annotation.tailrec import ai.rapids.cudf.{Scalar, Table} -import ai.rapids.cudf.ast.CompiledExpression +import ai.rapids.cudf.ast.{AstExpression, CompiledExpression} import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} import com.nvidia.spark.rapids.GpuMetric.OP_TIME_LEGACY import com.nvidia.spark.rapids.RapidsPluginImplicits._ @@ -27,110 +27,146 @@ 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.catalyst.expressions.{Expression, ExprId, NamedExpression} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.rapids.catalyst.expressions.GpuEquivalentExpressions import org.apache.spark.sql.types.DataType import org.apache.spark.sql.vectorized.ColumnarBatch -object GpuProjectAstExpression { - private def replaceChild(alias: GpuAlias, child: Expression): GpuAlias = { - if (child eq alias.child) { - alias - } else { - GpuAlias(child, alias.name)(alias.exprId, alias.qualifier, alias.explicitMetadata) +trait GpuProjectAstExpressionBase + extends ShimUnaryExpression with GpuExpression with GpuMetricsInjectable with AutoCloseable { + override def child: GpuExpression + + protected def backendName: String + protected def compileNvtxId: NvtxId + protected def computeNvtxId: NvtxId + protected def compileAst(ast: AstExpression): CompiledExpression = ast.compile() + + @transient private[this] var compiledExpression: CompiledExpression = _ + private[this] var opTime: GpuMetric = NoopMetric + + override final def dataType: DataType = child.dataType + + override final def nullable: Boolean = child.nullable + + override final def injectMetrics(metrics: Map[String, GpuMetric]): Unit = { + // OP_TIME_LEGACY is the owning operator's non-RDD timing metric, not a legacy AST metric. + opTime = metrics.getOrElse(OP_TIME_LEGACY, NoopMetric) + } + + override final def close(): Unit = { + val toClose = synchronized { + val current = compiledExpression + compiledExpression = null + current } + Option(toClose).foreach(_.safeClose()) } - private def asAst(child: GpuExpression): GpuProjectAstExpression = child match { - case astExpression: GpuProjectAstExpression => astExpression - case other => GpuProjectAstExpression(other) + override final def columnarEval(batch: ColumnarBatch): GpuColumnVector = { + withResource(GpuProjectAstExpressionBase.tableFromBatch(batch)) { table => + computeColumn(table) + } } - private[rapids] def wrap(expression: NamedExpression): NamedExpression = expression match { - case alias @ GpuAlias(child: GpuExpression, _) => - replaceChild(alias, asAst(child)) - case other => other + private[rapids] final def computeColumn(table: Table): GpuColumnVector = { + val compiled = getCompiledExpression + withComputeMetrics { + closeOnExcept(compiled.computeColumn(table)) { result => + GpuColumnVector.from(result, dataType) + } + } + } + + private[rapids] final def withComputeMetrics[T](body: => T): T = + NvtxIdWithMetrics(computeNvtxId, opTime)(body) + + private[rapids] final def getCompiledExpression: CompiledExpression = synchronized { + if (compiledExpression == null) { + val compiled = NvtxIdWithMetrics(compileNvtxId, opTime) { + // Force every bound reference to the left table; Project AST has one input table. + compileAst(child.convertToAst(Int.MaxValue)) + } + closeOnExcept(compiled) { _ => + var completed = false + Option(TaskContext.get()).foreach { taskContext => + onTaskCompletion(taskContext) { + completed = true + close() + } + } + if (completed) { + throw new IllegalStateException( + s"Task completed while registering the $backendName cleanup callback") + } + compiledExpression = compiled + } + } + compiledExpression } +} +object GpuProjectAstExpressionBase { @tailrec - private[rapids] def extractTopLevel(expression: Expression): Option[GpuProjectAstExpression] = { + private[rapids] def extractTopLevel( + expression: Expression): Option[GpuProjectAstExpressionBase] = { expression match { case alias: GpuAlias => extractTopLevel(alias.child) - case astExpression: GpuProjectAstExpression => Some(astExpression) + case astExpression: GpuProjectAstExpressionBase => Some(astExpression) case _ => None } } + private[rapids] def replaceChild(alias: GpuAlias, child: Expression): GpuAlias = { + if (child eq alias.child) { + alias + } else { + GpuAlias(child, alias.name)(alias.exprId, alias.qualifier, alias.explicitMetadata) + } + } + private def unwrap(expression: Expression): Expression = expression match { case alias: GpuAlias => replaceChild(alias, unwrap(alias.child)) case astExpression: GpuProjectAstExpression => astExpression.child + case jitExpression: GpuAstJitExpression => jitExpression.child case other => other } - private def rewrap(expression: Expression): Expression = expression match { - case namedExpression: NamedExpression => wrap(namedExpression) - case other => other - } - - private def rewrapAstTiers( - tiers: Seq[Seq[Expression]], - astOutputs: Seq[Boolean]): Seq[Seq[Expression]] = { - val finalTier = tiers.last - require(finalTier.size == astOutputs.size, - "The final expression tier must preserve the project output count") - val astReferences = finalTier.iterator.zip(astOutputs.iterator) - .collect { case (expression, true) => expression } - .flatMap(_.references.iterator) - .map(_.exprId) - .toSet - - // Tier aliases are the dataflow graph after CSE, so follow them backwards from AST outputs. - val (commonTiers, _) = tiers.dropRight(1).foldRight( - (List.empty[Seq[Expression]], astReferences)) { - case (tier, (rewrittenTiers, requiredExprIds)) => - val astAliases = tier.collect { - case alias: GpuAlias if requiredExprIds.contains(alias.exprId) => alias - } - val astAliasIds = astAliases.iterator.map(_.exprId).toSet - val dependencies = astAliases.iterator - .flatMap(_.references.iterator) - .map(_.exprId) - .toSet - val rewrittenTier = tier.map { - case alias: GpuAlias - if astAliasIds.contains(alias.exprId) && - GpuBatchUtils.isFixedWidth(alias.dataType) => - rewrap(alias) - case expression => expression - } - (rewrittenTier :: rewrittenTiers, requiredExprIds ++ dependencies) - } - - commonTiers :+ finalTier.zip(astOutputs).map { - case (expression, true) => rewrap(expression) - case (expression, false) => expression - } - } - private[rapids] def buildExprTiers( expressions: Seq[Expression], - conf: SQLConf): Seq[Seq[Expression]] = { - val astOutputs = expressions.map(extractTopLevel(_).isDefined) + conf: SQLConf, + enableProjectAstJit: Boolean = false): Seq[Seq[Expression]] = { + val astOutputs = expressions.map(GpuProjectAstExpression.extractTopLevel(_).isDefined) val hasAstOutputs = astOutputs.contains(true) - // CSE must see through the marker so AST and non-AST outputs can share the same tiers. - val unwrapped = if (hasAstOutputs) expressions.map(unwrap) else expressions - val replaced = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) { - GpuEquivalentExpressions.replaceMultiExpressions(unwrapped, conf) + val hasJitOutputs = expressions.exists(GpuAstJitExpression.extractTopLevel(_).isDefined) + // CSE must see through backend markers so all outputs can share the same tiers. + val unwrapped = if (hasAstOutputs || hasJitOutputs) { + expressions.map(unwrap) } else { - unwrapped + expressions } - val tiers = GpuEquivalentExpressions.getExprTiers(replaced) - if (hasAstOutputs) { - rewrapAstTiers(tiers, astOutputs) + val tiers = if (enableProjectAstJit) { + GpuAstJitProjectPlanner.buildExprTiers(unwrapped, conf) + } else { + val replaced = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) { + GpuEquivalentExpressions.replaceMultiExpressions(unwrapped, conf) + } else { + unwrapped + } + GpuEquivalentExpressions.getExprTiers(replaced) + } + val astTiers = if (hasAstOutputs) { + GpuProjectAstExpression.rewrapAstTiers(tiers, astOutputs) } else { tiers } + if (enableProjectAstJit) { + // Project binding selects JIT after CSE so newly exposed tiers are eligible. + astTiers.map(_.map(GpuAstJitExpression.wrapTierExpression)) + } else { + // Only the Project-specific binder selects JIT after CSE. + astTiers + } } private[rapids] def tableFromBatch(batch: ColumnarBatch): Table = { @@ -147,56 +183,76 @@ object GpuProjectAstExpression { } } -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)" +object GpuProjectAstExpression { - override def injectMetrics(metrics: Map[String, GpuMetric]): Unit = { - opTime = metrics.getOrElse(OP_TIME_LEGACY, NoopMetric) + private def asAst(child: GpuExpression): GpuProjectAstExpression = child match { + case astExpression: GpuProjectAstExpression => astExpression + case other => GpuProjectAstExpression(other) } - override def close(): Unit = synchronized { - Option(compiledExpression).foreach(_.safeClose()) - compiledExpression = null + private[rapids] def wrap(expression: NamedExpression): NamedExpression = expression match { + case alias @ GpuAlias(child: GpuExpression, _) => + GpuProjectAstExpressionBase.replaceChild(alias, asAst(child)) + case other => other } - override def columnarEval(batch: ColumnarBatch): GpuColumnVector = { - withResource(GpuProjectAstExpression.tableFromBatch(batch)) { table => - computeColumn(table) + /** Extracts a legacy Project AST wrapper after unwrapping any top-level aliases. */ + private[rapids] def extractTopLevel(expression: Expression): Option[GpuProjectAstExpression] = + GpuProjectAstExpressionBase.extractTopLevel(expression).collect { + case astExpression: GpuProjectAstExpression => astExpression } + + private def rewrap(expression: Expression): Expression = expression match { + case namedExpression: NamedExpression => wrap(namedExpression) + case other => other } - private[rapids] def computeColumn(table: Table): GpuColumnVector = { - val compiled = getCompiledExpression - NvtxIdWithMetrics(NvtxRegistry.PROJECT_AST, opTime) { - closeOnExcept(compiled.computeColumn(table)) { result => - GpuColumnVector.from(result, dataType) - } + private[rapids] def rewrapAstTiers( + tiers: Seq[Seq[Expression]], + astOutputs: Seq[Boolean]): Seq[Seq[Expression]] = { + val finalTier = tiers.last + require(finalTier.size == astOutputs.size, + "The final expression tier must preserve the project output count") + def referenceSet(taggedTier: Iterable[(Expression, Boolean)]): Set[ExprId] = { + taggedTier.collect { case (expression, true) => expression } + .flatMap(_.references.iterator) + .map(_.exprId) + .toSet } - } + val astReferences = referenceSet(finalTier.zip(astOutputs)) - private def getCompiledExpression: CompiledExpression = synchronized { - if (compiledExpression == null) { - val compiled = NvtxIdWithMetrics(NvtxRegistry.COMPILE_ASTS, opTime) { - // Force every bound reference to the left table; Project AST has one input table. - child.convertToAst(Int.MaxValue).compile() - } - closeOnExcept(compiled) { _ => - Option(TaskContext.get()).foreach { taskContext => - onTaskCompletion(taskContext) { - close() - } + // Tier aliases are the dataflow graph after CSE, so follow them backwards from AST outputs. + val (commonTiers, _) = tiers.dropRight(1).foldRight( + (List.empty[Seq[Expression]], astReferences)) { + case (tier, (rewrittenTiers, requiredExprIds)) => + val taggedTier = tier.map { + case alias: GpuAlias if requiredExprIds.contains(alias.exprId) => (alias, true) + case expression => (expression, false) } - compiledExpression = compiled - } + val dependencies = referenceSet(taggedTier) + val rewrittenTier = taggedTier.map { + case (alias, true) if GpuBatchUtils.isFixedWidth(alias.dataType) => rewrap(alias) + case (expression, _) => expression + } + (rewrittenTier :: rewrittenTiers, requiredExprIds ++ dependencies) + } + + commonTiers :+ finalTier.zip(astOutputs).map { + case (expression, true) => rewrap(expression) + case (expression, false) => expression } - compiledExpression } + +} + +case class GpuProjectAstExpression(child: GpuExpression) + extends GpuProjectAstExpressionBase { + + override protected def backendName: String = "Project AST" + + override protected def compileNvtxId: NvtxId = NvtxRegistry.COMPILE_ASTS + + override protected def computeNvtxId: NvtxId = NvtxRegistry.PROJECT_AST + + override def toString: String = s"AST($child)" } diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxRangeWithDoc.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxRangeWithDoc.scala index f5c9c4a137c..556fa571f9e 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxRangeWithDoc.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/NvtxRangeWithDoc.scala @@ -327,6 +327,12 @@ object NvtxRegistry { val COMPILE_ASTS: NvtxId = NvtxId("Compile ASTs", NvtxColor.ORANGE, "Compiling abstract syntax trees for expression evaluation") + val PROJECT_AST_JIT: NvtxId = NvtxId("Project AST JIT", NvtxColor.CYAN, + "Applying JIT-compiled AST projection to batch") + + val COMPILE_AST_JIT: NvtxId = NvtxId("Compile AST JIT", NvtxColor.ORANGE, + "Compiling an AST expression for JIT evaluation") + // Aggregate operations val COMPUTE_AGGREGATE: NvtxId = NvtxId("computeAggregate", NvtxColor.CYAN, "Computing aggregation on input batch") @@ -759,6 +765,8 @@ object NvtxRegistry { register(PROJECT_EXEC) register(PROJECT_AST) register(COMPILE_ASTS) + register(PROJECT_AST_JIT) + register(COMPILE_AST_JIT) register(COMPUTE_AGGREGATE) register(FINALIZE_AGG) register(POST_PROCESS_AGG) diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala index 1fbab1427d6..5cff2e9aa62 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala @@ -1247,6 +1247,21 @@ 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. " + + "When both project AST backends are enabled, JIT takes precedence for supported " + + "expressions within each projection tier.") + .internal() + .booleanConf + .createWithDefault(false) + + val ENABLE_PROJECT_AST_JIT_MULTI_OUTPUT = + conf("spark.rapids.sql.projectAstJitMultiOutputEnabled") + .doc("Evaluate multiple Project AST JIT expressions in one cuDF call when possible.") + .internal() + .booleanConf + .createWithDefault(true) + val ENABLE_TIERED_PROJECT = conf("spark.rapids.sql.tiered.project.enabled") .doc("Enable tiered projections.") .internal() @@ -2545,6 +2560,12 @@ val SHUFFLE_COMPRESSION_LZ4_CHUNK_SIZE = conf("spark.rapids.shuffle.compression. .stringConf .createWithDefault("NOT_ON_GPU") + private[rapids] def shouldExplain(explain: String): Boolean = + !explain.equalsIgnoreCase("NONE") + + private[rapids] def shouldExplainAll(explain: String): Boolean = + explain.equalsIgnoreCase("ALL") + val SHIMS_PROVIDER_OVERRIDE = conf("spark.rapids.shims-provider-override") .internal() .startupOnly() @@ -3538,9 +3559,9 @@ class RapidsConf(conf: Map[String, String]) extends Logging { lazy val explain: String = get(EXPLAIN) - lazy val shouldExplain: Boolean = !explain.equalsIgnoreCase("NONE") + lazy val shouldExplain: Boolean = RapidsConf.shouldExplain(explain) - lazy val shouldExplainAll: Boolean = explain.equalsIgnoreCase("ALL") + lazy val shouldExplainAll: Boolean = RapidsConf.shouldExplainAll(explain) lazy val chunkedReaderEnabled: Boolean = get(CHUNKED_READER) @@ -3634,6 +3655,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) diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala index 9f9c0ffdc17..7f199aea50d 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala @@ -16,6 +16,8 @@ package com.nvidia.spark.rapids +import java.util.{ArrayDeque, IdentityHashMap} + import scala.annotation.tailrec import scala.collection.mutable.ArrayBuffer @@ -62,35 +64,50 @@ 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() + val jitExprs = if (conf.isProjectAstJitEnabled) { + GpuAstJitExpression.wrapProjectExpressions(gpuExprs) + } else { + gpuExprs + } val projectList = if (conf.isProjectAstEnabled) { - val astExprs = childExprs.zip(gpuExprs).map { case (meta, expr) => + childExprs.zip(jitExprs).map { case (meta, expr) => // cuDF requires return column is fixed width // Regular projection can reuse its cached null vector across outputs. - if (GpuBatchUtils.isFixedWidth(expr.dataType) && meta.canThisBeAst && + if (GpuAstJitExpression.extractTopLevel(expr).isEmpty && + GpuBatchUtils.isFixedWidth(expr.dataType) && meta.canThisBeAst && !isTopLevelNullLiteral(expr)) { GpuProjectAstExpression.wrap(expr) } else { expr } }.toList - // explain AST because this is optional and it is sometimes hard to debug - if (conf.shouldExplain) { - val explain = (childExprs.iterator.map(_.explainAst(conf.shouldExplainAll)) - .filter(_.nonEmpty) ++ gpuExprs.iterator.collect { - case expr if !GpuBatchUtils.isFixedWidth(expr.dataType) => - s" $expr cannot be converted to AST because its return type " + - s"${expr.dataType} is not fixed-width\n" - case expr if isTopLevelNullLiteral(expr) => - s" $expr will use the regular GPU projection so null outputs can reuse " + - "the cached null vector\n" - }).mkString - if (explain.nonEmpty) { - logWarning(s"AST PROJECT\n$explain") - } - } - astExprs } else { - gpuExprs + jitExprs + } + // Legacy Project AST eligibility is decided here. JIT selection is reported after tiering. + if (conf.shouldExplain && conf.isProjectAstEnabled) { + val legacyExplain = childExprs.iterator.zip(projectList.iterator).flatMap { + case (meta, expression) if GpuAstJitExpression.extractTopLevel(expression).isEmpty => + Some(meta.explainAst(conf.shouldExplainAll)).filter(_.nonEmpty) + .map(explanation => s" Legacy Project AST eligibility:\n$explanation") + case _ => None + } + val regularExplain = gpuExprs.iterator.zip(projectList.iterator).collect { + case (expr, expression) + if GpuProjectAstExpressionBase.extractTopLevel(expression).isEmpty && + !GpuBatchUtils.isFixedWidth(expr.dataType) => + s" $expr cannot be converted to legacy Project AST because its return type " + + s"${expr.dataType} is not fixed-width\n" + case (expr, expression) + if GpuProjectAstExpressionBase.extractTopLevel(expression).isEmpty && + isTopLevelNullLiteral(expr) => + s" $expr will use the regular GPU projection instead of legacy Project AST so " + + "null outputs can reuse the cached null vector\n" + } + val explain = (legacyExplain ++ regularExplain).mkString + if (explain.nonEmpty) { + logWarning(s"LEGACY PROJECT AST\n$explain") + } } GpuProjectExec(projectList, gpuChild) } @@ -153,15 +170,45 @@ object GpuProjectExec { } } - val hasAstExpressions = boundExprs.exists { expression => - GpuProjectAstExpression.extractTopLevel(expression).isDefined - } - if (hasAstExpressions) { - withResource(GpuProjectAstExpression.tableFromBatch(cb)) { table => - projectWithEval { expression => - GpuProjectAstExpression.extractTopLevel(expression) match { - case Some(astExpression) => astExpression.computeColumn(table) - case None => expression.columnarEval(cb) + val astExpressions = boundExprs.flatMap( + GpuProjectAstExpressionBase.extractTopLevel) + if (astExpressions.nonEmpty) { + withResource(GpuProjectAstExpressionBase.tableFromBatch(cb)) { table => + val jitExpressions = astExpressions.collect { + case jitExpression: GpuAstJitExpression => jitExpression + } + if (jitExpressions.size > 1 && + RapidsConf.ENABLE_PROJECT_AST_JIT_MULTI_OUTPUT.get(SQLConf.get)) { + withResource(GpuAstJitExpression.computeColumns(jitExpressions, table)) { jitBatch => + val jitColumnIndexes = + new IdentityHashMap[GpuAstJitExpression, ArrayDeque[Int]]() + jitExpressions.zipWithIndex.foreach { case (jitExpression, index) => + var indexes = jitColumnIndexes.get(jitExpression) + if (indexes == null) { + indexes = new ArrayDeque[Int]() + jitColumnIndexes.put(jitExpression, indexes) + } + indexes.addLast(index) + } + projectWithEval { expression => + GpuProjectAstExpressionBase.extractTopLevel(expression) match { + case Some(jitExpression: GpuAstJitExpression) => + val indexes = jitColumnIndexes.get(jitExpression) + require(indexes != null && !indexes.isEmpty, + s"Missing multi-output AST JIT column for $jitExpression") + jitBatch.column(indexes.removeFirst()) + .asInstanceOf[GpuColumnVector].incRefCount() + case Some(astExpression) => astExpression.computeColumn(table) + case None => expression.columnarEval(cb) + } + } + } + } else { + projectWithEval { expression => + GpuProjectAstExpressionBase.extractTopLevel(expression) match { + case Some(astExpression) => astExpression.computeColumn(table) + case None => expression.columnarEval(cb) + } } } } @@ -902,8 +949,8 @@ case class GpuProjectExec( val opTime = gpuLongMetric(OP_TIME_LEGACY) val numPreSplit = gpuLongMetric(KEY_NUM_PRE_SPLIT) - val boundProjectList = GpuBindReferences.bindGpuReferencesTiered(projectList, child.output, - conf, allMetrics) + val boundProjectList = GpuBindReferences.bindGpuProjectReferencesTiered( + projectList, child.output, conf, allMetrics) val localEnablePreSplit = enablePreSplit val rdd = child.executeColumnar() @@ -1046,8 +1093,8 @@ case class GpuProjectExec( None } withResource(sbToClose) { _ => - retryables.foreach(_.checkpoint()) RmmRapidsRetryIterator.withRetryNoSplit { + retryables.foreach(_.checkpoint()) withResource(sb.getColumnarBatch()) { cb => withRestoreOnRetry(retryables) { project(cb) diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/literals.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/literals.scala index fc6566dc222..e471189701d 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/literals.scala +++ b/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 diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/arithmetic.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/arithmetic.scala index 7c823128e10..83a632b05a5 100644 --- a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/arithmetic.scala +++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/arithmetic.scala @@ -314,6 +314,9 @@ abstract class GpuAddBase extends CudfBinaryArithmetic with Serializable { override def binaryOp: BinaryOp = BinaryOp.ADD override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.ADD) + override protected def astJitCompatible: Boolean = + !failOnError && (dataType == IntegerType || dataType == LongType) + override def hasSideEffects: Boolean = (failOnError && GpuAnsi.needBasicOpOverflowCheck(dataType)) || super.hasSideEffects @@ -768,6 +771,9 @@ case class GpuMultiply( override def binaryOp: BinaryOp = BinaryOp.MUL override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.MUL) + override protected def astJitCompatible: Boolean = + !failOnError && (dataType == IntegerType || dataType == LongType) + private def multiplyOverflowError(msg: String): ArithmeticException = { RapidsErrorUtils.arithmeticOverflowError(msg, origin) } diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastHashJoinExecBase.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastHashJoinExecBase.scala index 13cee7f7f3c..569e1790959 100644 --- a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastHashJoinExecBase.scala +++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastHashJoinExecBase.scala @@ -184,14 +184,13 @@ abstract class GpuBroadcastHashJoinExecBase( val boundProjects = projects.map { project => // Match GpuProjectExec's tiered binding so build-side extraction has the same // splitting and retry behavior as a normal project. - GpuBindReferences.bindGpuReferencesTiered( + GpuBindReferences.bindGpuProjectReferencesTiered( project.projectList, project.child.output, conf, allMetrics) } Some((batch: ColumnarBatch) => boundProjects.foldLeft(batch) { case (currentBatch, boundProject) => - val spillableBatch = SpillableColumnarBatch( - currentBatch, SpillPriorities.ACTIVE_ON_DECK_PRIORITY) - boundProject.projectAndCloseWithRetrySingleBatch(spillableBatch) + boundProject.projectAndCloseWithRetrySingleBatch( + SpillableColumnarBatch(currentBatch, SpillPriorities.ACTIVE_ON_DECK_PRIORITY)) }) } else { None diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExecBase.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExecBase.scala index 398af19c4da..35a9354d43d 100644 --- a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExecBase.scala +++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExecBase.scala @@ -726,18 +726,18 @@ abstract class GpuBroadcastNestedLoopJoinExecBase( Some(buildDataSize)) } - private def buildSidePostProjection: Option[ColumnarBatch => ColumnarBatch] = { + private[execution] def buildSidePostProjection: Option[ColumnarBatch => ColumnarBatch] = { buildPlan match { case p: GpuProjectExec => // Need to manually do project columnar execution other than calling child's // internalDoExecuteColumnar. This is to workaround especial handle to build broadcast // batch. - val proj = GpuBindReferences.bindGpuReferencesTiered( + val proj = GpuBindReferences.bindGpuProjectReferencesTiered( postBuildCondition, p.child.output, conf, allMetrics) - val fn = (batch: ColumnarBatch) => { - withResource(batch)(proj.project) + Some { batch: ColumnarBatch => + proj.projectAndCloseWithRetrySingleBatch( + SpillableColumnarBatch(batch, SpillPriorities.ACTIVE_ON_DECK_PRIORITY)) } - Some(fn) case _ => None } diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuArrayHofFusionSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuArrayHofFusionSuite.scala index b6aaac3def3..3d5420dac96 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuArrayHofFusionSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuArrayHofFusionSuite.scala @@ -216,27 +216,27 @@ class GpuArrayHofFusionSuite extends GpuUnitTests { withResource(FuzzerUtils.createColumnarBatch(schema, 8))(check) } - test("fused HOF project preserves the shared AST input table") { + test("fused HOF project preserves the shared legacy and JIT AST input table") { val arrayType = ArrayType(IntegerType, containsNull = true) val schema = FuzzerUtils.createSchema(arrayType, LongType, LongType) - val firstAst = spy(GpuProjectAstExpression(GpuAdd( + val legacyAst = spy(GpuProjectAstExpression(GpuAdd( GpuBoundReference(1, LongType, nullable = true)(ExprId(400), "a"), GpuBoundReference(2, LongType, nullable = true)(ExprId(401), "b"), failOnError = false)())) - val secondAst = spy(GpuProjectAstExpression(GpuMultiply( + val jitAst = spy(GpuAstJitExpression(GpuMultiply( GpuBoundReference(1, LongType, nullable = true)(ExprId(402), "a"), GpuBoundReference(2, LongType, nullable = true)(ExprId(403), "b"), failOnError = false)())) // Prevent side effects in ANSI mode. val expressions = Seq( alias(executableTransform(404), "left"), - alias(firstAst, "sum"), - alias(secondAst, "product"), + alias(legacyAst, "sum"), + alias(jitAst, "product"), alias(executableTransform(405), "right")) assertResult(Seq(Seq(0, 3))) { GpuArrayHofFusion.findFusedGroupIndexes(expressions) } - withResource(Seq(firstAst, secondAst)) { _ => + withResource(Seq(legacyAst, jitAst)) { _ => withResource(FuzzerUtils.createColumnarBatch(schema, 8)) { batch => withResource(GpuProjectExec.project(batch, expressions)) { projected => assertResult(4)(projected.numCols()) @@ -245,10 +245,10 @@ class GpuArrayHofFusionSuite extends GpuUnitTests { } } - val firstTable = ArgumentCaptor.forClass(classOf[Table]) - val secondTable = ArgumentCaptor.forClass(classOf[Table]) - verify(firstAst).computeColumn(firstTable.capture()) - verify(secondAst).computeColumn(secondTable.capture()) - assert(firstTable.getValue eq secondTable.getValue) + val legacyTable = ArgumentCaptor.forClass(classOf[Table]) + val jitTable = ArgumentCaptor.forClass(classOf[Table]) + verify(legacyAst).computeColumn(legacyTable.capture()) + verify(jitAst).computeColumn(jitTable.capture()) + assert(legacyTable.getValue eq jitTable.getValue) } } diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala new file mode 100644 index 00000000000..95024aaa7e1 --- /dev/null +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala @@ -0,0 +1,453 @@ +/* + * 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.{AstExpression, CompiledExpression} +import com.nvidia.spark.rapids.ProjectAstTestUtils.collectExpressions +import org.mockito.Mockito.{doThrow, mock, times, verify, when} +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.TaskContext +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.rapids.{GpuAdd, GpuGreatest, GpuMultiply, GpuSubtract} +import org.apache.spark.sql.rapids.metrics.source.MockTaskContext +import org.apache.spark.sql.types.{FloatType, IntegerType, LongType} +import org.apache.spark.util.TaskCompletionListener + +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)() + + private def mockCompiledChild( + compiled: CompiledExpression, + jit: Boolean = false): GpuExpression = { + val child = mock(classOf[GpuExpression]) + val ast = mock(classOf[AstExpression]) + when(child.convertToAst(Int.MaxValue)).thenReturn(ast) + if (jit) { + when(ast.compileJit()).thenReturn(compiled) + } else { + when(ast.compile()).thenReturn(compiled) + } + child + } + + private def mockJitExpression(compiled: CompiledExpression): GpuAstJitExpression = + GpuAstJitExpression(mockCompiledChild(compiled, jit = true)) + + private def projectConf( + tiered: Boolean = true, + jit: Boolean = true, + legacy: Boolean = false): SQLConf = { + val conf = new SQLConf() + conf.setConfString(RapidsConf.ENABLE_TIERED_PROJECT.key, tiered.toString) + conf.setConfString(RapidsConf.ENABLE_PROJECT_AST_JIT.key, jit.toString) + conf.setConfString(RapidsConf.ENABLE_PROJECT_AST.key, legacy.toString) + conf + } + + 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] + val wrappedAgain = GpuAstJitExpression.wrapProjectExpressions(wrapped) + assert(jit.child.isInstanceOf[GpuMultiply]) + assert(jit.child.find(_.isInstanceOf[GpuAstJitExpression]).isEmpty) + assert(GpuAstJitExpression.extractTopLevel(wrapped.head).contains(jit)) + assert(wrappedAgain.head eq wrapped.head) + } + + test("project AST JIT only wraps a fully supported top-level expression") { + 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(wrapped.forall(GpuAstJitExpression.extractTopLevel(_).isEmpty)) + assert(subtract.left.isInstanceOf[GpuAdd]) + assert(subtract.right.isInstanceOf[GpuMultiply]) + } + + test("project wave exports a shared supported expression to JIT") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val third = reference(2, IntegerType) + val fourth = reference(3, IntegerType) + val shared = GpuAdd(left, right, failOnError = false)() + // [AST((left+right)-third) AS legacy, greatest(left+right, fourth) AS regular] + val expressions = Seq( + GpuProjectAstExpression.wrap( + alias(GpuSubtract(shared, third, failOnError = false)(), "legacy")), + alias(GpuGreatest(Seq(shared, fourth)), "regular")) + + val tiered = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + expressions, Seq(left, right, third, fourth), projectConf()) + + // after wave planning: + // tier 0: [AST_JIT(left+right) AS t1] + // tier 1: [AST(t1-third) AS legacy, greatest(t1, fourth) AS regular] + val jitExpressionTiers = tiered.exprTiers.map(collectExpressions[GpuAstJitExpression]) + assertResult(Seq(1, 0))(jitExpressionTiers.map(_.size)) + assert(jitExpressionTiers.head.head.child.isInstanceOf[GpuAdd]) + assert(GpuProjectAstExpression.extractTopLevel(tiered.exprTiers.last.head).isDefined) + assert(GpuProjectAstExpression.extractTopLevel(tiered.exprTiers.last(1)).isEmpty) + // Final references: [t1, t1] (distinct: {t1}). + val waveExprId = tiered.exprTiers.head.collectFirst { + case alias: GpuAlias if GpuAstJitExpression.extractTopLevel(alias).isDefined => alias.exprId + }.get + val finalTierReferences = collectExpressions[GpuBoundReference](tiered.exprTiers.last) + .filter(_.exprId == waveExprId) + assertResult(2)(finalTierReferences.size) + assertResult(1)(finalTierReferences.map(_.exprId).distinct.size) + } + + test("same-wave JIT roots keep their shared subtree inside one group") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val third = reference(2, IntegerType) + val fourth = reference(3, IntegerType) + def shared = GpuAdd(left, right, failOnError = false)() + val expressions = Seq( + alias(GpuMultiply(shared, third, failOnError = false)(), "first"), + alias(GpuMultiply(shared, fourth, failOnError = false)(), "second")) + + val tiered = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + expressions, Seq(left, right, third, fourth), projectConf()) + val jitTiers = tiered.exprTiers.map(_.flatMap(GpuAstJitExpression.extractTopLevel)) + + assertResult(1)(tiered.exprTiers.size) + assertResult(Seq(2))(jitTiers.map(_.size)) + assert(jitTiers.head.forall(_.child.isInstanceOf[GpuMultiply])) + assert(jitTiers.head.forall(_.child.find(_.isInstanceOf[GpuAdd]).isDefined)) + } + + test("JIT wave exports a shared subtree only for its regular consumer") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val third = reference(2, IntegerType) + val fourth = reference(3, IntegerType) + def shared = GpuAdd(left, right, failOnError = false)() + val expressions = Seq( + alias(GpuMultiply(shared, third, failOnError = false)(), "jit"), + alias(GpuGreatest(Seq(shared, fourth)), "regular")) + + val tiered = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + expressions, Seq(left, right, third, fourth), projectConf()) + val jitTiers = tiered.exprTiers.map(_.flatMap(GpuAstJitExpression.extractTopLevel)) + + assertResult(Seq(2, 0))(jitTiers.map(_.size)) + assertResult(1)(jitTiers.head.count(_.child.isInstanceOf[GpuMultiply])) + assertResult(1)(jitTiers.head.count(_.child.isInstanceOf[GpuAdd])) + assertResult(2)(tiered.exprTiers.last.size) + assert(collectExpressions[GpuGreatest](tiered.exprTiers.last).nonEmpty) + } + + test("JIT and regular dependencies form three waves") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val third = reference(2, IntegerType) + val fourth = reference(3, IntegerType) + def shared = GpuAdd(left, right, failOnError = false)() + def regular = GpuGreatest(Seq(shared, third)) + val expressions = Seq( + alias(GpuMultiply(shared, fourth, failOnError = false)(), "early"), + alias(regular, "regular"), + alias(GpuMultiply(regular, fourth, failOnError = false)(), "late")) + + val tiered = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + expressions, Seq(left, right, third, fourth), projectConf()) + val jitTiers = tiered.exprTiers.map(_.flatMap(GpuAstJitExpression.extractTopLevel)) + + assertResult(3)(tiered.exprTiers.size) + assertResult(Seq(2, 0, 1))(jitTiers.map(_.size)) + assertResult(1)(jitTiers.head.count(_.child.isInstanceOf[GpuAdd])) + assertResult(1)(jitTiers.head.count(_.child.isInstanceOf[GpuMultiply])) + assert(jitTiers.last.head.child.isInstanceOf[GpuMultiply]) + assert(collectExpressions[GpuGreatest](tiered.exprTiers(1)).nonEmpty) + } + + test("unsupported root uses maximal JIT children in an earlier wave") { + 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 tiered = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + Seq(expression), Seq(left, right), projectConf()) + val jitTiers = tiered.exprTiers.map(_.flatMap(GpuAstJitExpression.extractTopLevel)) + + assertResult(2)(tiered.exprTiers.size) + assertResult(Seq(2, 0))(jitTiers.map(_.size)) + assert(collectExpressions[GpuSubtract](tiered.exprTiers.last).nonEmpty) + } + + test("JIT planner keeps non-deterministic arithmetic on the regular backend") { + val expression = alias( + GpuAdd( + GpuMonotonicallyIncreasingID(), + GpuLiteral(1L, LongType), + failOnError = false)(), + "result") + + val tiered = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + Seq(expression), Seq.empty, projectConf()) + + assertResult(1)(tiered.exprTiers.size) + assert(collectExpressions[GpuAstJitExpression](tiered.exprTiers.head).isEmpty) + assert(collectExpressions[GpuMonotonicallyIncreasingID](tiered.exprTiers.head).nonEmpty) + } + + test("final JIT explanation includes a shared expression selected in an earlier tier") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val third = reference(2, IntegerType) + val fourth = reference(3, IntegerType) + val shared = GpuAdd(left, right, failOnError = false)() + val expressions = Seq( + alias(GpuSubtract(shared, third, failOnError = false)(), "first"), + alias(GpuSubtract(shared, fourth, failOnError = false)(), "second")) + + val tiered = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + expressions, Seq(left, right, third, fourth), projectConf()) + val all = GpuAstJitExpression.explainFinalSelections(tiered.exprTiers, all = true) + + assert(all.contains("TIER 0"), all) + assert(all.contains("AST_JIT"), all) + assert(all.contains("final backend: Project AST JIT"), all) + assert(all.contains("TIER 1"), all) + assert(all.contains("final backend: the regular GPU projection"), all) + assertResult("")( + GpuAstJitExpression.explainFinalSelections(tiered.exprTiers, all = false)) + } + + test("project JIT takes precedence while unsupported legacy AST falls back") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val jitCandidate = GpuProjectAstExpression.wrap( + alias(GpuAdd(left, right, failOnError = false)(), "jit")) + val legacyCandidate = GpuProjectAstExpression.wrap( + alias(GpuSubtract(left, right, failOnError = false)(), "legacy")) + + val tiered = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + Seq(jitCandidate, legacyCandidate), Seq(left, right), projectConf(legacy = true)) + val outputs = tiered.exprTiers.last + + assert(GpuAstJitExpression.extractTopLevel(outputs.head).isDefined) + assert(GpuProjectAstExpression.extractTopLevel(outputs(1)).isDefined) + } + + test("only the project binder selects the JIT backend") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val expression = alias(GpuAdd(left, right, failOnError = false)(), "result") + val conf = projectConf() + + val generic = GpuBindReferences.bindGpuReferencesTieredNoMetrics( + Seq(expression), Seq(left, right), conf) + val project = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + Seq(expression), Seq(left, right), conf) + + assert(collectExpressions[GpuAstJitExpression](generic.exprTiers.flatten).isEmpty) + assertResult(1)(collectExpressions[GpuAstJitExpression](project.exprTiers.flatten).size) + } + + test("the project binder respects a disabled JIT setting") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val expression = alias(GpuAdd(left, right, failOnError = false)(), "result") + + val project = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + Seq(expression), Seq(left, right), projectConf(jit = false)) + + assert(collectExpressions[GpuAstJitExpression](project.exprTiers.flatten).isEmpty) + } + + test("project JIT remains available when tiered projection is disabled") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val expressions = Seq( + alias(GpuAdd(left, right, failOnError = false)(), "sum"), + alias(GpuMultiply(left, right, failOnError = false)(), "product")) + + val project = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + expressions, Seq(left, right), projectConf(tiered = false)) + + assertResult(1)(project.exprTiers.size) + assertResult(2)(collectExpressions[GpuAstJitExpression](project.exprTiers.head).size) + } + + 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)) + } + + test("project AST JIT does not compile a literal or pass-through") { + val input = reference(0, IntegerType) + val expressions = List( + alias(GpuLiteral(7, IntegerType), "literal"), + GpuAlias(input, "pass_through")()) + + val wrapped = GpuAstJitExpression.wrapProjectExpressions(expressions) + + assert(wrapped.forall(GpuAstJitExpression.extractTopLevel(_).isEmpty)) + } + + test("final JIT explanation reports only unselected JIT candidates by default") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val third = reference(2, IntegerType) + val add = alias(GpuAdd(left, right, failOnError = false)(), "jit") + val subtract = alias( + GpuSubtract( + GpuAdd(left, right, failOnError = false)(), + third, + failOnError = false)(), + "regular") + val jit = GpuAstJitExpression.wrapProjectExpressions(List(add)).head + val legacy = GpuProjectAstExpression.wrap(subtract) + val selections = Seq(Seq(jit, legacy, subtract)) + + val all = GpuAstJitExpression.explainFinalSelections(selections, all = true) + assert(all.contains("final backend: Project AST JIT"), all) + assert(all.contains("final backend: legacy Project AST"), all) + assert(all.contains("final backend: the regular GPU projection"), all) + assertResult("")( + GpuAstJitExpression.explainFinalSelections(Seq(Seq(jit)), all = false)) + + val notOnGpu = GpuAstJitExpression.explainFinalSelections(selections, all = false) + assert(!notOnGpu.contains("final backend: Project AST JIT"), notOnGpu) + assert(notOnGpu.contains("final backend: legacy Project AST"), notOnGpu) + assert(notOnGpu.contains("final backend: the regular GPU projection"), notOnGpu) + } + + test("project AST JIT keeps its compiled expression across retry") { + val child = mock(classOf[GpuExpression]) + val ast = mock(classOf[AstExpression]) + val compiled = mock(classOf[CompiledExpression]) + when(child.convertToAst(Int.MaxValue)).thenReturn(ast) + when(ast.compileJit()).thenReturn(compiled) + val jit = GpuAstJitExpression(child) + + TestUtils.withMockTaskContext() { + jit.checkpoint() + jit.restore() + jit.checkpoint() + verify(ast, times(1)).compileJit() + verify(compiled, times(0)).close() + } + verify(compiled, times(1)).close() + } + + test("project AST JIT cleans up when task completion registration fails") { + val registrationFailure = new RuntimeException("task completion registration failed") + val closeFailure = new RuntimeException("compiled expression close failed") + val taskContext = new MockTaskContext(taskAttemptId = 1, partitionId = 0) { + override def addTaskCompletionListener(listener: TaskCompletionListener): TaskContext = + throw registrationFailure + } + val compiled = mock(classOf[CompiledExpression]) + doThrow(closeFailure).when(compiled).close() + val jit = mockJitExpression(compiled) + + TestUtils.withTaskContext(taskContext) { + val thrown = intercept[RuntimeException] { + jit.checkpoint() + } + assert(thrown eq registrationFailure) + assertResult(Seq(closeFailure))(thrown.getSuppressed.toSeq) + jit.close() + verify(compiled, times(1)).close() + } + } + + test("project AST JIT rejects a cleanup callback consumed during registration") { + val taskContext = new MockTaskContext(taskAttemptId = 1, partitionId = 0) { + override def addTaskCompletionListener(listener: TaskCompletionListener): TaskContext = { + listener.onTaskCompletion(this) + this + } + } + val compiled = mock(classOf[CompiledExpression]) + val jit = mockJitExpression(compiled) + + TestUtils.withTaskContext(taskContext) { + val thrown = intercept[IllegalStateException] { + jit.checkpoint() + } + assertResult("Task completed while registering the AST JIT cleanup callback")( + thrown.getMessage) + jit.close() + verify(compiled, times(1)).close() + } + } + + test("legacy project AST rejects a cleanup callback consumed during registration") { + val taskContext = new MockTaskContext(taskAttemptId = 1, partitionId = 0) { + override def addTaskCompletionListener(listener: TaskCompletionListener): TaskContext = { + listener.onTaskCompletion(this) + this + } + } + val compiled = mock(classOf[CompiledExpression]) + val astExpression = GpuProjectAstExpression(mockCompiledChild(compiled)) + + TestUtils.withTaskContext(taskContext) { + val thrown = intercept[IllegalStateException] { + astExpression.computeColumn(mock(classOf[Table])) + } + assertResult("Task completed while registering the Project AST cleanup callback")( + thrown.getMessage) + astExpression.close() + verify(compiled, times(1)).close() + } + } +} diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/ProjectAstTestUtils.scala b/tests/src/test/scala/com/nvidia/spark/rapids/ProjectAstTestUtils.scala new file mode 100644 index 00000000000..d2f8d501f5c --- /dev/null +++ b/tests/src/test/scala/com/nvidia/spark/rapids/ProjectAstTestUtils.scala @@ -0,0 +1,36 @@ +/* + * 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.reflect.ClassTag + +import org.apache.spark.sql.catalyst.expressions.Expression + +object ProjectAstTestUtils { + def collectExpressions[T <: Expression : ClassTag]( + expressions: Seq[Expression]): Seq[T] = { + expressions.flatMap(_.collect { + case expression: T => expression + }) + } + + def tierReferences(expression: Expression): Seq[GpuBoundReference] = { + expression.collect { + case reference: GpuBoundReference if reference.name.startsWith("tiered_input_") => reference + } + } +} diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/ProjectSplitRetrySuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/ProjectSplitRetrySuite.scala index 9f50c2aff90..fe6fc1e6e4a 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/ProjectSplitRetrySuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/ProjectSplitRetrySuite.scala @@ -17,9 +17,10 @@ package com.nvidia.spark.rapids import ai.rapids.cudf.ColumnVector +import com.nvidia.spark.Retryable import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} import com.nvidia.spark.rapids.RapidsPluginImplicits.AutoCloseableProducingSeq -import com.nvidia.spark.rapids.jni.{GpuSplitAndRetryOOM, RmmSpark} +import com.nvidia.spark.rapids.jni.{GpuRetryOOM, GpuSplitAndRetryOOM, RmmSpark} import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD @@ -98,6 +99,21 @@ class ProjectSplitRetrySuite extends RmmSparkRetrySuiteBase { batch.column(ordinal).asInstanceOf[GpuColumnVector].incRefCount() } + private case class GpuCheckpointRetryPassthrough(ordinal: Int, dataType: DataType) + extends GpuLeafExpression with Retryable { + var checkpointCount: Int = 0 + override def nullable: Boolean = false + override def checkpoint(): Unit = { + checkpointCount += 1 + if (checkpointCount == 1) { + throw new GpuRetryOOM("in checkpoint") + } + } + override def restore(): Unit = () + override def columnarEval(batch: ColumnarBatch): GpuColumnVector = + batch.column(ordinal).asInstanceOf[GpuColumnVector].incRefCount() + } + private def mixedNonRetryableExprs(): Seq[GpuExpression] = addOneExprs() :+ GpuAlias(GpuNonRetryablePassthrough(0, IntegerType), "non_retryable")() @@ -177,6 +193,20 @@ class ProjectSplitRetrySuite extends RmmSparkRetrySuiteBase { } } + test("tiered project retries checkpoint OOM when split retry is disabled") { + val sqlConf = new SQLConf() + sqlConf.setConfString(RapidsConf.PROJECT_SPLIT_RETRY_ENABLED.key, "false") + SQLConf.withExistingConf(sqlConf) { + val expression = GpuCheckpointRetryPassthrough(0, IntegerType) + val tier = GpuTieredProject(Seq(Seq(expression))) + withResource(tier.projectAndCloseWithRetrySingleBatch(newSpillable())) { output => + assertResult(NUM_ROWS)(output.numRows()) + assertResult((0 until NUM_ROWS).toArray)(collectInts(output, 0)) + } + assertResult(2)(expression.checkpointCount) + } + } + test("tiered project split-retry produces correct output") { val tier = GpuBindReferences.bindGpuReferencesTiered( addOneExprs(), batchAttrs, new SQLConf(), Map.empty) diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/TestUtils.scala b/tests/src/test/scala/com/nvidia/spark/rapids/TestUtils.scala index 9f453c8cdf6..383d0179456 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/TestUtils.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/TestUtils.scala @@ -29,6 +29,7 @@ import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.rapids.execution.TrampolineUtil +import org.apache.spark.sql.rapids.metrics.source.{MockTaskContext, MockTaskContextBase} import org.apache.spark.sql.vectorized.ColumnarBatch /** A collection of utility methods useful in tests. */ @@ -44,6 +45,28 @@ object TestUtils extends Assertions { System.getProperty("test.build.data", System.getProperty("java.io.tmpdir", "/tmp")), basename) + def withTaskContext[T]( + taskContext: MockTaskContextBase, + completesTask: Boolean = false)(body: => T): T = { + TrampolineUtil.setTaskContext(taskContext) + try { + body + } finally { + try { + if (completesTask) { + taskContext.markTaskComplete() + } + } finally { + TrampolineUtil.unsetTaskContext() + ScalableTaskCompletion.reset() + } + } + } + + def withMockTaskContext[T](completesTask: Boolean = false)(body: => T): T = { + withTaskContext(new MockTaskContext(taskAttemptId = 1, partitionId = 0), completesTask)(body) + } + // Spark caches the configured serializer in a JVM-global singleton, so suites that select a // different serializer must reset it at suite boundaries. def clearCachedBatchSerializer(): Unit = { diff --git a/tests/src/test/scala/org/apache/spark/sql/rapids/ProjectExprSuite.scala b/tests/src/test/scala/org/apache/spark/sql/rapids/ProjectExprSuite.scala index 33c8d5a3c4c..84a67685a06 100644 --- a/tests/src/test/scala/org/apache/spark/sql/rapids/ProjectExprSuite.scala +++ b/tests/src/test/scala/org/apache/spark/sql/rapids/ProjectExprSuite.scala @@ -22,34 +22,20 @@ import java.nio.file.Files import ai.rapids.cudf.Table import com.nvidia.spark.rapids._ import com.nvidia.spark.rapids.Arm.withResource +import com.nvidia.spark.rapids.ProjectAstTestUtils.{collectExpressions, tierReferences} import com.nvidia.spark.rapids.jni.RmmSpark import org.mockito.Mockito.{never, spy, verify} import org.apache.spark.SparkConf import org.apache.spark.sql.Row -import org.apache.spark.sql.catalyst.expressions.{ - AttributeReference, Expression, Literal, NamedExpression} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Literal, NamedExpression} import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.rapids.execution.TrampolineUtil -import org.apache.spark.sql.rapids.metrics.source.MockTaskContext import org.apache.spark.sql.rapids.shims.TrampolineConnectShims._ import org.apache.spark.sql.tests.datagen.DataGenExprShims import org.apache.spark.sql.types._ class ProjectExprSuite extends SparkQueryCompareTestSuite { - private def astExpressions(expressions: Seq[Expression]): Seq[GpuProjectAstExpression] = { - expressions.flatMap(_.collect { - case astExpression: GpuProjectAstExpression => astExpression - }) - } - - private def tierReferences(expression: Expression): Seq[GpuBoundReference] = { - expression.collect { - case reference: GpuBoundReference if reference.name.startsWith("tiered_input_") => reference - } - } - def forceHostColumnarToGpu(): SparkConf = { // turns off BatchScanExec, so we get a CPU BatchScanExec together with a HostColumnarToGpu new SparkConf().set("spark.rapids.sql.exec.BatchScanExec", "false") @@ -215,10 +201,12 @@ class ProjectExprSuite extends SparkQueryCompareTestSuite { // tier 1: [AST(t1*c) AS t2] // tier 2: [AST(t2+d) AS first, AST(t2+e) AS second, // t1 AS shared_first, t1 AS shared_second, greatest(t1, f) AS regular] - assertResult(Seq(1, 1, 2))(tiered.exprTiers.map(astExpressions(_).size)) - assert(astExpressions(tiered.exprTiers.head).head.child.isInstanceOf[GpuAdd]) - assert(astExpressions(tiered.exprTiers(1)).head.child.isInstanceOf[GpuMultiply]) - assertResult(1)(tierReferences(astExpressions(tiered.exprTiers(1)).head).size) + val astExpressionTiers = tiered.exprTiers.map( + collectExpressions[GpuProjectAstExpression]) + assertResult(Seq(1, 1, 2))(astExpressionTiers.map(_.size)) + assert(astExpressionTiers.head.head.child.isInstanceOf[GpuAdd]) + assert(astExpressionTiers(1).head.child.isInstanceOf[GpuMultiply]) + assertResult(1)(tierReferences(astExpressionTiers(1).head).size) // Final references: [t2, t2, t1, t1, t1] (distinct: {t2, t1}). val finalReferences = tiered.exprTiers.last.flatMap(tierReferences) assertResult(5)(finalReferences.size) @@ -226,28 +214,70 @@ class ProjectExprSuite extends SparkQueryCompareTestSuite { } test("AST compiled expression closes at task completion") { - val context = new MockTaskContext(taskAttemptId = 1, partitionId = 0) val astExpression = spy(GpuProjectAstExpression(GpuAdd( GpuBoundReference(0, LongType, true)(NamedExpression.newExprId, "a"), GpuBoundReference(1, LongType, true)(NamedExpression.newExprId, "b"), false)())) - TrampolineUtil.setTaskContext(context) try { - withResource(buildProjectBatch()) { spillableBatch => - withResource(spillableBatch.getColumnarBatch()) { inputBatch => - withResource(GpuProjectExec.project( - inputBatch, Seq(GpuAlias(astExpression, "sum")()))) { _ => } + TestUtils.withMockTaskContext(completesTask = true) { + withResource(buildProjectBatch()) { spillableBatch => + withResource(spillableBatch.getColumnarBatch()) { inputBatch => + withResource(GpuProjectExec.project( + inputBatch, Seq(GpuAlias(astExpression, "sum")()))) { _ => } + } } + verify(astExpression, never()).close() } - verify(astExpression, never()).close() - context.markTaskComplete() verify(astExpression).close() } finally { - TrampolineUtil.unsetTaskContext() - ScalableTaskCompletion.reset() astExpression.close() } } + test("multi-output AST JIT project retries and preserves output order and nulls") { + val left = GpuBoundReference(0, LongType, nullable = true)( + NamedExpression.newExprId, "a") + val right = GpuBoundReference(1, LongType, nullable = true)( + NamedExpression.newExprId, "b") + def shared = GpuAdd(left, GpuLiteral(1L, LongType), failOnError = false)() + val jitExpressions = Seq( + GpuAstJitExpression(GpuMultiply(shared, right, failOnError = false)()), + GpuAstJitExpression(GpuMultiply(shared, left, failOnError = false)())) + val expressions = Seq( + GpuAlias(jitExpressions.head, "right_product")(), + GpuAlias(jitExpressions(1), "left_product")()) + + RmmSpark.currentThreadIsDedicatedToTask(0) + try { + withResource(jitExpressions) { _ => + TestUtils.withMockTaskContext(completesTask = true) { + val spillableBatch = buildProjectBatch() + RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 1, + RmmSpark.OomInjectionType.GPU.ordinal, 0) + withResource(GpuProjectExec.projectAndCloseWithRetrySingleBatch( + spillableBatch, expressions)) { result => + assertResult(2)(result.numCols()) + val rightProduct = result.column(0).asInstanceOf[GpuColumnVector] + val leftProduct = result.column(1).asInstanceOf[GpuColumnVector] + withResource(rightProduct.getBase.copyToHost()) { hostProduct => + assertResult(36L)(hostProduct.getLong(0)) + assert(hostProduct.isNull(1)) + assertResult(32L)(hostProduct.getLong(2)) + assertResult(18L)(hostProduct.getLong(3)) + } + withResource(leftProduct.getBase.copyToHost()) { hostProduct => + assertResult(30L)(hostProduct.getLong(0)) + assert(hostProduct.isNull(1)) + assertResult(12L)(hostProduct.getLong(2)) + assertResult(2L)(hostProduct.getLong(3)) + } + } + } + } + } finally { + RmmSpark.removeCurrentDedicatedThreadAssociation(0) + } + } + testSparkResultsAreEqual("Test literal values in select", mixedFloatDf) { frame => frame.select(col("floats"), diff --git a/tests/src/test/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinRetrySuite.scala b/tests/src/test/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinRetrySuite.scala new file mode 100644 index 00000000000..065d04af737 --- /dev/null +++ b/tests/src/test/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinRetrySuite.scala @@ -0,0 +1,146 @@ +/* + * 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 org.apache.spark.sql.rapids.execution + +import ai.rapids.cudf.Table +import com.nvidia.spark.rapids._ +import com.nvidia.spark.rapids.Arm.withResource +import com.nvidia.spark.rapids.ProjectAstTestUtils.collectExpressions +import com.nvidia.spark.rapids.jni.RmmSpark + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, NamedExpression} +import org.apache.spark.sql.catalyst.plans.Inner +import org.apache.spark.sql.execution.{LeafExecNode, SparkPlan} +import org.apache.spark.sql.rapids.{GpuAdd, GpuMultiply, GpuSubtract} +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.vectorized.ColumnarBatch + +class GpuBroadcastNestedLoopJoinRetrySuite extends RmmSparkRetrySuiteBase { + private val taskId = 1 + private val x = AttributeReference("x", IntegerType, nullable = false)() + private val y = AttributeReference("y", IntegerType, nullable = false)() + private val z = AttributeReference("z", IntegerType, nullable = false)() + private val buildAttributes = Seq(x, y, z) + + private case class TestLeafExec(override val output: Seq[Attribute]) extends LeafExecNode { + override protected def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException("TestLeafExec is not executable") + } + + private case class TestBroadcastNestedLoopJoin( + left: SparkPlan, + right: SparkPlan, + postBuildCondition: List[NamedExpression]) + extends GpuBroadcastNestedLoopJoinExecBase( + left, + right, + Inner, + GpuBuildRight, + condition = None, + postBuildCondition, + targetSizeBytes = 1024L) { + override lazy val allMetrics: Map[String, GpuMetric] = Map.empty + } + + override def afterEach(): Unit = { + RmmSpark.getAndResetNumRetryThrow(taskId) + RmmSpark.getAndResetNumSplitRetryThrow(taskId) + super.afterEach() + } + + private def postBuildExpressions: List[NamedExpression] = { + val shared = GpuMultiply( + GpuAdd(x, y, failOnError = false)(), z, failOnError = false)() + List( + GpuAlias(GpuSubtract(shared, x, failOnError = false)(), "shared_minus_x")(), + GpuAlias(GpuSubtract(shared, y, failOnError = false)(), "shared_minus_y")(), + x, + y, + z) + } + + private def buildBatch(): ColumnarBatch = { + val table = new Table.TestBuilder() + .column(1.asInstanceOf[java.lang.Integer], 2, 3) + .column(4.asInstanceOf[java.lang.Integer], 5, 6) + .column(2.asInstanceOf[java.lang.Integer], 3, 4) + .build() + withResource(table) { tbl => + GpuColumnVector.from(tbl, Array(IntegerType, IntegerType, IntegerType)) + } + } + + private def collectInts(batch: ColumnarBatch, column: Int): Seq[Int] = { + val gpuColumn = batch.column(column).asInstanceOf[GpuColumnVector] + withResource(gpuColumn.getBase.copyToHost()) { hostColumn => + (0 until batch.numRows()).map(row => hostColumn.getInt(row)) + } + } + + test("BNLJ build-side shared JIT tier retries GpuRetryOOM") { + val spark = SparkSession.builder() + .master("local[1]") + .appName(this.getClass.getSimpleName) + .config(RapidsConf.ENABLE_TIERED_PROJECT.key, "true") + .config(RapidsConf.ENABLE_COMBINED_EXPRESSIONS.key, "true") + .config(RapidsConf.ENABLE_PROJECT_AST_JIT.key, "true") + .config(RapidsConf.PROJECT_SPLIT_RETRY_ENABLED.key, "true") + .getOrCreate() + val conf = spark.sessionState.conf + // x=[1,2,3], y=[4,5,6], z=[2,3,4] + // [((x+y)*z)-x AS shared_minus_x, ((x+y)*z)-y AS shared_minus_y, x, y, z] + val expressions = postBuildExpressions + val buildProject = GpuProjectExec(expressions, TestLeafExec(buildAttributes)) + val join = TestBroadcastNestedLoopJoin( + TestLeafExec(Seq.empty), buildProject, expressions) + + TestUtils.withMockTaskContext(completesTask = true) { + val boundProject = GpuBindReferences.bindGpuProjectReferencesTiered( + expressions, buildAttributes, conf, Map.empty) + // after CSE: + // tier 0: [x, y, z, AST_JIT((x+y)*z) AS t1] + // tier 1: [t1-x AS shared_minus_x, t1-y AS shared_minus_y, x, y, z] + val jitExpressions = collectExpressions[GpuAstJitExpression](boundProject.exprTiers.flatten) + assertResult(Seq(4, 5))(boundProject.exprTiers.map(_.size)) + assertResult(1)(jitExpressions.size) + assert(jitExpressions.head.child.isInstanceOf[GpuMultiply]) + assert(jitExpressions.head.child.find(_.isInstanceOf[GpuAdd]).isDefined) + + val projectBuildSide = join.buildSidePostProjection.get + // Warm up AST JIT so first-use setup cannot consume the injected OOM; the + // next call exercises retry during query execution. + withResource(projectBuildSide(buildBatch())) { _ => } + + val retryInput = buildBatch() + RmmSpark.getAndResetNumRetryThrow(taskId) + RmmSpark.getAndResetNumSplitRetryThrow(taskId) + RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 1, + RmmSpark.OomInjectionType.GPU.ordinal, 0) + withResource(projectBuildSide(retryInput)) { output => + assertResult(3)(output.numRows()) + assertResult(5)(output.numCols()) + assertResult(Seq(9, 19, 33))(collectInts(output, 0)) + assertResult(Seq(6, 16, 30))(collectInts(output, 1)) + } + assert(RmmSpark.getAndResetNumRetryThrow(taskId) > 0) + assertResult(0)(RmmSpark.getAndResetNumSplitRetryThrow(taskId)) + } + } +}