Keep the TableCacheQueryStageExec on the CPU but convert the child plan

integration_tests/src/main/python/cache_test.py

@@ -17,10 +17,11 @@
 from asserts import assert_gpu_and_cpu_are_equal_collect, assert_equal
 from conftest import is_not_utc
 from data_gen import *
+from pyspark import StorageLevel
 import pyspark.sql.functions as f
-from spark_session import with_cpu_session, with_gpu_session, is_before_spark_330
+from spark_session import with_cpu_session, with_gpu_session, is_before_spark_330, is_spark_350_or_351
 from join_test import create_df
-from marks import incompat, allow_non_gpu, ignore_order, disable_ansi_mode
+from marks import incompat, allow_non_gpu, allow_non_gpu_conditional, ignore_order, disable_ansi_mode
 import pyspark.mllib.linalg as mllib
 import pyspark.ml.linalg as ml
 
@@ -361,4 +362,61 @@ def test_inmem_cache_count():
 
 @pytest.mark.parametrize('with_x_session', [with_gpu_session, with_cpu_session])
 def test_batch_no_cols(with_x_session):
-    function_to_test_on_df(with_x_session, lambda spark: unary_op_df(spark, int_gen).drop("a"), lambda df: df.count(), test_conf={})
+    function_to_test_on_df(with_x_session, lambda spark: unary_op_df(spark, int_gen).drop("a"), lambda df: df.count(), test_conf={})
+
+@ignore_order(local=True)
+@allow_non_gpu("ShuffleExchangeExec", "ColumnarToRowExec")
+@allow_non_gpu_conditional(is_spark_350_or_351(), "InMemoryTableScanExec")
+@pytest.mark.parametrize("data_gen", integral_gens, ids=idfn)
+@pytest.mark.parametrize('enable_vectorized_conf', enable_vectorized_confs, ids=idfn)
+def test_aqe_cache_version_specific_behavior(data_gen, enable_vectorized_conf):
+    """
+    Test InMemoryTableScan + AQE behavior across Spark versions.
+    - Spark 3.2.0-3.4.x: InMemoryTableScan works on GPU
+    - Spark 3.5.0-3.5.1: InMemoryTableScan disabled by default due to missing InMemoryTableScanLike trait
+    - Spark 3.5.2+: InMemoryTableScan works on GPU with proper trait support
+    """
+
+    def do_it(spark):
+        df1 = unary_op_df(spark, data_gen).orderBy('a').cache()
+        df2 = unary_op_df(spark, data_gen).withColumnRenamed("a", "r_a").cache()
+        df1.count()
+        df2.count()
+        return df1.join(df2, df1.a == df2.r_a, 'Outer')
+
+    assert_gpu_and_cpu_are_equal_collect(do_it, conf=enable_vectorized_conf)
+
+@ignore_order(local=True)
+@allow_non_gpu("CollectLimitExec", "ShuffleExchangeExec", "ColumnarToRowExec")
+@pytest.mark.parametrize('enable_vectorized_conf', enable_vectorized_confs, ids=idfn)
+@allow_non_gpu_conditional(is_spark_350_or_351(), "InMemoryTableScanExec")
+def test_persist_with_groupby_join_version_specific(enable_vectorized_conf):
+    """
+    Expected behavior:
+    - Spark 3.2.0-3.4.x: InMemoryTableScan works on GPU
+    - Spark 3.5.0-3.5.1: InMemoryTableScan falls back to CPU due to missing InMemoryTableScanLike trait
+    - Spark 3.5.2+: InMemoryTableScan works on GPU with proper trait support
+    """
+
+    def do_it(spark):
+        df = spark.range(0, 1000, 1, 2).select(
+            f.col("id").alias("_1"),
+            f.col("id").alias("_2")
+        )
+
+        ee = df.select(
+            f.col("_1").alias("src"),
+            f.col("_2").alias("dst")
+        ).persist(StorageLevel.MEMORY_AND_DISK)
+        ee.count()
+
+        minNbrs1 = ee.groupBy("src").agg(
+            f.min(f.col("dst")).alias("min_number")
+        ).persist(StorageLevel.MEMORY_AND_DISK)
+        minNbrs1.count()
+
+        ee.join(minNbrs1, "src")
+
+        return ee.join(minNbrs1, "src")
+
+    assert_gpu_and_cpu_are_equal_collect(do_it, conf=enable_vectorized_conf)

