diff --git a/delta-lake/common/src/main/delta-33x-42x/scala/org/apache/spark/sql/delta/rapids/GpuCreateDeltaTableCommandBase.scala b/delta-lake/common/src/main/delta-33x-42x/scala/org/apache/spark/sql/delta/rapids/GpuCreateDeltaTableCommandBase.scala index b24f747fa1e..72b15a877d0 100644 --- a/delta-lake/common/src/main/delta-33x-42x/scala/org/apache/spark/sql/delta/rapids/GpuCreateDeltaTableCommandBase.scala +++ b/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 || diff --git a/delta-lake/common/src/main/delta-io/scala/org/apache/spark/sql/delta/rapids/DeltaRuntimeShim.scala b/delta-lake/common/src/main/delta-io/scala/org/apache/spark/sql/delta/rapids/DeltaRuntimeShim.scala index 6efd944e7fd..de7dde263e3 100644 --- a/delta-lake/common/src/main/delta-io/scala/org/apache/spark/sql/delta/rapids/DeltaRuntimeShim.scala +++ b/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) } diff --git a/delta-lake/delta-40x/src/main/scala/org/apache/spark/sql/delta/rapids/delta40x/Delta40xRuntimeShim.scala b/delta-lake/delta-40x/src/main/scala/org/apache/spark/sql/delta/rapids/delta40x/Delta40xRuntimeShim.scala index b1be29b6afc..7430c97ae4f 100644 --- a/delta-lake/delta-40x/src/main/scala/org/apache/spark/sql/delta/rapids/delta40x/Delta40xRuntimeShim.scala +++ b/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 = { diff --git a/delta-lake/delta-41x/src/main/scala/org/apache/spark/sql/delta/rapids/delta41x/Delta41xRuntimeShim.scala b/delta-lake/delta-41x/src/main/scala/org/apache/spark/sql/delta/rapids/delta41x/Delta41xRuntimeShim.scala index eb672505861..94ae14518d6 100644 --- a/delta-lake/delta-41x/src/main/scala/org/apache/spark/sql/delta/rapids/delta41x/Delta41xRuntimeShim.scala +++ b/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 = { diff --git a/delta-lake/delta-42x/src/main/scala/org/apache/spark/sql/delta/rapids/delta42x/Delta42xRuntimeShim.scala b/delta-lake/delta-42x/src/main/scala/org/apache/spark/sql/delta/rapids/delta42x/Delta42xRuntimeShim.scala index 21213c8f932..2f1f94fcaf7 100644 --- a/delta-lake/delta-42x/src/main/scala/org/apache/spark/sql/delta/rapids/delta42x/Delta42xRuntimeShim.scala +++ b/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 = { diff --git a/integration_tests/src/main/python/delta_lake_write_test.py b/integration_tests/src/main/python/delta_lake_write_test.py index 1db32d36539..7a8a1ed2113 100644 --- a/integration_tests/src/main/python/delta_lake_write_test.py +++ b/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)