From a3190f185db285f3c52001cee31cb3e4a58c6ec9 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Tue, 21 Jul 2026 14:09:51 +0800 Subject: [PATCH 01/20] Add experimental AST JIT for integral add and multiply --- integration_tests/src/main/python/ast_test.py | 32 +++++ .../spark/rapids/GpuAstJitExpression.scala | 115 ++++++++++++++++++ .../spark/rapids/GpuBoundAttribute.scala | 2 + .../nvidia/spark/rapids/GpuExpressions.scala | 15 +++ .../com/nvidia/spark/rapids/RapidsConf.scala | 8 ++ .../spark/rapids/basicPhysicalOperators.scala | 6 + .../com/nvidia/spark/rapids/literals.scala | 2 + .../apache/spark/sql/rapids/arithmetic.scala | 10 ++ .../spark/rapids/GpuProjectAstJitSuite.scala | 80 ++++++++++++ 9 files changed, 270 insertions(+) create mode 100644 sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala create mode 100644 tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index 5cc2d700f01..787856101fe 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/integration_tests/src/main/python/ast_test.py @@ -66,6 +66,7 @@ ast_acosh_descr = [(double_gen, not (is_spark_403() or is_spark_412_or_later()))] _project_ast_enabled_conf = {"spark.rapids.sql.projectAstEnabled": "true"} +_project_ast_jit_enabled_conf = {"spark.rapids.sql.projectAstJitEnabled": "true"} def assert_gpu_ast(is_supported, func, conf={}): exist = "GpuProjectAstExec" @@ -370,6 +371,37 @@ def test_multiplication(data_descr): f.lit(-12).cast(data_type) * f.col('b'), f.col('a') * f.col('b'))) +@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn) +@disable_ansi_mode +def test_jit_add_multiply(data_gen): + assert_cpu_and_gpu_are_equal_collect_with_capture( + lambda spark: binary_op_df(spark, data_gen).select( + f.col('a') + f.col('b'), + f.col('a') * f.col('b')), + exist_classes=r"GpuProject.*AST_JIT", + non_exist_classes="GpuProjectAst", + conf=_project_ast_jit_enabled_conf) + +@disable_ansi_mode +def test_jit_mixed_nested_subexpressions(): + assert_cpu_and_gpu_are_equal_collect_with_capture( + lambda spark: binary_op_df(spark, int_gen).select( + (f.col('a') + f.col('b')) - (f.col('a') * f.col('b'))), + exist_classes=r"GpuProject.*AST_JIT.*- AST_JIT", + non_exist_classes="GpuProjectAst", + conf=_project_ast_jit_enabled_conf) + +@disable_ansi_mode +def test_jit_mixed_project_expressions(): + assert_cpu_and_gpu_are_equal_collect_with_capture( + lambda spark: binary_op_df(spark, int_gen).select( + (f.col('a') + f.col('b')).alias('jit'), + (f.col('a') - f.col('b')).alias('gpu'), + ((f.col('a') * f.col('b')) - f.col('a')).alias('mixed')), + exist_classes=r"GpuProject.*AST_JIT.*AS jit.*AS gpu.*AST_JIT.*AS mixed", + non_exist_classes="GpuProjectAst", + conf=_project_ast_jit_enabled_conf) + # Each descriptor contains a list of data generators and a corresponding boolean # indicating whether that data type is supported by the AST # all the below desc are not supported by the AST because ANSI mode is on 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..4b20f02cd9d --- /dev/null +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala @@ -0,0 +1,115 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.spark.rapids + +import ai.rapids.cudf.{Scalar, Table} +import ai.rapids.cudf.ast.CompiledExpression +import com.nvidia.spark.Retryable +import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} +import com.nvidia.spark.rapids.RapidsPluginImplicits._ +import com.nvidia.spark.rapids.ScalableTaskCompletion.onTaskCompletion +import com.nvidia.spark.rapids.shims.ShimUnaryExpression + +import org.apache.spark.TaskContext +import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression} +import org.apache.spark.sql.types.DataType +import org.apache.spark.sql.vectorized.ColumnarBatch + +object GpuAstJitExpression { + private def wrapMaximalSubtrees(expression: Expression): Expression = expression match { + case gpuExpression: GpuExpression + if gpuExpression.supportsAstJit && gpuExpression.containsAstJitOperator => + GpuAstJitExpression(gpuExpression) + case gpuExpression: GpuExpression => + gpuExpression.mapChildren { + case child: GpuExpression => wrapMaximalSubtrees(child) + case child => child + } + case other => other + } + + private[rapids] def wrapProjectExpressions( + expressions: List[NamedExpression]): List[NamedExpression] = { + expressions.map(wrapMaximalSubtrees(_).asInstanceOf[NamedExpression]) + } +} + +case class GpuAstJitExpression(child: Expression) + extends ShimUnaryExpression with GpuExpression with Retryable with AutoCloseable { + require(child.isInstanceOf[GpuExpression], "AST JIT child must be a GPU expression") + + @transient private[this] var compiledExpression: CompiledExpression = _ + @transient private[this] var completionRegistered = false + + override def dataType: DataType = child.dataType + + override def nullable: Boolean = child.nullable + + override def disableTieredProjectCombine: Boolean = true + + override def toString: String = s"AST_JIT($child)" + + override def checkpoint(): Unit = { + getCompiledExpression + } + + override def restore(): Unit = closeCompiledExpression() + + override def close(): Unit = closeCompiledExpression() + + override def columnarEval(batch: ColumnarBatch): GpuColumnVector = { + withResource(tableFromBatch(batch)) { table => + closeOnExcept(getCompiledExpression.computeColumnJit(table)) { result => + GpuColumnVector.from(result, dataType) + } + } + } + + private def getCompiledExpression: CompiledExpression = synchronized { + if (compiledExpression == null) { + compiledExpression = child.asInstanceOf[GpuExpression] + .convertToAst(Int.MaxValue) + .compile() + } + if (!completionRegistered) { + Option(TaskContext.get()).foreach { taskContext => + onTaskCompletion(taskContext) { + close() + } + completionRegistered = true + } + } + compiledExpression + } + + private def closeCompiledExpression(): Unit = synchronized { + Option(compiledExpression).foreach(_.safeClose()) + compiledExpression = null + } + + private def tableFromBatch(batch: ColumnarBatch): Table = { + if (batch.numCols() != 0) { + GpuColumnVector.from(batch) + } else { + withResource(Scalar.fromBool(false)) { falseScalar => + withResource(ai.rapids.cudf.ColumnVector.fromScalar(falseScalar, batch.numRows())) { + falseColumn => new Table(falseColumn) + } + } + } + } +} 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 3c5b9f26872..2aa230704a6 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 @@ -268,6 +268,8 @@ case class GpuBoundReference(ordinal: Int, dataType: DataType, nullable: Boolean (val exprId: ExprId, val name: String) extends GpuLeafExpression with ShimExpression { + override def selfSupportsAstJit: Boolean = true + override def toString: String = s"input[$ordinal, ${dataType.simpleString}, $nullable]($name#${exprId.id})" 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..17f87eed041 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,21 @@ trait GpuExpression extends Expression { def convertToAst(numFirstTableColumns: Int): ast.AstExpression = throw new IllegalStateException(s"Cannot convert ${this.getClass.getSimpleName} to AST") + def selfSupportsAstJit: Boolean = false + + def selfIsAstJitOperator: Boolean = false + + final def supportsAstJit: Boolean = selfSupportsAstJit && children.forall { + case child: GpuExpression => child.supportsAstJit + case _: AttributeReference => true + case _ => false + } + + final def containsAstJitOperator: Boolean = selfIsAstJitOperator || children.exists { + case child: GpuExpression => child.containsAstJitOperator + case _ => false + } + /** Could evaluating this expression cause side-effects, such as throwing an exception? */ def hasSideEffects: Boolean = children.exists { 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 1bc3da5da96..31b4de07a23 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,12 @@ val GPU_COREDUMP_PIPE_PATTERN = conf("spark.rapids.gpu.coreDump.pipePattern") .booleanConf .createWithDefault(false) + val ENABLE_PROJECT_AST_JIT = conf("spark.rapids.sql.projectAstJitEnabled") + .doc("Enable the experimental cuDF JIT backend for supported project AST expressions.") + .internal() + .booleanConf + .createWithDefault(false) + val ENABLE_TIERED_PROJECT = conf("spark.rapids.sql.tiered.project.enabled") .doc("Enable tiered projections.") .internal() @@ -3639,6 +3645,8 @@ class RapidsConf(conf: Map[String, String]) extends Logging { lazy val isProjectAstEnabled: Boolean = get(ENABLE_PROJECT_AST) + lazy val isProjectAstJitEnabled: Boolean = get(ENABLE_PROJECT_AST_JIT) + lazy val isTieredProjectEnabled: Boolean = get(ENABLE_TIERED_PROJECT) lazy val isCombinedExpressionsEnabled: Boolean = get(ENABLE_COMBINED_EXPRESSIONS) 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 a2cc24b041a..0decb6373fb 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 @@ -59,6 +59,12 @@ class GpuProjectExecMeta( // Force list to avoid recursive Java serialization of lazy list Seq implementation val gpuExprs = childExprs.map(_.convertToGpu().asInstanceOf[NamedExpression]).toList val gpuChild = childPlans.head.convertIfNeeded() + if (conf.isProjectAstJitEnabled) { + val jitProjectList = GpuAstJitExpression.wrapProjectExpressions(gpuExprs) + if (jitProjectList.exists(_.find(_.isInstanceOf[GpuAstJitExpression]).isDefined)) { + return GpuProjectExec(jitProjectList, gpuChild) + } + } if (conf.isProjectAstEnabled) { // cuDF requires return column is fixed width val allReturnTypesFixedWidth = gpuExprs.forall(e => GpuBatchUtils.isFixedWidth(e.dataType)) 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..f113de4ebcb 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 @@ -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..0fcabb4537b 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,11 @@ abstract class GpuAddBase extends CudfBinaryArithmetic with Serializable { override def binaryOp: BinaryOp = BinaryOp.ADD override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.ADD) + override def selfSupportsAstJit: Boolean = + !failOnError && (dataType == IntegerType || dataType == LongType) + + override def selfIsAstJitOperator: Boolean = selfSupportsAstJit + override def hasSideEffects: Boolean = (failOnError && GpuAnsi.needBasicOpOverflowCheck(dataType)) || super.hasSideEffects @@ -768,6 +773,11 @@ case class GpuMultiply( override def binaryOp: BinaryOp = BinaryOp.MUL override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.MUL) + override def selfSupportsAstJit: Boolean = + !failOnError && (dataType == IntegerType || dataType == LongType) + + override def selfIsAstJitOperator: Boolean = selfSupportsAstJit + private def multiplyOverflowError(msg: String): ArithmeticException = { RapidsErrorUtils.arithmeticOverflowError(msg, origin) } 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..e4388250057 --- /dev/null +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala @@ -0,0 +1,80 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.spark.rapids + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.rapids.{GpuAdd, GpuMultiply, GpuSubtract} +import org.apache.spark.sql.types.{FloatType, IntegerType, LongType} + +class GpuProjectAstJitSuite extends AnyFunSuite { + private def reference(ordinal: Int, dataType: org.apache.spark.sql.types.DataType) = + AttributeReference(s"c$ordinal", dataType, nullable = true)() + + private def alias(expression: GpuExpression, name: String) = GpuAlias(expression, name)() + + test("project AST JIT is disabled by default") { + assert(!new RapidsConf(Map.empty[String, String]).isProjectAstJitEnabled) + } + + test("project AST JIT supports non-ANSI integral add and multiply") { + val left = reference(0, LongType) + val right = reference(1, LongType) + val expression = alias( + GpuMultiply( + GpuAdd(left, right, failOnError = false)(), + right, + failOnError = false)(), + "result") + + val wrapped = GpuAstJitExpression.wrapProjectExpressions(List(expression)) + val jit = wrapped.head.asInstanceOf[GpuAlias].child.asInstanceOf[GpuAstJitExpression] + assert(jit.child.isInstanceOf[GpuMultiply]) + assert(jit.child.find(_.isInstanceOf[GpuAstJitExpression]).isEmpty) + } + + test("project AST JIT wraps maximal nested subtrees independently") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val expression = alias( + GpuSubtract( + GpuAdd(left, right, failOnError = false)(), + GpuMultiply(left, right, failOnError = false)(), + failOnError = false)(), + "result") + + val wrapped = GpuAstJitExpression.wrapProjectExpressions(List(expression)) + val subtract = wrapped.head.asInstanceOf[GpuAlias].child.asInstanceOf[GpuSubtract] + assert(subtract.left.asInstanceOf[GpuAstJitExpression].child.isInstanceOf[GpuAdd]) + assert(subtract.right.asInstanceOf[GpuAstJitExpression].child.isInstanceOf[GpuMultiply]) + } + + test("project AST JIT excludes ANSI and floating point arithmetic") { + val intLeft = reference(0, IntegerType) + val intRight = reference(1, IntegerType) + val floatLeft = reference(0, FloatType) + val floatRight = reference(1, FloatType) + + val ansiAdd = alias(GpuAdd(intLeft, intRight, failOnError = true)(), "ansi_sum") + val floatMultiply = alias( + GpuMultiply(floatLeft, floatRight, failOnError = false)(), "float_product") + + val wrapped = GpuAstJitExpression.wrapProjectExpressions(List(ansiAdd, floatMultiply)) + assert(wrapped.forall(_.find(_.isInstanceOf[GpuAstJitExpression]).isEmpty)) + } +} From 4f6972ab655c0d7008eec1f6912eba00505c6f55 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Tue, 21 Jul 2026 15:06:10 +0800 Subject: [PATCH 02/20] Update copyright years Signed-off-by: Haoyang Li --- .../main/scala/com/nvidia/spark/rapids/GpuBoundAttribute.scala | 2 +- .../src/main/scala/com/nvidia/spark/rapids/literals.scala | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 2aa230704a6..925fcd72c38 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 @@ -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. 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 f113de4ebcb..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. From f75b0c1179d5229c935c8ef82064fe819d7a1850 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Tue, 21 Jul 2026 15:41:38 +0800 Subject: [PATCH 03/20] Address AST JIT review feedback Signed-off-by: Haoyang Li --- integration_tests/src/main/python/ast_test.py | 10 ++++++---- .../com/nvidia/spark/rapids/GpuAstJitExpression.scala | 5 ++++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index 435fa28112c..4d8fdce7ed7 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/integration_tests/src/main/python/ast_test.py @@ -383,19 +383,21 @@ def test_jit_add_multiply(data_gen): 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_mixed_nested_subexpressions(): +def test_jit_mixed_nested_subexpressions(data_gen): assert_cpu_and_gpu_are_equal_collect_with_capture( - lambda spark: binary_op_df(spark, int_gen).select( + 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.*- 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_mixed_project_expressions(): +def test_jit_mixed_project_expressions(data_gen): assert_cpu_and_gpu_are_equal_collect_with_capture( - lambda spark: binary_op_df(spark, int_gen).select( + 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')), 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 index 4b20f02cd9d..ab8c03ced98 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala @@ -67,7 +67,10 @@ case class GpuAstJitExpression(child: Expression) getCompiledExpression } - override def restore(): Unit = closeCompiledExpression() + override def restore(): Unit = { + // The existing task callback closes the expression recompiled after a retry. + closeCompiledExpression() + } override def close(): Unit = closeCompiledExpression() From b4a1182791e6903c95f4d88e5d373f5f01dac451 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Fri, 24 Jul 2026 15:37:22 +0800 Subject: [PATCH 04/20] Enable per-expression legacy AST projection --- integration_tests/src/main/python/ast_test.py | 42 +++-- .../spark/rapids/GpuBoundAttribute.scala | 9 +- .../rapids/GpuProjectAstExpression.scala | 166 ++++++++++++++++++ .../spark/rapids/basicPhysicalOperators.scala | 143 ++++----------- .../spark/rapids/higherOrderFunctions.scala | 10 +- .../GpuEquivalentExpressions.scala | 53 ++++-- .../spark/rapids/GpuArrayHofFusionSuite.scala | 43 ++++- .../spark/sql/rapids/ProjectExprSuite.scala | 133 ++++++++++---- .../rapids/IntervalArithmeticSuite.scala | 3 +- 9 files changed, 419 insertions(+), 183 deletions(-) create mode 100644 sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuProjectAstExpression.scala diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index efb4f5502cc..cae95fd09f9 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/integration_tests/src/main/python/ast_test.py @@ -68,11 +68,12 @@ _project_ast_enabled_conf = {"spark.rapids.sql.projectAstEnabled": "true"} def assert_gpu_ast(is_supported, func, conf={}): - exist = "GpuProjectAstExec" - non_exist = "GpuProjectExec" + ast_expression = "GpuProjectAstExpression" + exist = ast_expression + non_exist = '' if not is_supported: exist = "GpuProjectExec" - non_exist = "GpuProjectAstExec" + non_exist = ast_expression ast_conf = copy_and_update(conf, _project_ast_enabled_conf) assert_cpu_and_gpu_are_equal_collect_with_capture( func, @@ -428,14 +429,29 @@ def test_multi_tier_ast(): func=lambda spark: spark.range(10).withColumn("x", f.col("id")).repartition(1)\ .selectExpr("x", "(id < x) == (id < (id + x))")) - -# MUST NOT use GPU AST when project refers to string type(non-fixed-width), -# or cudf::compute_column will throw error: Invalid, non-fixed-width type -# ANSI mode is disabled here due to an overflow issue with integer multiplication on Spark 4.0.0. @disable_ansi_mode -@ignore_order(local=True) -def test_refer_to_non_fixed_width_column(): - gens = [('col_int', int_gen), ('col_string', string_gen)] - assert_gpu_and_cpu_are_equal_collect( - lambda spark: gen_df(spark, gens).selectExpr("col_int * col_int", "col_string"), - conf=_project_ast_enabled_conf) +@pytest.mark.parametrize( + 'tiered_project_enabled', ['true', 'false'], ids=['tiered', 'single_tier']) +def test_project_ast_mixed_expressions(tiered_project_enabled): + def project(spark): + df = gen_df(spark, [ + ('a', int_gen), + ('b', int_gen), + ('c', int_gen), + ('d', int_gen), + ('col_string', string_gen) + ]) + shared = f.col('a') + f.col('b') + return df.select( + (shared * f.col('c')).alias('ast_first'), + (shared * f.col('d')).alias('ast_second'), + f.greatest(shared, f.col('c')).alias('gpu_shared'), + f.length(f.col('col_string')).alias('gpu_string'), + f.col('col_string').alias('raw_string')) + + assert_gpu_ast( + is_supported=True, + func=project, + conf={ + 'spark.rapids.sql.tiered.project.enabled': tiered_project_enabled + }) 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 3c5b9f26872..5c6058cdafa 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 @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2025, NVIDIA CORPORATION. + * Copyright (c) 2019-2026, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -135,12 +135,7 @@ object GpuBindReferences extends Logging { conf: SQLConf): GpuTieredProject = { if (RapidsConf.ENABLE_TIERED_PROJECT.get(conf)) { - val replaced = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) { - GpuEquivalentExpressions.replaceMultiExpressions(expressions, conf) - } else { - expressions - } - val exprTiers = GpuEquivalentExpressions.getExprTiers(replaced) + val exprTiers = GpuProjectAstExpression.buildExprTiers(expressions, conf) val inputTiers = GpuEquivalentExpressions.getInputTiers(exprTiers, input) // Update ExprTiers to include the columns that are pass through and drop unneeded columns val newExprTiers = exprTiers.zipWithIndex.map { 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 new file mode 100644 index 00000000000..756286dc666 --- /dev/null +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuProjectAstExpression.scala @@ -0,0 +1,166 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.nvidia.spark.rapids + +import scala.annotation.tailrec + +import ai.rapids.cudf.{Scalar, Table} +import ai.rapids.cudf.ast.CompiledExpression +import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} +import com.nvidia.spark.rapids.GpuMetric.OP_TIME_LEGACY +import com.nvidia.spark.rapids.RapidsPluginImplicits._ +import com.nvidia.spark.rapids.ScalableTaskCompletion.onTaskCompletion +import com.nvidia.spark.rapids.shims.ShimUnaryExpression + +import org.apache.spark.TaskContext +import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.rapids.catalyst.expressions.{ + GpuEquivalentExpressions, GpuExpressionEquals} +import org.apache.spark.sql.types.DataType +import org.apache.spark.sql.vectorized.ColumnarBatch + +object GpuProjectAstExpression { + private[rapids] def wrap(expression: NamedExpression): NamedExpression = { + expression match { + case alias @ GpuAlias(child: GpuExpression, name) => + GpuAlias(GpuProjectAstExpression(child), name)( + alias.exprId, alias.qualifier, alias.explicitMetadata) + case other => other + } + } + + @tailrec + private[rapids] def extractTopLevel(expression: Expression): Option[GpuProjectAstExpression] = { + expression match { + case alias: GpuAlias => extractTopLevel(alias.child) + case astExpression: GpuProjectAstExpression => Some(astExpression) + case _ => None + } + } + + private def unwrap(expression: Expression): Expression = expression match { + case alias @ GpuAlias(astExpression: GpuProjectAstExpression, name) => + GpuAlias(astExpression.child, name)( + alias.exprId, alias.qualifier, alias.explicitMetadata) + case astExpression: GpuProjectAstExpression => astExpression.child + case other => other + } + + private def rewrap(expression: Expression): Expression = expression match { + case namedExpression: NamedExpression => wrap(namedExpression) + case gpuExpression: GpuExpression => GpuProjectAstExpression(gpuExpression) + case other => other + } + + private[rapids] def buildExprTiers( + expressions: Seq[Expression], + conf: SQLConf): Seq[Seq[Expression]] = { + val astOutputs = expressions.map(extractTopLevel) + val astSubexpressions = astOutputs.flatten.flatMap { astExpression => + astExpression.child.collect { + case gpuExpression: GpuExpression => GpuExpressionEquals(gpuExpression) + } + }.toSet + // CSE must see through the marker so AST and non-AST outputs can share the same tiers. + val unwrapped = expressions.map(unwrap) + val replaced = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) { + GpuEquivalentExpressions.replaceMultiExpressions(unwrapped, conf) + } else { + unwrapped + } + val tiers = GpuEquivalentExpressions.getExprTiers( + replaced, + (original, substituted) => (original, substituted) match { + case (gpuOriginal: GpuExpression, gpuSubstituted: GpuExpression) + if GpuBatchUtils.isFixedWidth(gpuOriginal.dataType) && + astSubexpressions.contains(GpuExpressionEquals(gpuOriginal)) => + GpuProjectAstExpression(gpuSubstituted) + case _ => substituted + }) + val finalTier = tiers.last.zip(astOutputs).map { + case (expression, Some(_)) => rewrap(expression) + case (expression, None) => expression + } + tiers.dropRight(1) :+ finalTier + } + + private[rapids] def tableFromBatch(batch: ColumnarBatch): Table = { + if (batch.numCols() != 0) { + GpuColumnVector.from(batch) + } else { + withResource(Scalar.fromBool(false)) { falseScalar => + withResource(ai.rapids.cudf.ColumnVector.fromScalar(falseScalar, batch.numRows())) { + falseColumn => new Table(falseColumn) + } + } + } + } +} + +case class GpuProjectAstExpression(child: GpuExpression) + extends ShimUnaryExpression with GpuExpression with GpuMetricsInjectable with AutoCloseable { + @transient private[this] var compiledExpression: CompiledExpression = _ + private[this] var opTime: GpuMetric = NoopMetric + + override def dataType: DataType = child.dataType + + override def nullable: Boolean = child.nullable + + override def toString: String = s"AST($child)" + + override def injectMetrics(metrics: Map[String, GpuMetric]): Unit = { + opTime = metrics.getOrElse(OP_TIME_LEGACY, NoopMetric) + } + + override def close(): Unit = synchronized { + Option(compiledExpression).foreach(_.safeClose()) + compiledExpression = null + } + + override def columnarEval(batch: ColumnarBatch): GpuColumnVector = { + withResource(GpuProjectAstExpression.tableFromBatch(batch)) { table => + computeColumn(table) + } + } + + def computeColumn(table: Table): GpuColumnVector = { + NvtxIdWithMetrics(NvtxRegistry.PROJECT_AST, opTime) { + closeOnExcept(getCompiledExpression.computeColumn(table)) { result => + GpuColumnVector.from(result, dataType) + } + } + } + + private def getCompiledExpression: CompiledExpression = synchronized { + if (compiledExpression == null) { + val compiled = NvtxIdWithMetrics(NvtxRegistry.COMPILE_ASTS, opTime) { + // Project AST has a single input table. + child.convertToAst(Int.MaxValue).compile() + } + closeOnExcept(compiled) { _ => + Option(TaskContext.get()).foreach { taskContext => + onTaskCompletion(taskContext) { + close() + } + } + compiledExpression = compiled + } + } + compiledExpression + } +} 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 d70e96b1dae..44684e18ffb 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 @@ -59,12 +59,16 @@ class GpuProjectExecMeta( // Force list to avoid recursive Java serialization of lazy list Seq implementation val gpuExprs = childExprs.map(_.convertToGpu().asInstanceOf[NamedExpression]).toList val gpuChild = childPlans.head.convertIfNeeded() - if (conf.isProjectAstEnabled) { - // cuDF requires return column is fixed width + val projectList = if (conf.isProjectAstEnabled) { val allReturnTypesFixedWidth = gpuExprs.forall(e => GpuBatchUtils.isFixedWidth(e.dataType)) - if (allReturnTypesFixedWidth && childExprs.forall(_.canThisBeAst)) { - return GpuProjectAstExec(gpuExprs, gpuChild) - } + val astExprs = childExprs.zip(gpuExprs).map { case (meta, expr) => + // cuDF requires return column is fixed width + if (GpuBatchUtils.isFixedWidth(expr.dataType) && meta.canThisBeAst) { + 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.map(_.explainAst(conf.shouldExplainAll)) @@ -77,8 +81,11 @@ class GpuProjectExecMeta( s"return types: ${gpuExprs.map(_.dataType)}") } } + astExprs + } else { + gpuExprs } - GpuProjectExec(gpuExprs, gpuChild) + GpuProjectExec(projectList, gpuChild) } } @@ -132,9 +139,27 @@ object GpuProjectExec { // different vector length, thus not able to reuse cached vectors. GpuExpressionsUtils.cachedNullVectors.get.clear() - GpuArrayHofFusion.project(cb, boundExprs).getOrElse { - val newColumns = boundExprs.safeMap(_.columnarEval(cb)).toArray[ColumnVector] - new ColumnarBatch(newColumns, cb.numRows()) + def projectWithEvaluator(evaluateExpression: Expression => ColumnVector): ColumnarBatch = { + GpuArrayHofFusion.project(cb, boundExprs, evaluateExpression).getOrElse { + val newColumns = boundExprs.safeMap(evaluateExpression).toArray[ColumnVector] + new ColumnarBatch(newColumns, cb.numRows()) + } + } + + val hasAstExpressions = boundExprs.exists { expression => + GpuProjectAstExpression.extractTopLevel(expression).isDefined + } + if (hasAstExpressions) { + withResource(GpuProjectAstExpression.tableFromBatch(cb)) { table => + projectWithEvaluator { expression => + GpuProjectAstExpression.extractTopLevel(expression) match { + case Some(astExpression) => astExpression.computeColumn(table) + case None => expression.columnarEval(cb) + } + } + } + } else { + projectWithEvaluator(_.columnarEval(cb)) } } finally { GpuExpressionsUtils.cachedNullVectors.get.clear() @@ -913,106 +938,6 @@ case class GpuProjectExec( } } -/** Use cudf AST expressions to project columnar batches */ -case class GpuProjectAstExec( - // NOTE for Scala 2.12.x and below we enforce usage of (eager) List to prevent running - // into a deep recursion during serde of lazy lists. See - // https://github.com/NVIDIA/spark-rapids/issues/2036 - // - // Whereas a similar issue https://issues.apache.org/jira/browse/SPARK-27100 is resolved - // using an Array, we opt in for List because it implements Seq while having non-recursive - // serde: https://github.com/scala/scala/blob/2.12.x/src/library/scala/collection/ - // immutable/List.scala#L516 - projectList: List[Expression], - child: SparkPlan -) extends GpuProjectExecLike { - - override def output: Seq[Attribute] = { - projectList.collect { case ne: NamedExpression => ne.toAttribute } - } - - override def internalDoExecuteColumnar(): RDD[ColumnarBatch] = { - child.executeColumnar().mapPartitions(buildRetryableAstIterator) - } - - def buildRetryableAstIterator( - input: Iterator[ColumnarBatch]): GpuColumnarBatchIterator = { - val numOutputRows = gpuLongMetric(NUM_OUTPUT_ROWS) - val numOutputBatches = gpuLongMetric(NUM_OUTPUT_BATCHES) - val opTime = gpuLongMetric(OP_TIME_LEGACY) - val boundProjectList = GpuBindReferences.bindGpuReferences(projectList, child.output, - allMetrics) - val outputTypes = output.map(_.dataType).toArray - new GpuColumnarBatchIterator(true) { - private[this] var maybeSplittedItr: Iterator[ColumnarBatch] = Iterator.empty - private[this] var compiledAstExprs = - NvtxIdWithMetrics(NvtxRegistry.COMPILE_ASTS, opTime) { - boundProjectList.safeMap { expr => - // Use intmax for the left table column count since there's only one input table here. - expr.convertToAst(Int.MaxValue).compile() - } - } - - override def hasNext: Boolean = maybeSplittedItr.hasNext || { - if (input.hasNext) { - true - } else { - close() - false - } - } - - override def next(): ColumnarBatch = { - if (!maybeSplittedItr.hasNext) { - val spillable = SpillableColumnarBatch( - input.next(), SpillPriorities.ACTIVE_ON_DECK_PRIORITY) - // AST currently doesn't support non-deterministic expressions so it's not needed - // to check whether compiled expressions are retryable. - maybeSplittedItr = withRetry(spillable, splitSpillableInHalfByRows) { spillable => - NvtxIdWithMetrics(NvtxRegistry.PROJECT_AST, opTime) { - withResource(spillable.getColumnarBatch()) { cb => - val projectedTable = withResource(tableFromBatch(cb)) { table => - withResource( - compiledAstExprs.safeMap(_.computeColumn(table))) { projectedColumns => - new Table(projectedColumns: _*) - } - } - withResource(projectedTable) { _ => - GpuColumnVector.from(projectedTable, outputTypes) - } - } - } - } - } - - val ret = maybeSplittedItr.next() - numOutputBatches += 1 - numOutputRows += ret.numRows() - ret - } - - override def doClose(): Unit = { - compiledAstExprs.safeClose() - compiledAstExprs = Nil - } - - private def tableFromBatch(cb: ColumnarBatch): Table = { - if (cb.numCols != 0) { - GpuColumnVector.from(cb) - } else { - // Count-only batch but cudf Table cannot be created with no columns. - // Create the cheapest table we can to evaluate the AST expression. - withResource(Scalar.fromBool(false)) { falseScalar => - withResource(cudf.ColumnVector.fromScalar(falseScalar, cb.numRows())) { falseColumn => - new Table(falseColumn) - } - } - } - } - } - } -} - /** * Do projections in a tiered fashion, where earlier tiers contain sub-expressions that are * referenced in later tiers. Each tier adds columns to the original batch corresponding diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/higherOrderFunctions.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/higherOrderFunctions.scala index 0626f110361..e592b0408a3 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/higherOrderFunctions.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/higherOrderFunctions.scala @@ -375,13 +375,14 @@ private[rapids] object GpuArrayHofFusion { private[rapids] def project( batch: ColumnarBatch, - boundExprs: Seq[Expression]): Option[ColumnarBatch] = { + boundExprs: Seq[Expression], + evaluateExpression: Expression => ColumnVector): Option[ColumnarBatch] = { val fusedGroups = findFusedGroups(boundExprs) if (fusedGroups.isEmpty) { None } else { val groupsByStartIndex = fusedGroups.map(group => group.startIndex -> group).toMap - Some(projectWithFusedGroups(batch, boundExprs, groupsByStartIndex)) + Some(projectWithFusedGroups(batch, boundExprs, groupsByStartIndex, evaluateExpression)) } } @@ -481,7 +482,8 @@ private[rapids] object GpuArrayHofFusion { private def projectWithFusedGroups( batch: ColumnarBatch, boundExprs: Seq[Expression], - groupsByStartIndex: Map[Int, HofGroup]): ColumnarBatch = { + groupsByStartIndex: Map[Int, HofGroup], + evaluateExpression: Expression => ColumnVector): ColumnarBatch = { val outputColumns = new Array[ColumnVector](boundExprs.length) closeOnExcept(outputColumns) { _ => boundExprs.indices.foreach { index => @@ -490,7 +492,7 @@ private[rapids] object GpuArrayHofFusion { case Some(group) => evaluateFusedGroup(batch, group, outputColumns) case None => - outputColumns(index) = boundExprs(index).columnarEval(batch) + outputColumns(index) = evaluateExpression(boundExprs(index)) } } } diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/catalyst/expressions/GpuEquivalentExpressions.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/catalyst/expressions/GpuEquivalentExpressions.scala index e076bb28a81..0270f9f9ee9 100644 --- a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/catalyst/expressions/GpuEquivalentExpressions.scala +++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/catalyst/expressions/GpuEquivalentExpressions.scala @@ -275,6 +275,8 @@ class GpuEquivalentExpressions { } object GpuEquivalentExpressions { + private case class TierExpression(originalExpression: Expression, expression: Expression) + /** * Recursively replaces semantic equal expression with its proxy expression in `substitutionMap`. */ @@ -328,16 +330,22 @@ object GpuEquivalentExpressions { /** * Applies substitutions to all expression tiers. */ - private def doSubstitutions(exprTiers: Seq[Seq[Expression]], currentTier: Seq[Expression], - substitutionMap: mutable.HashMap[Expression, Expression]): Seq[Seq[Expression]] = { + private def doSubstitutions( + exprTiers: Seq[Seq[TierExpression]], + currentTier: Seq[TierExpression], + substitutionMap: mutable.HashMap[Expression, Expression]): Seq[Seq[TierExpression]] = { // Make substitutions in given tiers, filtering out matches from original current tier, // but don't filter the last tier - it needs to match original size val subTiers = exprTiers.dropRight(1) val lastTier = exprTiers.last + val currentExpressions = currentTier.map(_.expression) val updatedSubTiers = subTiers.map { - t => t.filter(e => !currentTier.contains(e)).map(replaceWithCommonRef(_, substitutionMap)) + t => t.filter(e => !currentExpressions.contains(e.expression)) + .map(e => e.copy(expression = replaceWithCommonRef(e.expression, substitutionMap))) + } + val updatedLastTier = lastTier.map { + e => e.copy(expression = replaceWithCommonRef(e.expression, substitutionMap)) } - val updatedLastTier = lastTier.map(replaceWithCommonRef(_, substitutionMap)) updatedSubTiers ++ Seq(updatedLastTier) } @@ -345,34 +353,37 @@ object GpuEquivalentExpressions { * Apply subexpression substitutions to all tiers. */ @tailrec - private def recurseUpdateTiers(exprTiers: Seq[Seq[Expression]], + private def recurseUpdateTiers(exprTiers: Seq[Seq[TierExpression]], updatedTiers: Seq[Seq[Expression]], substitutionMap: mutable.HashMap[Expression, Expression], - startIndex: Int): Seq[Seq[Expression]] = { + startIndex: Int, + rewriteCommonExpression: (Expression, Expression) => Expression): Seq[Seq[Expression]] = { exprTiers match { case Nil => updatedTiers case tier :: tail => { // Last tier should already be updated. if (tail.isEmpty) { - updatedTiers ++ Seq(tier) + updatedTiers ++ Seq(tier.map(_.expression)) } else { // Replace expressions in this tier with GpuAlias val aliasedTier = tier.zipWithIndex.map { case (e, i) => - GpuAlias(e, s"tiered_input_${startIndex + i}")() + val tierExpression = + rewriteCommonExpression(e.originalExpression, e.expression) + GpuAlias(tierExpression, s"tiered_input_${startIndex + i}")() } // Add them to the map tier.zip(aliasedTier).foreach { case (expr, alias) => { - substitutionMap.get(expr) match { - case None => substitutionMap.put(expr, alias.toAttribute) + substitutionMap.get(expr.expression) match { + case None => substitutionMap.put(expr.expression, alias.toAttribute) case Some(_) => } } } val newUpdatedTiers = doSubstitutions(tail, tier, substitutionMap) recurseUpdateTiers(newUpdatedTiers, updatedTiers ++ Seq(aliasedTier), - substitutionMap, startIndex + aliasedTier.size) + substitutionMap, startIndex + aliasedTier.size, rewriteCommonExpression) } } } @@ -430,11 +441,29 @@ object GpuEquivalentExpressions { } def getExprTiers(expressions: Seq[Expression]): Seq[Seq[Expression]] = { + getExprTiers(expressions, (_, expression) => expression) + } + + /** + * Creates expression tiers and rewrites common expressions after substitutions. + * + * The callback receives both the original common expression and its potentially substituted + * form. + */ + def getExprTiers( + expressions: Seq[Expression], + rewriteCommonExpression: (Expression, Expression) => Expression): Seq[Seq[Expression]] = { // Get tiers of common expressions val expressionTiers = recurseCommonExpressions(expressions, Seq(expressions)) + val tierExpressions = expressionTiers.map { tier => + tier.map { expression => + TierExpression(expression, expression) + } + } val substitutionMap = mutable.HashMap.empty[Expression, Expression] // Update expression with common expressions from previous tiers - recurseUpdateTiers(expressionTiers, Seq.empty, substitutionMap, 0) + recurseUpdateTiers( + tierExpressions, Seq.empty, substitutionMap, 0, rewriteCommonExpression) } // Determine which of the inputAttrs are needed for remaining tiers 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 344a104d1a2..9d75f80d009 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuArrayHofFusionSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuArrayHofFusionSuite.scala @@ -16,9 +16,14 @@ package com.nvidia.spark.rapids +import ai.rapids.cudf.Table import com.nvidia.spark.rapids.Arm.withResource +import com.nvidia.spark.rapids.RapidsPluginImplicits._ +import org.mockito.ArgumentCaptor +import org.mockito.Mockito.{spy, verify} import org.apache.spark.sql.catalyst.expressions.{Expression, ExprId} +import org.apache.spark.sql.rapids.{GpuAdd, GpuMultiply} import org.apache.spark.sql.types.{ArrayType, BooleanType, DataType, IntegerType, LongType} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -195,7 +200,8 @@ class GpuArrayHofFusionSuite extends GpuUnitTests { alias(executableTransform(301), "right")) def check(batch: ColumnarBatch): Unit = { - val fused = GpuArrayHofFusion.project(batch, exprs) + val fused = GpuArrayHofFusion.project( + batch, exprs, _.columnarEval(batch)) assert(fused.isDefined) withResource(fused.get) { projected => assertResult(3)(projected.numCols()) @@ -209,4 +215,39 @@ class GpuArrayHofFusionSuite extends GpuUnitTests { withResource(GpuColumnVector.emptyBatchFromTypes(Array(arrayType)))(check) withResource(FuzzerUtils.createColumnarBatch(schema, 8))(check) } + + test("fused HOF project preserves the shared AST input table") { + val arrayType = ArrayType(IntegerType, containsNull = true) + val schema = FuzzerUtils.createSchema(arrayType, LongType, LongType) + val firstAst = 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( + GpuBoundReference(1, LongType, nullable = true)(ExprId(402), "a"), + GpuBoundReference(2, LongType, nullable = true)(ExprId(403), "b"))())) + val expressions = Seq( + alias(executableTransform(404), "left"), + alias(firstAst, "sum"), + alias(secondAst, "product"), + alias(executableTransform(405), "right")) + + assertResult(Seq(Seq(0, 3))) { + GpuArrayHofFusion.findFusedGroupIndexes(expressions) + } + withResource(Seq(firstAst, secondAst)) { _ => + withResource(FuzzerUtils.createColumnarBatch(schema, 8)) { batch => + withResource(GpuProjectExec.project(batch, expressions)) { projected => + assertResult(4)(projected.numCols()) + assertResult(batch.numRows())(projected.numRows()) + } + } + } + + 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) + } } 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 e3a5207d887..bd0bc77a12b 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 @@ -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. @@ -23,19 +23,33 @@ import ai.rapids.cudf.Table import com.nvidia.spark.rapids._ import com.nvidia.spark.rapids.Arm.withResource import com.nvidia.spark.rapids.jni.RmmSpark -import org.mockito.Mockito.{mock, spy, when} +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.execution.SparkPlan +import org.apache.spark.sql.catalyst.expressions.{ + AttributeReference, Expression, 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") @@ -134,49 +148,98 @@ class ProjectExprSuite extends SparkQueryCompareTestSuite { test("AST retry with split") { RmmSpark.currentThreadIsDedicatedToTask(0) try { - val a = AttributeReference("a", LongType)() - val b = AttributeReference("b", LongType)() val sb = buildProjectBatch() - val expr = GpuAlias(GpuAdd( + val astExpression = GpuProjectAstExpression(GpuAdd( GpuBoundReference(0, LongType, true)(NamedExpression.newExprId, "a"), - GpuBoundReference(1, LongType, true)(NamedExpression.newExprId, "b"), false)(), - "ret")() - val mockPlan = mock(classOf[SparkPlan]) - when(mockPlan.output).thenReturn(Seq(a, b)) - val ast = GpuProjectAstExec(List(expr.asInstanceOf[Expression]), mockPlan) - RmmSpark.forceSplitAndRetryOOM(RmmSpark.getCurrentThreadId, 1, - RmmSpark.OomInjectionType.GPU.ordinal, 0) - withResource(sb) { sb => - withResource(ast.buildRetryableAstIterator(Seq(sb.getColumnarBatch).iterator)) { result => - withResource(result.next()) { cb => - assertResult(2)(cb.numRows) - assertResult(1)(cb.numCols) - val gcv = cb.column(0).asInstanceOf[GpuColumnVector] - withResource(gcv.getBase.copyToHost()) { hcv => - assert(!hcv.isNull(0)) - assertResult(11L)(hcv.getLong(0)) - assert(hcv.isNull(1)) - } + GpuBoundReference(1, LongType, true)(NamedExpression.newExprId, "b"), false)()) + val expr = GpuAlias(astExpression, "ret")() + val tieredProject = GpuTieredProject(Seq(Seq(expr))) + withResource(astExpression) { _ => + RmmSpark.forceSplitAndRetryOOM(RmmSpark.getCurrentThreadId, 1, + RmmSpark.OomInjectionType.GPU.ordinal, 0) + val result = tieredProject.projectAndCloseStreamingWithSplitRetry(sb) + withResource(result.next()) { cb => + assertResult(2)(cb.numRows) + assertResult(1)(cb.numCols) + val gcv = cb.column(0).asInstanceOf[GpuColumnVector] + withResource(gcv.getBase.copyToHost()) { hcv => + assert(!hcv.isNull(0)) + assertResult(11L)(hcv.getLong(0)) + assert(hcv.isNull(1)) } + } - withResource(result.next()) { cb => - assertResult(2)(cb.numRows) - assertResult(1)(cb.numCols) - val gcv = cb.column(0).asInstanceOf[GpuColumnVector] - withResource(gcv.getBase.copyToHost()) { hcv => - assert(!hcv.isNull(0)) - assertResult(11L)(hcv.getLong(0)) - assert(!hcv.isNull(1)) - assertResult(10L)(hcv.getLong(1)) - } + withResource(result.next()) { cb => + assertResult(2)(cb.numRows) + assertResult(1)(cb.numCols) + val gcv = cb.column(0).asInstanceOf[GpuColumnVector] + withResource(gcv.getBase.copyToHost()) { hcv => + assert(!hcv.isNull(0)) + assertResult(11L)(hcv.getLong(0)) + assert(!hcv.isNull(1)) + assertResult(10L)(hcv.getLong(1)) } } + assert(!result.hasNext) } } finally { RmmSpark.removeCurrentDedicatedThreadAssociation(0) } } + test("tiered project preserves AST across multi-level shared expressions") { + val a = AttributeReference("a", LongType)() + val b = AttributeReference("b", LongType)() + val c = AttributeReference("c", LongType)() + val d = AttributeReference("d", LongType)() + val e = AttributeReference("e", LongType)() + val f = AttributeReference("f", LongType)() + def shared: GpuAdd = GpuAdd(a, b, failOnError = false)() + def intermediate: GpuMultiply = GpuMultiply(shared, c)() + def ast(expression: GpuExpression, name: String): GpuAlias = + GpuAlias(GpuProjectAstExpression(expression), name)() + val expressions = Seq( + ast(GpuAdd(intermediate, d, failOnError = false)(), "first"), + ast(GpuAdd(intermediate, e, failOnError = false)(), "second"), + ast(shared, "shared_first"), + ast(shared, "shared_second"), + GpuAlias(GpuGreatest(Seq(shared, f)), "regular")()) + + val tiered = GpuBindReferences.bindGpuReferencesTieredNoMetrics( + expressions, Seq(a, b, c, d, e, f), new SQLConf()) + + 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 finalReferences = tiered.exprTiers.last.flatMap(tierReferences) + assertResult(5)(finalReferences.size) + assertResult(2)(finalReferences.map(_.exprId).distinct.size) + } + + 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")()))) { _ => } + } + } + verify(astExpression, never()).close() + context.markTaskComplete() + verify(astExpression).close() + } finally { + TrampolineUtil.unsetTaskContext() + ScalableTaskCompletion.reset() + astExpression.close() + } + } + testSparkResultsAreEqual("Test literal values in select", mixedFloatDf) { frame => frame.select(col("floats"), diff --git a/tests/src/test/spark330/scala/com/nvidia/spark/rapids/IntervalArithmeticSuite.scala b/tests/src/test/spark330/scala/com/nvidia/spark/rapids/IntervalArithmeticSuite.scala index c21b7311d47..7e4a6a276de 100644 --- a/tests/src/test/spark330/scala/com/nvidia/spark/rapids/IntervalArithmeticSuite.scala +++ b/tests/src/test/spark330/scala/com/nvidia/spark/rapids/IntervalArithmeticSuite.scala @@ -204,8 +204,7 @@ class IntervalArithmeticSuite extends SparkQueryCompareTestSuite { spark.createDataFrame(spark.sparkContext.parallelize(data), schema) }, new SparkConf().set(RapidsConf.ENABLE_PROJECT_AST.key, "true"), - existClasses = "GpuProjectAstExec", - nonExistClasses = "GpuProjectExec" + existClasses = "GpuProjectAstExpression" ) { df => { df.selectExpr("+c_year_month1") From 6cea0cfba6055e9fdf1e69c3228eadfd2f7171d5 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Fri, 24 Jul 2026 15:40:05 +0800 Subject: [PATCH 05/20] Add signoff Signed-off-by: Haoyang Li From f3ebe60eb67f40c30882ae849480290ec88bac44 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Fri, 24 Jul 2026 18:29:39 +0800 Subject: [PATCH 06/20] small refactor Signed-off-by: Haoyang Li --- integration_tests/src/main/python/ast_test.py | 25 +++-- .../rapids/GpuProjectAstExpression.scala | 96 +++++++++++++------ .../GpuEquivalentExpressions.scala | 53 +++------- .../rapids/IntervalArithmeticSuite.scala | 5 +- 4 files changed, 91 insertions(+), 88 deletions(-) diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index 8f57f38901b..5298d60bace 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/integration_tests/src/main/python/ast_test.py @@ -81,6 +81,9 @@ def assert_gpu_ast(is_supported, func, conf={}): non_exist_classes=non_exist, conf=ast_conf) +def assert_gpu_project_without_ast(func, conf={}): + assert_gpu_ast(False, func, conf) + def assert_unary_ast(data_descr, func, conf={}): (data_gen, is_supported) = data_descr assert_gpu_ast(is_supported, lambda spark: func(unary_op_df(spark, data_gen)), conf=conf) @@ -119,22 +122,16 @@ def test_isnotnull(data_descr): def test_bitwise_not(data_descr): assert_unary_ast(data_descr, lambda df: df.selectExpr('~a')) -# This just ends up being a pass through. There is no good way to force -# a unary positive into a plan, because it gets optimized out, but this -# verifies that we can handle it. -@pytest.mark.parametrize('data_descr', [ - (byte_gen, True), - (short_gen, True), - (int_gen, True), - (long_gen, True), - (float_gen, True), - (double_gen, True)], ids=idfn) -def test_unary_positive(data_descr): - assert_unary_ast(data_descr, lambda df: df.selectExpr('+a')) +# Unary positive is optimized to a pass-through, so per-expression AST has nothing to compile. +@pytest.mark.parametrize( + 'data_gen', [byte_gen, short_gen, int_gen, long_gen, float_gen, double_gen], ids=idfn) +def test_unary_positive(data_gen): + assert_gpu_project_without_ast( + lambda spark: unary_op_df(spark, data_gen).selectExpr('+a')) def test_unary_positive_for_daytime_interval(): - data_descr = (DayTimeIntervalGen(), True) - assert_unary_ast(data_descr, lambda df: df.selectExpr('+a')) + assert_gpu_project_without_ast( + lambda spark: unary_op_df(spark, DayTimeIntervalGen()).selectExpr('+a')) @pytest.mark.parametrize('data_descr', ast_arithmetic_descrs, ids=idfn) @disable_ansi_mode 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 756286dc666..a8f1c2bdc64 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 @@ -29,17 +29,20 @@ import com.nvidia.spark.rapids.shims.ShimUnaryExpression import org.apache.spark.TaskContext import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.rapids.catalyst.expressions.{ - GpuEquivalentExpressions, GpuExpressionEquals} +import org.apache.spark.sql.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 = { + GpuAlias(child, alias.name)(alias.exprId, alias.qualifier, alias.explicitMetadata) + } + private[rapids] def wrap(expression: NamedExpression): NamedExpression = { expression match { - case alias @ GpuAlias(child: GpuExpression, name) => - GpuAlias(GpuProjectAstExpression(child), name)( - alias.exprId, alias.qualifier, alias.explicitMetadata) + case alias @ GpuAlias(_: GpuProjectAstExpression, _) => alias + case alias @ GpuAlias(child: GpuExpression, _) => + replaceChild(alias, GpuProjectAstExpression(child)) case other => other } } @@ -54,55 +57,85 @@ object GpuProjectAstExpression { } private def unwrap(expression: Expression): Expression = expression match { - case alias @ GpuAlias(astExpression: GpuProjectAstExpression, name) => - GpuAlias(astExpression.child, name)( - alias.exprId, alias.qualifier, alias.explicitMetadata) + case alias: GpuAlias => + val child = unwrap(alias.child) + if (child eq alias.child) alias else replaceChild(alias, child) case astExpression: GpuProjectAstExpression => astExpression.child case other => other } private def rewrap(expression: Expression): Expression = expression match { + case astExpression: GpuProjectAstExpression => astExpression case namedExpression: NamedExpression => wrap(namedExpression) case gpuExpression: GpuExpression => GpuProjectAstExpression(gpuExpression) 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) - val astSubexpressions = astOutputs.flatten.flatMap { astExpression => - astExpression.child.collect { - case gpuExpression: GpuExpression => GpuExpressionEquals(gpuExpression) - } - }.toSet + val hasAstOutputs = expressions.exists(extractTopLevel(_).isDefined) // CSE must see through the marker so AST and non-AST outputs can share the same tiers. - val unwrapped = expressions.map(unwrap) + val unwrapped = if (hasAstOutputs) expressions.map(unwrap) else expressions val replaced = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) { GpuEquivalentExpressions.replaceMultiExpressions(unwrapped, conf) } else { unwrapped } - val tiers = GpuEquivalentExpressions.getExprTiers( - replaced, - (original, substituted) => (original, substituted) match { - case (gpuOriginal: GpuExpression, gpuSubstituted: GpuExpression) - if GpuBatchUtils.isFixedWidth(gpuOriginal.dataType) && - astSubexpressions.contains(GpuExpressionEquals(gpuOriginal)) => - GpuProjectAstExpression(gpuSubstituted) - case _ => substituted - }) - val finalTier = tiers.last.zip(astOutputs).map { - case (expression, Some(_)) => rewrap(expression) - case (expression, None) => expression + val tiers = GpuEquivalentExpressions.getExprTiers(replaced) + if (hasAstOutputs) { + val astOutputs = expressions.map(extractTopLevel(_).isDefined) + rewrapAstTiers(tiers, astOutputs) + } else { + tiers } - tiers.dropRight(1) :+ finalTier } private[rapids] def tableFromBatch(batch: ColumnarBatch): Table = { if (batch.numCols() != 0) { GpuColumnVector.from(batch) } else { + // cuDF cannot represent a row-count-only table, so use a dummy fixed-width column. withResource(Scalar.fromBool(false)) { falseScalar => withResource(ai.rapids.cudf.ColumnVector.fromScalar(falseScalar, batch.numRows())) { falseColumn => new Table(falseColumn) @@ -138,9 +171,10 @@ case class GpuProjectAstExpression(child: GpuExpression) } } - def computeColumn(table: Table): GpuColumnVector = { + private[rapids] def computeColumn(table: Table): GpuColumnVector = { + val compiled = getCompiledExpression NvtxIdWithMetrics(NvtxRegistry.PROJECT_AST, opTime) { - closeOnExcept(getCompiledExpression.computeColumn(table)) { result => + closeOnExcept(compiled.computeColumn(table)) { result => GpuColumnVector.from(result, dataType) } } @@ -149,7 +183,7 @@ case class GpuProjectAstExpression(child: GpuExpression) private def getCompiledExpression: CompiledExpression = synchronized { if (compiledExpression == null) { val compiled = NvtxIdWithMetrics(NvtxRegistry.COMPILE_ASTS, opTime) { - // Project AST has a single input table. + // Force every bound reference to the left table; Project AST has one input table. child.convertToAst(Int.MaxValue).compile() } closeOnExcept(compiled) { _ => diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/catalyst/expressions/GpuEquivalentExpressions.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/catalyst/expressions/GpuEquivalentExpressions.scala index 0270f9f9ee9..e076bb28a81 100644 --- a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/catalyst/expressions/GpuEquivalentExpressions.scala +++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/catalyst/expressions/GpuEquivalentExpressions.scala @@ -275,8 +275,6 @@ class GpuEquivalentExpressions { } object GpuEquivalentExpressions { - private case class TierExpression(originalExpression: Expression, expression: Expression) - /** * Recursively replaces semantic equal expression with its proxy expression in `substitutionMap`. */ @@ -330,22 +328,16 @@ object GpuEquivalentExpressions { /** * Applies substitutions to all expression tiers. */ - private def doSubstitutions( - exprTiers: Seq[Seq[TierExpression]], - currentTier: Seq[TierExpression], - substitutionMap: mutable.HashMap[Expression, Expression]): Seq[Seq[TierExpression]] = { + private def doSubstitutions(exprTiers: Seq[Seq[Expression]], currentTier: Seq[Expression], + substitutionMap: mutable.HashMap[Expression, Expression]): Seq[Seq[Expression]] = { // Make substitutions in given tiers, filtering out matches from original current tier, // but don't filter the last tier - it needs to match original size val subTiers = exprTiers.dropRight(1) val lastTier = exprTiers.last - val currentExpressions = currentTier.map(_.expression) val updatedSubTiers = subTiers.map { - t => t.filter(e => !currentExpressions.contains(e.expression)) - .map(e => e.copy(expression = replaceWithCommonRef(e.expression, substitutionMap))) - } - val updatedLastTier = lastTier.map { - e => e.copy(expression = replaceWithCommonRef(e.expression, substitutionMap)) + t => t.filter(e => !currentTier.contains(e)).map(replaceWithCommonRef(_, substitutionMap)) } + val updatedLastTier = lastTier.map(replaceWithCommonRef(_, substitutionMap)) updatedSubTiers ++ Seq(updatedLastTier) } @@ -353,37 +345,34 @@ object GpuEquivalentExpressions { * Apply subexpression substitutions to all tiers. */ @tailrec - private def recurseUpdateTiers(exprTiers: Seq[Seq[TierExpression]], + private def recurseUpdateTiers(exprTiers: Seq[Seq[Expression]], updatedTiers: Seq[Seq[Expression]], substitutionMap: mutable.HashMap[Expression, Expression], - startIndex: Int, - rewriteCommonExpression: (Expression, Expression) => Expression): Seq[Seq[Expression]] = { + startIndex: Int): Seq[Seq[Expression]] = { exprTiers match { case Nil => updatedTiers case tier :: tail => { // Last tier should already be updated. if (tail.isEmpty) { - updatedTiers ++ Seq(tier.map(_.expression)) + updatedTiers ++ Seq(tier) } else { // Replace expressions in this tier with GpuAlias val aliasedTier = tier.zipWithIndex.map { case (e, i) => - val tierExpression = - rewriteCommonExpression(e.originalExpression, e.expression) - GpuAlias(tierExpression, s"tiered_input_${startIndex + i}")() + GpuAlias(e, s"tiered_input_${startIndex + i}")() } // Add them to the map tier.zip(aliasedTier).foreach { case (expr, alias) => { - substitutionMap.get(expr.expression) match { - case None => substitutionMap.put(expr.expression, alias.toAttribute) + substitutionMap.get(expr) match { + case None => substitutionMap.put(expr, alias.toAttribute) case Some(_) => } } } val newUpdatedTiers = doSubstitutions(tail, tier, substitutionMap) recurseUpdateTiers(newUpdatedTiers, updatedTiers ++ Seq(aliasedTier), - substitutionMap, startIndex + aliasedTier.size, rewriteCommonExpression) + substitutionMap, startIndex + aliasedTier.size) } } } @@ -441,29 +430,11 @@ object GpuEquivalentExpressions { } def getExprTiers(expressions: Seq[Expression]): Seq[Seq[Expression]] = { - getExprTiers(expressions, (_, expression) => expression) - } - - /** - * Creates expression tiers and rewrites common expressions after substitutions. - * - * The callback receives both the original common expression and its potentially substituted - * form. - */ - def getExprTiers( - expressions: Seq[Expression], - rewriteCommonExpression: (Expression, Expression) => Expression): Seq[Seq[Expression]] = { // Get tiers of common expressions val expressionTiers = recurseCommonExpressions(expressions, Seq(expressions)) - val tierExpressions = expressionTiers.map { tier => - tier.map { expression => - TierExpression(expression, expression) - } - } val substitutionMap = mutable.HashMap.empty[Expression, Expression] // Update expression with common expressions from previous tiers - recurseUpdateTiers( - tierExpressions, Seq.empty, substitutionMap, 0, rewriteCommonExpression) + recurseUpdateTiers(expressionTiers, Seq.empty, substitutionMap, 0) } // Determine which of the inputAttrs are needed for remaining tiers diff --git a/tests/src/test/spark330/scala/com/nvidia/spark/rapids/IntervalArithmeticSuite.scala b/tests/src/test/spark330/scala/com/nvidia/spark/rapids/IntervalArithmeticSuite.scala index 7e4a6a276de..b563cbff998 100644 --- a/tests/src/test/spark330/scala/com/nvidia/spark/rapids/IntervalArithmeticSuite.scala +++ b/tests/src/test/spark330/scala/com/nvidia/spark/rapids/IntervalArithmeticSuite.scala @@ -197,14 +197,15 @@ class IntervalArithmeticSuite extends SparkQueryCompareTestSuite { } testSparkResultsAreEqual( - "test year month interval arithmetic: Positive, AST", + "test year month interval arithmetic: Positive, AST config", spark => { val data = Seq(Row(Period.ofYears(100))) val schema = StructType(Seq(StructField("c_year_month1", YearMonthIntervalType()))) spark.createDataFrame(spark.sparkContext.parallelize(data), schema) }, new SparkConf().set(RapidsConf.ENABLE_PROJECT_AST.key, "true"), - existClasses = "GpuProjectAstExpression" + existClasses = "GpuProjectExec", + nonExistClasses = "GpuProjectAstExpression" ) { df => { df.selectExpr("+c_year_month1") From b242c733611d207edb8807734c3cd7eff0c58a55 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Wed, 29 Jul 2026 17:20:59 +0800 Subject: [PATCH 07/20] address comments Signed-off-by: Haoyang Li --- integration_tests/src/main/python/ast_test.py | 1 + .../rapids/GpuProjectAstExpression.scala | 32 ++++++++++--------- .../spark/rapids/basicPhysicalOperators.scala | 23 +++++++------ .../spark/rapids/higherOrderFunctions.scala | 8 ++--- .../spark/sql/rapids/ProjectExprSuite.scala | 8 +++++ 5 files changed, 41 insertions(+), 31 deletions(-) diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index 5298d60bace..7205ac23dbf 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/integration_tests/src/main/python/ast_test.py @@ -426,6 +426,7 @@ def test_multi_tier_ast(): func=lambda spark: spark.range(10).withColumn("x", f.col("id")).repartition(1)\ .selectExpr("x", "(id < x) == (id < (id + x))")) +# ANSI mode is disabled here due to an overflow issue with integer multiplication on Spark 4.0.0. @disable_ansi_mode @pytest.mark.parametrize( 'tiered_project_enabled', ['true', 'false'], ids=['tiered', 'single_tier']) 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 a8f1c2bdc64..79b5f8ed13a 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 @@ -35,16 +35,22 @@ import org.apache.spark.sql.vectorized.ColumnarBatch object GpuProjectAstExpression { private def replaceChild(alias: GpuAlias, child: Expression): GpuAlias = { - GpuAlias(child, alias.name)(alias.exprId, alias.qualifier, alias.explicitMetadata) + if (child eq alias.child) { + alias + } else { + GpuAlias(child, alias.name)(alias.exprId, alias.qualifier, alias.explicitMetadata) + } } - private[rapids] def wrap(expression: NamedExpression): NamedExpression = { - expression match { - case alias @ GpuAlias(_: GpuProjectAstExpression, _) => alias - case alias @ GpuAlias(child: GpuExpression, _) => - replaceChild(alias, GpuProjectAstExpression(child)) - case other => other - } + private def asAst(child: GpuExpression): GpuProjectAstExpression = child match { + case astExpression: GpuProjectAstExpression => astExpression + case other => GpuProjectAstExpression(other) + } + + private[rapids] def wrap(expression: NamedExpression): NamedExpression = expression match { + case alias @ GpuAlias(child: GpuExpression, _) => + replaceChild(alias, asAst(child)) + case other => other } @tailrec @@ -57,17 +63,13 @@ object GpuProjectAstExpression { } private def unwrap(expression: Expression): Expression = expression match { - case alias: GpuAlias => - val child = unwrap(alias.child) - if (child eq alias.child) alias else replaceChild(alias, child) + case alias: GpuAlias => replaceChild(alias, unwrap(alias.child)) case astExpression: GpuProjectAstExpression => astExpression.child case other => other } private def rewrap(expression: Expression): Expression = expression match { - case astExpression: GpuProjectAstExpression => astExpression case namedExpression: NamedExpression => wrap(namedExpression) - case gpuExpression: GpuExpression => GpuProjectAstExpression(gpuExpression) case other => other } @@ -114,7 +116,8 @@ object GpuProjectAstExpression { private[rapids] def buildExprTiers( expressions: Seq[Expression], conf: SQLConf): Seq[Seq[Expression]] = { - val hasAstOutputs = expressions.exists(extractTopLevel(_).isDefined) + val astOutputs = expressions.map(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)) { @@ -124,7 +127,6 @@ object GpuProjectAstExpression { } val tiers = GpuEquivalentExpressions.getExprTiers(replaced) if (hasAstOutputs) { - val astOutputs = expressions.map(extractTopLevel(_).isDefined) rewrapAstTiers(tiers, astOutputs) } else { tiers 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 44684e18ffb..57e73d062e0 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 @@ -60,7 +60,6 @@ class GpuProjectExecMeta( val gpuExprs = childExprs.map(_.convertToGpu().asInstanceOf[NamedExpression]).toList val gpuChild = childPlans.head.convertIfNeeded() val projectList = if (conf.isProjectAstEnabled) { - val allReturnTypesFixedWidth = gpuExprs.forall(e => GpuBatchUtils.isFixedWidth(e.dataType)) val astExprs = childExprs.zip(gpuExprs).map { case (meta, expr) => // cuDF requires return column is fixed width if (GpuBatchUtils.isFixedWidth(expr.dataType) && meta.canThisBeAst) { @@ -71,15 +70,15 @@ class GpuProjectExecMeta( }.toList // explain AST because this is optional and it is sometimes hard to debug if (conf.shouldExplain) { - val explain = childExprs.map(_.explainAst(conf.shouldExplainAll)) - .filter(_.nonEmpty) + 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" + }).mkString if (explain.nonEmpty) { logWarning(s"AST PROJECT\n$explain") } - if (!allReturnTypesFixedWidth) { - logWarning(s"AST PROJECT\n have non fixed return column, " + - s"return types: ${gpuExprs.map(_.dataType)}") - } } astExprs } else { @@ -139,9 +138,9 @@ object GpuProjectExec { // different vector length, thus not able to reuse cached vectors. GpuExpressionsUtils.cachedNullVectors.get.clear() - def projectWithEvaluator(evaluateExpression: Expression => ColumnVector): ColumnarBatch = { - GpuArrayHofFusion.project(cb, boundExprs, evaluateExpression).getOrElse { - val newColumns = boundExprs.safeMap(evaluateExpression).toArray[ColumnVector] + def projectWithEval(evalColumn: Expression => ColumnVector): ColumnarBatch = { + GpuArrayHofFusion.project(cb, boundExprs, evalColumn).getOrElse { + val newColumns = boundExprs.safeMap(evalColumn).toArray[ColumnVector] new ColumnarBatch(newColumns, cb.numRows()) } } @@ -151,7 +150,7 @@ object GpuProjectExec { } if (hasAstExpressions) { withResource(GpuProjectAstExpression.tableFromBatch(cb)) { table => - projectWithEvaluator { expression => + projectWithEval { expression => GpuProjectAstExpression.extractTopLevel(expression) match { case Some(astExpression) => astExpression.computeColumn(table) case None => expression.columnarEval(cb) @@ -159,7 +158,7 @@ object GpuProjectExec { } } } else { - projectWithEvaluator(_.columnarEval(cb)) + projectWithEval(_.columnarEval(cb)) } } finally { GpuExpressionsUtils.cachedNullVectors.get.clear() diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/higherOrderFunctions.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/higherOrderFunctions.scala index e592b0408a3..7171b869ab4 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/higherOrderFunctions.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/higherOrderFunctions.scala @@ -376,13 +376,13 @@ private[rapids] object GpuArrayHofFusion { private[rapids] def project( batch: ColumnarBatch, boundExprs: Seq[Expression], - evaluateExpression: Expression => ColumnVector): Option[ColumnarBatch] = { + evalColumn: Expression => ColumnVector): Option[ColumnarBatch] = { val fusedGroups = findFusedGroups(boundExprs) if (fusedGroups.isEmpty) { None } else { val groupsByStartIndex = fusedGroups.map(group => group.startIndex -> group).toMap - Some(projectWithFusedGroups(batch, boundExprs, groupsByStartIndex, evaluateExpression)) + Some(projectWithFusedGroups(batch, boundExprs, groupsByStartIndex, evalColumn)) } } @@ -483,7 +483,7 @@ private[rapids] object GpuArrayHofFusion { batch: ColumnarBatch, boundExprs: Seq[Expression], groupsByStartIndex: Map[Int, HofGroup], - evaluateExpression: Expression => ColumnVector): ColumnarBatch = { + evalColumn: Expression => ColumnVector): ColumnarBatch = { val outputColumns = new Array[ColumnVector](boundExprs.length) closeOnExcept(outputColumns) { _ => boundExprs.indices.foreach { index => @@ -492,7 +492,7 @@ private[rapids] object GpuArrayHofFusion { case Some(group) => evaluateFusedGroup(batch, group, outputColumns) case None => - outputColumns(index) = evaluateExpression(boundExprs(index)) + outputColumns(index) = evalColumn(boundExprs(index)) } } } 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 bd0bc77a12b..33c8d5a3c4c 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 @@ -198,6 +198,8 @@ class ProjectExprSuite extends SparkQueryCompareTestSuite { def intermediate: GpuMultiply = GpuMultiply(shared, c)() def ast(expression: GpuExpression, name: String): GpuAlias = GpuAlias(GpuProjectAstExpression(expression), name)() + // [AST((a+b)*c+d) AS first, AST((a+b)*c+e) AS second, + // AST(a+b) AS shared_first, AST(a+b) AS shared_second, greatest(a+b, f) AS regular] val expressions = Seq( ast(GpuAdd(intermediate, d, failOnError = false)(), "first"), ast(GpuAdd(intermediate, e, failOnError = false)(), "second"), @@ -208,10 +210,16 @@ class ProjectExprSuite extends SparkQueryCompareTestSuite { val tiered = GpuBindReferences.bindGpuReferencesTieredNoMetrics( expressions, Seq(a, b, c, d, e, f), new SQLConf()) + // After CSE: + // tier 0: [AST(a+b) AS t1] + // 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) + // Final references: [t2, t2, t1, t1, t1] (distinct: {t2, t1}). val finalReferences = tiered.exprTiers.last.flatMap(tierReferences) assertResult(5)(finalReferences.size) assertResult(2)(finalReferences.map(_.exprId).distinct.size) From 8ffede9171b3d9802a0b5edb3010dc28d6d2b02b Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Wed, 29 Jul 2026 22:10:28 +0800 Subject: [PATCH 08/20] fix tests Signed-off-by: Haoyang Li --- integration_tests/src/main/python/ast_test.py | 34 +++++++++++----- .../com/nvidia/spark/rapids/RapidsMeta.scala | 13 +++++-- .../spark/rapids/basicPhysicalOperators.scala | 7 +++- .../nvidia/spark/rapids/AstUtilSuite.scala | 39 ++++++++++++++++++- 4 files changed, 76 insertions(+), 17 deletions(-) diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index 7205ac23dbf..3b1174d793c 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/integration_tests/src/main/python/ast_test.py @@ -98,8 +98,8 @@ def test_literal(spark_tmp_path, data_gen): data_path = spark_tmp_path + '/AST_TEST_DATA' with_cpu_session(lambda spark: gen_df(spark, [("a", IntegerGen())]).write.parquet(data_path)) scalar = with_cpu_session(lambda spark: gen_scalar(data_gen, force_no_nulls=True)) - assert_gpu_ast(is_supported=True, - func=lambda spark: spark.read.parquet(data_path).select(scalar)) + assert_gpu_project_without_ast( + func=lambda spark: spark.read.parquet(data_path).select(scalar)) @pytest.mark.parametrize('data_gen', [boolean_gen, byte_gen, short_gen, int_gen, long_gen, float_gen, double_gen, timestamp_gen, date_gen], ids=idfn) def test_null_literal(spark_tmp_path, data_gen): @@ -107,8 +107,8 @@ def test_null_literal(spark_tmp_path, data_gen): data_path = spark_tmp_path + '/AST_TEST_DATA' with_cpu_session(lambda spark: gen_df(spark, [("a", IntegerGen())]).write.parquet(data_path)) data_type = data_gen.data_type - assert_gpu_ast(is_supported=True, - func=lambda spark: spark.read.parquet(data_path).select(f.lit(None).cast(data_type))) + assert_gpu_project_without_ast( + func=lambda spark: spark.read.parquet(data_path).select(f.lit(None).cast(data_type))) @pytest.mark.parametrize('data_descr', ast_descrs, ids=idfn) def test_isnull(data_descr): @@ -258,9 +258,18 @@ def test_exp(data_descr): def test_expm1(data_descr): assert_unary_ast(data_descr, lambda df: df.selectExpr('expm1(a)')) +@pytest.mark.parametrize('data_gen', [float_gen, double_gen], ids=idfn) +def test_folded_null_literal_stays_on_regular_project(data_gen): + assert_gpu_project_without_ast( + lambda spark: binary_op_df(spark, data_gen).select( + f.col('a') == f.lit(None).cast(data_gen.data_type), + f.col('a') == f.col('b'))) + +# Keep null scalars from folding unsupported comparisons into AST-compatible null literals. @pytest.mark.parametrize('data_descr', ast_comparable_descrs, ids=idfn) def test_eq(data_descr): - (s1, s2) = with_cpu_session(lambda spark: gen_scalars(data_descr[0], 2)) + (s1, s2) = with_cpu_session( + lambda spark: gen_scalars(data_descr[0], 2, force_no_nulls=True)) assert_binary_ast(data_descr, lambda df: df.select( f.col('a') == s1, @@ -269,7 +278,8 @@ def test_eq(data_descr): @pytest.mark.parametrize('data_descr', ast_comparable_descrs, ids=idfn) def test_ne(data_descr): - (s1, s2) = with_cpu_session(lambda spark: gen_scalars(data_descr[0], 2)) + (s1, s2) = with_cpu_session( + lambda spark: gen_scalars(data_descr[0], 2, force_no_nulls=True)) assert_binary_ast(data_descr, lambda df: df.select( f.col('a') != s1, @@ -278,7 +288,8 @@ def test_ne(data_descr): @pytest.mark.parametrize('data_descr', ast_comparable_descrs, ids=idfn) def test_lt(data_descr): - (s1, s2) = with_cpu_session(lambda spark: gen_scalars(data_descr[0], 2)) + (s1, s2) = with_cpu_session( + lambda spark: gen_scalars(data_descr[0], 2, force_no_nulls=True)) assert_binary_ast(data_descr, lambda df: df.select( f.col('a') < s1, @@ -287,7 +298,8 @@ def test_lt(data_descr): @pytest.mark.parametrize('data_descr', ast_comparable_descrs, ids=idfn) def test_lte(data_descr): - (s1, s2) = with_cpu_session(lambda spark: gen_scalars(data_descr[0], 2)) + (s1, s2) = with_cpu_session( + lambda spark: gen_scalars(data_descr[0], 2, force_no_nulls=True)) assert_binary_ast(data_descr, lambda df: df.select( f.col('a') <= s1, @@ -296,7 +308,8 @@ def test_lte(data_descr): @pytest.mark.parametrize('data_descr', ast_comparable_descrs, ids=idfn) def test_gt(data_descr): - (s1, s2) = with_cpu_session(lambda spark: gen_scalars(data_descr[0], 2)) + (s1, s2) = with_cpu_session( + lambda spark: gen_scalars(data_descr[0], 2, force_no_nulls=True)) assert_binary_ast(data_descr, lambda df: df.select( f.col('a') > s1, @@ -305,7 +318,8 @@ def test_gt(data_descr): @pytest.mark.parametrize('data_descr', ast_comparable_descrs, ids=idfn) def test_gte(data_descr): - (s1, s2) = with_cpu_session(lambda spark: gen_scalars(data_descr[0], 2)) + (s1, s2) = with_cpu_session( + lambda spark: gen_scalars(data_descr[0], 2, force_no_nulls=True)) assert_binary_ast(data_descr, lambda df: df.select( f.col('a') >= s1, diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala index 3931867fe05..417d2d69553 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala @@ -1350,10 +1350,14 @@ abstract class BaseExprMeta[INPUT <: Expression]( } protected def willWorkInAstInfo: String = { - if (cannotBeAstReasons.isEmpty) { - "will run in AST" - } else { + if (!canThisBeReplaced) { + "cannot be converted to GPU AST because it cannot run on GPU" + } else if (willUseGpuCpuBridge) { + "cannot be converted to GPU AST because it uses the CPU bridge" + } else if (cannotBeAstReasons.nonEmpty) { s"cannot be converted to GPU AST because ${cannotBeAstReasons.mkString(";")}" + } else { + "is AST-compatible" } } @@ -1364,7 +1368,8 @@ abstract class BaseExprMeta[INPUT <: Expression]( * @param all should all the data be printed or just what does not work in the AST? */ protected def printAst(strBuilder: StringBuilder, depth: Int, all: Boolean): Unit = { - if (all || !canThisBeAst) { + val selfAstCompatible = canSelfBeAst + if (all || !selfAstCompatible) { indent(strBuilder, depth) strBuilder.append(operationName) .append(" <") 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 57e73d062e0..8ae2edb8bcc 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 @@ -62,7 +62,9 @@ class GpuProjectExecMeta( val projectList = if (conf.isProjectAstEnabled) { val astExprs = childExprs.zip(gpuExprs).map { case (meta, expr) => // cuDF requires return column is fixed width - if (GpuBatchUtils.isFixedWidth(expr.dataType) && meta.canThisBeAst) { + // Top-level literals are cheaper on the regular projection path. + if (GpuBatchUtils.isFixedWidth(expr.dataType) && meta.canThisBeAst && + GpuExpressionsUtils.extractGpuLit(expr).isEmpty) { GpuProjectAstExpression.wrap(expr) } else { expr @@ -75,6 +77,9 @@ class GpuProjectExecMeta( 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 GpuExpressionsUtils.extractGpuLit(expr).isDefined => + s" $expr will use the regular GPU projection because top-level literals " + + "are cheaper there\n" }).mkString if (explain.nonEmpty) { logWarning(s"AST PROJECT\n$explain") diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/AstUtilSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/AstUtilSuite.scala index c4f71df7906..ed41ca8f3fa 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/AstUtilSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/AstUtilSuite.scala @@ -18,13 +18,48 @@ package com.nvidia.spark.rapids import org.mockito.Mockito.{mock, when} -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Expression} +import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, AttributeSet, EqualTo, + Expression, Literal} import org.apache.spark.sql.rapids.{GpuAnd, GpuGreaterThan, GpuLength, GpuLessThan, GpuStringTrim} -import org.apache.spark.sql.types.{BooleanType, DataType, IntegerType, LongType, StringType} +import org.apache.spark.sql.types.{BooleanType, DataType, FloatType, IntegerType, LongType, + StringType} class AstUtilSuite extends GpuUnitTests { + private def floatComparisonAliasMeta(): BaseExprMeta[_] = { + val attr = AttributeReference("a", FloatType, nullable = false)() + val expr = Alias(EqualTo(attr, Literal(1.0f)), "result")() + val meta = GpuOverrides.wrapExpr( + expr, new RapidsConf(Map.empty[String, String]), None) + meta.tagForGpu() + meta + } + + test("explainAst only prints node-local AST blockers") { + val meta = floatComparisonAliasMeta() + + assert(meta.canSelfBeAst) + assert(!meta.canThisBeAst) + assert(!meta.childExprs.head.canSelfBeAst) + + val explain = meta.explainAst(all = false) + assert(!explain.contains(""), explain) + assert(explain.contains(""), explain) + assert(explain.contains("cannot be converted to GPU AST"), explain) + } + + test("explainAst all reports node-local AST compatibility") { + val meta = floatComparisonAliasMeta() + val explain = meta.explainAst(all = true) + val lines = explain.split("\n") + + assert(lines.find(_.contains("")) + .exists(_.contains("is AST-compatible")), explain) + assert(lines.find(_.contains("")) + .exists(_.contains("cannot be converted to GPU AST")), explain) + } + private[this] def testSingleNode(containsNonAstAble: Boolean, crossMultiChildPlan: Boolean) : Boolean = { val l1 = AttributeReference("l1", StringType)() From 8fc3fb5bce289bcf9f0d2a7f275529feb7b81e56 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Thu, 30 Jul 2026 16:36:32 +0800 Subject: [PATCH 09/20] address comments Signed-off-by: Haoyang Li --- integration_tests/src/main/python/ast_test.py | 7 ++-- .../com/nvidia/spark/rapids/RapidsMeta.scala | 33 +++++++++---------- .../spark/rapids/basicPhysicalOperators.scala | 13 +++++--- 3 files changed, 27 insertions(+), 26 deletions(-) diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index 3b1174d793c..a980c3811dc 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/integration_tests/src/main/python/ast_test.py @@ -98,8 +98,8 @@ def test_literal(spark_tmp_path, data_gen): data_path = spark_tmp_path + '/AST_TEST_DATA' with_cpu_session(lambda spark: gen_df(spark, [("a", IntegerGen())]).write.parquet(data_path)) scalar = with_cpu_session(lambda spark: gen_scalar(data_gen, force_no_nulls=True)) - assert_gpu_project_without_ast( - func=lambda spark: spark.read.parquet(data_path).select(scalar)) + assert_gpu_ast(is_supported=True, + func=lambda spark: spark.read.parquet(data_path).select(scalar)) @pytest.mark.parametrize('data_gen', [boolean_gen, byte_gen, short_gen, int_gen, long_gen, float_gen, double_gen, timestamp_gen, date_gen], ids=idfn) def test_null_literal(spark_tmp_path, data_gen): @@ -265,7 +265,8 @@ def test_folded_null_literal_stays_on_regular_project(data_gen): f.col('a') == f.lit(None).cast(data_gen.data_type), f.col('a') == f.col('b'))) -# Keep null scalars from folding unsupported comparisons into AST-compatible null literals. +# Use non-null scalars here because NullPropagation otherwise folds the comparisons into null +# literals, bypassing the comparison AST compatibility these tests are intended to verify. @pytest.mark.parametrize('data_descr', ast_comparable_descrs, ids=idfn) def test_eq(data_descr): (s1, s2) = with_cpu_session( diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala index 417d2d69553..98fc0aaf02a 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala @@ -1291,13 +1291,7 @@ abstract class BaseExprMeta[INPUT <: Expression]( */ final def mustBeAstExpression: Boolean = mustBeAst - final def canThisBeAst: Boolean = { - tagForAst() - // An expression cannot be AST if it cannot be replaced (disabled), uses CPU bridge, - // or has AST-specific issues - canThisBeReplaced && !willUseGpuCpuBridge && - childExprs.forall(_.canThisBeAst) && cannotBeAstReasons.isEmpty - } + final def canThisBeAst: Boolean = canSelfBeAst && childExprs.forall(_.canThisBeAst) /** * Check whether this node itself can be converted to AST. It will not recursively check its @@ -1307,8 +1301,8 @@ abstract class BaseExprMeta[INPUT <: Expression]( // undoBridgeOptimization() after a first read, so caching would return a stale answer. final def canSelfBeAst: Boolean = { tagForAst() - // Not AST-able if disabled, bridged (a GpuCpuBridgeExpression has no AST form), or it has - // AST-specific issues. + // An expression cannot be AST if it cannot be replaced (disabled), uses CPU bridge + // (a GpuCpuBridgeExpression has no AST form), or has AST-specific issues. canThisBeReplaced && !willUseGpuCpuBridge && cannotBeAstReasons.isEmpty } @@ -1350,14 +1344,18 @@ abstract class BaseExprMeta[INPUT <: Expression]( } protected def willWorkInAstInfo: String = { - if (!canThisBeReplaced) { - "cannot be converted to GPU AST because it cannot run on GPU" - } else if (willUseGpuCpuBridge) { - "cannot be converted to GPU AST because it uses the CPU bridge" - } else if (cannotBeAstReasons.nonEmpty) { - s"cannot be converted to GPU AST because ${cannotBeAstReasons.mkString(";")}" - } else { + if (canSelfBeAst) { "is AST-compatible" + } else { + val reason = if (!canThisBeReplaced) { + "it cannot run on GPU" + } else if (willUseGpuCpuBridge) { + "it uses the CPU bridge" + } else { + assert(cannotBeAstReasons.nonEmpty) + cannotBeAstReasons.mkString(";") + } + s"cannot be converted to GPU AST because $reason" } } @@ -1368,8 +1366,7 @@ abstract class BaseExprMeta[INPUT <: Expression]( * @param all should all the data be printed or just what does not work in the AST? */ protected def printAst(strBuilder: StringBuilder, depth: Int, all: Boolean): Unit = { - val selfAstCompatible = canSelfBeAst - if (all || !selfAstCompatible) { + if (all || !canSelfBeAst) { indent(strBuilder, depth) strBuilder.append(operationName) .append(" <") 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 8ae2edb8bcc..9f9c0ffdc17 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 @@ -55,6 +55,9 @@ class GpuProjectExecMeta( p: Option[RapidsMeta[_, _, _]], r: DataFromReplacementRule) extends SparkPlanMeta[ProjectExec](proj, conf, p, r) with Logging { + private def isTopLevelNullLiteral(expr: Expression): Boolean = + GpuExpressionsUtils.extractGpuLit(expr).exists(_.value == null) + override def convertToGpu(): GpuExec = { // Force list to avoid recursive Java serialization of lazy list Seq implementation val gpuExprs = childExprs.map(_.convertToGpu().asInstanceOf[NamedExpression]).toList @@ -62,9 +65,9 @@ class GpuProjectExecMeta( val projectList = if (conf.isProjectAstEnabled) { val astExprs = childExprs.zip(gpuExprs).map { case (meta, expr) => // cuDF requires return column is fixed width - // Top-level literals are cheaper on the regular projection path. + // Regular projection can reuse its cached null vector across outputs. if (GpuBatchUtils.isFixedWidth(expr.dataType) && meta.canThisBeAst && - GpuExpressionsUtils.extractGpuLit(expr).isEmpty) { + !isTopLevelNullLiteral(expr)) { GpuProjectAstExpression.wrap(expr) } else { expr @@ -77,9 +80,9 @@ class GpuProjectExecMeta( 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 GpuExpressionsUtils.extractGpuLit(expr).isDefined => - s" $expr will use the regular GPU projection because top-level literals " + - "are cheaper there\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") From ce448ef431d4510e02da145851504cfc6da9d2eb Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Fri, 31 Jul 2026 13:39:18 +0800 Subject: [PATCH 10/20] Document AST compatibility reason invariant --- .../src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala | 1 + 1 file changed, 1 insertion(+) diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala index 98fc0aaf02a..4a311f25143 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsMeta.scala @@ -1347,6 +1347,7 @@ abstract class BaseExprMeta[INPUT <: Expression]( if (canSelfBeAst) { "is AST-compatible" } else { + // These reasons must enumerate exactly the conditions checked by canSelfBeAst. val reason = if (!canThisBeReplaced) { "it cannot run on GPU" } else if (willUseGpuCpuBridge) { From 3cc28a889dc2f6ea70f5dc14e19c8408920b5a88 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Mon, 3 Aug 2026 01:03:01 +0800 Subject: [PATCH 11/20] fix scala 2.13 ci Signed-off-by: Haoyang Li --- .../scala/com/nvidia/spark/rapids/GpuArrayHofFusionSuite.scala | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 9d75f80d009..b6aaac3def3 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuArrayHofFusionSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuArrayHofFusionSuite.scala @@ -225,7 +225,8 @@ class GpuArrayHofFusionSuite extends GpuUnitTests { failOnError = false)())) val secondAst = spy(GpuProjectAstExpression(GpuMultiply( GpuBoundReference(1, LongType, nullable = true)(ExprId(402), "a"), - GpuBoundReference(2, LongType, nullable = true)(ExprId(403), "b"))())) + 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"), From fabcb01e96636cec32d5e6b76defb79689e60f7d Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Wed, 5 Aug 2026 13:24:51 +0800 Subject: [PATCH 12/20] Make Project AST JIT follow expression tiers Signed-off-by: Haoyang Li --- integration_tests/src/main/python/ast_test.py | 28 +++- .../spark/rapids/GpuAstJitExpression.scala | 21 +-- .../spark/rapids/GpuBoundAttribute.scala | 42 ++++- .../rapids/GpuProjectAstExpression.scala | 64 +++++--- .../com/nvidia/spark/rapids/RapidsConf.scala | 2 +- .../spark/rapids/basicPhysicalOperators.scala | 4 +- .../GpuBroadcastHashJoinExecBase.scala | 2 +- .../GpuBroadcastNestedLoopJoinExecBase.scala | 2 +- .../spark/rapids/GpuProjectAstJitSuite.scala | 152 +++++++++++++++++- 9 files changed, 265 insertions(+), 52 deletions(-) diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index 6de92a80f86..eb1adf9efb7 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/integration_tests/src/main/python/ast_test.py @@ -405,11 +405,28 @@ def test_jit_add_multiply(data_gen): @pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn) @disable_ansi_mode -def test_jit_mixed_nested_subexpressions(data_gen): +def test_jit_does_not_split_unique_unsupported_expression(data_gen): assert_cpu_and_gpu_are_equal_collect_with_capture( lambda spark: binary_op_df(spark, data_gen).select( (f.col('a') + f.col('b')) - (f.col('a') * f.col('b'))), - exist_classes=r"GpuProject.*AST_JIT.*- AST_JIT", + exist_classes="GpuProject", + non_exist_classes="AST_JIT,GpuProjectAst", + conf=_project_ast_jit_enabled_conf) + +@pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn) +@disable_ansi_mode +def test_jit_cse_shared_subexpression(data_gen): + def project_shared_expression(spark): + df = binary_op_df(spark, data_gen) + shared = f.col('a') + f.col('b') + return df.select( + shared.alias('shared'), + (shared - f.col('a')).alias('left'), + (shared - f.col('b')).alias('right')) + + assert_cpu_and_gpu_are_equal_collect_with_capture( + project_shared_expression, + exist_classes=r"GpuProject.*AST_JIT.*AS shared.*AS left.*AS right", non_exist_classes="GpuProjectAst", conf=_project_ast_jit_enabled_conf) @@ -421,8 +438,8 @@ def test_jit_mixed_project_expressions(data_gen): (f.col('a') + f.col('b')).alias('jit'), (f.col('a') - f.col('b')).alias('gpu'), ((f.col('a') * f.col('b')) - f.col('a')).alias('mixed')), - exist_classes=r"GpuProject.*AST_JIT.*AS jit.*AS gpu.*AST_JIT.*AS mixed", - non_exist_classes="GpuProjectAst", + exist_classes=r"GpuProject.*AST_JIT.*AS jit.*AS gpu.*AS mixed", + non_exist_classes=r"GpuProjectAst,AS gpu.*AST_JIT", conf=_project_ast_jit_enabled_conf) @pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn) @@ -436,7 +453,8 @@ def test_jit_and_legacy_ast_mixed_project_expressions(data_gen): (f.col('a') * f.col('b'))).alias('mixed')), exist_classes=( r"GpuProject.*AST_JIT.*AS jit.*AST\(.*AS legacy.*" - r"AST_JIT.*- AST_JIT.*AS mixed,GpuProjectAstExpression"), + r"AST\(.*AS mixed,GpuProjectAstExpression"), + non_exist_classes=r"AS legacy.*AST_JIT", conf=_project_ast_jit_and_legacy_enabled_conf) 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 index 38d9d50916d..341e5d0496e 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala @@ -30,21 +30,22 @@ import org.apache.spark.sql.types.DataType import org.apache.spark.sql.vectorized.ColumnarBatch object GpuAstJitExpression { - private def wrapMaximalSubtrees(expression: Expression): Expression = expression match { - case gpuExpression: GpuExpression - if gpuExpression.supportsAstJit && gpuExpression.containsAstJitOperator => - GpuAstJitExpression(gpuExpression) - case gpuExpression: GpuExpression => - gpuExpression.mapChildren { - case child: GpuExpression => wrapMaximalSubtrees(child) - case child => child - } + private def canUseAstJit(expression: GpuExpression): Boolean = + GpuBatchUtils.isFixedWidth(expression.dataType) && + expression.supportsAstJit && expression.containsAstJitOperator + + private[rapids] def wrapTierExpression(expression: Expression): Expression = expression match { + case alias @ GpuAlias(astExpression: GpuProjectAstExpression, _) + if canUseAstJit(astExpression.child) => + GpuProjectAstExpression.replaceChild(alias, GpuAstJitExpression(astExpression.child)) + case alias @ GpuAlias(child: GpuExpression, _) if canUseAstJit(child) => + GpuProjectAstExpression.replaceChild(alias, GpuAstJitExpression(child)) case other => other } private[rapids] def wrapProjectExpressions( expressions: List[NamedExpression]): List[NamedExpression] = { - expressions.map(wrapMaximalSubtrees(_).asInstanceOf[NamedExpression]) + expressions.map(wrapTierExpression(_).asInstanceOf[NamedExpression]) } private[rapids] def contains(expression: Expression): Boolean = 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 070a7a1a9b5..2d87a74e590 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 @@ -129,13 +129,15 @@ object GpuBindReferences extends Logging { * from SparkPlan nodes. Use the public API that requires metrics instead, except * when absolutely needed. */ - def bindGpuReferencesTieredNoMetrics[A <: Expression]( + private def bindGpuReferencesTieredNoMetricsInternal[A <: Expression]( expressions: Seq[A], input: AttributeSeq, - conf: SQLConf): GpuTieredProject = { + conf: SQLConf, + enableProjectAstJit: Boolean): GpuTieredProject = { if (RapidsConf.ENABLE_TIERED_PROJECT.get(conf)) { - val exprTiers = GpuProjectAstExpression.buildExprTiers(expressions, conf) + val exprTiers = GpuProjectAstExpression.buildExprTiers( + expressions, conf, enableProjectAstJit) val inputTiers = GpuEquivalentExpressions.getInputTiers(exprTiers, input) // Update ExprTiers to include the columns that are pass through and drop unneeded columns val newExprTiers = exprTiers.zipWithIndex.map { @@ -174,10 +176,32 @@ 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))) } } + def bindGpuReferencesTieredNoMetrics[A <: Expression]( + expressions: Seq[A], + input: AttributeSeq, + conf: SQLConf): GpuTieredProject = { + bindGpuReferencesTieredNoMetricsInternal( + expressions, input, conf, enableProjectAstJit = false) + } + + private[rapids] def bindGpuProjectReferencesTieredNoMetrics[A <: Expression]( + expressions: Seq[A], + input: AttributeSeq, + conf: SQLConf): GpuTieredProject = { + bindGpuReferencesTieredNoMetricsInternal( + expressions, input, conf, RapidsConf.ENABLE_PROJECT_AST_JIT.get(conf)) + } + // ========== Public "Front Door" APIs (for use by SparkPlan nodes) ========== // These methods require metrics and inject them after binding @@ -257,6 +281,16 @@ object GpuBindReferences extends Logging { bound.injectMetrics(metrics) bound } + + 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) 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 8202dd32ab1..85ccef6789a 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 @@ -50,7 +50,7 @@ object GpuProjectAstExpressionBase { } object GpuProjectAstExpression { - private def replaceChild(alias: GpuAlias, child: Expression): GpuAlias = { + private[rapids] def replaceChild(alias: GpuAlias, child: Expression): GpuAlias = { if (child eq alias.child) { alias } else { @@ -78,9 +78,11 @@ object GpuProjectAstExpression { } } - private def unwrap(expression: Expression): Expression = expression match { - case alias: GpuAlias => replaceChild(alias, unwrap(alias.child)) - case astExpression: GpuProjectAstExpression => astExpression.child + private def unwrap(expression: Expression, unwrapJit: Boolean): Expression = expression match { + case alias: GpuAlias => replaceChild(alias, unwrap(alias.child, unwrapJit)) + case astExpression: GpuProjectAstExpression => unwrap(astExpression.child, unwrapJit) + case jitExpression: GpuAstJitExpression if unwrapJit => + unwrap(jitExpression.child, unwrapJit) case other => other } @@ -89,64 +91,82 @@ object GpuProjectAstExpression { case other => other } - private def rewrapAstTiers( + private def rewrapBackendTiers( tiers: Seq[Seq[Expression]], - astOutputs: Seq[Boolean]): Seq[Seq[Expression]] = { + backendOutputs: Seq[Boolean], + wrapExpression: Expression => Expression): Seq[Seq[Expression]] = { val finalTier = tiers.last - require(finalTier.size == astOutputs.size, + require(finalTier.size == backendOutputs.size, "The final expression tier must preserve the project output count") - val astReferences = finalTier.iterator.zip(astOutputs.iterator) + val backendReferences = finalTier.iterator.zip(backendOutputs.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. + // Tier aliases are the dataflow graph after CSE, so follow them backwards from the outputs. val (commonTiers, _) = tiers.dropRight(1).foldRight( - (List.empty[Seq[Expression]], astReferences)) { + (List.empty[Seq[Expression]], backendReferences)) { case (tier, (rewrittenTiers, requiredExprIds)) => - val astAliases = tier.collect { + val backendAliases = tier.collect { case alias: GpuAlias if requiredExprIds.contains(alias.exprId) => alias } - val astAliasIds = astAliases.iterator.map(_.exprId).toSet - val dependencies = astAliases.iterator + val backendAliasIds = backendAliases.iterator.map(_.exprId).toSet + val dependencies = backendAliases.iterator .flatMap(_.references.iterator) .map(_.exprId) .toSet val rewrittenTier = tier.map { case alias: GpuAlias - if astAliasIds.contains(alias.exprId) && + if backendAliasIds.contains(alias.exprId) && GpuBatchUtils.isFixedWidth(alias.dataType) => - rewrap(alias) + wrapExpression(alias) case expression => expression } (rewrittenTier :: rewrittenTiers, requiredExprIds ++ dependencies) } - commonTiers :+ finalTier.zip(astOutputs).map { - case (expression, true) => rewrap(expression) + commonTiers :+ finalTier.zip(backendOutputs).map { + case (expression, true) => wrapExpression(expression) case (expression, false) => expression } } private[rapids] def buildExprTiers( expressions: Seq[Expression], - conf: SQLConf): Seq[Seq[Expression]] = { + conf: SQLConf, + enableProjectAstJit: Boolean = false): Seq[Seq[Expression]] = { val astOutputs = expressions.map(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 jitOutputs = expressions.map { expression => + GpuProjectAstExpressionBase.extractTopLevel(expression) + .exists(_.isInstanceOf[GpuAstJitExpression]) + } + val hasJitOutputs = jitOutputs.contains(true) + // CSE must see through backend markers so all outputs can share the same tiers. + val unwrapped = if (hasAstOutputs || hasJitOutputs) { + expressions.map(unwrap(_, unwrapJit = true)) + } else { + expressions + } val replaced = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) { GpuEquivalentExpressions.replaceMultiExpressions(unwrapped, conf) } else { unwrapped } val tiers = GpuEquivalentExpressions.getExprTiers(replaced) - if (hasAstOutputs) { - rewrapAstTiers(tiers, astOutputs) + val astTiers = if (hasAstOutputs) { + rewrapBackendTiers(tiers, astOutputs, rewrap) } else { tiers } + if (enableProjectAstJit) { + astTiers.map(_.map(GpuAstJitExpression.wrapTierExpression)) + } else if (hasJitOutputs) { + rewrapBackendTiers(astTiers, jitOutputs, GpuAstJitExpression.wrapTierExpression) + } else { + astTiers + } } private[rapids] def tableFromBatch(batch: ColumnarBatch): Table = { 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 63d0b8e1d51..36d1c91aa19 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 @@ -1250,7 +1250,7 @@ val GPU_COREDUMP_PIPE_PATTERN = conf("spark.rapids.gpu.coreDump.pipePattern") 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 " + - "subexpressions.") + "expressions within each projection tier.") .internal() .booleanConf .createWithDefault(false) 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 c169d98c368..d8d4e65fcb2 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 @@ -908,8 +908,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() 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..aed990542de 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,7 +184,7 @@ 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) { 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..26d251da8c2 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 @@ -732,7 +732,7 @@ abstract class GpuBroadcastNestedLoopJoinExecBase( // 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) diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala index de061cac599..49b31b6fbad 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala @@ -18,8 +18,9 @@ package com.nvidia.spark.rapids import org.scalatest.funsuite.AnyFunSuite -import org.apache.spark.sql.catalyst.expressions.AttributeReference -import org.apache.spark.sql.rapids.{GpuAdd, GpuMultiply, GpuSubtract} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.rapids.{GpuAdd, GpuGreatest, GpuMultiply, GpuSubtract} import org.apache.spark.sql.types.{FloatType, IntegerType, LongType} class GpuProjectAstJitSuite extends AnyFunSuite { @@ -28,6 +29,23 @@ class GpuProjectAstJitSuite extends AnyFunSuite { private def alias(expression: GpuExpression, name: String) = GpuAlias(expression, name)() + private def jitExpressions(expressions: Seq[Expression]): Seq[GpuAstJitExpression] = { + expressions.flatMap(_.collect { + case expression: GpuAstJitExpression => expression + }) + } + + 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) } @@ -49,7 +67,7 @@ class GpuProjectAstJitSuite extends AnyFunSuite { assert(GpuProjectAstExpressionBase.extractTopLevel(wrapped.head).contains(jit)) } - test("project AST JIT wraps maximal nested subtrees independently") { + 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( @@ -61,9 +79,120 @@ class GpuProjectAstJitSuite extends AnyFunSuite { val wrapped = GpuAstJitExpression.wrapProjectExpressions(List(expression)) val subtract = wrapped.head.asInstanceOf[GpuAlias].child.asInstanceOf[GpuSubtract] - assert(GpuAstJitExpression.contains(wrapped.head)) - assert(subtract.left.asInstanceOf[GpuAstJitExpression].child.isInstanceOf[GpuAdd]) - assert(subtract.right.asInstanceOf[GpuAstJitExpression].child.isInstanceOf[GpuMultiply]) + assert(jitExpressions(wrapped).isEmpty) + assert(subtract.left.isInstanceOf[GpuAdd]) + assert(subtract.right.isInstanceOf[GpuMultiply]) + } + + test("project CSE exposes 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)() + 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()) + + assertResult(2)(tiered.exprTiers.size) + assertResult(Seq(1, 0))(tiered.exprTiers.map(jitExpressions(_).size)) + assert(jitExpressions(tiered.exprTiers.head).head.child.isInstanceOf[GpuAdd]) + assert(GpuProjectAstExpressionBase.extractTopLevel(tiered.exprTiers.last.head) + .exists(_.isInstanceOf[GpuProjectAstExpression])) + assert(GpuProjectAstExpressionBase.extractTopLevel(tiered.exprTiers.last(1)).isEmpty) + val finalTierReferences = tiered.exprTiers.last.flatMap(_.collect { + case reference: GpuBoundReference if reference.name.startsWith("tiered_input_") => reference + }) + assertResult(2)(finalTierReferences.size) + assertResult(1)(finalTierReferences.map(_.exprId).distinct.size) + } + + 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(GpuProjectAstExpressionBase.extractTopLevel(outputs.head) + .exists(_.isInstanceOf[GpuAstJitExpression])) + assert(GpuProjectAstExpressionBase.extractTopLevel(outputs(1)) + .exists(_.isInstanceOf[GpuProjectAstExpression])) + } + + 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(jitExpressions(generic.exprTiers.flatten).isEmpty) + assertResult(1)(jitExpressions(project.exprTiers.flatten).size) + } + + test("the generic binder preserves only the preselected JIT dataflow") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val shared = GpuAdd(left, right, failOnError = false)() + val preselected = GpuAstJitExpression.wrapProjectExpressions(List( + alias(shared, "jit"))).head + val firstUse = alias(GpuSubtract(shared, left, failOnError = false)(), "first_use") + val secondUse = alias(GpuSubtract(shared, right, failOnError = false)(), "second_use") + val unrelated = alias(GpuMultiply(left, right, failOnError = false)(), "unrelated") + + val generic = GpuBindReferences.bindGpuReferencesTieredNoMetrics( + Seq(preselected, firstUse, secondUse, unrelated), Seq(left, right), projectConf()) + + assertResult(2)(generic.exprTiers.size) + withClue(generic.exprTiers.mkString("\n")) { + assertResult(Seq(1, 0))(generic.exprTiers.map(jitExpressions(_).size)) + } + assert(jitExpressions(generic.exprTiers.head).head.child.isInstanceOf[GpuAdd]) + assert(generic.exprTiers.last.forall( + GpuProjectAstExpressionBase.extractTopLevel(_).isEmpty)) + val sharedReferences = generic.exprTiers.last.flatMap(_.collect { + case reference: GpuBoundReference if reference.name.startsWith("tiered_input_") => + reference + }) + assertResult(3)(sharedReferences.size) + assertResult(1)(sharedReferences.map(_.exprId).distinct.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(jitExpressions(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 expression = alias(GpuAdd(left, right, failOnError = false)(), "result") + + val project = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( + Seq(expression), Seq(left, right), projectConf(tiered = false)) + + assertResult(1)(project.exprTiers.size) + assertResult(1)(jitExpressions(project.exprTiers.head).size) } test("project AST JIT excludes ANSI and floating point arithmetic") { @@ -79,4 +208,15 @@ class GpuProjectAstJitSuite extends AnyFunSuite { 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(jitExpressions(wrapped).isEmpty) + } } From 93b73a867f2efef3de0098f6e7e3f08dd56af0ec Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Wed, 5 Aug 2026 19:23:53 +0800 Subject: [PATCH 13/20] address local comments Signed-off-by: Haoyang Li --- .../GpuBroadcastNestedLoopJoinExecBase.scala | 6 +- ...GpuBroadcastNestedLoopJoinRetrySuite.scala | 152 ++++++++++++++++++ 2 files changed, 156 insertions(+), 2 deletions(-) create mode 100644 tests/src/test/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinRetrySuite.scala 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 26d251da8c2..807ad3ad92c 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,7 +726,7 @@ 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 @@ -735,7 +735,9 @@ abstract class GpuBroadcastNestedLoopJoinExecBase( val proj = GpuBindReferences.bindGpuProjectReferencesTiered( postBuildCondition, p.child.output, conf, allMetrics) val fn = (batch: ColumnarBatch) => { - withResource(batch)(proj.project) + val spillableBatch = SpillableColumnarBatch( + batch, SpillPriorities.ACTIVE_ON_DECK_PRIORITY) + proj.projectAndCloseWithRetrySingleBatch(spillableBatch) } Some(fn) case _ => 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..c894ad47e61 --- /dev/null +++ b/tests/src/test/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinRetrySuite.scala @@ -0,0 +1,152 @@ +/* + * 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.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.internal.SQLConf +import org.apache.spark.sql.rapids.{GpuAdd, GpuMultiply, GpuSubtract} +import org.apache.spark.sql.rapids.metrics.source.MockTaskContext +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.vectorized.ColumnarBatch + +class GpuBroadcastNestedLoopJoinRetrySuite extends RmmSparkRetrySuiteBase { + 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 */ 1) + RmmSpark.getAndResetNumSplitRetryThrow(/* taskId */ 1) + 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 conf = new SQLConf() + conf.setConfString(RapidsConf.ENABLE_TIERED_PROJECT.key, "true") + conf.setConfString(RapidsConf.ENABLE_COMBINED_EXPRESSIONS.key, "true") + conf.setConfString(RapidsConf.ENABLE_PROJECT_AST_JIT.key, "true") + conf.setConfString(RapidsConf.PROJECT_SPLIT_RETRY_ENABLED.key, "true") + val expressions = postBuildExpressions + val buildProject = GpuProjectExec(expressions, TestLeafExec(buildAttributes)) + val join = TestBroadcastNestedLoopJoin( + TestLeafExec(Seq.empty), buildProject, expressions) + val taskContext = new MockTaskContext(taskAttemptId = 1, partitionId = 0) + val spark = SparkSession.builder() + .master("local[1]") + .appName("GpuBroadcastNestedLoopJoinRetrySuite") + .getOrCreate() + + TrampolineUtil.setTaskContext(taskContext) + try { + SQLConf.withExistingConf(conf) { + val boundProject = GpuBindReferences.bindGpuProjectReferencesTiered( + expressions, buildAttributes, conf, Map.empty) + val jitExpressions = boundProject.exprTiers.flatten.flatMap(_.collect { + case expression: GpuAstJitExpression => expression + }) + assertResult(2)(boundProject.exprTiers.size) + assertResult(1)(jitExpressions.size) + assert(jitExpressions.head.child.isInstanceOf[GpuMultiply]) + assert(jitExpressions.head.child.find(_.isInstanceOf[GpuAdd]).isDefined) + + val projectBuildSide = join.buildSidePostProjection.get + // Compile before arming the OOM so the retry is exercised by computeColumnJit. + withResource(projectBuildSide(buildBatch())) { _ => } + + val retryInput = buildBatch() + RmmSpark.getAndResetNumRetryThrow(/* taskId */ 1) + RmmSpark.getAndResetNumSplitRetryThrow(/* taskId */ 1) + RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 1, + RmmSpark.OomInjectionType.GPU.ordinal, 0) + withResource(projectBuildSide(retryInput)) { output => + assertResult(Seq(9, 19, 33))(collectInts(output, 0)) + assertResult(Seq(6, 16, 30))(collectInts(output, 1)) + } + assert(RmmSpark.getAndResetNumRetryThrow(/* taskId */ 1) > 0) + assertResult(0)(RmmSpark.getAndResetNumSplitRetryThrow(/* taskId */ 1)) + } + } finally { + try { + taskContext.markTaskComplete() + } finally { + TrampolineUtil.unsetTaskContext() + ScalableTaskCompletion.reset() + spark.stop() + } + } + } +} From 116528e1063a23ee2c2c3a43b6007512b789c169 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Wed, 5 Aug 2026 20:26:18 +0800 Subject: [PATCH 14/20] Fix AST JIT cleanup registration failure --- .../spark/rapids/GpuAstJitExpression.scala | 27 ++++++-- .../spark/rapids/GpuProjectAstJitSuite.scala | 67 +++++++++++++++++++ 2 files changed, 89 insertions(+), 5 deletions(-) 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 index 341e5d0496e..954d07bacf2 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala @@ -96,17 +96,34 @@ case class GpuAstJitExpression(child: GpuExpression) } if (!completionRegistered) { Option(TaskContext.get()).foreach { taskContext => - onTaskCompletion(taskContext) { - close() - } completionRegistered = true + try { + onTaskCompletion(taskContext) { + closeAtTaskCompletion() + } + if (!completionRegistered) { + throw new IllegalStateException( + "Task completed while registering the AST JIT cleanup callback") + } + } catch { + case t: Throwable => + completionRegistered = false + closeCompiledExpression(t) + throw t + } } } compiledExpression } - private def closeCompiledExpression(): Unit = synchronized { - Option(compiledExpression).foreach(_.safeClose()) + private def closeAtTaskCompletion(): Unit = synchronized { + completionRegistered = false + closeCompiledExpression() + } + + private def closeCompiledExpression(error: Throwable = null): Unit = synchronized { + val toClose = compiledExpression compiledExpression = null + Option(toClose).foreach(_.safeClose(error)) } } diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala index 49b31b6fbad..7fbb4c4a3a9 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala @@ -16,12 +16,18 @@ package com.nvidia.spark.rapids +import ai.rapids.cudf.ast.{AstExpression, CompiledExpression} +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, Expression} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.rapids.{GpuAdd, GpuGreatest, GpuMultiply, GpuSubtract} +import org.apache.spark.sql.rapids.execution.TrampolineUtil +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) = @@ -29,6 +35,24 @@ class GpuProjectAstJitSuite extends AnyFunSuite { private def alias(expression: GpuExpression, name: String) = GpuAlias(expression, name)() + private def mockJitExpression(compiled: CompiledExpression): GpuAstJitExpression = { + val child = mock(classOf[GpuExpression]) + val ast = mock(classOf[AstExpression]) + when(child.convertToAst(Int.MaxValue)).thenReturn(ast) + when(ast.compile()).thenReturn(compiled) + GpuAstJitExpression(child) + } + + private def withTaskContext[T](taskContext: TaskContext)(body: => T): T = { + TrampolineUtil.setTaskContext(taskContext) + try { + body + } finally { + TrampolineUtil.unsetTaskContext() + ScalableTaskCompletion.reset() + } + } + private def jitExpressions(expressions: Seq[Expression]): Seq[GpuAstJitExpression] = { expressions.flatMap(_.collect { case expression: GpuAstJitExpression => expression @@ -219,4 +243,47 @@ class GpuProjectAstJitSuite extends AnyFunSuite { assert(jitExpressions(wrapped).isEmpty) } + + 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) + + 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) + + 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() + } + } } From 59e26f226cd6aeaab8a80ffe0624a8e970270d13 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Mon, 10 Aug 2026 22:22:06 +0800 Subject: [PATCH 15/20] address comments Signed-off-by: Haoyang Li --- integration_tests/src/main/python/ast_test.py | 8 +- .../spark/rapids/GpuAstJitExpression.scala | 61 ++++++++------ .../spark/rapids/GpuBoundAttribute.scala | 25 +++++- .../nvidia/spark/rapids/GpuExpressions.scala | 12 +++ .../rapids/GpuProjectAstExpression.scala | 21 ++--- .../spark/rapids/basicPhysicalOperators.scala | 60 ++++++++++---- .../apache/spark/sql/rapids/arithmetic.scala | 4 - .../GpuBroadcastHashJoinExecBase.scala | 5 +- .../GpuBroadcastNestedLoopJoinExecBase.scala | 8 +- .../spark/rapids/GpuProjectAstJitSuite.scala | 81 +++++++++--------- .../com/nvidia/spark/rapids/TestUtils.scala | 19 +++++ ...GpuBroadcastNestedLoopJoinRetrySuite.scala | 82 +++++++++---------- 12 files changed, 229 insertions(+), 157 deletions(-) diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index eb1adf9efb7..88ecf187d58 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/integration_tests/src/main/python/ast_test.py @@ -431,8 +431,10 @@ def project_shared_expression(spark): 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): +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'), @@ -440,7 +442,9 @@ def test_jit_mixed_project_expressions(data_gen): ((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=_project_ast_jit_enabled_conf) + 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 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 index 954d07bacf2..52f74c02346 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala @@ -20,6 +20,7 @@ import ai.rapids.cudf.Table import ai.rapids.cudf.ast.CompiledExpression import com.nvidia.spark.Retryable import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} +import com.nvidia.spark.rapids.GpuMetric.OP_TIME_LEGACY import com.nvidia.spark.rapids.RapidsPluginImplicits._ import com.nvidia.spark.rapids.ScalableTaskCompletion.onTaskCompletion import com.nvidia.spark.rapids.shims.ShimUnaryExpression @@ -35,6 +36,7 @@ object GpuAstJitExpression { expression.supportsAstJit && expression.containsAstJitOperator private[rapids] def wrapTierExpression(expression: Expression): Expression = expression match { + case alias @ GpuAlias(_: GpuAstJitExpression, _) => alias case alias @ GpuAlias(astExpression: GpuProjectAstExpression, _) if canUseAstJit(astExpression.child) => GpuProjectAstExpression.replaceChild(alias, GpuAstJitExpression(astExpression.child)) @@ -49,32 +51,34 @@ object GpuAstJitExpression { } private[rapids] def contains(expression: Expression): Boolean = - expression.find(_.isInstanceOf[GpuAstJitExpression]).isDefined + GpuProjectAstExpressionBase.extractTopLevel(expression) + .exists(_.isInstanceOf[GpuAstJitExpression]) } case class GpuAstJitExpression(child: GpuExpression) extends ShimUnaryExpression with GpuProjectAstExpressionBase - with Retryable with AutoCloseable { + with GpuMetricsInjectable with Retryable with AutoCloseable { @transient private[this] var compiledExpression: CompiledExpression = _ @transient private[this] var completionRegistered = false + private[this] var opTime: GpuMetric = NoopMetric override def dataType: DataType = child.dataType override def nullable: Boolean = child.nullable - override def disableTieredProjectCombine: Boolean = true - override def toString: String = s"AST_JIT($child)" + override def injectMetrics(metrics: Map[String, GpuMetric]): Unit = { + opTime = metrics.getOrElse(OP_TIME_LEGACY, NoopMetric) + } + override def checkpoint(): Unit = { getCompiledExpression } - override def restore(): Unit = { - // The existing task callback closes the expression recompiled after a retry. - closeCompiledExpression() - } + // Compiled ASTs are immutable and remain valid across retry attempts. + override def restore(): Unit = () override def close(): Unit = closeCompiledExpression() @@ -84,22 +88,28 @@ case class GpuAstJitExpression(child: GpuExpression) } } - private[rapids] override def computeColumn(table: Table): GpuColumnVector = - closeOnExcept(getCompiledExpression.computeColumnJit(table)) { result => - GpuColumnVector.from(result, dataType) + private[rapids] override def computeColumn(table: Table): GpuColumnVector = { + val compiled = getCompiledExpression + NvtxIdWithMetrics(NvtxRegistry.PROJECT_AST, opTime) { + closeOnExcept(compiled.computeColumnJit(table)) { result => + GpuColumnVector.from(result, dataType) + } } + } private def getCompiledExpression: CompiledExpression = synchronized { if (compiledExpression == null) { - compiledExpression = child.convertToAst(Int.MaxValue) - .compile() + compiledExpression = NvtxIdWithMetrics(NvtxRegistry.COMPILE_ASTS, opTime) { + // Force every bound reference to the left table; Project AST has one input table. + child.convertToAst(Int.MaxValue).compile() + } } if (!completionRegistered) { Option(TaskContext.get()).foreach { taskContext => completionRegistered = true try { onTaskCompletion(taskContext) { - closeAtTaskCompletion() + clearRegistrationAndClose() } if (!completionRegistered) { throw new IllegalStateException( @@ -107,8 +117,7 @@ case class GpuAstJitExpression(child: GpuExpression) } } catch { case t: Throwable => - completionRegistered = false - closeCompiledExpression(t) + clearRegistrationAndClose(t) throw t } } @@ -116,14 +125,20 @@ case class GpuAstJitExpression(child: GpuExpression) compiledExpression } - private def closeAtTaskCompletion(): Unit = synchronized { - completionRegistered = false - closeCompiledExpression() - } + private def clearRegistrationAndClose(error: Throwable = null): Unit = + closeCompiledExpression(error, clearRegistration = true) - private def closeCompiledExpression(error: Throwable = null): Unit = synchronized { - val toClose = compiledExpression - compiledExpression = null + private def closeCompiledExpression( + error: Throwable = null, + clearRegistration: Boolean = false): Unit = { + val toClose = synchronized { + if (clearRegistration) { + completionRegistered = false + } + val current = compiledExpression + compiledExpression = null + current + } Option(toClose).foreach(_.safeClose(error)) } } 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 2d87a74e590..d180b94e2c7 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 @@ -124,10 +124,9 @@ 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 */ private def bindGpuReferencesTieredNoMetricsInternal[A <: Expression]( expressions: Seq[A], @@ -186,6 +185,12 @@ 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. + */ def bindGpuReferencesTieredNoMetrics[A <: Expression]( expressions: Seq[A], input: AttributeSeq, @@ -194,6 +199,10 @@ object GpuBindReferences extends Logging { 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, @@ -282,6 +291,14 @@ object GpuBindReferences extends Logging { 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, 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 17f87eed041..d5bb09beb46 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,16 +197,26 @@ trait GpuExpression extends Expression { def convertToAst(numFirstTableColumns: Int): ast.AstExpression = throw new IllegalStateException(s"Cannot convert ${this.getClass.getSimpleName} to AST") + /** + * Whether this node supports AST JIT for its current semantics and types, excluding its + * children. Operator overrides must validate their execution modes and local input/output types. + */ def selfSupportsAstJit: Boolean = false + /** + * Whether this node is an operation, rather than an AST-compatible leaf. Literals and references + * leave this false so they do not trigger compilation without useful work to JIT. + */ def selfIsAstJitOperator: Boolean = false + /** Whether this node and its complete expression subtree support AST JIT. */ final def supportsAstJit: Boolean = selfSupportsAstJit && children.forall { case child: GpuExpression => child.supportsAstJit case _: AttributeReference => true case _ => false } + /** Whether this expression subtree contains an operation that makes AST JIT useful. */ final def containsAstJitOperator: Boolean = selfIsAstJitOperator || children.exists { case child: GpuExpression => child.containsAstJitOperator case _ => false @@ -406,6 +416,8 @@ trait CudfBinaryExpression extends GpuBinaryExpression { def castOutputAtEnd: Boolean = false def astOperator: Option[ast.BinaryOperator] = None + override def selfIsAstJitOperator: Boolean = selfSupportsAstJit + def outputType(l: BinaryOperable, r: BinaryOperable): DType = { val over = outputTypeOverride if (over == null) { 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 85ccef6789a..e81376c6cbc 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 @@ -78,11 +78,10 @@ object GpuProjectAstExpression { } } - private def unwrap(expression: Expression, unwrapJit: Boolean): Expression = expression match { - case alias: GpuAlias => replaceChild(alias, unwrap(alias.child, unwrapJit)) - case astExpression: GpuProjectAstExpression => unwrap(astExpression.child, unwrapJit) - case jitExpression: GpuAstJitExpression if unwrapJit => - unwrap(jitExpression.child, unwrapJit) + 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 } @@ -138,14 +137,10 @@ object GpuProjectAstExpression { enableProjectAstJit: Boolean = false): Seq[Seq[Expression]] = { val astOutputs = expressions.map(extractTopLevel(_).isDefined) val hasAstOutputs = astOutputs.contains(true) - val jitOutputs = expressions.map { expression => - GpuProjectAstExpressionBase.extractTopLevel(expression) - .exists(_.isInstanceOf[GpuAstJitExpression]) - } - val hasJitOutputs = jitOutputs.contains(true) + val hasJitOutputs = expressions.exists(GpuAstJitExpression.contains) // CSE must see through backend markers so all outputs can share the same tiers. val unwrapped = if (hasAstOutputs || hasJitOutputs) { - expressions.map(unwrap(_, unwrapJit = true)) + expressions.map(unwrap) } else { expressions } @@ -161,10 +156,10 @@ object GpuProjectAstExpression { tiers } if (enableProjectAstJit) { + // Project binding selects JIT after CSE so newly exposed tiers are eligible. astTiers.map(_.map(GpuAstJitExpression.wrapTierExpression)) - } else if (hasJitOutputs) { - rewrapBackendTiers(astTiers, jitOutputs, GpuAstJitExpression.wrapTierExpression) } else { + // Only the Project-specific binder selects JIT after CSE. astTiers } } 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 d8d4e65fcb2..aad8b881939 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 @@ -49,6 +49,16 @@ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.random.BernoulliCellSampler +object GpuProjectExecMeta { + private[rapids] def explainBackend(expression: Expression): String = { + GpuProjectAstExpressionBase.extractTopLevel(expression) match { + case Some(_: GpuAstJitExpression) => "Project AST JIT" + case Some(_: GpuProjectAstExpression) => "legacy Project AST" + case _ => "the regular GPU projection" + } + } +} + class GpuProjectExecMeta( proj: ProjectExec, conf: RapidsConf, @@ -68,7 +78,7 @@ class GpuProjectExecMeta( gpuExprs } val projectList = if (conf.isProjectAstEnabled) { - val astExprs = childExprs.zip(jitExprs).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 (!GpuAstJitExpression.contains(expr) && @@ -79,25 +89,41 @@ class GpuProjectExecMeta( 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 { jitExprs } + // Explain the planned top-level backend because both AST backends are optional and can be + // difficult to distinguish from the regular projection in a transformed plan. + if (conf.shouldExplain && (conf.isProjectAstEnabled || conf.isProjectAstJitEnabled)) { + val backendExplain = projectList.iterator.map { expression => + s" $expression will use ${GpuProjectExecMeta.explainBackend(expression)}\n" + } + val legacyExplain = if (conf.isProjectAstEnabled) { + childExprs.iterator.zip(projectList.iterator).collect { + case (meta, expression) + if GpuProjectAstExpressionBase.extractTopLevel(expression).isEmpty => + meta.explainAst(conf.shouldExplainAll) + }.filter(_.nonEmpty) + } else { + Iterator.empty + } + 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 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 so null outputs can reuse " + + "the cached null vector\n" + } + val explain = (backendExplain ++ legacyExplain ++ regularExplain).mkString + if (explain.nonEmpty) { + logWarning(s"AST PROJECT\n$explain") + } + } GpuProjectExec(projectList, gpuChild) } } 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 0fcabb4537b..e67681c4f66 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 @@ -317,8 +317,6 @@ abstract class GpuAddBase extends CudfBinaryArithmetic with Serializable { override def selfSupportsAstJit: Boolean = !failOnError && (dataType == IntegerType || dataType == LongType) - override def selfIsAstJitOperator: Boolean = selfSupportsAstJit - override def hasSideEffects: Boolean = (failOnError && GpuAnsi.needBasicOpOverflowCheck(dataType)) || super.hasSideEffects @@ -776,8 +774,6 @@ case class GpuMultiply( override def selfSupportsAstJit: Boolean = !failOnError && (dataType == IntegerType || dataType == LongType) - override def selfIsAstJitOperator: Boolean = selfSupportsAstJit - 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 aed990542de..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 @@ -189,9 +189,8 @@ abstract class GpuBroadcastHashJoinExecBase( } 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 807ad3ad92c..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 @@ -734,12 +734,10 @@ abstract class GpuBroadcastNestedLoopJoinExecBase( // batch. val proj = GpuBindReferences.bindGpuProjectReferencesTiered( postBuildCondition, p.child.output, conf, allMetrics) - val fn = (batch: ColumnarBatch) => { - val spillableBatch = SpillableColumnarBatch( - batch, SpillPriorities.ACTIVE_ON_DECK_PRIORITY) - proj.projectAndCloseWithRetrySingleBatch(spillableBatch) + 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/GpuProjectAstJitSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala index 7fbb4c4a3a9..26d8ef4e7ee 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala @@ -24,7 +24,6 @@ import org.apache.spark.TaskContext import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.rapids.{GpuAdd, GpuGreatest, GpuMultiply, GpuSubtract} -import org.apache.spark.sql.rapids.execution.TrampolineUtil import org.apache.spark.sql.rapids.metrics.source.MockTaskContext import org.apache.spark.sql.types.{FloatType, IntegerType, LongType} import org.apache.spark.util.TaskCompletionListener @@ -43,16 +42,6 @@ class GpuProjectAstJitSuite extends AnyFunSuite { GpuAstJitExpression(child) } - private def withTaskContext[T](taskContext: TaskContext)(body: => T): T = { - TrampolineUtil.setTaskContext(taskContext) - try { - body - } finally { - TrampolineUtil.unsetTaskContext() - ScalableTaskCompletion.reset() - } - } - private def jitExpressions(expressions: Seq[Expression]): Seq[GpuAstJitExpression] = { expressions.flatMap(_.collect { case expression: GpuAstJitExpression => expression @@ -86,9 +75,11 @@ class GpuProjectAstJitSuite extends AnyFunSuite { 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(GpuProjectAstExpressionBase.extractTopLevel(wrapped.head).contains(jit)) + assert(GpuAstJitExpression.contains(wrapped.head)) + assert(wrappedAgain.head eq wrapped.head) } test("project AST JIT only wraps a fully supported top-level expression") { @@ -119,6 +110,8 @@ class GpuProjectAstJitSuite extends AnyFunSuite { alias(GpuSubtract(shared, third, failOnError = false)(), "legacy")), alias(GpuGreatest(Seq(shared, fourth)), "regular")) + // Inputs: AST((left + right) - third), greatest(left + right, fourth). + // Tier 0 computes AST_JIT(left + right); tier 1 consumes that shared result in both outputs. val tiered = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( expressions, Seq(left, right, third, fourth), projectConf()) @@ -168,34 +161,6 @@ class GpuProjectAstJitSuite extends AnyFunSuite { assertResult(1)(jitExpressions(project.exprTiers.flatten).size) } - test("the generic binder preserves only the preselected JIT dataflow") { - val left = reference(0, IntegerType) - val right = reference(1, IntegerType) - val shared = GpuAdd(left, right, failOnError = false)() - val preselected = GpuAstJitExpression.wrapProjectExpressions(List( - alias(shared, "jit"))).head - val firstUse = alias(GpuSubtract(shared, left, failOnError = false)(), "first_use") - val secondUse = alias(GpuSubtract(shared, right, failOnError = false)(), "second_use") - val unrelated = alias(GpuMultiply(left, right, failOnError = false)(), "unrelated") - - val generic = GpuBindReferences.bindGpuReferencesTieredNoMetrics( - Seq(preselected, firstUse, secondUse, unrelated), Seq(left, right), projectConf()) - - assertResult(2)(generic.exprTiers.size) - withClue(generic.exprTiers.mkString("\n")) { - assertResult(Seq(1, 0))(generic.exprTiers.map(jitExpressions(_).size)) - } - assert(jitExpressions(generic.exprTiers.head).head.child.isInstanceOf[GpuAdd]) - assert(generic.exprTiers.last.forall( - GpuProjectAstExpressionBase.extractTopLevel(_).isEmpty)) - val sharedReferences = generic.exprTiers.last.flatMap(_.collect { - case reference: GpuBoundReference if reference.name.startsWith("tiered_input_") => - reference - }) - assertResult(3)(sharedReferences.size) - assertResult(1)(sharedReferences.map(_.exprId).distinct.size) - } - test("the project binder respects a disabled JIT setting") { val left = reference(0, IntegerType) val right = reference(1, IntegerType) @@ -244,6 +209,38 @@ class GpuProjectAstJitSuite extends AnyFunSuite { assert(jitExpressions(wrapped).isEmpty) } + test("project backend explanation follows the final wrapper") { + val left = reference(0, IntegerType) + val right = reference(1, IntegerType) + val add = alias(GpuAdd(left, right, failOnError = false)(), "jit") + val subtract = alias(GpuSubtract(left, right, failOnError = false)(), "regular") + val jit = GpuAstJitExpression.wrapProjectExpressions(List(add)).head + val legacy = GpuProjectAstExpression.wrap(subtract) + + assertResult("Project AST JIT")(GpuProjectExecMeta.explainBackend(jit)) + assertResult("legacy Project AST")(GpuProjectExecMeta.explainBackend(legacy)) + assertResult("the regular GPU projection")(GpuProjectExecMeta.explainBackend(subtract)) + } + + test("project AST JIT keeps its compiled expression across retry") { + val taskContext = new MockTaskContext(taskAttemptId = 1, partitionId = 0) + val child = mock(classOf[GpuExpression]) + val ast = mock(classOf[AstExpression]) + val compiled = mock(classOf[CompiledExpression]) + when(child.convertToAst(Int.MaxValue)).thenReturn(ast) + when(ast.compile()).thenReturn(compiled) + val jit = GpuAstJitExpression(child) + + TestUtils.withTaskContext(taskContext) { + jit.checkpoint() + jit.restore() + jit.checkpoint() + verify(ast, times(1)).compile() + 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") @@ -255,7 +252,7 @@ class GpuProjectAstJitSuite extends AnyFunSuite { doThrow(closeFailure).when(compiled).close() val jit = mockJitExpression(compiled) - withTaskContext(taskContext) { + TestUtils.withTaskContext(taskContext) { val thrown = intercept[RuntimeException] { jit.checkpoint() } @@ -276,7 +273,7 @@ class GpuProjectAstJitSuite extends AnyFunSuite { val compiled = mock(classOf[CompiledExpression]) val jit = mockJitExpression(compiled) - withTaskContext(taskContext) { + TestUtils.withTaskContext(taskContext) { val thrown = intercept[IllegalStateException] { jit.checkpoint() } 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..ed88e40ff92 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.MockTaskContextBase import org.apache.spark.sql.vectorized.ColumnarBatch /** A collection of utility methods useful in tests. */ @@ -44,6 +45,24 @@ object TestUtils extends Assertions { System.getProperty("test.build.data", System.getProperty("java.io.tmpdir", "/tmp")), basename) + def withTaskContext[T]( + taskContext: MockTaskContextBase, + markTaskComplete: Boolean = false)(body: => T): T = { + TrampolineUtil.setTaskContext(taskContext) + try { + body + } finally { + try { + if (markTaskComplete) { + taskContext.markTaskComplete() + } + } finally { + TrampolineUtil.unsetTaskContext() + ScalableTaskCompletion.reset() + } + } + } + // 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/execution/GpuBroadcastNestedLoopJoinRetrySuite.scala b/tests/src/test/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinRetrySuite.scala index c894ad47e61..6ddc2b3851d 100644 --- 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 @@ -27,7 +27,6 @@ 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.internal.SQLConf import org.apache.spark.sql.rapids.{GpuAdd, GpuMultiply, GpuSubtract} import org.apache.spark.sql.rapids.metrics.source.MockTaskContext import org.apache.spark.sql.types.IntegerType @@ -95,58 +94,53 @@ class GpuBroadcastNestedLoopJoinRetrySuite extends RmmSparkRetrySuiteBase { } test("BNLJ build-side shared JIT tier retries GpuRetryOOM") { - val conf = new SQLConf() - conf.setConfString(RapidsConf.ENABLE_TIERED_PROJECT.key, "true") - conf.setConfString(RapidsConf.ENABLE_COMBINED_EXPRESSIONS.key, "true") - conf.setConfString(RapidsConf.ENABLE_PROJECT_AST_JIT.key, "true") - conf.setConfString(RapidsConf.PROJECT_SPLIT_RETRY_ENABLED.key, "true") + val spark = SparkSession.builder() + .master("local[1]") + .appName("GpuBroadcastNestedLoopJoinRetrySuite") + .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 val expressions = postBuildExpressions val buildProject = GpuProjectExec(expressions, TestLeafExec(buildAttributes)) val join = TestBroadcastNestedLoopJoin( TestLeafExec(Seq.empty), buildProject, expressions) val taskContext = new MockTaskContext(taskAttemptId = 1, partitionId = 0) - val spark = SparkSession.builder() - .master("local[1]") - .appName("GpuBroadcastNestedLoopJoinRetrySuite") - .getOrCreate() - TrampolineUtil.setTaskContext(taskContext) - try { - SQLConf.withExistingConf(conf) { - val boundProject = GpuBindReferences.bindGpuProjectReferencesTiered( - expressions, buildAttributes, conf, Map.empty) - val jitExpressions = boundProject.exprTiers.flatten.flatMap(_.collect { - case expression: GpuAstJitExpression => expression - }) - assertResult(2)(boundProject.exprTiers.size) - assertResult(1)(jitExpressions.size) - assert(jitExpressions.head.child.isInstanceOf[GpuMultiply]) - assert(jitExpressions.head.child.find(_.isInstanceOf[GpuAdd]).isDefined) + TestUtils.withTaskContext(taskContext, markTaskComplete = true) { + // Input: x=[1,2,3], y=[4,5,6], z=[2,3,4]. CSE produces a first tier that + // passes through x/y/z and computes AST_JIT((x+y)*z), followed by the five final outputs + // [shared-x, shared-y, x, y, z]. + val boundProject = GpuBindReferences.bindGpuProjectReferencesTiered( + expressions, buildAttributes, conf, Map.empty) + val jitExpressions = boundProject.exprTiers.flatten.flatMap(_.collect { + case expression: GpuAstJitExpression => expression + }) + 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 - // Compile before arming the OOM so the retry is exercised by computeColumnJit. - withResource(projectBuildSide(buildBatch())) { _ => } + val projectBuildSide = join.buildSidePostProjection.get + // Warm up computeColumnJit so first-use JIT setup cannot consume the injected OOM; the + // next call exercises retry during query execution. + withResource(projectBuildSide(buildBatch())) { _ => } - val retryInput = buildBatch() - RmmSpark.getAndResetNumRetryThrow(/* taskId */ 1) - RmmSpark.getAndResetNumSplitRetryThrow(/* taskId */ 1) - RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 1, - RmmSpark.OomInjectionType.GPU.ordinal, 0) - withResource(projectBuildSide(retryInput)) { output => - assertResult(Seq(9, 19, 33))(collectInts(output, 0)) - assertResult(Seq(6, 16, 30))(collectInts(output, 1)) - } - assert(RmmSpark.getAndResetNumRetryThrow(/* taskId */ 1) > 0) - assertResult(0)(RmmSpark.getAndResetNumSplitRetryThrow(/* taskId */ 1)) - } - } finally { - try { - taskContext.markTaskComplete() - } finally { - TrampolineUtil.unsetTaskContext() - ScalableTaskCompletion.reset() - spark.stop() + val retryInput = buildBatch() + RmmSpark.getAndResetNumRetryThrow(/* taskId */ 1) + RmmSpark.getAndResetNumSplitRetryThrow(/* taskId */ 1) + 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 */ 1) > 0) + assertResult(0)(RmmSpark.getAndResetNumSplitRetryThrow(/* taskId */ 1)) } } } From 1058db8bb78d033570c267726d8943d9690a2962 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Thu, 13 Aug 2026 01:03:37 +0800 Subject: [PATCH 16/20] address comments Signed-off-by: Haoyang Li --- docs/dev/nvtx_ranges.md | 2 + .../spark/rapids/GpuAstJitExpression.scala | 110 +++++++++++------- .../spark/rapids/GpuBoundAttribute.scala | 19 ++- .../nvidia/spark/rapids/GpuExpressions.scala | 18 ++- .../rapids/GpuProjectAstExpression.scala | 70 ++++++----- .../spark/rapids/NvtxRangeWithDoc.scala | 8 ++ .../spark/rapids/basicPhysicalOperators.scala | 48 +++----- .../apache/spark/sql/rapids/arithmetic.scala | 4 +- .../spark/rapids/GpuProjectAstJitSuite.scala | 104 +++++++++++------ .../spark/rapids/ProjectAstTestUtils.scala | 37 ++++++ .../com/nvidia/spark/rapids/TestUtils.scala | 10 +- .../spark/sql/rapids/ProjectExprSuite.scala | 45 +++---- ...GpuBroadcastNestedLoopJoinRetrySuite.scala | 32 ++--- 13 files changed, 309 insertions(+), 198 deletions(-) create mode 100644 tests/src/test/scala/com/nvidia/spark/rapids/ProjectAstTestUtils.scala diff --git a/docs/dev/nvtx_ranges.md b/docs/dev/nvtx_ranges.md index f3c31cf4324..caf31095174 100644 --- a/docs/dev/nvtx_ranges.md +++ b/docs/dev/nvtx_ranges.md @@ -38,6 +38,7 @@ DoubleBatchedWindow_PRE|Pre-processing for double-batched window operation GpuGenerate project split|Splitting projection in generate operation parquet parse filter footer|Parsing and filtering Parquet footer by range Compile ASTs|Compiling abstract syntax trees for expression evaluation +Compile AST JIT|Compiling an AST expression for JIT evaluation parquet filter blocks|Filtering Parquet row group blocks based on predicates PageableH2D|Copying from pageable host memory to device TOP N|Computing top N rows @@ -134,6 +135,7 @@ single build batch concat|Concatenating batches for single build batch sort copy boundaries|Copying boundary data for sort operation WaitingForWrites|Rapids Shuffle Manager (multi threaded) is waiting for any queued writes to finish before finalizing the map output writer Project AST|Applying AST-based projection to batch +Project AST JIT|Applying JIT-compiled AST projection to batch ORC readBatches|Reading ORC batches Shuffle Transfer Request|Handling shuffle data transfer request consumeWindow|Consuming transfer window 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 index 52f74c02346..1356f2ef32d 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala @@ -35,13 +35,16 @@ object GpuAstJitExpression { GpuBatchUtils.isFixedWidth(expression.dataType) && expression.supportsAstJit && expression.containsAstJitOperator + private def asAstJit(child: GpuExpression): Option[GpuAstJitExpression] = child match { + case jitExpression: GpuAstJitExpression => Some(jitExpression) + case astExpression: GpuProjectAstExpression => asAstJit(astExpression.child) + case other if canUseAstJit(other) => Some(GpuAstJitExpression(other)) + case _ => None + } + private[rapids] def wrapTierExpression(expression: Expression): Expression = expression match { - case alias @ GpuAlias(_: GpuAstJitExpression, _) => alias - case alias @ GpuAlias(astExpression: GpuProjectAstExpression, _) - if canUseAstJit(astExpression.child) => - GpuProjectAstExpression.replaceChild(alias, GpuAstJitExpression(astExpression.child)) - case alias @ GpuAlias(child: GpuExpression, _) if canUseAstJit(child) => - GpuProjectAstExpression.replaceChild(alias, GpuAstJitExpression(child)) + case alias @ GpuAlias(child: GpuExpression, _) => + asAstJit(child).map(GpuProjectAstExpression.replaceChild(alias, _)).getOrElse(alias) case other => other } @@ -50,9 +53,42 @@ object GpuAstJitExpression { expressions.map(wrapTierExpression(_).asInstanceOf[NamedExpression]) } - private[rapids] def contains(expression: Expression): Boolean = - GpuProjectAstExpressionBase.extractTopLevel(expression) - .exists(_.isInstanceOf[GpuAstJitExpression]) + /** 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) @@ -60,7 +96,6 @@ case class GpuAstJitExpression(child: GpuExpression) with GpuMetricsInjectable with Retryable with AutoCloseable { @transient private[this] var compiledExpression: CompiledExpression = _ - @transient private[this] var completionRegistered = false private[this] var opTime: GpuMetric = NoopMetric override def dataType: DataType = child.dataType @@ -70,6 +105,7 @@ case class GpuAstJitExpression(child: GpuExpression) override def toString: String = s"AST_JIT($child)" override 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) } @@ -80,7 +116,14 @@ case class GpuAstJitExpression(child: GpuExpression) // Compiled ASTs are immutable and remain valid across retry attempts. override def restore(): Unit = () - override def close(): Unit = closeCompiledExpression() + override def close(): Unit = { + val toClose = synchronized { + val current = compiledExpression + compiledExpression = null + current + } + Option(toClose).foreach(_.safeClose()) + } override def columnarEval(batch: ColumnarBatch): GpuColumnVector = { withResource(GpuProjectAstExpression.tableFromBatch(batch)) { table => @@ -90,7 +133,7 @@ case class GpuAstJitExpression(child: GpuExpression) private[rapids] override def computeColumn(table: Table): GpuColumnVector = { val compiled = getCompiledExpression - NvtxIdWithMetrics(NvtxRegistry.PROJECT_AST, opTime) { + NvtxIdWithMetrics(NvtxRegistry.PROJECT_AST_JIT, opTime) { closeOnExcept(compiled.computeColumnJit(table)) { result => GpuColumnVector.from(result, dataType) } @@ -99,46 +142,25 @@ case class GpuAstJitExpression(child: GpuExpression) private def getCompiledExpression: CompiledExpression = synchronized { if (compiledExpression == null) { - compiledExpression = NvtxIdWithMetrics(NvtxRegistry.COMPILE_ASTS, opTime) { + val compiled = NvtxIdWithMetrics(NvtxRegistry.COMPILE_AST_JIT, opTime) { // Force every bound reference to the left table; Project AST has one input table. child.convertToAst(Int.MaxValue).compile() } - } - if (!completionRegistered) { - Option(TaskContext.get()).foreach { taskContext => - completionRegistered = true - try { + closeOnExcept(compiled) { _ => + var completed = false + Option(TaskContext.get()).foreach { taskContext => onTaskCompletion(taskContext) { - clearRegistrationAndClose() + completed = true + close() } - if (!completionRegistered) { - throw new IllegalStateException( - "Task completed while registering the AST JIT cleanup callback") - } - } catch { - case t: Throwable => - clearRegistrationAndClose(t) - throw t } + if (completed) { + throw new IllegalStateException( + "Task completed while registering the AST JIT cleanup callback") + } + compiledExpression = compiled } } compiledExpression } - - private def clearRegistrationAndClose(error: Throwable = null): Unit = - closeCompiledExpression(error, clearRegistration = true) - - private def closeCompiledExpression( - error: Throwable = null, - clearRegistration: Boolean = false): Unit = { - val toClose = synchronized { - if (clearRegistration) { - completionRegistered = false - } - val current = compiledExpression - compiledExpression = null - current - } - Option(toClose).foreach(_.safeClose(error)) - } } 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 d180b94e2c7..1e520fabe8a 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 (!explain.equalsIgnoreCase("NONE")) { + val explanation = GpuAstJitExpression.explainFinalSelections( + tieredProject.exprTiers, explain.equalsIgnoreCase("ALL")) + 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 @@ -134,7 +147,7 @@ object GpuBindReferences extends Logging { conf: SQLConf, enableProjectAstJit: Boolean): GpuTieredProject = { - if (RapidsConf.ENABLE_TIERED_PROJECT.get(conf)) { + val tieredProject = if (RapidsConf.ENABLE_TIERED_PROJECT.get(conf)) { val exprTiers = GpuProjectAstExpression.buildExprTiers( expressions, conf, enableProjectAstJit) val inputTiers = GpuEquivalentExpressions.getInputTiers(exprTiers, input) @@ -183,6 +196,10 @@ object GpuBindReferences extends Logging { GpuTieredProject(Seq( GpuBindReferences.bindGpuReferencesNoMetrics(projectExpressions, input))) } + if (enableProjectAstJit) { + explainFinalProjectAstJitSelection(tieredProject, conf) + } + tieredProject } /** 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 d5bb09beb46..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 @@ -199,7 +199,8 @@ trait GpuExpression extends Expression { /** * Whether this node supports AST JIT for its current semantics and types, excluding its - * children. Operator overrides must validate their execution modes and local input/output types. + * children. Returning true requires `convertToAst` to work for the same execution modes and + * input/output types. */ def selfSupportsAstJit: Boolean = false @@ -346,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 = { @@ -416,7 +425,12 @@ trait CudfBinaryExpression extends GpuBinaryExpression { def castOutputAtEnd: Boolean = false def astOperator: Option[ast.BinaryOperator] = None - override def selfIsAstJitOperator: Boolean = selfSupportsAstJit + 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 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 e81376c6cbc..48faf8db64c 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 @@ -27,7 +27,7 @@ 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 @@ -69,6 +69,7 @@ object GpuProjectAstExpression { case other => other } + /** Extracts a legacy Project AST wrapper after unwrapping any top-level aliases. */ @tailrec private[rapids] def extractTopLevel(expression: Expression): Option[GpuProjectAstExpression] = { expression match { @@ -90,43 +91,38 @@ object GpuProjectAstExpression { case other => other } - private def rewrapBackendTiers( + private def rewrapAstTiers( tiers: Seq[Seq[Expression]], - backendOutputs: Seq[Boolean], - wrapExpression: Expression => Expression): Seq[Seq[Expression]] = { + astOutputs: Seq[Boolean]): Seq[Seq[Expression]] = { val finalTier = tiers.last - require(finalTier.size == backendOutputs.size, + require(finalTier.size == astOutputs.size, "The final expression tier must preserve the project output count") - val backendReferences = finalTier.iterator.zip(backendOutputs.iterator) - .collect { case (expression, true) => expression } - .flatMap(_.references.iterator) - .map(_.exprId) - .toSet + 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)) - // Tier aliases are the dataflow graph after CSE, so follow them backwards from the outputs. + // 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]], backendReferences)) { + (List.empty[Seq[Expression]], astReferences)) { case (tier, (rewrittenTiers, requiredExprIds)) => - val backendAliases = tier.collect { - case alias: GpuAlias if requiredExprIds.contains(alias.exprId) => alias + val taggedTier = tier.map { + case alias: GpuAlias if requiredExprIds.contains(alias.exprId) => (alias, true) + case expression => (expression, false) } - val backendAliasIds = backendAliases.iterator.map(_.exprId).toSet - val dependencies = backendAliases.iterator - .flatMap(_.references.iterator) - .map(_.exprId) - .toSet - val rewrittenTier = tier.map { - case alias: GpuAlias - if backendAliasIds.contains(alias.exprId) && - GpuBatchUtils.isFixedWidth(alias.dataType) => - wrapExpression(alias) - case expression => expression + 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(backendOutputs).map { - case (expression, true) => wrapExpression(expression) + commonTiers :+ finalTier.zip(astOutputs).map { + case (expression, true) => rewrap(expression) case (expression, false) => expression } } @@ -137,7 +133,7 @@ object GpuProjectAstExpression { enableProjectAstJit: Boolean = false): Seq[Seq[Expression]] = { val astOutputs = expressions.map(extractTopLevel(_).isDefined) val hasAstOutputs = astOutputs.contains(true) - val hasJitOutputs = expressions.exists(GpuAstJitExpression.contains) + 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) @@ -151,7 +147,7 @@ object GpuProjectAstExpression { } val tiers = GpuEquivalentExpressions.getExprTiers(replaced) val astTiers = if (hasAstOutputs) { - rewrapBackendTiers(tiers, astOutputs, rewrap) + rewrapAstTiers(tiers, astOutputs) } else { tiers } @@ -194,9 +190,13 @@ case class GpuProjectAstExpression(child: GpuExpression) opTime = metrics.getOrElse(OP_TIME_LEGACY, NoopMetric) } - override def close(): Unit = synchronized { - Option(compiledExpression).foreach(_.safeClose()) - compiledExpression = null + override def close(): Unit = { + val toClose = synchronized { + val current = compiledExpression + compiledExpression = null + current + } + Option(toClose).foreach(_.safeClose()) } override def columnarEval(batch: ColumnarBatch): GpuColumnVector = { @@ -221,11 +221,17 @@ case class GpuProjectAstExpression(child: GpuExpression) child.convertToAst(Int.MaxValue).compile() } closeOnExcept(compiled) { _ => + var completed = false Option(TaskContext.get()).foreach { taskContext => onTaskCompletion(taskContext) { + completed = true close() } } + if (completed) { + throw new IllegalStateException( + "Task completed while registering the Project AST cleanup callback") + } compiledExpression = compiled } } 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/basicPhysicalOperators.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala index aad8b881939..4d4a4b40d72 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 @@ -49,16 +49,6 @@ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.random.BernoulliCellSampler -object GpuProjectExecMeta { - private[rapids] def explainBackend(expression: Expression): String = { - GpuProjectAstExpressionBase.extractTopLevel(expression) match { - case Some(_: GpuAstJitExpression) => "Project AST JIT" - case Some(_: GpuProjectAstExpression) => "legacy Project AST" - case _ => "the regular GPU projection" - } - } -} - class GpuProjectExecMeta( proj: ProjectExec, conf: RapidsConf, @@ -81,7 +71,7 @@ class GpuProjectExecMeta( 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 (!GpuAstJitExpression.contains(expr) && + if (GpuAstJitExpression.extractTopLevel(expr).isEmpty && GpuBatchUtils.isFixedWidth(expr.dataType) && meta.canThisBeAst && !isTopLevelNullLiteral(expr)) { GpuProjectAstExpression.wrap(expr) @@ -92,36 +82,32 @@ class GpuProjectExecMeta( } else { jitExprs } - // Explain the planned top-level backend because both AST backends are optional and can be - // difficult to distinguish from the regular projection in a transformed plan. - if (conf.shouldExplain && (conf.isProjectAstEnabled || conf.isProjectAstJitEnabled)) { - val backendExplain = projectList.iterator.map { expression => - s" $expression will use ${GpuProjectExecMeta.explainBackend(expression)}\n" - } - val legacyExplain = if (conf.isProjectAstEnabled) { - childExprs.iterator.zip(projectList.iterator).collect { - case (meta, expression) - if GpuProjectAstExpressionBase.extractTopLevel(expression).isEmpty => - meta.explainAst(conf.shouldExplainAll) - }.filter(_.nonEmpty) - } else { - Iterator.empty - } + // 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).collect { + case (meta, expression) if GpuAstJitExpression.extractTopLevel(expression).isEmpty => + val explanation = meta.explainAst(conf.shouldExplainAll) + if (explanation.nonEmpty) { + s" Legacy Project AST eligibility:\n$explanation" + } else { + "" + } + }.filter(_.nonEmpty) 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 AST because its return type " + + 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 so null outputs can reuse " + - "the cached null vector\n" + 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 = (backendExplain ++ legacyExplain ++ regularExplain).mkString + val explain = (legacyExplain ++ regularExplain).mkString if (explain.nonEmpty) { - logWarning(s"AST PROJECT\n$explain") + logWarning(s"LEGACY PROJECT AST\n$explain") } } GpuProjectExec(projectList, gpuChild) 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 e67681c4f66..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,7 +314,7 @@ abstract class GpuAddBase extends CudfBinaryArithmetic with Serializable { override def binaryOp: BinaryOp = BinaryOp.ADD override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.ADD) - override def selfSupportsAstJit: Boolean = + override protected def astJitCompatible: Boolean = !failOnError && (dataType == IntegerType || dataType == LongType) override def hasSideEffects: Boolean = @@ -771,7 +771,7 @@ case class GpuMultiply( override def binaryOp: BinaryOp = BinaryOp.MUL override def astOperator: Option[BinaryOperator] = Some(ast.BinaryOperator.MUL) - override def selfSupportsAstJit: Boolean = + override protected def astJitCompatible: Boolean = !failOnError && (dataType == IntegerType || dataType == LongType) private def multiplyOverflowError(msg: String): ArithmeticException = { diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala index 26d8ef4e7ee..a52455e865b 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala @@ -17,11 +17,12 @@ package com.nvidia.spark.rapids import ai.rapids.cudf.ast.{AstExpression, CompiledExpression} +import com.nvidia.spark.rapids.ProjectAstTestUtils.{collectExpressions, tierReferences} 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, Expression} +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 @@ -42,12 +43,6 @@ class GpuProjectAstJitSuite extends AnyFunSuite { GpuAstJitExpression(child) } - private def jitExpressions(expressions: Seq[Expression]): Seq[GpuAstJitExpression] = { - expressions.flatMap(_.collect { - case expression: GpuAstJitExpression => expression - }) - } - private def projectConf( tiered: Boolean = true, jit: Boolean = true, @@ -78,7 +73,7 @@ class GpuProjectAstJitSuite extends AnyFunSuite { val wrappedAgain = GpuAstJitExpression.wrapProjectExpressions(wrapped) assert(jit.child.isInstanceOf[GpuMultiply]) assert(jit.child.find(_.isInstanceOf[GpuAstJitExpression]).isEmpty) - assert(GpuAstJitExpression.contains(wrapped.head)) + assert(GpuAstJitExpression.extractTopLevel(wrapped.head).contains(jit)) assert(wrappedAgain.head eq wrapped.head) } @@ -94,7 +89,7 @@ class GpuProjectAstJitSuite extends AnyFunSuite { val wrapped = GpuAstJitExpression.wrapProjectExpressions(List(expression)) val subtract = wrapped.head.asInstanceOf[GpuAlias].child.asInstanceOf[GpuSubtract] - assert(jitExpressions(wrapped).isEmpty) + assert(wrapped.forall(GpuAstJitExpression.extractTopLevel(_).isEmpty)) assert(subtract.left.isInstanceOf[GpuAdd]) assert(subtract.right.isInstanceOf[GpuMultiply]) } @@ -105,29 +100,52 @@ class GpuProjectAstJitSuite extends AnyFunSuite { 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")) - // Inputs: AST((left + right) - third), greatest(left + right, fourth). - // Tier 0 computes AST_JIT(left + right); tier 1 consumes that shared result in both outputs. val tiered = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( expressions, Seq(left, right, third, fourth), projectConf()) - assertResult(2)(tiered.exprTiers.size) - assertResult(Seq(1, 0))(tiered.exprTiers.map(jitExpressions(_).size)) - assert(jitExpressions(tiered.exprTiers.head).head.child.isInstanceOf[GpuAdd]) - assert(GpuProjectAstExpressionBase.extractTopLevel(tiered.exprTiers.last.head) - .exists(_.isInstanceOf[GpuProjectAstExpression])) - assert(GpuProjectAstExpressionBase.extractTopLevel(tiered.exprTiers.last(1)).isEmpty) - val finalTierReferences = tiered.exprTiers.last.flatMap(_.collect { - case reference: GpuBoundReference if reference.name.startsWith("tiered_input_") => reference - }) + // after CSE: + // 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 finalTierReferences = tiered.exprTiers.last.flatMap(tierReferences) assertResult(2)(finalTierReferences.size) assertResult(1)(finalTierReferences.map(_.exprId).distinct.size) } + 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) @@ -140,10 +158,8 @@ class GpuProjectAstJitSuite extends AnyFunSuite { Seq(jitCandidate, legacyCandidate), Seq(left, right), projectConf(legacy = true)) val outputs = tiered.exprTiers.last - assert(GpuProjectAstExpressionBase.extractTopLevel(outputs.head) - .exists(_.isInstanceOf[GpuAstJitExpression])) - assert(GpuProjectAstExpressionBase.extractTopLevel(outputs(1)) - .exists(_.isInstanceOf[GpuProjectAstExpression])) + assert(GpuAstJitExpression.extractTopLevel(outputs.head).isDefined) + assert(GpuProjectAstExpression.extractTopLevel(outputs(1)).isDefined) } test("only the project binder selects the JIT backend") { @@ -157,8 +173,8 @@ class GpuProjectAstJitSuite extends AnyFunSuite { val project = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( Seq(expression), Seq(left, right), conf) - assert(jitExpressions(generic.exprTiers.flatten).isEmpty) - assertResult(1)(jitExpressions(project.exprTiers.flatten).size) + assert(collectExpressions[GpuAstJitExpression](generic.exprTiers.flatten).isEmpty) + assertResult(1)(collectExpressions[GpuAstJitExpression](project.exprTiers.flatten).size) } test("the project binder respects a disabled JIT setting") { @@ -169,7 +185,7 @@ class GpuProjectAstJitSuite extends AnyFunSuite { val project = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( Seq(expression), Seq(left, right), projectConf(jit = false)) - assert(jitExpressions(project.exprTiers.flatten).isEmpty) + assert(collectExpressions[GpuAstJitExpression](project.exprTiers.flatten).isEmpty) } test("project JIT remains available when tiered projection is disabled") { @@ -181,7 +197,7 @@ class GpuProjectAstJitSuite extends AnyFunSuite { Seq(expression), Seq(left, right), projectConf(tiered = false)) assertResult(1)(project.exprTiers.size) - assertResult(1)(jitExpressions(project.exprTiers.head).size) + assertResult(1)(collectExpressions[GpuAstJitExpression](project.exprTiers.head).size) } test("project AST JIT excludes ANSI and floating point arithmetic") { @@ -206,24 +222,38 @@ class GpuProjectAstJitSuite extends AnyFunSuite { val wrapped = GpuAstJitExpression.wrapProjectExpressions(expressions) - assert(jitExpressions(wrapped).isEmpty) + assert(wrapped.forall(GpuAstJitExpression.extractTopLevel(_).isEmpty)) } - test("project backend explanation follows the final wrapper") { + 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(left, right, failOnError = false)(), "regular") + 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) - - assertResult("Project AST JIT")(GpuProjectExecMeta.explainBackend(jit)) - assertResult("legacy Project AST")(GpuProjectExecMeta.explainBackend(legacy)) - assertResult("the regular GPU projection")(GpuProjectExecMeta.explainBackend(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 taskContext = new MockTaskContext(taskAttemptId = 1, partitionId = 0) val child = mock(classOf[GpuExpression]) val ast = mock(classOf[AstExpression]) val compiled = mock(classOf[CompiledExpression]) @@ -231,7 +261,7 @@ class GpuProjectAstJitSuite extends AnyFunSuite { when(ast.compile()).thenReturn(compiled) val jit = GpuAstJitExpression(child) - TestUtils.withTaskContext(taskContext) { + TestUtils.withMockTaskContext() { jit.checkpoint() jit.restore() jit.checkpoint() 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..b1ddb9d10ee --- /dev/null +++ b/tests/src/test/scala/com/nvidia/spark/rapids/ProjectAstTestUtils.scala @@ -0,0 +1,37 @@ +/* + * 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] = { + val runtimeClass = implicitly[ClassTag[T]].runtimeClass + expressions.flatMap(_.collect { + case expression if runtimeClass.isInstance(expression) => expression.asInstanceOf[T] + }) + } + + 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/TestUtils.scala b/tests/src/test/scala/com/nvidia/spark/rapids/TestUtils.scala index ed88e40ff92..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,7 +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.MockTaskContextBase +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. */ @@ -47,13 +47,13 @@ object TestUtils extends Assertions { def withTaskContext[T]( taskContext: MockTaskContextBase, - markTaskComplete: Boolean = false)(body: => T): T = { + completesTask: Boolean = false)(body: => T): T = { TrampolineUtil.setTaskContext(taskContext) try { body } finally { try { - if (markTaskComplete) { + if (completesTask) { taskContext.markTaskComplete() } } finally { @@ -63,6 +63,10 @@ object TestUtils extends Assertions { } } + 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..d377740cc1c 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,24 +214,21 @@ 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() } } 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 index 6ddc2b3851d..909f0db394a 100644 --- 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 @@ -19,6 +19,7 @@ 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 @@ -28,11 +29,11 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, 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.rapids.metrics.source.MockTaskContext 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)() @@ -59,8 +60,8 @@ class GpuBroadcastNestedLoopJoinRetrySuite extends RmmSparkRetrySuiteBase { } override def afterEach(): Unit = { - RmmSpark.getAndResetNumRetryThrow(/* taskId */ 1) - RmmSpark.getAndResetNumSplitRetryThrow(/* taskId */ 1) + RmmSpark.getAndResetNumRetryThrow(taskId) + RmmSpark.getAndResetNumSplitRetryThrow(taskId) super.afterEach() } @@ -96,28 +97,27 @@ class GpuBroadcastNestedLoopJoinRetrySuite extends RmmSparkRetrySuiteBase { test("BNLJ build-side shared JIT tier retries GpuRetryOOM") { val spark = SparkSession.builder() .master("local[1]") - .appName("GpuBroadcastNestedLoopJoinRetrySuite") + .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) - val taskContext = new MockTaskContext(taskAttemptId = 1, partitionId = 0) - TestUtils.withTaskContext(taskContext, markTaskComplete = true) { - // Input: x=[1,2,3], y=[4,5,6], z=[2,3,4]. CSE produces a first tier that - // passes through x/y/z and computes AST_JIT((x+y)*z), followed by the five final outputs - // [shared-x, shared-y, x, y, z]. + TestUtils.withMockTaskContext(completesTask = true) { val boundProject = GpuBindReferences.bindGpuProjectReferencesTiered( expressions, buildAttributes, conf, Map.empty) - val jitExpressions = boundProject.exprTiers.flatten.flatMap(_.collect { - case expression: GpuAstJitExpression => expression - }) + // 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]) @@ -129,8 +129,8 @@ class GpuBroadcastNestedLoopJoinRetrySuite extends RmmSparkRetrySuiteBase { withResource(projectBuildSide(buildBatch())) { _ => } val retryInput = buildBatch() - RmmSpark.getAndResetNumRetryThrow(/* taskId */ 1) - RmmSpark.getAndResetNumSplitRetryThrow(/* taskId */ 1) + RmmSpark.getAndResetNumRetryThrow(taskId) + RmmSpark.getAndResetNumSplitRetryThrow(taskId) RmmSpark.forceRetryOOM(RmmSpark.getCurrentThreadId, 1, RmmSpark.OomInjectionType.GPU.ordinal, 0) withResource(projectBuildSide(retryInput)) { output => @@ -139,8 +139,8 @@ class GpuBroadcastNestedLoopJoinRetrySuite extends RmmSparkRetrySuiteBase { assertResult(Seq(9, 19, 33))(collectInts(output, 0)) assertResult(Seq(6, 16, 30))(collectInts(output, 1)) } - assert(RmmSpark.getAndResetNumRetryThrow(/* taskId */ 1) > 0) - assertResult(0)(RmmSpark.getAndResetNumSplitRetryThrow(/* taskId */ 1)) + assert(RmmSpark.getAndResetNumRetryThrow(taskId) > 0) + assertResult(0)(RmmSpark.getAndResetNumSplitRetryThrow(taskId)) } } } From 084c57f25e23adefe8ba55a77b39fce74eb7cdd0 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Mon, 17 Aug 2026 15:46:02 +0800 Subject: [PATCH 17/20] address comments Signed-off-by: Haoyang Li --- .../spark/rapids/GpuAstJitExpression.scala | 79 +----- .../spark/rapids/GpuBoundAttribute.scala | 6 +- .../nvidia/spark/rapids/GpuOverrides.scala | 3 +- .../rapids/GpuProjectAstExpression.scala | 251 ++++++++++-------- .../com/nvidia/spark/rapids/RapidsConf.scala | 10 +- .../spark/rapids/basicPhysicalOperators.scala | 15 +- .../spark/rapids/GpuProjectAstJitSuite.scala | 29 +- .../spark/rapids/ProjectAstTestUtils.scala | 3 +- 8 files changed, 189 insertions(+), 207 deletions(-) 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 index 1356f2ef32d..a36ce52c948 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala @@ -16,19 +16,11 @@ package com.nvidia.spark.rapids -import ai.rapids.cudf.Table +import ai.rapids.cudf.{ColumnVector, Table} import ai.rapids.cudf.ast.CompiledExpression import com.nvidia.spark.Retryable -import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} -import com.nvidia.spark.rapids.GpuMetric.OP_TIME_LEGACY -import com.nvidia.spark.rapids.RapidsPluginImplicits._ -import com.nvidia.spark.rapids.ScalableTaskCompletion.onTaskCompletion -import com.nvidia.spark.rapids.shims.ShimUnaryExpression -import org.apache.spark.TaskContext import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression} -import org.apache.spark.sql.types.DataType -import org.apache.spark.sql.vectorized.ColumnarBatch object GpuAstJitExpression { private def canUseAstJit(expression: GpuExpression): Boolean = @@ -44,7 +36,7 @@ object GpuAstJitExpression { private[rapids] def wrapTierExpression(expression: Expression): Expression = expression match { case alias @ GpuAlias(child: GpuExpression, _) => - asAstJit(child).map(GpuProjectAstExpression.replaceChild(alias, _)).getOrElse(alias) + asAstJit(child).map(GpuProjectAstExpressionBase.replaceChild(alias, _)).getOrElse(alias) case other => other } @@ -92,22 +84,19 @@ object GpuAstJitExpression { } case class GpuAstJitExpression(child: GpuExpression) - extends ShimUnaryExpression with GpuProjectAstExpressionBase - with GpuMetricsInjectable with Retryable with AutoCloseable { + extends GpuProjectAstExpressionBase with Retryable { - @transient private[this] var compiledExpression: CompiledExpression = _ - private[this] var opTime: GpuMetric = NoopMetric + override protected def backendName: String = "AST JIT" - override def dataType: DataType = child.dataType + override protected def compileNvtxId: NvtxId = NvtxRegistry.COMPILE_AST_JIT - override def nullable: Boolean = child.nullable + override protected def computeNvtxId: NvtxId = NvtxRegistry.PROJECT_AST_JIT - override def toString: String = s"AST_JIT($child)" + override protected def evaluate( + compiled: CompiledExpression, + table: Table): ColumnVector = compiled.computeColumnJit(table) - override 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 def toString: String = s"AST_JIT($child)" override def checkpoint(): Unit = { getCompiledExpression @@ -115,52 +104,4 @@ case class GpuAstJitExpression(child: GpuExpression) // Compiled ASTs are immutable and remain valid across retry attempts. override def restore(): Unit = () - - override def close(): Unit = { - val toClose = synchronized { - val current = compiledExpression - compiledExpression = null - current - } - Option(toClose).foreach(_.safeClose()) - } - - override def columnarEval(batch: ColumnarBatch): GpuColumnVector = { - withResource(GpuProjectAstExpression.tableFromBatch(batch)) { table => - computeColumn(table) - } - } - - private[rapids] override def computeColumn(table: Table): GpuColumnVector = { - val compiled = getCompiledExpression - NvtxIdWithMetrics(NvtxRegistry.PROJECT_AST_JIT, opTime) { - closeOnExcept(compiled.computeColumnJit(table)) { result => - GpuColumnVector.from(result, dataType) - } - } - } - - private def getCompiledExpression: CompiledExpression = synchronized { - if (compiledExpression == null) { - val compiled = NvtxIdWithMetrics(NvtxRegistry.COMPILE_AST_JIT, opTime) { - // Force every bound reference to the left table; Project AST has one input table. - child.convertToAst(Int.MaxValue).compile() - } - closeOnExcept(compiled) { _ => - var completed = false - Option(TaskContext.get()).foreach { taskContext => - onTaskCompletion(taskContext) { - completed = true - close() - } - } - if (completed) { - throw new IllegalStateException( - "Task completed while registering the AST JIT cleanup callback") - } - compiledExpression = compiled - } - } - compiledExpression - } } 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 1e520fabe8a..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 @@ -46,9 +46,9 @@ object GpuBindReferences extends Logging { tieredProject: GpuTieredProject, conf: SQLConf): Unit = { val explain = RapidsConf.EXPLAIN.get(conf) - if (!explain.equalsIgnoreCase("NONE")) { + if (RapidsConf.shouldExplain(explain)) { val explanation = GpuAstJitExpression.explainFinalSelections( - tieredProject.exprTiers, explain.equalsIgnoreCase("ALL")) + tieredProject.exprTiers, RapidsConf.shouldExplainAll(explain)) if (explanation.nonEmpty) { logWarning(s"FINAL PROJECT AST JIT SELECTION\n$explanation") } @@ -148,7 +148,7 @@ object GpuBindReferences extends Logging { enableProjectAstJit: Boolean): GpuTieredProject = { val tieredProject = if (RapidsConf.ENABLE_TIERED_PROJECT.get(conf)) { - val exprTiers = GpuProjectAstExpression.buildExprTiers( + 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 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 48faf8db64c..738a2f000f6 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 @@ -18,7 +18,7 @@ package com.nvidia.spark.rapids import scala.annotation.tailrec -import ai.rapids.cudf.{Scalar, Table} +import ai.rapids.cudf.{ColumnVector, Scalar, Table} import ai.rapids.cudf.ast.CompiledExpression import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource} import com.nvidia.spark.rapids.GpuMetric.OP_TIME_LEGACY @@ -33,8 +33,74 @@ import org.apache.spark.sql.rapids.catalyst.expressions.GpuEquivalentExpressions import org.apache.spark.sql.types.DataType import org.apache.spark.sql.vectorized.ColumnarBatch -trait GpuProjectAstExpressionBase extends GpuExpression { - private[rapids] def computeColumn(table: Table): GpuColumnVector +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 evaluate(compiled: CompiledExpression, table: Table): ColumnVector + + @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()) + } + + override final def columnarEval(batch: ColumnarBatch): GpuColumnVector = { + withResource(GpuProjectAstExpressionBase.tableFromBatch(batch)) { table => + computeColumn(table) + } + } + + private[rapids] final def computeColumn(table: Table): GpuColumnVector = { + val compiled = getCompiledExpression + NvtxIdWithMetrics(computeNvtxId, opTime) { + closeOnExcept(evaluate(compiled, table)) { result => + GpuColumnVector.from(result, dataType) + } + } + } + + protected 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. + child.convertToAst(Int.MaxValue).compile() + } + 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 { @@ -47,9 +113,7 @@ object GpuProjectAstExpressionBase { case _ => None } } -} -object GpuProjectAstExpression { private[rapids] def replaceChild(alias: GpuAlias, child: Expression): GpuAlias = { if (child eq alias.child) { alias @@ -58,27 +122,6 @@ object GpuProjectAstExpression { } } - private def asAst(child: GpuExpression): GpuProjectAstExpression = child match { - case astExpression: GpuProjectAstExpression => astExpression - case other => GpuProjectAstExpression(other) - } - - private[rapids] def wrap(expression: NamedExpression): NamedExpression = expression match { - case alias @ GpuAlias(child: GpuExpression, _) => - replaceChild(alias, asAst(child)) - case other => other - } - - /** Extracts a legacy Project AST wrapper after unwrapping any top-level aliases. */ - @tailrec - private[rapids] def extractTopLevel(expression: Expression): Option[GpuProjectAstExpression] = { - expression match { - case alias: GpuAlias => extractTopLevel(alias.child) - case astExpression: GpuProjectAstExpression => Some(astExpression) - case _ => None - } - } - private def unwrap(expression: Expression): Expression = expression match { case alias: GpuAlias => replaceChild(alias, unwrap(alias.child)) case astExpression: GpuProjectAstExpression => astExpression.child @@ -86,52 +129,11 @@ object GpuProjectAstExpression { 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") - 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)) - - // 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) - } - 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 - } - } - private[rapids] def buildExprTiers( expressions: Seq[Expression], conf: SQLConf, enableProjectAstJit: Boolean = false): Seq[Seq[Expression]] = { - val astOutputs = expressions.map(extractTopLevel(_).isDefined) + val astOutputs = expressions.map(GpuProjectAstExpression.extractTopLevel(_).isDefined) val hasAstOutputs = astOutputs.contains(true) val hasJitOutputs = expressions.exists(GpuAstJitExpression.extractTopLevel(_).isDefined) // CSE must see through backend markers so all outputs can share the same tiers. @@ -147,7 +149,7 @@ object GpuProjectAstExpression { } val tiers = GpuEquivalentExpressions.getExprTiers(replaced) val astTiers = if (hasAstOutputs) { - rewrapAstTiers(tiers, astOutputs) + GpuProjectAstExpression.rewrapAstTiers(tiers, astOutputs) } else { tiers } @@ -174,67 +176,80 @@ object GpuProjectAstExpression { } } -case class GpuProjectAstExpression(child: GpuExpression) - extends ShimUnaryExpression with GpuProjectAstExpressionBase - 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 = { - val toClose = synchronized { - val current = compiledExpression - compiledExpression = null - current - } - Option(toClose).foreach(_.safeClose()) + 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] override 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) { _ => - var completed = false - Option(TaskContext.get()).foreach { taskContext => - onTaskCompletion(taskContext) { - completed = true - 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) } - if (completed) { - throw new IllegalStateException( - "Task completed while registering the Project AST cleanup callback") + val dependencies = referenceSet(taggedTier) + val rewrittenTier = taggedTier.map { + case (alias, true) if GpuBatchUtils.isFixedWidth(alias.dataType) => rewrap(alias) + case (expression, _) => expression } - compiledExpression = compiled - } + (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 protected def evaluate( + compiled: CompiledExpression, + table: Table): ColumnVector = compiled.computeColumn(table) + + override def toString: String = s"AST($child)" } 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 36d1c91aa19..d2e326c4dbd 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 @@ -2553,6 +2553,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() @@ -3546,9 +3552,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) 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 4d4a4b40d72..be14a2e68e0 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 @@ -84,15 +84,12 @@ class GpuProjectExecMeta( } // 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).collect { + val legacyExplain = childExprs.iterator.zip(projectList.iterator).flatMap { case (meta, expression) if GpuAstJitExpression.extractTopLevel(expression).isEmpty => - val explanation = meta.explainAst(conf.shouldExplainAll) - if (explanation.nonEmpty) { - s" Legacy Project AST eligibility:\n$explanation" - } else { - "" - } - }.filter(_.nonEmpty) + 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 && @@ -175,7 +172,7 @@ object GpuProjectExec { GpuProjectAstExpressionBase.extractTopLevel(expression).isDefined } if (hasAstExpressions) { - withResource(GpuProjectAstExpression.tableFromBatch(cb)) { table => + withResource(GpuProjectAstExpressionBase.tableFromBatch(cb)) { table => projectWithEval { expression => GpuProjectAstExpressionBase.extractTopLevel(expression) match { case Some(astExpression) => astExpression.computeColumn(table) diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala index a52455e865b..a539aebcebf 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala @@ -16,6 +16,7 @@ package com.nvidia.spark.rapids +import ai.rapids.cudf.Table import ai.rapids.cudf.ast.{AstExpression, CompiledExpression} import com.nvidia.spark.rapids.ProjectAstTestUtils.{collectExpressions, tierReferences} import org.mockito.Mockito.{doThrow, mock, times, verify, when} @@ -35,14 +36,17 @@ class GpuProjectAstJitSuite extends AnyFunSuite { private def alias(expression: GpuExpression, name: String) = GpuAlias(expression, name)() - private def mockJitExpression(compiled: CompiledExpression): GpuAstJitExpression = { + private def mockCompiledChild(compiled: CompiledExpression): GpuExpression = { val child = mock(classOf[GpuExpression]) val ast = mock(classOf[AstExpression]) when(child.convertToAst(Int.MaxValue)).thenReturn(ast) when(ast.compile()).thenReturn(compiled) - GpuAstJitExpression(child) + child } + private def mockJitExpression(compiled: CompiledExpression): GpuAstJitExpression = + GpuAstJitExpression(mockCompiledChild(compiled)) + private def projectConf( tiered: Boolean = true, jit: Boolean = true, @@ -313,4 +317,25 @@ class GpuProjectAstJitSuite extends AnyFunSuite { 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 index b1ddb9d10ee..d2f8d501f5c 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/ProjectAstTestUtils.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/ProjectAstTestUtils.scala @@ -23,9 +23,8 @@ import org.apache.spark.sql.catalyst.expressions.Expression object ProjectAstTestUtils { def collectExpressions[T <: Expression : ClassTag]( expressions: Seq[Expression]): Seq[T] = { - val runtimeClass = implicitly[ClassTag[T]].runtimeClass expressions.flatMap(_.collect { - case expression if runtimeClass.isInstance(expression) => expression.asInstanceOf[T] + case expression: T => expression }) } From c3a1f6da4a0cbb997343c41684b95ac2e31cb9f5 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Wed, 19 Aug 2026 20:33:36 +0800 Subject: [PATCH 18/20] add nvtx docs Signed-off-by: Haoyang Li --- docs/dev/nvtx_ranges.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/dev/nvtx_ranges.md b/docs/dev/nvtx_ranges.md index caf31095174..187622cd3f1 100644 --- a/docs/dev/nvtx_ranges.md +++ b/docs/dev/nvtx_ranges.md @@ -38,7 +38,6 @@ DoubleBatchedWindow_PRE|Pre-processing for double-batched window operation GpuGenerate project split|Splitting projection in generate operation parquet parse filter footer|Parsing and filtering Parquet footer by range Compile ASTs|Compiling abstract syntax trees for expression evaluation -Compile AST JIT|Compiling an AST expression for JIT evaluation parquet filter blocks|Filtering Parquet row group blocks based on predicates PageableH2D|Copying from pageable host memory to device TOP N|Computing top N rows @@ -116,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 @@ -135,7 +135,6 @@ single build batch concat|Concatenating batches for single build batch sort copy boundaries|Copying boundary data for sort operation WaitingForWrites|Rapids Shuffle Manager (multi threaded) is waiting for any queued writes to finish before finalizing the map output writer Project AST|Applying AST-based projection to batch -Project AST JIT|Applying JIT-compiled AST projection to batch ORC readBatches|Reading ORC batches Shuffle Transfer Request|Handling shuffle data transfer request consumeWindow|Consuming transfer window @@ -192,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 From 3372a30b8c3bd73f06b1481e234f9bb20333c1db Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Wed, 26 Aug 2026 16:00:55 +0800 Subject: [PATCH 19/20] Add multi-output AST JIT project waves Signed-off-by: Haoyang Li --- integration_tests/src/main/python/ast_test.py | 38 ++- .../spark/rapids/GpuAstJitExpression.scala | 25 +- .../rapids/GpuAstJitProjectPlanner.scala | 270 ++++++++++++++++++ .../rapids/GpuProjectAstExpression.scala | 19 +- .../com/nvidia/spark/rapids/RapidsConf.scala | 7 + .../spark/rapids/basicPhysicalOperators.scala | 48 +++- .../spark/rapids/GpuProjectAstJitSuite.scala | 120 +++++++- .../spark/sql/rapids/ProjectExprSuite.scala | 45 +++ 8 files changed, 546 insertions(+), 26 deletions(-) create mode 100644 sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitProjectPlanner.scala diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index 88ecf187d58..15b81c2fa2c 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/integration_tests/src/main/python/ast_test.py @@ -405,12 +405,46 @@ def test_jit_add_multiply(data_gen): @pytest.mark.parametrize('data_gen', [int_gen, long_gen], ids=idfn) @disable_ansi_mode -def test_jit_does_not_split_unique_unsupported_expression(data_gen): +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="AST_JIT,GpuProjectAst", + 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) 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 index a36ce52c948..100ee94d434 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala @@ -19,13 +19,18 @@ package com.nvidia.spark.rapids import ai.rapids.cudf.{ColumnVector, Table} import ai.rapids.cudf.ast.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 def canUseAstJit(expression: GpuExpression): Boolean = - GpuBatchUtils.isFixedWidth(expression.dataType) && - expression.supportsAstJit && expression.containsAstJitOperator + 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) @@ -45,6 +50,18 @@ object GpuAstJitExpression { 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 { @@ -86,6 +103,8 @@ object GpuAstJitExpression { 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 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/GpuProjectAstExpression.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuProjectAstExpression.scala index 738a2f000f6..ddd1073cf03 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 @@ -71,14 +71,17 @@ trait GpuProjectAstExpressionBase private[rapids] final def computeColumn(table: Table): GpuColumnVector = { val compiled = getCompiledExpression - NvtxIdWithMetrics(computeNvtxId, opTime) { + withComputeMetrics { closeOnExcept(evaluate(compiled, table)) { result => GpuColumnVector.from(result, dataType) } } } - protected final def getCompiledExpression: CompiledExpression = synchronized { + 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. @@ -142,12 +145,16 @@ object GpuProjectAstExpressionBase { } else { expressions } - val replaced = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) { - GpuEquivalentExpressions.replaceMultiExpressions(unwrapped, conf) + val tiers = if (enableProjectAstJit) { + GpuAstJitProjectPlanner.buildExprTiers(unwrapped, conf) } else { - unwrapped + val replaced = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) { + GpuEquivalentExpressions.replaceMultiExpressions(unwrapped, conf) + } else { + unwrapped + } + GpuEquivalentExpressions.getExprTiers(replaced) } - val tiers = GpuEquivalentExpressions.getExprTiers(replaced) val astTiers = if (hasAstOutputs) { GpuProjectAstExpression.rewrapAstTiers(tiers, astOutputs) } else { 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 d2e326c4dbd..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 @@ -1255,6 +1255,13 @@ val GPU_COREDUMP_PIPE_PATTERN = conf("spark.rapids.gpu.coreDump.pipePattern") .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() 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 be14a2e68e0..b5a5506d51e 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 @@ -168,15 +170,45 @@ object GpuProjectExec { } } - val hasAstExpressions = boundExprs.exists { expression => - GpuProjectAstExpressionBase.extractTopLevel(expression).isDefined - } - if (hasAstExpressions) { + val astExpressions = boundExprs.flatMap( + GpuProjectAstExpressionBase.extractTopLevel) + if (astExpressions.nonEmpty) { withResource(GpuProjectAstExpressionBase.tableFromBatch(cb)) { table => - projectWithEval { expression => - GpuProjectAstExpressionBase.extractTopLevel(expression) match { - case Some(astExpression) => astExpression.computeColumn(table) - case None => expression.columnarEval(cb) + 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) + } } } } diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala index a539aebcebf..c4e59d73377 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala @@ -18,7 +18,7 @@ package com.nvidia.spark.rapids import ai.rapids.cudf.Table import ai.rapids.cudf.ast.{AstExpression, CompiledExpression} -import com.nvidia.spark.rapids.ProjectAstTestUtils.{collectExpressions, tierReferences} +import com.nvidia.spark.rapids.ProjectAstTestUtils.collectExpressions import org.mockito.Mockito.{doThrow, mock, times, verify, when} import org.scalatest.funsuite.AnyFunSuite @@ -98,7 +98,7 @@ class GpuProjectAstJitSuite extends AnyFunSuite { assert(subtract.right.isInstanceOf[GpuMultiply]) } - test("project CSE exposes a shared supported expression to JIT") { + 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) @@ -113,7 +113,7 @@ class GpuProjectAstJitSuite extends AnyFunSuite { val tiered = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( expressions, Seq(left, right, third, fourth), projectConf()) - // after CSE: + // 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]) @@ -122,11 +122,115 @@ class GpuProjectAstJitSuite extends AnyFunSuite { assert(GpuProjectAstExpression.extractTopLevel(tiered.exprTiers.last.head).isDefined) assert(GpuProjectAstExpression.extractTopLevel(tiered.exprTiers.last(1)).isEmpty) // Final references: [t1, t1] (distinct: {t1}). - val finalTierReferences = tiered.exprTiers.last.flatMap(tierReferences) + 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) @@ -195,13 +299,15 @@ class GpuProjectAstJitSuite extends AnyFunSuite { test("project JIT remains available when tiered projection is disabled") { val left = reference(0, IntegerType) val right = reference(1, IntegerType) - val expression = alias(GpuAdd(left, right, failOnError = false)(), "result") + val expressions = Seq( + alias(GpuAdd(left, right, failOnError = false)(), "sum"), + alias(GpuMultiply(left, right, failOnError = false)(), "product")) val project = GpuBindReferences.bindGpuProjectReferencesTieredNoMetrics( - Seq(expression), Seq(left, right), projectConf(tiered = false)) + expressions, Seq(left, right), projectConf(tiered = false)) assertResult(1)(project.exprTiers.size) - assertResult(1)(collectExpressions[GpuAstJitExpression](project.exprTiers.head).size) + assertResult(2)(collectExpressions[GpuAstJitExpression](project.exprTiers.head).size) } test("project AST JIT excludes ANSI and floating point arithmetic") { 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 d377740cc1c..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 @@ -233,6 +233,51 @@ class ProjectExprSuite extends SparkQueryCompareTestSuite { } } + 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"), From 7dc09923f786d8d32a7ecedcdd1cc0ae48754ba7 Mon Sep 17 00:00:00 2001 From: Haoyang Li Date: Wed, 26 Aug 2026 18:04:55 +0800 Subject: [PATCH 20/20] use new api Signed-off-by: Haoyang Li --- .../spark/rapids/GpuAstJitExpression.scala | 8 ++--- .../rapids/GpuProjectAstExpression.scala | 14 +++----- .../spark/rapids/basicPhysicalOperators.scala | 2 +- .../spark/rapids/GpuProjectAstJitSuite.scala | 16 +++++++--- .../spark/rapids/ProjectSplitRetrySuite.scala | 32 ++++++++++++++++++- ...GpuBroadcastNestedLoopJoinRetrySuite.scala | 2 +- 6 files changed, 52 insertions(+), 22 deletions(-) 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 index 100ee94d434..cf6a10a678e 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAstJitExpression.scala @@ -16,8 +16,8 @@ package com.nvidia.spark.rapids -import ai.rapids.cudf.{ColumnVector, Table} -import ai.rapids.cudf.ast.CompiledExpression +import ai.rapids.cudf.Table +import ai.rapids.cudf.ast.{AstExpression, CompiledExpression} import com.nvidia.spark.Retryable import com.nvidia.spark.rapids.Arm.withResource @@ -111,9 +111,7 @@ case class GpuAstJitExpression(child: GpuExpression) override protected def computeNvtxId: NvtxId = NvtxRegistry.PROJECT_AST_JIT - override protected def evaluate( - compiled: CompiledExpression, - table: Table): ColumnVector = compiled.computeColumnJit(table) + override protected def compileAst(ast: AstExpression): CompiledExpression = ast.compileJit() override def toString: String = s"AST_JIT($child)" 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 ddd1073cf03..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 @@ -18,8 +18,8 @@ package com.nvidia.spark.rapids import scala.annotation.tailrec -import ai.rapids.cudf.{ColumnVector, Scalar, Table} -import ai.rapids.cudf.ast.CompiledExpression +import ai.rapids.cudf.{Scalar, Table} +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._ @@ -40,7 +40,7 @@ trait GpuProjectAstExpressionBase protected def backendName: String protected def compileNvtxId: NvtxId protected def computeNvtxId: NvtxId - protected def evaluate(compiled: CompiledExpression, table: Table): ColumnVector + protected def compileAst(ast: AstExpression): CompiledExpression = ast.compile() @transient private[this] var compiledExpression: CompiledExpression = _ private[this] var opTime: GpuMetric = NoopMetric @@ -72,7 +72,7 @@ trait GpuProjectAstExpressionBase private[rapids] final def computeColumn(table: Table): GpuColumnVector = { val compiled = getCompiledExpression withComputeMetrics { - closeOnExcept(evaluate(compiled, table)) { result => + closeOnExcept(compiled.computeColumn(table)) { result => GpuColumnVector.from(result, dataType) } } @@ -85,7 +85,7 @@ trait GpuProjectAstExpressionBase if (compiledExpression == null) { val compiled = NvtxIdWithMetrics(compileNvtxId, opTime) { // Force every bound reference to the left table; Project AST has one input table. - child.convertToAst(Int.MaxValue).compile() + compileAst(child.convertToAst(Int.MaxValue)) } closeOnExcept(compiled) { _ => var completed = false @@ -254,9 +254,5 @@ case class GpuProjectAstExpression(child: GpuExpression) override protected def computeNvtxId: NvtxId = NvtxRegistry.PROJECT_AST - override protected def evaluate( - compiled: CompiledExpression, - table: Table): ColumnVector = compiled.computeColumn(table) - override def toString: String = s"AST($child)" } 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 b5a5506d51e..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 @@ -1093,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/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala index c4e59d73377..95024aaa7e1 100644 --- a/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala +++ b/tests/src/test/scala/com/nvidia/spark/rapids/GpuProjectAstJitSuite.scala @@ -36,16 +36,22 @@ class GpuProjectAstJitSuite extends AnyFunSuite { private def alias(expression: GpuExpression, name: String) = GpuAlias(expression, name)() - private def mockCompiledChild(compiled: CompiledExpression): GpuExpression = { + 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) - when(ast.compile()).thenReturn(compiled) + if (jit) { + when(ast.compileJit()).thenReturn(compiled) + } else { + when(ast.compile()).thenReturn(compiled) + } child } private def mockJitExpression(compiled: CompiledExpression): GpuAstJitExpression = - GpuAstJitExpression(mockCompiledChild(compiled)) + GpuAstJitExpression(mockCompiledChild(compiled, jit = true)) private def projectConf( tiered: Boolean = true, @@ -368,14 +374,14 @@ class GpuProjectAstJitSuite extends AnyFunSuite { val ast = mock(classOf[AstExpression]) val compiled = mock(classOf[CompiledExpression]) when(child.convertToAst(Int.MaxValue)).thenReturn(ast) - when(ast.compile()).thenReturn(compiled) + when(ast.compileJit()).thenReturn(compiled) val jit = GpuAstJitExpression(child) TestUtils.withMockTaskContext() { jit.checkpoint() jit.restore() jit.checkpoint() - verify(ast, times(1)).compile() + verify(ast, times(1)).compileJit() verify(compiled, times(0)).close() } verify(compiled, times(1)).close() 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/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinRetrySuite.scala b/tests/src/test/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinRetrySuite.scala index 909f0db394a..065d04af737 100644 --- 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 @@ -124,7 +124,7 @@ class GpuBroadcastNestedLoopJoinRetrySuite extends RmmSparkRetrySuiteBase { assert(jitExpressions.head.child.find(_.isInstanceOf[GpuAdd]).isDefined) val projectBuildSide = join.buildSidePostProjection.get - // Warm up computeColumnJit so first-use JIT setup cannot consume the injected OOM; the + // Warm up AST JIT so first-use setup cannot consume the injected OOM; the // next call exercises retry during query execution. withResource(projectBuildSide(buildBatch())) { _ => }