integration_tests/src/main/python/spark_session.py

@@ -213,6 +213,9 @@ def is_before_spark_351():
 def is_before_spark_353():
     return spark_version() < "3.5.3"
 
+def is_spark_350_or_351():
+    return spark_version() >= "3.5.0" and spark_version() <= "3.5.1"
+
 def is_before_spark_400():
     return spark_version() < "4.0.0"
 

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

@@ -235,7 +235,10 @@ class GpuTransitionOverrides extends Rule[SparkPlan] {
       p.withNewChildren(Array(newChild))
 
     case p =>
-      p.withNewChildren(p.children.map(c => optimizeAdaptiveTransitions(c, Some(p))))
+      SparkShimImpl.handleTableCacheInOptimizeAdaptiveTransitions(p, parent) match {
+        case Some(handledPlan) => handledPlan
+        case None => p.withNewChildren(p.children.map(c => optimizeAdaptiveTransitions(c, Some(p))))
+      }
   }
 
   /**
@@ -869,7 +872,8 @@ object GpuTransitionOverrides {
         } else {
           sqse.plan
         }
-      case _ => plan
+      case _ =>
+        SparkShimImpl.getTableCacheNonQueryStagePlan(plan).getOrElse(plan)
     }
   }
 

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

@@ -196,4 +196,17 @@ trait SparkShims {
    * Handle regexp_replace inconsistency from https://issues.apache.org/jira/browse/SPARK-39107
    */
   def reproduceEmptyStringBug: Boolean
+
+  /**
+   * Handle TableCacheQueryStageExec for optimizeAdaptiveTransitions.
+   * Returns the original plan for versions where TableCacheQueryStageExec doesn't exist.
+   */
+  def handleTableCacheInOptimizeAdaptiveTransitions(plan: SparkPlan,
+      parent: Option[SparkPlan]): Option[SparkPlan] = None
+
+  /**
+   * Handle TableCacheQueryStageExec for getNonQueryStagePlan.
+   * Returns None for versions where TableCacheQueryStageExec doesn't exist.
+   */
+  def getTableCacheNonQueryStagePlan(plan: SparkPlan): Option[SparkPlan] = None
 }

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/GpuInMemoryTableScanExec.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,7 +18,7 @@ package org.apache.spark.sql.rapids
 
 import com.nvidia.spark.ParquetCachedBatchSerializer
 import com.nvidia.spark.rapids.{DataFromReplacementRule, ExecChecks, GpuExec, GpuMetric, RapidsConf, RapidsMeta, SparkPlanMeta}
