diff --git a/delta-lake/delta-spark350db143/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuWriteIntoDeltaCommand.scala b/delta-lake/delta-spark350db143/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuWriteIntoDeltaCommand.scala new file mode 100644 index 00000000000..f3bd275987f --- /dev/null +++ b/delta-lake/delta-spark350db143/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuWriteIntoDeltaCommand.scala @@ -0,0 +1,322 @@ +/* + * 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.databricks.sql.transaction.tahoe.rapids + +import com.databricks.sql.io.skipping.liquid.ClusteredTableUtils +import com.databricks.sql.transaction.tahoe.{DeltaColumnMapping, DeltaParquetFileFormat} +import com.databricks.sql.transaction.tahoe.commands.WriteIntoDeltaCommand +import com.databricks.sql.transaction.tahoe.files.TahoeBatchFileIndex +import com.databricks.sql.transaction.tahoe.schema.InnerInvariantViolationException +import com.databricks.sql.transaction.tahoe.stats.{DeltaJobStatisticsTracker, + StatisticsOnLoadJobTracker} +import com.nvidia.spark.rapids.{DataFromReplacementRule, DataWritingCommandMeta, + GpuDataWritingCommand, GpuMetric, GpuParquetFileFormat, RapidsConf, RapidsMeta} +import com.nvidia.spark.rapids.delta.RapidsDeltaUtils + +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.{BasicWriteJobStatsTracker, GpuWriteFiles, + HadoopFsRelation, LogicalRelation} +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.rapids.{BasicColumnarWriteJobStatsTracker, ColumnarWriteJobStatsTracker, + GpuFileFormatWriter} +import org.apache.spark.sql.rapids.BasicColumnarWriteJobStatsTracker.TASK_COMMIT_TIME +import org.apache.spark.sql.rapids.shims.TrampolineConnectShims +import org.apache.spark.sql.rapids.shims.TrampolineConnectShims.SparkSession +import org.apache.spark.sql.types.StructField +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.SerializableConfiguration + +private object DeltaWriteAttributeMapping { + private def structurallyMatches(left: Attribute, right: Attribute): Boolean = { + left.name == right.name && left.dataType == right.dataType + } + + private def structurallyMatches(left: Attribute, right: StructField): Boolean = { + left.name == right.name && left.dataType == right.dataType + } + + def validateQueryAttributeAtOrdinal( + queryOutput: Seq[Attribute], + logicalField: StructField, + physicalField: StructField, + attribute: Attribute, + ordinal: Int, + description: String): Either[String, Unit] = { + if (ordinal >= queryOutput.size) { + Left(s"Delta $description ${attribute.exprId} maps to missing query output ordinal $ordinal") + } else { + val queryAttribute = queryOutput(ordinal) + val exprIdOrdinal = queryOutput.indexWhere(_.exprId == attribute.exprId) + val matchesExpectedSchema = structurallyMatches(attribute, queryAttribute) || + (structurallyMatches(queryAttribute, logicalField) && + structurallyMatches(attribute, physicalField)) + if (attribute.exprId == queryAttribute.exprId || + (exprIdOrdinal == -1 && matchesExpectedSchema)) { + Right(()) + } else if (exprIdOrdinal != -1) { + Left(s"Delta $description ${attribute.exprId} maps to query output ordinal " + + s"$exprIdOrdinal instead of native output ordinal $ordinal") + } else { + Left(s"Delta $description ${attribute.exprId} has no query ExprId match and does not " + + s"match the logical or physical Delta schema at query output ordinal $ordinal") + } + } + } + + def remapToDataAttribute( + outputAttribute: Attribute, + queryAttribute: Attribute, + dataAttribute: Attribute, + description: String, + ordinal: Int): Either[String, Attribute] = { + if (structurallyMatches(queryAttribute, dataAttribute)) { + Right(outputAttribute.withExprId(dataAttribute.exprId)) + } else { + Left(s"Delta $description at ordinal $ordinal does not structurally match the data plan: " + + s"query=${queryAttribute.name}:${queryAttribute.dataType.catalogString}, " + + s"data=${dataAttribute.name}:${dataAttribute.dataType.catalogString}") + } + } +} + +class GpuWriteIntoDeltaCommandMeta( + cmd: WriteIntoDeltaCommand, + conf: RapidsConf, + parent: Option[RapidsMeta[_, _, _]], + rule: DataFromReplacementRule) + extends DataWritingCommandMeta[WriteIntoDeltaCommand](cmd, conf, parent, rule) { + + private var fileFormat: Option[GpuParquetFileFormat] = None + private lazy val logicalDataFields = cmd.metadata.dataSchema.fields + private lazy val physicalDataFields = DeltaColumnMapping.createPhysicalSchema( + cmd.metadata.dataSchema, + cmd.metadata.schema, + cmd.metadata.columnMappingMode).fields + + private def tagAttributeMapping( + attribute: Attribute, + ordinal: Int, + description: String): Unit = { + DeltaWriteAttributeMapping.validateQueryAttributeAtOrdinal( + cmd.query.output, + logicalDataFields(ordinal), + physicalDataFields(ordinal), + attribute, + ordinal, + description) match { + case Left(reason) => willNotWorkOnGpu(reason) + case Right(_) => () + } + } + + override protected def tagSelfForGpuInternal(): Unit = { + val isNativeOptimize = cmd.query.collectLeaves() match { + case Seq(LogicalRelation(HadoopFsRelation( + index: TahoeBatchFileIndex, _, _, _, _, _), _, _, _)) => + index.actionType.equalsIgnoreCase("Optimize") + case _ => false + } + val isSupportedNativeOptimize = + isNativeOptimize && !ClusteredTableUtils.isSupported(cmd.protocol) + if (!isSupportedNativeOptimize) { + willNotWorkOnGpu( + "DBR WriteIntoDeltaCommand GPU support is limited to native non-clustered OPTIMIZE") + } + if (!conf.isDeltaWriteEnabled) { + willNotWorkOnGpu("Delta Lake output acceleration has been disabled") + } + val spark = TrampolineConnectShims.getActiveSession + RapidsDeltaUtils.tagForDeltaWrite( + this, cmd.query.schema, Some(cmd.deltaLog), cmd.options, spark) + if (cmd.fileFormat.getClass != classOf[DeltaParquetFileFormat]) { + willNotWorkOnGpu(s"Delta file format ${cmd.fileFormat.getClass.getName} is not supported") + } else { + fileFormat = GpuParquetFileFormat.tagGpuSupport( + this, spark, cmd.options, cmd.hadoopConf, cmd.query.schema) + } + if (cmd.bucketSpec.nonEmpty) { + willNotWorkOnGpu("Bucketed Delta writes are not supported") + } + if (cmd.staticPartitions.nonEmpty) { + willNotWorkOnGpu("Static partition Delta writes are not supported by this command path") + } + if (cmd.partitionColExprIds.nonEmpty) { + willNotWorkOnGpu("Partitioned DBR Delta writes are not supported by this command path " + + "until native partition materialization and partition-evolution semantics are validated") + } + val columnCounts = Seq( + "output specification" -> cmd.outputSpec.outputColumns.size, + "query" -> cmd.query.output.size, + "logical data schema" -> logicalDataFields.length, + "physical data schema" -> physicalDataFields.length) + if (columnCounts.map(_._2).distinct.size != 1) { + willNotWorkOnGpu(s"Delta column count mismatch: ${columnCounts.map { + case (description, count) => s"$description=$count" + }.mkString(", ")}") + } else { + cmd.outputSpec.outputColumns.zipWithIndex.foreach { case (attribute, ordinal) => + tagAttributeMapping(attribute, ordinal, "output column") + } + cmd.partitionColExprIds.foreach { exprId => + val matches = cmd.outputSpec.outputColumns.zipWithIndex.filter(_._1.exprId == exprId) + matches match { + case Seq((attribute, ordinal)) => + tagAttributeMapping(attribute, ordinal, "partition column") + case _ => + willNotWorkOnGpu( + s"Delta partition column $exprId has ${matches.size} output specification matches") + } + } + } + cmd.statsTrackers.foreach { + case tracker: BasicWriteJobStatsTracker => + val metrics = tracker.driverSideMetrics ++ cmd.writeJobMetrics + if (!metrics.contains(TASK_COMMIT_TIME)) { + willNotWorkOnGpu(s"Delta basic statistics tracker is missing $TASK_COMMIT_TIME") + } + case _: DeltaJobStatisticsTracker => + case _: StatisticsOnLoadJobTracker => + willNotWorkOnGpu("DBR StatisticsOnLoadJobTracker is not supported on GPU") + case tracker => + willNotWorkOnGpu(s"Delta write statistics tracker ${tracker.getClass.getName} " + + "is not supported on GPU") + } + } + + override def convertToGpu(): GpuDataWritingCommand = { + val gpuFileFormat = fileFormat.getOrElse( + throw new IllegalStateException("fileFormat missing, tagSelfForGpu not called?")) + GpuWriteIntoDeltaCommand(cmd, conf, gpuFileFormat) + } +} + +case class GpuWriteIntoDeltaCommand( + cpuCmd: WriteIntoDeltaCommand, + @transient rapidsConf: RapidsConf, + fileFormat: GpuParquetFileFormat) extends GpuDataWritingCommand { + + private lazy val logicalDataFields = cpuCmd.metadata.dataSchema.fields + private lazy val physicalDataFields = DeltaColumnMapping.createPhysicalSchema( + cpuCmd.metadata.dataSchema, + cpuCmd.metadata.schema, + cpuCmd.metadata.columnMappingMode).fields + + override def query: LogicalPlan = cpuCmd.query + + override def outputColumnNames: Seq[String] = cpuCmd.outputColumnNames + + override lazy val metrics: Map[String, SQLMetric] = cpuCmd.writeJobMetrics + + override def requireSingleBatch: Boolean = false + + private def columnarStatsTrackers( + sparkSession: SparkSession): Seq[ColumnarWriteJobStatsTracker] = { + val serializableConf = new SerializableConfiguration(cpuCmd.hadoopConf) + cpuCmd.statsTrackers.map { + case tracker: BasicWriteJobStatsTracker => + val metrics = tracker.driverSideMetrics ++ cpuCmd.writeJobMetrics + new BasicColumnarWriteJobStatsTracker( + serializableConf, GpuMetric.wrap(metrics)) + case tracker: DeltaJobStatisticsTracker => + GpuWriteIntoDeltaCommandStats(cpuCmd, tracker, sparkSession) + case tracker => + throw new IllegalStateException( + s"Unsupported Delta write statistics tracker ${tracker.getClass.getName}") + } + } + + override def runColumnar( + sparkSession: SparkSession, + child: SparkPlan): Seq[ColumnarBatch] = { + val dataPlan = GpuWriteFiles.getWriteFilesOpt(child).map(_.child).getOrElse(child) + + if (cpuCmd.outputSpec.outputColumns.size != cpuCmd.query.output.size || + dataPlan.output.size != cpuCmd.query.output.size) { + throw new IllegalStateException( + s"Delta output size mismatch: outputSpec=${cpuCmd.outputSpec.outputColumns.size}, " + + s"query=${cpuCmd.query.output.size}, dataPlan=${dataPlan.output.size}") + } + + def resolveAttribute( + attribute: Attribute, + ordinal: Int, + description: String): Attribute = { + DeltaWriteAttributeMapping.validateQueryAttributeAtOrdinal( + cpuCmd.query.output, + logicalDataFields(ordinal), + physicalDataFields(ordinal), + attribute, + ordinal, + description) match { + case Right(_) => () + case Left(reason) => throw new IllegalStateException(reason) + } + DeltaWriteAttributeMapping.remapToDataAttribute( + attribute, + cpuCmd.query.output(ordinal), + dataPlan.output(ordinal), + description, + ordinal) match { + case Right(value) => value + case Left(reason) => throw new IllegalStateException(reason) + } + } + + val outputColumns = cpuCmd.outputSpec.outputColumns.zipWithIndex.map { + case (attribute, ordinal) => + resolveAttribute(attribute, ordinal, "output column") + } + val partitionColumns = cpuCmd.partitionColExprIds.map { exprId => + val matches = cpuCmd.outputSpec.outputColumns.zipWithIndex.filter(_._1.exprId == exprId) + matches match { + case Seq((attribute, ordinal)) => + resolveAttribute(attribute, ordinal, "partition column") + case _ => throw new IllegalStateException( + s"Delta partition column $exprId has ${matches.size} output specification matches") + } + } + val outputSpec = cpuCmd.outputSpec.copy(outputColumns = outputColumns) + val writePartitionColumns = WriteIntoDeltaCommand.writePartitionColumns( + cpuCmd.metadata, sparkSession) + if (partitionColumns.nonEmpty && writePartitionColumns) { + throw new IllegalStateException( + "Writing partition columns into Delta Parquet data files is not supported") + } + + try { + GpuFileFormatWriter.write( + sparkSession = sparkSession, + plan = child, + fileFormat = fileFormat, + committer = cpuCmd.committer, + outputSpec = outputSpec, + hadoopConf = cpuCmd.hadoopConf, + partitionColumns = partitionColumns, + bucketSpec = None, + statsTrackers = columnarStatsTrackers(sparkSession), + options = cpuCmd.options, + useStableSort = rapidsConf.stableSort, + concurrentWriterPartitionFlushSize = rapidsConf.concurrentWriterPartitionFlushSize, + baseDebugOutputPath = rapidsConf.outputDebugDumpPrefix) + } catch { + case InnerInvariantViolationException(violation) => throw violation + } + Seq.empty + } +} diff --git a/delta-lake/delta-spark350db143/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuWriteIntoDeltaCommandStats.scala b/delta-lake/delta-spark350db143/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuWriteIntoDeltaCommandStats.scala new file mode 100644 index 00000000000..c7d841bbda2 --- /dev/null +++ b/delta-lake/delta-spark350db143/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuWriteIntoDeltaCommandStats.scala @@ -0,0 +1,76 @@ +/* + * 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.databricks.sql.transaction.tahoe.rapids + +import com.databricks.sql.transaction.tahoe.{DeltaColumnMapping, DeltaConfigs} +import com.databricks.sql.transaction.tahoe.commands.{DeletionVectorUtils, WriteIntoDeltaCommand} +import com.databricks.sql.transaction.tahoe.sources.DeltaSQLConf +import com.databricks.sql.transaction.tahoe.stats.DeltaJobStatisticsTracker +import com.nvidia.spark.rapids.delta.{GpuDeltaJobStatisticsTracker, GpuStatisticsCollection} + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.execution.datasources.WriteTaskStats +import org.apache.spark.sql.rapids.{ColumnarWriteJobStatsTracker, ColumnarWriteTaskStatsTracker} +import org.apache.spark.sql.rapids.shims.TrampolineConnectShims.SparkSession +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.ColumnarBatch + +private[rapids] object GpuWriteIntoDeltaCommandStats { + def apply( + cpuCmd: WriteIntoDeltaCommand, + nativeTracker: DeltaJobStatisticsTracker, + sparkSession: SparkSession): ColumnarWriteJobStatsTracker = { + val useTableSchema = sparkSession.sessionState.conf.getConf( + DeltaSQLConf.DELTA_COLLECT_STATS_USING_TABLE_SCHEMA) + val statsCollection = new GpuStatisticsCollection { + override val spark = sparkSession + override val deletionVectorsSupported: Boolean = + DeletionVectorUtils.deletionVectorsWritable( + cpuCmd.deltaLog.unsafeVolatileSnapshot, Some(cpuCmd.protocol), Some(cpuCmd.metadata)) + override val tableDataSchema: StructType = if (useTableSchema) { + DeltaColumnMapping.createPhysicalSchema( + cpuCmd.metadata.dataSchema, + cpuCmd.metadata.schema, + cpuCmd.metadata.columnMappingMode) + } else { + nativeTracker.dataCols.toStructType + } + override val dataSchema: StructType = nativeTracker.dataCols.toStructType + override val numIndexedCols: Int = + DeltaConfigs.DATA_SKIPPING_NUM_INDEXED_COLS.fromMetaData(cpuCmd.metadata) + override val stringPrefixLength: Int = + spark.sessionState.conf.getConf(DeltaSQLConf.DATA_SKIPPING_STRING_PREFIX_LENGTH) + } + val statsSchema = statsCollection.statCollectionSchema + val explodedDataSchema = statsCollection.explodedDataSchema + val batchStatsToRow = (batch: ColumnarBatch, row: InternalRow) => { + GpuStatisticsCollection.batchStatsToRow(statsSchema, explodedDataSchema, batch, row) + } + val gpuTracker = new GpuDeltaJobStatisticsTracker( + nativeTracker.dataCols, nativeTracker.statsColExpr, batchStatsToRow) + + new ColumnarWriteJobStatsTracker { + override def newTaskInstance(): ColumnarWriteTaskStatsTracker = + gpuTracker.newTaskInstance() + + override def processStats(stats: Seq[WriteTaskStats], jobCommitTime: Long): Unit = { + gpuTracker.processStats(stats, jobCommitTime) + nativeTracker.recordedStats = gpuTracker.recordedStats + } + } + } +} diff --git a/delta-lake/delta-spark350db143/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark350DB143Provider.scala b/delta-lake/delta-spark350db143/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark350DB143Provider.scala index 0e712debd10..90f6e2d3a06 100644 --- a/delta-lake/delta-spark350db143/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark350DB143Provider.scala +++ b/delta-lake/delta-spark350db143/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark350DB143Provider.scala @@ -22,14 +22,16 @@ package com.nvidia.spark.rapids.delta import com.databricks.sql.transaction.tahoe.DeltaOptions -import com.databricks.sql.transaction.tahoe.commands.WriteIntoDeltaEdge -import com.databricks.sql.transaction.tahoe.rapids.{GpuDeltaCatalog, GpuDeltaLog, GpuDeltaV1Write, GpuWriteIntoDelta} +import com.databricks.sql.transaction.tahoe.commands.{WriteIntoDeltaCommand, WriteIntoDeltaEdge} +import com.databricks.sql.transaction.tahoe.rapids.{GpuDeltaCatalog, GpuDeltaLog, GpuDeltaV1Write, + GpuWriteIntoDelta, GpuWriteIntoDeltaCommandMeta} import com.nvidia.spark.rapids._ import com.nvidia.spark.rapids.delta.shims.DeltaLogShim import org.apache.spark.sql.{DataFrame, SaveMode} import org.apache.spark.sql.SaveMode import org.apache.spark.sql.connector.write.V1Write +import org.apache.spark.sql.execution.command.DataWritingCommand import org.apache.spark.sql.execution.datasources.LogicalRelation import org.apache.spark.sql.execution.datasources.v2.{AtomicCreateTableAsSelectExec, AtomicReplaceTableAsSelectExec} import org.apache.spark.sql.execution.datasources.v2.rapids.{GpuAtomicCreateTableAsSelectExec, GpuAtomicReplaceTableAsSelectExec} @@ -37,6 +39,15 @@ import org.apache.spark.sql.sources.InsertableRelation object DeltaSpark350DB143Provider extends DatabricksDeltaProviderBase { + override def getDataWritingCommandRules: Map[Class[_ <: DataWritingCommand], + DataWritingCommandRule[_ <: DataWritingCommand]] = { + Seq( + GpuOverrides.dataWriteCmd[WriteIntoDeltaCommand]( + "Write files for a DBR Delta OPTIMIZE transaction", + (a, conf, p, r) => new GpuWriteIntoDeltaCommandMeta(a, conf, p, r)) + ).map(r => (r.getClassFor.asSubclass(classOf[DataWritingCommand]), r)).toMap + } + override protected def toGpuWrite( writeConfig: DeltaWriteV1Config, rapidsConf: RapidsConf): V1Write = new GpuDeltaV1Write { diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuWriteIntoDeltaCommand.scala b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuWriteIntoDeltaCommand.scala index 1c973296070..c9af80a875c 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuWriteIntoDeltaCommand.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuWriteIntoDeltaCommand.scala @@ -16,8 +16,10 @@ package com.databricks.sql.transaction.tahoe.rapids +import com.databricks.sql.io.skipping.liquid.ClusteredTableUtils import com.databricks.sql.transaction.tahoe.{DeltaColumnMapping, DeltaParquetFileFormat} import com.databricks.sql.transaction.tahoe.commands.WriteIntoDeltaCommand +import com.databricks.sql.transaction.tahoe.files.TahoeBatchFileIndex import com.databricks.sql.transaction.tahoe.schema.InnerInvariantViolationException import com.databricks.sql.transaction.tahoe.stats.{DeltaJobStatisticsTracker, StatisticsOnLoadJobTracker} @@ -28,7 +30,8 @@ import com.nvidia.spark.rapids.delta.RapidsDeltaUtils import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.datasources.{BasicWriteJobStatsTracker, GpuWriteFiles} +import org.apache.spark.sql.execution.datasources.{BasicWriteJobStatsTracker, GpuWriteFiles, + HadoopFsRelation, LogicalRelation} import org.apache.spark.sql.execution.datasources.v2.rapids.GpuAtomicDeltaWriteContext import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.rapids.{BasicColumnarWriteJobStatsTracker, ColumnarWriteJobStatsTracker, @@ -124,9 +127,18 @@ class GpuWriteIntoDeltaCommandMeta( } override protected def tagSelfForGpuInternal(): Unit = { - if (!GpuAtomicDeltaWriteContext.isActive) { + val isNativeOptimize = cmd.query.collectLeaves() match { + case Seq(LogicalRelation(HadoopFsRelation( + index: TahoeBatchFileIndex, _, _, _, _, _), _, _, _, _, _, _)) => + index.actionType.equalsIgnoreCase("Optimize") + case _ => false + } + val isSupportedNativeOptimize = + isNativeOptimize && !ClusteredTableUtils.isSupported(cmd.protocol) + if (!GpuAtomicDeltaWriteContext.isActive && !isSupportedNativeOptimize) { willNotWorkOnGpu( - "DBR WriteIntoDeltaCommand GPU support is limited to atomic CTAS/RTAS") + "DBR WriteIntoDeltaCommand GPU support is limited to atomic CTAS/RTAS or " + + "native non-clustered OPTIMIZE") } if (!conf.isDeltaWriteEnabled) { willNotWorkOnGpu("Delta Lake output acceleration has been disabled") diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormat.scala b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormat.scala index 1cd4d12cf86..480ddf7710b 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormat.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormat.scala @@ -146,7 +146,7 @@ case class GpuDeltaParquetFileFormat( tahoeFileIndexOpt.exists { tahoeFileIndex => tahoeFileIndex.rowIndexFilters.exists(_.nonEmpty) || tahoeFileIndex - .matchingFiles(partitionFilters = Seq(TrueLiteral), dataFilters = Seq(TrueLiteral)) + .matchingFiles(partitionFilters = Seq.empty, dataFilters = Seq(TrueLiteral)) .exists(_.deletionVector != null) } } @@ -287,7 +287,7 @@ object GpuDeltaParquetFileFormat { filters.values.exists(_.getRowIndexFilterType != RowIndexFilterType.IF_CONTAINED) } hasUnsupportedCachedFilter || tahoeFileIndex.matchingFiles( - Seq(TrueLiteral), Seq(TrueLiteral)).exists { addFile => + Seq.empty, Seq(TrueLiteral)).exists { addFile => val provider = try { tahoeFileIndex.getRowIndexFilterForFile(addFile.path) } catch { diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/RapidsDeletionVectors.scala b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/RapidsDeletionVectors.scala index faa0ad8ff76..a6e871610c2 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/RapidsDeletionVectors.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/RapidsDeletionVectors.scala @@ -130,7 +130,7 @@ object RapidsDeletionVectors extends Logging { val filterTypes = tahoeFileIndex.rowIndexFilters.getOrElse(Map.empty) .map(kv => kv._1 -> kv._2.getRowIndexFilterType) val matchingFiles = tahoeFileIndex - .matchingFiles(partitionFilters = Seq(TrueLiteral), dataFilters = Seq(TrueLiteral)) + .matchingFiles(partitionFilters = Seq.empty, dataFilters = Seq(TrueLiteral)) def fileKeys(relativePath: String): Seq[String] = { val absolute = absolutePath(tahoeFileIndex.path.toString, relativePath) diff --git a/integration_tests/src/main/python/delta_lake_optimize_table_test.py b/integration_tests/src/main/python/delta_lake_optimize_table_test.py index 206337a3b35..67dfb009453 100644 --- a/integration_tests/src/main/python/delta_lake_optimize_table_test.py +++ b/integration_tests/src/main/python/delta_lake_optimize_table_test.py @@ -18,7 +18,8 @@ from delta_lake_utils import * from marks import * from spark_session import with_cpu_session, with_gpu_session, with_spark_session, is_before_spark_353, \ - supports_delta_lake_deletion_vectors, is_databricks_runtime, is_databricks173_or_later + supports_delta_lake_deletion_vectors, is_databricks_runtime, is_databricks143, \ + is_databricks173_or_later from pyspark.sql.types import IntegerType, StringType _optimize_conf = copy_and_update(delta_writes_enabled_conf, { @@ -27,6 +28,14 @@ "spark.databricks.delta.autoCompact.enabled": "false" }) +_native_optimize_write_conf = copy_and_update(delta_writes_enabled_conf, { + "spark.rapids.sql.command.OptimizeTableCommand": "false", + "spark.rapids.sql.command.OptimizeTableCommandEdge": "false", + "spark.databricks.delta.autoCompact.enabled": "false", + "spark.databricks.delta.optimizeWrite.enabled": "false", + "spark.sql.adaptive.enabled": "false" +}) + _liquid_optimize_dv_conf = copy_and_update(_optimize_conf, { "spark.databricks.delta.delete.deletionVectors.persistent": "true" }) @@ -406,6 +415,45 @@ def assert_data_and_log_parity(): assert_data_and_log_parity() +def _assert_native_optimize_gpu_write_parity(spark_tmp_path): + data_path = spark_tmp_path + "/DELTA_NATIVE_OPTIMIZE_WRITE" + cpu_path = data_path + "/CPU" + gpu_path = data_path + "/GPU" + conf = _native_optimize_write_conf + _setup_tables(False, cpu_path, gpu_path, None, None, conf) + files_before = { + path: with_cpu_session( + lambda spark: len(spark.read.format("delta").load(path).inputFiles()), conf=conf) + for path in [cpu_path, gpu_path] + } + + cpu_result = with_cpu_session( + lambda spark: spark.sql(_optimize_sql(cpu_path)).collect(), conf=conf) + gpu_result = assert_rapids_delta_write( + lambda spark: spark.sql(_optimize_sql(gpu_path)).collect(), conf=conf, + required_gpu_classes=["GpuDataWritingCommandExec", "GpuWriteFilesExec"], + require_same_plan=True, + forbidden_cpu_fallback_classes=["DataWritingCommandExec", "WriteFilesExec"]) + + assert str(cpu_result[0][0]).rstrip('/').endswith('/CPU') + assert str(gpu_result[0][0]).rstrip('/').endswith('/GPU') + cpu_data = with_cpu_session(lambda spark: _read_sorted(spark, cpu_path).collect(), conf=conf) + gpu_data = with_cpu_session(lambda spark: _read_sorted(spark, gpu_path).collect(), conf=conf) + assert_equal(cpu_data, gpu_data) + for path in [cpu_path, gpu_path]: + files_after = with_cpu_session( + lambda spark: len(spark.read.format("delta").load(path).inputFiles()), conf=conf) + assert files_after < files_before[path], \ + f"OPTIMIZE did not reduce the file count for {path}" + optimize_count = with_cpu_session(lambda spark: spark.sql( + f"DESCRIBE HISTORY delta.`{path}`").filter("operation = 'OPTIMIZE'").count(), + conf=conf) + assert_equal(optimize_count, 1) + with_cpu_session( + lambda spark: assert_gpu_and_cpu_latest_delta_log_equivalent(spark, data_path), + conf=conf) + + @allow_non_gpu(*delta_meta_allow) @delta_lake @ignore_order @@ -417,6 +465,14 @@ def test_delta_optimize_unpartitioned_table(spark_tmp_path, enable_deletion_vect _assert_optimize_parity(enable_deletion_vectors, spark_tmp_path, partition_columns=None) +@allow_non_gpu('ExecutedCommandExec', 'HashAggregateExec', *delta_meta_allow) +@delta_lake +@pytest.mark.skipif(not (is_databricks143() or is_databricks173_or_later()), + reason="Native DBR OPTIMIZE write coverage is for DBR 14.3 and 17.3+") +def test_delta_native_optimize_gpu_write(spark_tmp_path): + _assert_native_optimize_gpu_write_parity(spark_tmp_path) + + @allow_non_gpu(*delta_meta_allow) @delta_lake @ignore_order diff --git a/integration_tests/src/main/python/delta_lake_utils.py b/integration_tests/src/main/python/delta_lake_utils.py index 7116269bd45..0943f58fb48 100644 --- a/integration_tests/src/main/python/delta_lake_utils.py +++ b/integration_tests/src/main/python/delta_lake_utils.py @@ -401,7 +401,9 @@ def read_sorted_delta_path(spark, path): with_cpu_session(lambda spark: assert_gpu_and_cpu_latest_delta_log_equivalent(spark, data_path), conf=conf) -def assert_rapids_delta_write(do_test, conf): +def assert_rapids_delta_write( + do_test, conf, required_gpu_classes=delta_write, require_same_plan=False, + forbidden_cpu_fallback_classes=None): """ Validates that a Delta write operation executed on the GPU produces the expected execution plans. This function starts a plan capture mechanism using the Spark JVM's ExecutionPlanCaptureCallback, @@ -416,28 +418,46 @@ def assert_rapids_delta_write(do_test, conf): conf : dict A dictionary of configuration options to be passed to the GPU session. + required_gpu_classes : list[str] + GPU class names that must occur in the captured plans. + require_same_plan : bool + Whether all required GPU classes must occur in the same captured plan. + forbidden_cpu_fallback_classes : list[str] or None + CPU class names that must not be reported as falling back in any captured plan. + Returns ------- result : Any The result returned by the `do_test` function. """ jvm = spark_jvm() - jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback.startCapture() + callback = jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback + callback.startCapture() try: result = with_gpu_session(do_test, conf=conf) - captured_plans = jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback.getResultsWithTimeout(10000) + captured_plans = callback.getResultsWithTimeout(10000) # Some write functions are no-op. We may not capture any GPU plan. - if len(captured_plans) > 0: - for cls in delta_write: + if require_same_plan: + found = any( + all(callback.contains(plan, cls) for cls in required_gpu_classes) + for plan in captured_plans) + assert found, \ + f"No captured plan contains all required GPU classes: {required_gpu_classes}" + elif len(captured_plans) > 0: + for cls in required_gpu_classes: found = False for plan in captured_plans: - found = jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback.contains(plan, cls) + found = callback.contains(plan, cls) if found: break assert found, f"{cls} is not found in any captured plan" + for plan in captured_plans: + for cls in forbidden_cpu_fallback_classes or []: + assert not callback.didFallBack(plan, cls), \ + f"Captured Delta write plan fell back to CPU {cls}" return result finally: - jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback.endCapture() + callback.endCapture() def assert_db173_gpu_data_writing_command( do_test, conf, optimized_write, expected_atomic_gpu_class, aqe_enabled=None, 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 78fd9d448cf..6e3fc60efc4 100644 --- a/integration_tests/src/main/python/delta_lake_write_test.py +++ b/integration_tests/src/main/python/delta_lake_write_test.py @@ -1800,9 +1800,8 @@ def test_delta_write_column_name_mapping(spark_tmp_path, mapping): # Hash aggregate can be used in a metadata query for compaction which completely falls back compaction_allow = "HashAggregateExec" if is_databricks_runtime(): - # compaction can fallback due to unsupported WriteIntoDeltaCommand - # tracked by https://github.com/NVIDIA/spark-rapids/issues/11169 - compaction_allow += "," + delta_write_fallback_allow + # The native OPTIMIZE command stays on CPU while its nested write runs on GPU. + compaction_allow += ",ExecutedCommandExec" @allow_non_gpu(compaction_allow, *delta_meta_allow) @delta_lake @ignore_order