Fix Delta v1 writer detection for spark 4.0+

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

@@ -28,7 +28,7 @@ import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FileSystem, Path}
 
 import org.apache.spark.SparkContext
-import org.apache.spark.sql.{DataFrame, DataFrameWriter, Row, SaveMode, SparkSession}
+import org.apache.spark.sql.{DataFrame, Row, SaveMode, SparkSession}
 import org.apache.spark.sql.catalyst.catalog.{CatalogTable, CatalogTableType}
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
@@ -364,6 +364,8 @@ abstract class GpuCreateDeltaTableCommandBase(
      tableWithLocation: CatalogTable): Unit = {
     val isManagedTable = tableWithLocation.tableType == CatalogTableType.MANAGED
     val options = new DeltaOptions(table.storage.properties, sparkSession.sessionState.conf)
+    val isV1WriterSaveAsTableOverwrite =
+      DeltaRuntimeShim.isV1WriterSaveAsTableOverwrite(options, mode)
 
     // Execute write command for `deltaWriter` by
     //   - replacing the metadata new target table for DataFrameWriterV2 writer if it is a
@@ -376,7 +378,7 @@ abstract class GpuCreateDeltaTableCommandBase(
       schema: StructType): (TaggedCommitData[Action], DeltaOperations.Operation) = {
       // In the V2 Writer, methods like "replace" and "createOrReplace" implicitly mean that
       // the metadata should be changed. This wasn't the behavior for DataFrameWriterV1.
-      if (!isV1Writer) {
+      if (!isV1WriterSaveAsTableOverwrite) {
         replaceMetadataIfNecessary(
           txn,
           tableWithLocation,
@@ -394,13 +396,13 @@ abstract class GpuCreateDeltaTableCommandBase(
         // saveAsTable() command uses this same code path and is marked as a V1 writer.
         // We do not want saveAsTable() to be treated as a REPLACE command wrt dynamic partition
         // overwrite.
-        isTableReplace = isReplace && !isV1Writer
+        isTableReplace = isReplace && !isV1WriterSaveAsTableOverwrite
       )
       // Metadata updates for creating table (with any writer) and replacing table
       // (only with V1 writer) will be handled inside WriteIntoDelta.
       // For createOrReplace operation, metadata updates are handled here if the table already
       // exists (replacing table), otherwise it is handled inside WriteIntoDelta (creating table).
-      if (!isV1Writer && isReplace && txn.readVersion > -1L) {
+      if (!isV1WriterSaveAsTableOverwrite && isReplace && txn.readVersion > -1L) {
         val newDomainMetadata = Seq.empty[DomainMetadata] ++
           ClusteredTableUtils.getDomainMetadataFromTransaction(
             ClusteredTableUtils.getClusterBySpecOptional(table), txn)
@@ -413,7 +415,7 @@ abstract class GpuCreateDeltaTableCommandBase(
       val op = getOperation(txn.metadata, isManagedTable, Some(options),
         clusterBy = ClusteredTableUtils.getLogicalClusteringColumnNames(
           txn, taggedCommitData.actions),
-        isV1SaveAsTableOverwrite = if (isV1Writer) Some(true) else None
+        isV1SaveAsTableOverwrite = if (isV1WriterSaveAsTableOverwrite) Some(true) else None
       )
       (taggedCommitData, op)
     }
@@ -868,19 +870,6 @@ abstract class GpuCreateDeltaTableCommandBase(
     }
   }
 
-  /**
-   * Horrible hack to differentiate between DataFrameWriterV1 and V2 so that we can decide
-   * what to do with table metadata. In DataFrameWriterV1, mode("overwrite").saveAsTable,
-   * behaves as a CreateOrReplace table, but we have asked for "overwriteSchema" as an
-   * explicit option to overwrite partitioning or schema information. With DataFrameWriterV2,
-   * the behavior asked for by the user is clearer: .createOrReplace(), which means that we
-   * should overwrite schema and/or partitioning. Therefore we have this hack.
-   */
-  private def isV1Writer: Boolean = {
-    Thread.currentThread().getStackTrace.exists(_.toString.contains(
-      classOf[DataFrameWriter[_]].getCanonicalName + "."))
-  }
-
   /** Returns true if the current operation could be replacing a table. */
   private def isReplace: Boolean = {
     operation == TableCreationModes.CreateOrReplace ||

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

@@ -22,7 +22,7 @@ import com.nvidia.spark.rapids.{RapidsConf, ShimLoader, ShimReflectionUtils, Ver
 import com.nvidia.spark.rapids.delta.{DeltaConfigChecker, DeltaProvider}
 
 import org.apache.spark.SPARK_VERSION
-import org.apache.spark.sql.{SaveMode, SparkSession}
+import org.apache.spark.sql.{DataFrameWriter, SaveMode, SparkSession}
 import org.apache.spark.sql.catalyst.catalog.CatalogTable
 import org.apache.spark.sql.connector.catalog.StagingTableCatalog
 import org.apache.spark.sql.delta.{DeltaLog, DeltaOperations, DeltaOptions, DeltaUDF, Snapshot}
@@ -73,6 +73,15 @@ trait DeltaRuntimeShim {
   def getTightBoundColumnOnFileInitDisabled(spark: SparkSession): Boolean
 
   def getGpuDeltaCatalog(cpuCatalog: DeltaCatalog, rapidsConf: RapidsConf): StagingTableCatalog
+
+  /**
+   * Detect a DataFrameWriter V1 mode("overwrite").saveAsTable operation so it retains the
+   * existing table metadata. Delta versions before 4.1 require stack-trace inspection.
+   */
+  def isV1WriterSaveAsTableOverwrite(options: DeltaOptions, mode: SaveMode): Boolean = {
+    mode == SaveMode.Overwrite && Thread.currentThread().getStackTrace.exists(_.toString.contains(
+      classOf[DataFrameWriter[_]].getCanonicalName + "."))
+  }
 }
 
 object DeltaRuntimeShim {
@@ -192,6 +201,9 @@ object DeltaRuntimeShim {
   def getTightBoundColumnOnFileInitDisabled(spark: SparkSession): Boolean =
     shimInstance.getTightBoundColumnOnFileInitDisabled(spark)
 
+  def isV1WriterSaveAsTableOverwrite(options: DeltaOptions, mode: SaveMode): Boolean =
+    shimInstance.isV1WriterSaveAsTableOverwrite(options, mode)
+
   def getGpuDeltaCatalog(cpuCatalog: DeltaCatalog, rapidsConf: RapidsConf): StagingTableCatalog = {
     shimInstance.getGpuDeltaCatalog(cpuCatalog, rapidsConf)
   }

delta-lake/delta-40x/src/main/scala/org/apache/spark/sql/delta/rapids/delta40x/Delta40xRuntimeShim.scala

@@ -22,6 +22,7 @@ import com.nvidia.spark.rapids.delta.delta40x.Delta40xProvider
 import com.nvidia.spark.rapids.delta.delta40x.GpuDeltaCatalog
 
 import org.apache.spark.sql.SaveMode
+import org.apache.spark.sql.classic.DataFrameWriter
 import org.apache.spark.sql.connector.catalog.StagingTableCatalog
 import org.apache.spark.sql.delta.{DeltaOperations, DeltaOptions}
 import org.apache.spark.sql.delta.actions.Metadata
@@ -41,6 +42,13 @@ class Delta40xRuntimeShim extends DeltaRuntimeShimBase {
 
   override def getDeltaProvider: DeltaProvider = Delta40xProvider
 
+  override def isV1WriterSaveAsTableOverwrite(
+      options: DeltaOptions,
+      mode: SaveMode): Boolean = {
+    mode == SaveMode.Overwrite && Thread.currentThread().getStackTrace.exists(_.toString.contains(
+      classOf[DataFrameWriter[_]].getCanonicalName + "."))
+  }
+
   override def createGpuWrite(
       gpuDeltaLog: GpuDeltaLog,
       cpuWrite: WriteIntoDelta): GpuWriteIntoDeltaLike = {

delta-lake/delta-41x/src/main/scala/org/apache/spark/sql/delta/rapids/delta41x/Delta41xRuntimeShim.scala

@@ -29,7 +29,7 @@ import org.apache.spark.sql.connector.catalog.StagingTableCatalog
 import org.apache.spark.sql.delta.{DeltaOperations, DeltaOptions}
 import org.apache.spark.sql.delta.actions.Metadata
 import org.apache.spark.sql.delta.catalog.DeltaCatalog
-import org.apache.spark.sql.delta.commands.WriteIntoDelta
+import org.apache.spark.sql.delta.commands.{CreateDeltaTableLikeShims, WriteIntoDelta}
 import org.apache.spark.sql.delta.hooks.GpuAutoCompact41x
 import org.apache.spark.sql.delta.rapids.{
   DeltaRuntimeShimBase,
@@ -45,6 +45,12 @@ class Delta41xRuntimeShim extends DeltaRuntimeShimBase {
 
   override def getDeltaProvider: DeltaProvider = Delta41xProvider
 
+  override def isV1WriterSaveAsTableOverwrite(
+      options: DeltaOptions,
+      mode: SaveMode): Boolean = {
+    CreateDeltaTableLikeShims.isV1WriterSaveAsTableOverwrite(options, mode)
+  }
+
   override def createGpuWrite(
       gpuDeltaLog: GpuDeltaLog,
       cpuWrite: WriteIntoDelta): GpuWriteIntoDeltaLike = {

delta-lake/delta-42x/src/main/scala/org/apache/spark/sql/delta/rapids/delta42x/Delta42xRuntimeShim.scala

@@ -29,7 +29,7 @@ import org.apache.spark.sql.connector.catalog.StagingTableCatalog
 import org.apache.spark.sql.delta.{DeltaOperations, DeltaOptions}
 import org.apache.spark.sql.delta.actions.Metadata
 import org.apache.spark.sql.delta.catalog.DeltaCatalog
-import org.apache.spark.sql.delta.commands.WriteIntoDelta
+import org.apache.spark.sql.delta.commands.{CreateDeltaTableLikeShims, WriteIntoDelta}
 import org.apache.spark.sql.delta.hooks.GpuAutoCompact42x
 import org.apache.spark.sql.delta.rapids.{DeltaRuntimeShimBase, GpuDeltaLog, GpuOptimisticTransaction,
   GpuOptimisticTransactionBase, GpuWriteIntoDeltaLike, StartTransactionArg}
@@ -40,6 +40,12 @@ class Delta42xRuntimeShim extends DeltaRuntimeShimBase {
 
   override def getDeltaProvider: DeltaProvider = Delta42xProvider
 
+  override def isV1WriterSaveAsTableOverwrite(
+      options: DeltaOptions,
+      mode: SaveMode): Boolean = {
+    CreateDeltaTableLikeShims.isV1WriterSaveAsTableOverwrite(options, mode)
+  }
+
   override def getGpuDeltaCatalog(
       cpuCatalog: DeltaCatalog,
       rapidsConf: RapidsConf): StagingTableCatalog = {

integration_tests/src/main/python/delta_lake_write_test.py

@@ -868,6 +868,75 @@ def read_ids(spark, table):
     assert [row.id for row in gpu_rows] == list(range(10, 20))
 
 
+@allow_non_gpu('DataWritingCommandExec', 'WriteFilesExec', *delta_meta_allow)
+@delta_lake
+@ignore_order(local=True)
+@pytest.mark.xfail(is_databricks_runtime(),
+                   reason="https://github.com/NVIDIA/spark-rapids/issues/11169")
+def test_delta_replace_where_save_as_table_preserves_partitioning(spark_tmp_table_factory):
+    cpu_table = spark_tmp_table_factory.get()
+    gpu_table = spark_tmp_table_factory.get()
+    confs = copy_and_update(writer_confs, delta_writes_enabled_conf)
+
+    def create_initial_tables(spark):
+        initial_values = ", ".join(
+            f"({record_id}L, '{region}', {record_id}.0D)"
+            for record_id, region in enumerate(
+                ["NA", "EMEA", "APAC", "LATAM", "NA", "EMEA", "APAC", "LATAM"]))
+        for table in [cpu_table, gpu_table]:
+            spark.sql(
+                f"CREATE TABLE {table} (record_id BIGINT, region STRING, amount DOUBLE) "
+                f"USING DELTA PARTITIONED BY (region)")
+            spark.sql(f"INSERT INTO {table} VALUES {initial_values}")
+
+    def replace_na_partition(spark, table):
+        replacement = spark.sql(
+            "SELECT 100L AS record_id, 'NA' AS region, 100.0D AS amount")
+        (replacement.write.format("delta").mode("overwrite")
+         .option("replaceWhere", "region = 'NA'")
+         .saveAsTable(table))
+
+    with_cpu_session(create_initial_tables, conf=confs)
+    with_cpu_session(lambda spark: replace_na_partition(spark, cpu_table), conf=confs)
+
+    callback = spark_jvm().org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback
+    callback.startCapture()
+    try:
+        with_gpu_session(lambda spark: replace_na_partition(spark, gpu_table), conf=confs)
+        plans = callback.getResultsWithTimeout(10000)
+        assert any(callback.contains(plan, "GpuAtomicReplaceTableAsSelectExec")
+                   for plan in plans), "GpuAtomicReplaceTableAsSelectExec was not executed"
+        # The RTAS data write runs as a nested query execution: Spark 4.0+ issues it as
+        # OverwriteByExpression, earlier versions as AppendData.
+        v1_write_node = ("GpuOverwriteByExpressionExecV1" if is_spark_400_or_later()
+                         else "GpuAppendDataExecV1")
+        assert any(callback.contains(plan, v1_write_node)
+                   for plan in plans), f"{v1_write_node} was not executed"
+    finally:
+        callback.endCapture()
+
+    def partition_counts(spark, table):
+        return spark.sql(
+            f"SELECT region, COUNT(*) AS count FROM {table} "
+            f"GROUP BY region ORDER BY region").collect()
+
+    cpu_counts = with_cpu_session(lambda spark: partition_counts(spark, cpu_table), conf=confs)
+    gpu_counts = with_cpu_session(lambda spark: partition_counts(spark, gpu_table), conf=confs)
+    assert_equal(cpu_counts, gpu_counts)
+    assert [(row.region, row["count"]) for row in gpu_counts] == [
+        ("APAC", 2), ("EMEA", 2), ("LATAM", 2), ("NA", 1)]
+
+    def partition_columns(spark, table):
+        return spark.sql(f"DESCRIBE DETAIL {table}").select("partitionColumns").head()[0]
+
+    cpu_partition_columns = with_cpu_session(
+        lambda spark: partition_columns(spark, cpu_table), conf=confs)
+    gpu_partition_columns = with_cpu_session(
+        lambda spark: partition_columns(spark, gpu_table), conf=confs)
+    assert cpu_partition_columns == ["region"]
+    assert gpu_partition_columns == cpu_partition_columns
+
+
 @allow_non_gpu(*delta_meta_allow)
 @delta_lake
 @ignore_order(local=True)