Keep the TableCacheQueryStageExec on the CPU but convert the child plan

delta-lake/common/src/main/delta-33x/scala/org/apache/spark/sql/delta/rapids/GpuOptimizeWriteExchangeExec.scala

@@ -108,27 +108,31 @@ case class GpuOptimizeWriteExchangeExec(
 
   @transient lazy val inputRDD: RDD[ColumnarBatch] = child.executeColumnar()
 
+  @transient private lazy val childNumPartitions = inputRDD.getNumPartitions
+
   @transient lazy val mapOutputStatisticsFuture: Future[MapOutputStatistics] = {
-    if (inputRDD.getNumPartitions == 0) {
+    if (childNumPartitions == 0) {
       Future.successful(null)
     } else {
       sparkContext.submitMapStage(shuffleDependency)
     }
   }
 
-  private lazy val childNumPartitions = inputRDD.getNumPartitions
-
-  private lazy val actualNumPartitions: Int = {
-    val targetShuffleBlocks = conf.getConf(DeltaSQLConf.DELTA_OPTIMIZE_WRITE_SHUFFLE_BLOCKS)
-    math.min(
-      math.max(targetShuffleBlocks / childNumPartitions, 1),
-      conf.getConf(DeltaSQLConf.DELTA_OPTIMIZE_WRITE_MAX_SHUFFLE_PARTITIONS))
+  @transient private lazy val actualNumPartitions: Int = {
+    if (childNumPartitions == 0) {
+      0
+    } else {
+      val targetShuffleBlocks = conf.getConf(DeltaSQLConf.DELTA_OPTIMIZE_WRITE_SHUFFLE_BLOCKS)
+      math.min(
+        math.max(targetShuffleBlocks / childNumPartitions, 1),
+        conf.getConf(DeltaSQLConf.DELTA_OPTIMIZE_WRITE_MAX_SHUFFLE_PARTITIONS))
+    }
   }
 
   // The actual partitioning to use for the shuffle exchange. The input partition count can be
   // adjusted based on the number of partitions in the input RDD and the target number of shuffle
   // blocks.
-  private lazy val actualPartitioning: GpuPartitioning = partitioning match {
+  @transient private lazy val actualPartitioning: GpuPartitioning = partitioning match {
     // Currently only hash and round-robin partitioning are supported.
     // See DeltaShufflePartitionsUtil.partitioningForRebalance() for more details.
     case p: GpuHashPartitioning => p.copy(numPartitions = actualNumPartitions)

docs/additional-functionality/advanced_configs.md

@@ -309,6 +309,7 @@ Name | SQL Function(s) | Description | Default Value | Notes
 <a name="sql.expression.MapEntries"></a>spark.rapids.sql.expression.MapEntries|`map_entries`|Returns an unordered array of all entries in the given map|true|None|
 <a name="sql.expression.MapFilter"></a>spark.rapids.sql.expression.MapFilter|`map_filter`|Filters entries in a map using the function|true|None|
 <a name="sql.expression.MapFromArrays"></a>spark.rapids.sql.expression.MapFromArrays|`map_from_arrays`|Creates a new map from two arrays|true|None|
+<a name="sql.expression.MapFromEntries"></a>spark.rapids.sql.expression.MapFromEntries|`map_from_entries`|Creates a map from an array of entries (structs of key-value pairs)|true|None|
 <a name="sql.expression.MapKeys"></a>spark.rapids.sql.expression.MapKeys|`map_keys`|Returns an unordered array containing the keys of the map|true|None|
 <a name="sql.expression.MapValues"></a>spark.rapids.sql.expression.MapValues|`map_values`|Returns an unordered array containing the values of the map|true|None|
 <a name="sql.expression.MapZipWith"></a>spark.rapids.sql.expression.MapZipWith|`map_zip_with`|Filters entries in a map using the function|true|None|