Fix incorrect GPU partitioning and ordering metadata [databricks]

delta-lake/common/src/main/db-350db143-410db183/scala/com/databricks/sql/transaction/tahoe/rapids/GpuOptimisticTransactionBaseCommon.scala

@@ -107,7 +107,7 @@ abstract class GpuOptimisticTransactionBaseCommon
         GpuAlias(GpuEmpty2Null(p), p.name)()
       case attr => attr
     }
-    if (needConvert) GpuProjectExec(projectList.toList, plan) else plan
+    if (needConvert) GpuProjectExec(projectList.toList, plan)() else plan
   }
 
   /**

delta-lake/common/src/main/db-400db173-410db183/scala/com/nvidia/spark/rapids/delta/DatabricksDeltaProvider173183Base.scala

@@ -193,7 +193,7 @@ abstract class DatabricksDeltaProvider173183Base extends DatabricksDeltaProvider
     def pruneMetadataProject(
         project: GpuProjectExec,
         scan: GpuFileSourceScanExec): SparkPlan = {
-      project.copy(projectList = project.projectList.filterNot(_.name == "_metadata"))
+      project.withProjectList(project.projectList.filterNot(_.name == "_metadata"))
         .withNewChildren(Seq(
           scan.copy(
             originalOutput = scan.originalOutput.filterNot(_.name == "_tmp_metadata_row_index"),
@@ -441,7 +441,7 @@ private object DatabricksDVPredicatePushdown extends ShimPredicateHelper {
         case project @ ProjectExec(projectList, _) =>
           project.copy(projectList = projectList.filterNot(isDeletionVectorSkipRowColumnRef))
         case project @ GpuProjectExec(projectList, _, _) =>
-          project.copy(projectList = projectList.filterNot(isDeletionVectorSkipRowColumnRef))
+          project.withProjectList(projectList.filterNot(isDeletionVectorSkipRowColumnRef))
         case fsse: FileSourceScanExec =>
           fsse.copy(
             output = fsse.output.filterNot(attr => isDeletionVectorSkipRowColumn(attr.name)),
@@ -531,7 +531,8 @@ private object DatabricksDVPredicatePushdown extends ShimPredicateHelper {
         // identical projections: merging over an alias-computing child would drop the
         // alias's only producer ("Couldn't find <attr>"). Merge only pure pass-throughs.
         if (projSet1 == projSet2 && projList2.forall(_.isInstanceOf[AttributeReference])) {
-          GpuProjectExec(projList1, child, enablePreSplit1 && enablePreSplit2)
+          p.withProjectList(projList1, newChild = child,
+            preSplit = enablePreSplit1 && enablePreSplit2)
         } else {
           p
         }

delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/DeltaProviderBase.scala

@@ -178,7 +178,7 @@ abstract class DeltaProviderBase extends DeltaIOProvider {
       // Data and partition columns precede metadata columns in the scan output.
       val dataAndPartitionOutput = scan.originalOutput.take(
         scan.requiredSchema.length + scan.relation.partitionSchema.length)
-      project.copy(projectList = project.projectList.filterNot(_.name == "_metadata"))
+      project.withProjectList(project.projectList.filterNot(_.name == "_metadata"))
         .withNewChildren(Seq(
           scan.copy(
             // Drop the temporary row index along with the unused metadata struct.
@@ -333,7 +333,7 @@ object DVPredicatePushdown extends ShimPredicateHelper {
       plan.transformUp {
         case project @ GpuProjectExec(projectList, _, _) =>
           val newProjList = projectList.filterNot(isRowDeletedColumnRef(_))
-          project.copy(projectList = newProjList)
+          project.withProjectList(newProjList)
         case fsse: GpuFileSourceScanExec =>
           fsse.copy(originalOutput = fsse.originalOutput.filterNot(
             _.name == IS_ROW_DELETED_COLUMN_NAME),
@@ -430,7 +430,8 @@ object DVPredicatePushdown extends ShimPredicateHelper {
         // identical projections: merging over an alias-computing child would drop the
         // alias's only producer ("Couldn't find <attr>"). Merge only pure pass-throughs.
         if (projSet1 == projSet2 && projList2.forall(_.isInstanceOf[AttributeReference])) {
-          GpuProjectExec(projList1, child, enablePreSplit1 && enablePreSplit2)
+          p.withProjectList(projList1, newChild = child,
+            preSplit = enablePreSplit1 && enablePreSplit2)
         } else {
           p
         }

delta-lake/common/src/main/delta-io/scala/org/apache/spark/sql/delta/rapids/AbstractGpuOptimisticTransactionBase.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * Copyright (c) 2023-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.
@@ -92,7 +92,7 @@ abstract class AbstractGpuOptimisticTransactionBase(
         GpuAlias(GpuEmpty2Null(p), p.name)()
       case attr => attr
     }
-    if (needConvert) GpuProjectExec(projectList.toList, plan) else plan
+    if (needConvert) GpuProjectExec(projectList.toList, plan)() else plan
   }
 
   /**

delta-lake/delta-spark330db/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuOptimisticTransactionBase.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.
  *
  * This file was derived from OptimisticTransaction.scala and TransactionalWrite.scala
  * in the Delta Lake project at https://github.com/delta-io/delta.
@@ -101,7 +101,7 @@ abstract class GpuOptimisticTransactionBase
         GpuAlias(GpuEmpty2Null(p), p.name)()
       case attr => attr
     }
-    if (needConvert) GpuProjectExec(projectList.toList, plan) else plan
+    if (needConvert) GpuProjectExec(projectList.toList, plan)() else plan
   }
 
   /**

delta-lake/delta-spark332db/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuOptimisticTransactionBase.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.
  *
  * This file was derived from OptimisticTransaction.scala and TransactionalWrite.scala
  * in the Delta Lake project at https://github.com/delta-io/delta.
@@ -101,7 +101,7 @@ abstract class GpuOptimisticTransactionBase
         GpuAlias(GpuEmpty2Null(p), p.name)()
       case attr => attr
     }
-    if (needConvert) GpuProjectExec(projectList.toList, plan) else plan
+    if (needConvert) GpuProjectExec(projectList.toList, plan)() else plan
   }
 
   /**

integration_tests/src/main/python/cache_test.py

@@ -291,6 +291,121 @@ def helper(spark):
 
     assert_equal(reg_result, cached_result)
 
+@pytest.mark.parametrize('aqe_enabled', ['true', 'false'])
+@allow_non_gpu('InMemoryTableScanExec', 'ColumnarToRowExec')
+def test_cache_project_partitioning(aqe_enabled):
+    cached_frames = []
+
+    def query(spark):
+        cached = (spark.range(0, 100, 1, 16).selectExpr('id % 10 AS k', 'id AS v')
+            .persist(StorageLevel.MEMORY_AND_DISK))
+        cached_frames.append(cached)
+        return cached.groupBy('k').agg(f.max('v').alias('max_v')).orderBy('k')
+
+    conf = {
+        'spark.sql.adaptive.enabled': aqe_enabled,
+        'spark.sql.shuffle.partitions': '4'}
+    try:
+        assert_gpu_and_cpu_are_equal_collect(query, conf=conf)
+    finally:
+        for cached in cached_frames:
+            cached.unpersist(blocking=True)
+
+
+def _assert_cached_project_aggregate_equal(projected_input, aqe_enabled, conf=None):
+    # Each input hashes by id, then projects id % 10 as k. The cached aggregate must
+    # shuffle again to combine equal k values from different input partitions.
+    cached_frames = []
+
+    def query(spark):
+        cached = projected_input(spark).persist(StorageLevel.MEMORY_AND_DISK)
+        cached_frames.append(cached)
+        return cached.groupBy('k').agg(f.max('v').alias('max_v')).orderBy('k')
+
+    test_conf = copy_and_update({
+        'spark.sql.adaptive.enabled': aqe_enabled,
+        # Keep multiple partitions under AQE so coalescing cannot hide a missing shuffle.
+        'spark.sql.adaptive.coalescePartitions.enabled': 'false',
+        'spark.sql.shuffle.partitions': '4'}, conf or {})
+    try:
+        assert_gpu_and_cpu_are_equal_collect(query, conf=test_conf)
+    finally:
+        for cached in cached_frames:
+            cached.unpersist(blocking=True)
+
+
+@pytest.mark.parametrize('aqe_enabled', ['true', 'false'])
+@allow_non_gpu('InMemoryTableScanExec', 'ColumnarToRowExec')
+def test_cache_project_hash_partitioning(aqe_enabled):
+    def projected_input(spark):
+        source = gen_df(spark, [('id', LongGen(min_val=0, max_val=99, nullable=False))],
+            length=256, seed=0)
+        # Compute id on the GPU to keep the hash shuffle on the GPU.
+        partitioned = source.selectExpr('id + 1 AS id').repartition(4, 'id')
+        projected = partitioned.selectExpr('id % 10 AS k', 'id AS v')
+        if spark.conf.get('spark.rapids.sql.enabled') == 'true':
+            # AQE converts query stages to GPU during execution; inspect the completed plan.
+            assert projected.collect(), 'Expected non-empty projected input'
+            capture = spark._jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback
+            capture.assertContains(projected._jdf, 'GpuProjectExec')
+            capture.assertContains(projected._jdf, 'GpuShuffleExchangeExec')
+        return projected
+
+    _assert_cached_project_aggregate_equal(projected_input, aqe_enabled)
+
+
+@pytest.mark.parametrize('aqe_enabled', ['true', 'false'])
+@allow_non_gpu('InMemoryTableScanExec', 'ColumnarToRowExec')
+def test_cache_project_after_broadcast_join(aqe_enabled):
+    def projected_input(spark):
+        left = gen_df(spark, [('id', LongGen(min_val=0, max_val=99, nullable=False))],
+            length=256, seed=0).selectExpr('id + 1 AS id').repartition(4, 'id')
+        right = gen_df(spark, [('rid', LongGen(min_val=1, max_val=100, nullable=False))],
+            length=100, seed=1)
+        joined = left.join(f.broadcast(right), left.id == right.rid)
+        projected = joined.selectExpr('id % 10 AS k', 'id AS v')
+        if spark.conf.get('spark.rapids.sql.enabled') == 'true':
+            # AQE converts query stages to GPU during execution; inspect the completed plan.
+            assert projected.collect(), 'Expected non-empty projected input'
+            capture = spark._jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback
+            capture.assertContains(projected._jdf, 'GpuProjectExec')
+            capture.assertContains(projected._jdf, 'GpuShuffleExchangeExec')
+            capture.assertContains(projected._jdf, 'GpuBroadcastHashJoinExec')
+        return projected
+
+    _assert_cached_project_aggregate_equal(projected_input, aqe_enabled)
+
+
+@pytest.mark.parametrize('aqe_enabled', ['true', 'false'])
+@allow_non_gpu('InMemoryTableScanExec', 'ColumnarToRowExec')
+def test_cache_project_after_split_condition_join(aqe_enabled):
+    def projected_input(spark):
+        left = gen_df(spark, [
+            ('id', LongGen(min_val=0, max_val=99, nullable=False)),
+            ('probe', LongGen(min_val=0, max_val=99, nullable=False))],
+            length=256, seed=0).selectExpr('id', 'probe + 1 AS probe').repartition(4, 'id')
+        right = gen_df(spark, [
+            ('rid', LongGen(min_val=0, max_val=99, nullable=False)),
+            ('offset', LongGen(min_val=-1, max_val=1, nullable=False))], length=100, seed=1)
+        # Force Add out of AST so the condition requires a pre-projection on the right input.
+        condition = (left.id == right.rid) & (left.probe < right.rid + right['offset'])
+        joined = left.join(f.broadcast(right), condition)
+        projected = joined.selectExpr('id % 10 AS k', 'id AS v')
+        if spark.conf.get('spark.rapids.sql.enabled') == 'true':
+            # AQE converts query stages to GPU during execution; inspect the completed plan.
+            assert projected.collect(), 'Expected non-empty projected input'
+            capture = spark._jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback
+            capture.assertContains(projected._jdf, 'GpuProjectExec')
+            capture.assertContains(projected._jdf, 'GpuShuffleExchangeExec')
+            capture.assertContains(projected._jdf, 'GpuBroadcastHashJoinExec')
+            plan = projected._jdf.queryExecution().executedPlan().toString()
+            assert '_agpu_non_ast' in plan, 'Expected join-condition pre-projection:\n' + plan
+        return projected
+
+    _assert_cached_project_aggregate_equal(projected_input, aqe_enabled,
+        conf={'spark.rapids.sql.test.forceNonAst': 'Add'})
+
+
 @pytest.mark.parametrize('data_gen', all_gen, ids=idfn)
 @pytest.mark.parametrize('with_x_session', [with_gpu_session, with_cpu_session])
 @pytest.mark.parametrize('enable_vectorized_conf', enable_vectorized_confs, ids=idfn)

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

@@ -28,7 +28,7 @@ import com.nvidia.spark.rapids.GpuOverrides.pluginSupportedOrderableSig
 import com.nvidia.spark.rapids.RapidsPluginImplicits._
 import com.nvidia.spark.rapids.RmmRapidsRetryIterator.{splitSpillableInHalfByRows, withRetry, withRetryNoSplit}
 import com.nvidia.spark.rapids.ScalableTaskCompletion.onTaskCompletion
-import com.nvidia.spark.rapids.shims.AggregationTagging
+import com.nvidia.spark.rapids.shims.{AggregationTagging, GpuPartitioningPreservingUnaryExecNodeShim}
 
 import org.apache.spark.TaskContext
 import org.apache.spark.internal.Logging
@@ -1425,7 +1425,7 @@ abstract class GpuBaseAggregateMeta[INPUT <: SparkPlan](
       allowSinglePassAgg,
       allowNonFullyAggregatedOutput,
       conf.skipAggPassReductionRatio
-    )
+    )(Some(agg.resultExpressions))
   }
 }
 
@@ -1515,7 +1515,7 @@ abstract class GpuTypedImperativeSupportedAggregateExecMeta[INPUT <: BaseAggrega
         forceSinglePassAgg = false,
         allowSinglePassAgg = false,
         allowNonFullyAggregatedOutput = false,
-        1)
+        1)(Some(retExpressions.map(_.wrapped.asInstanceOf[NamedExpression])))
 
     } else {
       super.convertToGpu()
@@ -2023,8 +2023,11 @@ case class GpuHashAggregateExec(
     forceSinglePassAgg: Boolean,
     allowSinglePassAgg: Boolean,
     allowNonFullyAggregatedOutput: Boolean,
-    skipAggPassReductionRatio: Double
-) extends GpuPartitioningPreservingUnaryExecNode with GpuExec with Logging {
+    skipAggPassReductionRatio: Double)(
+    val cpuResultExpressions: Option[Seq[NamedExpression]] = None
+) extends GpuPartitioningPreservingUnaryExecNodeShim with GpuExec with Logging {
+
+  override def otherCopyArgs: Seq[AnyRef] = cpuResultExpressions :: Nil
 
   // lifted directly from `BaseAggregateExec.inputAttributes`, edited comment.
   def inputAttributes: Seq[Attribute] =
@@ -2124,7 +2127,8 @@ case class GpuHashAggregateExec(
     }
   }
 
-  protected def outputExpressions: Seq[NamedExpression] = resultExpressions
+  override protected def cpuOutputExpressions: Seq[NamedExpression] =
+    cpuResultExpressions.getOrElse(GpuAliasAwareOutputExpressions.toCpuAliases(resultExpressions))
 
   //
   // This section is derived (copied in most cases) from HashAggregateExec
@@ -2133,32 +2137,6 @@ case class GpuHashAggregateExec(
     aggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes)
   }
 
-  /**
-   * Override buildAttributeMap to handle aggregate output semantics.
-   */
-  override protected def buildAttributeMap(): Map[Attribute, Attribute] = {
-    resultExpressions.collect {
-      case a @ GpuAlias(child: Attribute, _) => child -> a.toAttribute
-      case a @ Alias(child: Attribute, _) => child -> a.toAttribute
-    }.toMap
-  }
-
-  protected def hasAlias: Boolean = outputExpressions.collectFirst { case _: Alias => }.isDefined
-
-  protected def replaceAliases(exprs: Seq[Expression]): Seq[Expression] = {
-    exprs.map {
-      case a: AttributeReference => replaceAlias(a).getOrElse(a)
-      case other => other
-    }
-  }
-
-  protected def replaceAlias(attr: AttributeReference): Option[Attribute] = {
-    outputExpressions.collectFirst {
-      case a@Alias(child: AttributeReference, _) if child.semanticEquals(attr) =>
-        a.toAttribute
-    }
-  }
-
   // Used in de-duping and optimizer rules.
   // Final/PartialMerge aggregates read input buffer attributes from their child. If those
   // attrs share exprIds with child output but have different names, keep them out of this

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

