diff --git a/integration_tests/src/main/python/ast_test.py b/integration_tests/src/main/python/ast_test.py index 09689ac1cb7..d96b836efda 100644 --- a/integration_tests/src/main/python/ast_test.py +++ b/integration_tests/src/main/python/ast_test.py @@ -70,11 +70,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, @@ -82,6 +83,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) @@ -105,8 +109,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): @@ -120,22 +124,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 @@ -262,9 +260,19 @@ 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'))) + +# 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(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, @@ -273,7 +281,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, @@ -282,7 +291,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, @@ -291,7 +301,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, @@ -300,7 +311,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, @@ -309,7 +321,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, @@ -430,14 +443,30 @@ 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..79b5f8ed13a --- /dev/null +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuProjectAstExpression.scala @@ -0,0 +1,202 @@ +/* + * 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 +import org.apache.spark.sql.types.DataType +import org.apache.spark.sql.vectorized.ColumnarBatch + +object GpuProjectAstExpression { + private def replaceChild(alias: GpuAlias, child: Expression): GpuAlias = { + if (child eq alias.child) { + alias + } else { + GpuAlias(child, alias.name)(alias.exprId, alias.qualifier, alias.explicitMetadata) + } + } + + 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 + 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 + case other => other + } + + private def rewrap(expression: Expression): Expression = expression match { + case namedExpression: NamedExpression => wrap(namedExpression) + case other => other + } + + private def rewrapAstTiers( + tiers: Seq[Seq[Expression]], + astOutputs: Seq[Boolean]): Seq[Seq[Expression]] = { + val finalTier = tiers.last + require(finalTier.size == astOutputs.size, + "The final expression tier must preserve the project output count") + val astReferences = finalTier.iterator.zip(astOutputs.iterator) + .collect { case (expression, true) => expression } + .flatMap(_.references.iterator) + .map(_.exprId) + .toSet + + // Tier aliases are the dataflow graph after CSE, so follow them backwards from AST outputs. + val (commonTiers, _) = tiers.dropRight(1).foldRight( + (List.empty[Seq[Expression]], astReferences)) { + case (tier, (rewrittenTiers, requiredExprIds)) => + val astAliases = tier.collect { + case alias: GpuAlias if requiredExprIds.contains(alias.exprId) => alias + } + val astAliasIds = astAliases.iterator.map(_.exprId).toSet + val dependencies = astAliases.iterator + .flatMap(_.references.iterator) + .map(_.exprId) + .toSet + val rewrittenTier = tier.map { + case alias: GpuAlias + if astAliasIds.contains(alias.exprId) && + GpuBatchUtils.isFixedWidth(alias.dataType) => + rewrap(alias) + case expression => expression + } + (rewrittenTier :: rewrittenTiers, requiredExprIds ++ dependencies) + } + + commonTiers :+ finalTier.zip(astOutputs).map { + case (expression, true) => rewrap(expression) + case (expression, false) => expression + } + } + + private[rapids] def buildExprTiers( + expressions: Seq[Expression], + conf: SQLConf): Seq[Seq[Expression]] = { + val astOutputs = expressions.map(extractTopLevel(_).isDefined) + val hasAstOutputs = astOutputs.contains(true) + // CSE must see through the marker so AST and non-AST outputs can share the same tiers. + val unwrapped = if (hasAstOutputs) expressions.map(unwrap) else expressions + val replaced = if (RapidsConf.ENABLE_COMBINED_EXPRESSIONS.get(conf)) { + GpuEquivalentExpressions.replaceMultiExpressions(unwrapped, conf) + } else { + unwrapped + } + val tiers = GpuEquivalentExpressions.getExprTiers(replaced) + if (hasAstOutputs) { + rewrapAstTiers(tiers, astOutputs) + } else { + tiers + } + } + + 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) + } + } + } + } +} + +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) + } + } + + private[rapids] def computeColumn(table: Table): GpuColumnVector = { + val compiled = getCompiledExpression + NvtxIdWithMetrics(NvtxRegistry.PROJECT_AST, opTime) { + closeOnExcept(compiled.computeColumn(table)) { result => + GpuColumnVector.from(result, dataType) + } + } + } + + private def getCompiledExpression: CompiledExpression = synchronized { + if (compiledExpression == null) { + val compiled = NvtxIdWithMetrics(NvtxRegistry.COMPILE_ASTS, opTime) { + // Force every bound reference to the left table; Project AST has one input table. + child.convertToAst(Int.MaxValue).compile() + } + closeOnExcept(compiled) { _ => + Option(TaskContext.get()).foreach { taskContext => + onTaskCompletion(taskContext) { + close() + } + } + compiledExpression = compiled + } + } + compiledExpression + } +} 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 58a232c564c..30d10dbee8d 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 @@ -1289,13 +1289,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 @@ -1305,8 +1299,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 } @@ -1348,10 +1342,19 @@ abstract class BaseExprMeta[INPUT <: Expression]( } protected def willWorkInAstInfo: String = { - if (cannotBeAstReasons.isEmpty) { - "will run in AST" + if (canSelfBeAst) { + "is AST-compatible" } else { - s"cannot be converted to GPU AST because ${cannotBeAstReasons.mkString(";")}" + // These reasons must enumerate exactly the conditions checked by canSelfBeAst. + 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" } } @@ -1362,7 +1365,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 = { - if (all || !canThisBeAst) { + 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 d70e96b1dae..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,30 +55,44 @@ 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 val gpuChild = childPlans.head.convertIfNeeded() - if (conf.isProjectAstEnabled) { - // cuDF requires return column is fixed width - val allReturnTypesFixedWidth = gpuExprs.forall(e => GpuBatchUtils.isFixedWidth(e.dataType)) - if (allReturnTypesFixedWidth && childExprs.forall(_.canThisBeAst)) { - return GpuProjectAstExec(gpuExprs, gpuChild) - } + val projectList = if (conf.isProjectAstEnabled) { + val astExprs = childExprs.zip(gpuExprs).map { case (meta, expr) => + // cuDF requires return column is fixed width + // Regular projection can reuse its cached null vector across outputs. + if (GpuBatchUtils.isFixedWidth(expr.dataType) && meta.canThisBeAst && + !isTopLevelNullLiteral(expr)) { + GpuProjectAstExpression.wrap(expr) + } else { + expr + } + }.toList // explain AST because this is optional and it is sometimes hard to debug if (conf.shouldExplain) { - val explain = childExprs.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" + 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") } - if (!allReturnTypesFixedWidth) { - logWarning(s"AST PROJECT\n have non fixed return column, " + - s"return types: ${gpuExprs.map(_.dataType)}") - } } + astExprs + } else { + gpuExprs } - GpuProjectExec(gpuExprs, gpuChild) + GpuProjectExec(projectList, gpuChild) } } @@ -132,9 +146,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 projectWithEval(evalColumn: Expression => ColumnVector): ColumnarBatch = { + GpuArrayHofFusion.project(cb, boundExprs, evalColumn).getOrElse { + val newColumns = boundExprs.safeMap(evalColumn).toArray[ColumnVector] + new ColumnarBatch(newColumns, cb.numRows()) + } + } + + val hasAstExpressions = boundExprs.exists { expression => + GpuProjectAstExpression.extractTopLevel(expression).isDefined + } + if (hasAstExpressions) { + withResource(GpuProjectAstExpression.tableFromBatch(cb)) { table => + projectWithEval { expression => + GpuProjectAstExpression.extractTopLevel(expression) match { + case Some(astExpression) => astExpression.computeColumn(table) + case None => expression.columnarEval(cb) + } + } + } + } else { + projectWithEval(_.columnarEval(cb)) } } finally { GpuExpressionsUtils.cachedNullVectors.get.clear() @@ -913,106 +945,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..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 @@ -375,13 +375,14 @@ private[rapids] object GpuArrayHofFusion { private[rapids] def project( batch: ColumnarBatch, - boundExprs: Seq[Expression]): Option[ColumnarBatch] = { + boundExprs: Seq[Expression], + 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)) + Some(projectWithFusedGroups(batch, boundExprs, groupsByStartIndex, evalColumn)) } } @@ -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], + evalColumn: 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) = evalColumn(boundExprs(index)) } } } 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)() 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..b6aaac3def3 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,40 @@ 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"), + failOnError = false)())) // Prevent side effects in ANSI mode. + 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..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 @@ -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,106 @@ 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)() + // [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"), + 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()) + + // 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) + } + + 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 ffd6240200a..63ad3022993 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 @@ -198,15 +198,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 = "GpuProjectAstExec", - nonExistClasses = "GpuProjectExec" + existClasses = "GpuProjectExec", + nonExistClasses = "GpuProjectAstExpression" ) { df => { df.selectExpr("+c_year_month1")