-import com.nvidia.spark.rapids.shims.ShimLeafExecNode
+import com.nvidia.spark.rapids.shims.{InMemoryTableScanExecLikeShim, ShimLeafExecNode}
 
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
@@ -76,7 +76,8 @@ class InMemoryTableScanMeta(
 case class GpuInMemoryTableScanExec(
    attributes: Seq[Attribute],
    predicates: Seq[Expression],
-   @transient relation: InMemoryRelation) extends ShimLeafExecNode with GpuExec {
+   @transient relation: InMemoryRelation) extends ShimLeafExecNode with
+   GpuExec with InMemoryTableScanExecLikeShim {
 
   override val nodeName: String = {
     relation.cacheBuilder.tableName match {

sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/InMemoryTableScanExecLikeShim.scala

@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) 2025, 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.
+ */
+
+/*** spark-rapids-shim-json-lines
+{"spark": "320"}
+{"spark": "321"}
+{"spark": "321cdh"}
+{"spark": "322"}
+{"spark": "323"}
+{"spark": "324"}
+{"spark": "330"}
+{"spark": "330cdh"}
+{"spark": "330db"}
+{"spark": "331"}
+{"spark": "332"}
+{"spark": "332cdh"}
+{"spark": "332db"}
+{"spark": "333"}
+{"spark": "334"}
+{"spark": "340"}
+{"spark": "341"}
+{"spark": "341db"}
+{"spark": "342"}
+{"spark": "343"}
+{"spark": "344"}
+{"spark": "350"}
+{"spark": "350db143"}
+{"spark": "351"}
+spark-rapids-shim-json-lines ***/
+package com.nvidia.spark.rapids.shims
+
+trait InMemoryTableScanExecLikeShim {
+   // No additional implementations needed for pre-3.5.2 Spark versions
+   // since InMemoryTableScanLike doesn't exist
+}

sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/shims/InMemoryTableScanUtils.scala

@@ -0,0 +1,59 @@
+/*
+ * Copyright (c) 2025, 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.
+ */
+
+/*** spark-rapids-shim-json-lines
+{"spark": "350"}
+{"spark": "351"}
+spark-rapids-shim-json-lines ***/
+package com.nvidia.spark.rapids.shims
+
+import com.nvidia.spark.rapids.{ExecRule, GpuOverrides}
+
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.TableCacheQueryStageExec
+import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
+
+/**
+ * Utility object for handling InMemoryTableScan version differences.
+ * For Spark 3.5.0 and 3.5.1, we disable InMemoryTableScan by default due to 
+ * missing InMemoryTableScanLike trait.
+ */
+object InMemoryTableScanUtils {
+
+  /**
+   * Modifies the InMemoryTableScan rule to be disabled by default for Spark 3.5.0-3.5.1.
+   */
+  def getInMemoryTableScanExecRule: ExecRule[_ <: SparkPlan] = {
+    val imtsKey = classOf[InMemoryTableScanExec].asSubclass(classOf[SparkPlan])
+    GpuOverrides.commonExecs.getOrElse(imtsKey,
+        throw new IllegalStateException("InMemoryTableScan should be overridden by default before" +
+        " Spark 3.5.0")).
+      disabledByDefault(
+        """there could be complications when using it with AQE with Spark-3.5.0 and Spark-3.5.1.
+          |For more details please check
+          |https://github.com/NVIDIA/spark-rapids/issues/10603""".stripMargin.replaceAll("\n", " "))
+  }
+
+  def canTableCacheWrapGpuInMemoryTableScan: Boolean = false
+
+  /**
+   * Gets the TableCacheQueryStageExec rule for this Spark version.
+   * For Spark 3.5.0-3.5.1: neverReplaceExec because of missing InMemoryTableScanLike trait
+   */
+  def getTableCacheQueryStageExecRule: ExecRule[_ <: SparkPlan] = {
+    GpuOverrides.neverReplaceExec[TableCacheQueryStageExec]("Table cache query stage")
+  }
+}

sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/shims/Spark350PlusNonDBShims.scala

@@ -34,12 +34,51 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression, PythonUDAF, ToPrettyString}
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.adaptive.TableCacheQueryStageExec
-import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
 import org.apache.spark.sql.execution.datasources.{FileFormat, FilePartition, FileScanRDD, PartitionedFile}
 import org.apache.spark.sql.execution.window.WindowGroupLimitExec
 import org.apache.spark.sql.rapids.execution.python.GpuPythonUDAF
 import org.apache.spark.sql.types.{StringType, StructType}
 
+class TableCacheQueryStageExecMeta(
+    tcqs: TableCacheQueryStageExec,
+    conf: RapidsConf,
+    parent: Option[RapidsMeta[_, _, _]],
+    rule: DataFromReplacementRule)
+    extends SparkPlanMeta[TableCacheQueryStageExec](tcqs, conf, parent, rule) {
+
+  override val childPlans: Seq[SparkPlanMeta[SparkPlan]] =
+    Seq(GpuOverrides.wrapPlan(tcqs.plan, conf, Some(this)))
+
+  override def tagPlanForGpu(): Unit = {
+    willNotWorkOnGpu("TableCacheQueryStageExec wrapper stays on CPU for Spark AQE compatibility; " +
+      "child plan may run on GPU")
+  }
+
+  override def convertToGpu(): GpuExec = {
+    throw new IllegalStateException("TableCacheQueryStageExec should not be converted to GPU")
+  }
+
+  override def convertToCpu(): SparkPlan = {
+    val wrappedPlan = childPlans.head.convertIfNeeded()
+
+    // If the wrapped plan wasn't converted, return the original TableCacheQueryStageExec
+    if (wrappedPlan == tcqs.plan) {
+      return tcqs
+    }
+
+    // The wrapped plan was converted to GPU - check if we can safely wrap it
+    if (InMemoryTableScanUtils.canTableCacheWrapGpuInMemoryTableScan) {
+      // For Spark 3.5.2+: GPU InMemoryTableScan implements InMemoryTableScanLike,
+      // so TableCacheQueryStageExec can safely wrap it and pass Spark's validation
+      tcqs.copy(plan = wrappedPlan)
+    } else {
+      // For Spark 3.5.0-3.5.1: Missing InMemoryTableScanLike trait causes validation issues.
+      // Keep the original CPU plan to avoid AQE complications.
+      tcqs
+    }
+  }
+}
+
 trait Spark350PlusNonDBShims extends Spark340PlusNonDBShims {
   override def getFileScanRDD(
       sparkSession: SparkSession,
@@ -103,19 +142,9 @@ trait Spark350PlusNonDBShims extends Spark340PlusNonDBShims {
   }
 
   override def getExecs: Map[Class[_ <: SparkPlan], ExecRule[_ <: SparkPlan]] = {
-    val imtsKey = classOf[InMemoryTableScanExec].asSubclass(classOf[SparkPlan])
-    // To avoid code duplication we are reusing the rule from GpuOverrides
-    // but we disable it by default
-    val imtsRule = GpuOverrides.commonExecs.getOrElse(imtsKey,
-        throw new IllegalStateException("InMemoryTableScan should be overridden by default before" +
-        " Spark 3.5.0")).
-      disabledByDefault(
-        """there could be complications when using it with AQE with Spark-3.5.0 and Spark-3.5.1.
-          |For more details please check
-          |https://github.com/NVIDIA/spark-rapids/issues/10603""".stripMargin.replaceAll("\n", " "))
-
     val shimExecs: Map[Class[_ <: SparkPlan], ExecRule[_ <: SparkPlan]] = Seq(
-      imtsRule,
+      // Use version-specific InMemoryTableScan rule (disabledByDefault for 3.5.0-3.5.1)
+      InMemoryTableScanUtils.getInMemoryTableScanExecRule,
       GpuOverrides.exec[WindowGroupLimitExec](
         "Apply group-limits for row groups destined for rank-based window functions like " +
           "row_number(), rank(), and dense_rank()",
@@ -124,8 +153,23 @@ trait Spark350PlusNonDBShims extends Spark340PlusNonDBShims {
             TypeSig.STRUCT + TypeSig.ARRAY + TypeSig.MAP).nested(),
           TypeSig.all),
         (limit, conf, p, r) => new GpuWindowGroupLimitExecMeta(limit, conf, p, r)),
-      GpuOverrides.neverReplaceExec[TableCacheQueryStageExec]("Table cache query stage")
+      InMemoryTableScanUtils.getTableCacheQueryStageExecRule
     ).map(r => (r.getClassFor.asSubclass(classOf[SparkPlan]), r)).toMap
     super.getExecs ++ shimExecs
   }
+
+  override def handleTableCacheInOptimizeAdaptiveTransitions(plan: SparkPlan,
+      parent: Option[SparkPlan]): Option[SparkPlan] = {
+    plan match {
+      case tcqs: TableCacheQueryStageExec => Some(tcqs)
+      case _ => None
+    }
+  }
+
+  override def getTableCacheNonQueryStagePlan(plan: SparkPlan): Option[SparkPlan] = {
+    plan match {
+      case tcqs: TableCacheQueryStageExec => Some(tcqs.plan)
+      case _ => None
+    }
+  }
 }