@@ -0,0 +1,34 @@
+/*
+ * 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.apache.spark.sql.catalyst.expressions.{Alias, NamedExpression}
+
+private[rapids] object GpuAliasAwareOutputExpressions {
+  // Internally generated projections have no original CPU plan. Spark's metadata traits
+  // recognize Alias rather than GpuAlias; preserve the output identity and metadata.
+  def toCpuAliases(expressions: Seq[NamedExpression]): List[NamedExpression] = {
+    expressions.map { expr =>
+      expr.transformDown {
+        case alias: GpuAlias =>
+          Alias(alias.child, alias.name)(
+            alias.exprId, alias.qualifier, Some(alias.metadata))
+        case ast: GpuProjectAstExpression => ast.child
+      }.asInstanceOf[NamedExpression]
+    }.toList
+  }
+}

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

@@ -474,6 +474,8 @@ class RapidsConf(conf: Map[String, String]) extends Logging {
 
   lazy val isTestEnabled: Boolean = get(TEST_CONF)
 
+  lazy val testingForceNonAst: Set[String] = get(TEST_FORCE_NON_AST).toSet
+
   lazy val isRetryContextCheckEnabled: Boolean = get(TEST_RETRY_CONTEXT_CHECK_ENABLED)
 
   lazy val isFoldableNonLitAllowed: Boolean = get(FOLDABLE_NON_LIT_ALLOWED)

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

@@ -76,6 +76,14 @@ private[rapids] trait RapidsConfEntries extends RapidsConfSqlEntries {
       .booleanConf
       .createWithDefault(false)
 
+  val TEST_FORCE_NON_AST = conf("spark.rapids.sql.test.forceNonAst")
+    .doc("Only for tests. Comma-separated Spark expression class names forced out of AST " +
+      s"when ${TEST_CONF.key} is enabled, while retaining GPU expression support.")
+    .internal()
+    .stringConf
+    .toSequence
+    .createWithDefault(Nil)
+
   val TEST_ALLOWED_NONGPU = conf("spark.rapids.sql.test.allowedNonGpu")
     .doc("Comma separate string of exec or expression class names that are allowed " +
       "to not be GPU accelerated for testing.")

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

@@ -1321,6 +1321,9 @@ abstract class BaseExprMeta[INPUT <: Expression](
   private var taggedForAst = false
   private final def tagForAst(): Unit = {
     if (!taggedForAst) {
+      if (conf.isTestEnabled && conf.testingForceNonAst.contains(wrapped.getClass.getSimpleName)) {
+        willNotWorkInAst("AST disabled by test configuration")
+      }
       if (wrapped.foldable && !GpuOverrides.isLit(wrapped)) {
         willNotWorkInAst(s"Cannot convert to AST. Is ConstantFolding excluded? Expression " +
             s"$wrapped is foldable and operates on non literals")