From 26ed958de9bb2f34f0b597a6675a67783b34b72a Mon Sep 17 00:00:00 2001 From: Ray Liu Date: Mon, 7 Sep 2026 05:13:44 +0000 Subject: [PATCH 01/10] [FEA] Add DBR 17.3 GPU low shuffle merge Signed-off-by: Ray Liu --- .../GpuDeltaParquetFileFormatUtils.scala | 52 +- .../rapids/GpuCheckOverflowInTableWrite.scala | 15 +- .../rapids/GpuLowShuffleMergeCommand.scala | 1193 +++++++++++++++++ .../delta/DeltaSpark400DB173Provider.scala | 4 +- .../delta/GpuDeltaParquetFileFormat.scala | 146 +- .../rapids/delta/GpuLowShuffleMergeScan.scala | 50 + .../shims/MergeIntoCommandMetaShim.scala | 96 +- .../advanced_configs.md | 2 +- .../delta_lake_low_shuffle_merge_test.py | 92 +- .../com/nvidia/spark/rapids/RapidsConf.scala | 5 +- 10 files changed, 1573 insertions(+), 82 deletions(-) create mode 100644 delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala create mode 100644 delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuLowShuffleMergeScan.scala diff --git a/delta-lake/common/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatUtils.scala b/delta-lake/common/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatUtils.scala index 1ade53b21b9..5a92b821fd2 100644 --- a/delta-lake/common/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatUtils.scala +++ b/delta-lake/common/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatUtils.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-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. @@ -28,8 +28,8 @@ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} object GpuDeltaParquetFileFormatUtils { /** * Row number of the row in the file. When used with [[FILE_PATH_COL]] together, it can be used - * as unique id of a row in file. Currently to correctly calculate this, the caller needs to - * set both [[isSplitable]] to false, and [[RapidsConf.PARQUET_READER_TYPE]] to "PERFILE". + * as unique id of a row in file. To calculate this correctly, the caller needs to make each file + * unsplittable and reset the row offset when a multi-file reader advances to the next file. */ val METADATA_ROW_IDX_COL: String = "__metadata_row_index" val METADATA_ROW_IDX_FIELD: StructField = StructField(METADATA_ROW_IDX_COL, LongType, @@ -63,23 +63,37 @@ object GpuDeltaParquetFileFormatUtils { } var rowIndex = 0L input.map { batch => - withResource(batch) { _ => - val rowIdxCol = if (metadataRowIndexCol == -1) { - None - } else { - Some(metadataRowIndexCol) - } + val numRows = batch.numRows() + val newBatch = addMetadataColumnsToBatch(schema, delVector, batch, maxBatchSize, + rowIndex, delVectorScatterTimeMetric) + rowIndex += numRows + newBatch + } + } - val delRowIdx2 = if (delRowIdx == -1) { - None - } else { - Some(delRowIdx) - } - val newBatch = addMetadataColumns(rowIdxCol, delRowIdx2, delVector,maxBatchSize, - rowIndex, batch, delVectorScatterTimeMetric) - rowIndex += batch.numRows() - newBatch - } + /** + * Add low-shuffle metadata columns to one batch at the specified file-global row offset. + * This entry point is used by multi-file readers, which reset the offset when the input file + * changes. + */ + def addMetadataColumnsToBatch( + schema: StructType, + delVector: Option[Roaring64Bitmap], + batch: ColumnarBatch, + maxBatchSize: Int, + rowIndex: Long, + delVectorScatterTimeMetric: GpuMetric): ColumnarBatch = { + val metadataRowIndexCol = schema.fieldNames.indexOf(METADATA_ROW_IDX_COL) + val delRowIdx = schema.fieldNames.indexOf(METADATA_ROW_DEL_COL) + withResource(batch) { _ => + addMetadataColumns( + if (metadataRowIndexCol == -1) None else Some(metadataRowIndexCol), + if (delRowIdx == -1) None else Some(delRowIdx), + delVector, + maxBatchSize, + rowIndex, + batch, + delVectorScatterTimeMetric) } } diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuCheckOverflowInTableWrite.scala b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuCheckOverflowInTableWrite.scala index 13cd5593d03..647bf56c30f 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuCheckOverflowInTableWrite.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuCheckOverflowInTableWrite.scala @@ -30,7 +30,10 @@ import org.apache.spark.sql.types.DataType import org.apache.spark.sql.vectorized.ColumnarBatch /** GPU version of Delta's CheckOverflowInTableWrite expression. */ -case class GpuCheckOverflowInTableWrite(child: GpuCast, columnName: String) +case class GpuCheckOverflowInTableWrite( + child: GpuExpression, + columnName: String, + sourceType: DataType) extends ShimUnaryExpression with GpuExpression { override def dataType: DataType = child.dataType @@ -41,7 +44,7 @@ case class GpuCheckOverflowInTableWrite(child: GpuCast, columnName: String) } catch { case _: ArithmeticException => throw DeltaErrors.castingCauseOverflowErrorInTableWrite( - child.child.dataType, + sourceType, dataType, columnName) } @@ -60,9 +63,13 @@ object GpuCheckOverflowInTableWrite { (check, conf, parent, rule) => new UnaryExprMeta[CheckOverflowInTableWrite](check, conf, parent, rule) { override def convertToGpu(child: Expression): GpuExpression = child match { - case cast: GpuCast => GpuCheckOverflowInTableWrite(cast, check.columnName) + case gpuChild: GpuExpression => + val sourceType = check.child.children.headOption + .map(_.dataType) + .getOrElse(check.child.dataType) + GpuCheckOverflowInTableWrite(gpuChild, check.columnName, sourceType) case _ => - throw new IllegalStateException("Expression child is not of type GpuCast") + throw new IllegalStateException("Expression child cannot run on the GPU") } }) } diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala new file mode 100644 index 00000000000..3b3678e3f9c --- /dev/null +++ b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala @@ -0,0 +1,1193 @@ +/* + * Copyright (c) 2024-2026, NVIDIA CORPORATION. + * + * This file was derived from MergeIntoCommand.scala + * in the Delta Lake project at https://github.com/delta-io/delta. + * + * Copyright (2021) The Delta Lake Project Authors. + * + * 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 java.net.URI +import java.util.concurrent.TimeUnit + +import scala.annotation.nowarn +import scala.collection.mutable + +import com.databricks.sql.transaction.tahoe._ +import com.databricks.sql.transaction.tahoe.DeltaOperations.MergePredicate +import com.databricks.sql.transaction.tahoe.actions.{AddCDCFile, AddFile, FileAction} +import com.databricks.sql.transaction.tahoe.commands.DeltaCommand +import com.databricks.sql.transaction.tahoe.commands.merge.MergeIntoMaterializeSource +import com.databricks.sql.transaction.tahoe.files.TahoeFileIndex +import com.databricks.sql.transaction.tahoe.rapids.MergeExecutor.{ + totalBytesAndDistinctPartitionValues, + FILE_PATH_COL, + INCR_METRICS_COL, + INCR_METRICS_FIELD, + ROW_DROPPED_COL, + ROW_DROPPED_FIELD, + SOURCE_ROW_PRESENT_COL, + SOURCE_ROW_PRESENT_FIELD, + TARGET_ROW_PRESENT_COL, + TARGET_ROW_PRESENT_FIELD} +import com.databricks.sql.transaction.tahoe.schema.ImplicitMetadataOperation +import com.databricks.sql.transaction.tahoe.sources.DeltaSQLConf +import com.databricks.sql.transaction.tahoe.util.{AnalysisHelper, DeltaFileOperations} +import com.nvidia.spark.rapids.{GpuOverrides, RapidsConf, SparkPlanMeta} +import com.nvidia.spark.rapids.RapidsConf.DELTA_LOW_SHUFFLE_MERGE_DEL_VECTOR_BROADCAST_THRESHOLD +import com.nvidia.spark.rapids.delta._ +import com.nvidia.spark.rapids.delta.GpuDeltaParquetFileFormatUtils.{ + METADATA_ROW_DEL_COL, + METADATA_ROW_DEL_FIELD, + METADATA_ROW_IDX_COL, + METADATA_ROW_IDX_FIELD} +import com.nvidia.spark.rapids.shims.FileSourceScanExecMeta +import org.roaringbitmap.longlong.Roaring64Bitmap + +import org.apache.spark.SparkContext +import org.apache.spark.internal.Logging +import org.apache.spark.sql._ +import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute +import org.apache.spark.sql.catalyst.catalog.CatalogTable +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeReference, + CaseWhen, Expression, IsNull, Literal, NamedExpression, PredicateHelper} +import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral +import org.apache.spark.sql.catalyst.plans.logical.{DeltaMergeAction, DeltaMergeIntoClause, + DeltaMergeIntoMatchedClause, DeltaMergeIntoMatchedDeleteClause, + DeltaMergeIntoMatchedUpdateClause, DeltaMergeIntoNotMatchedBySourceClause, + DeltaMergeIntoNotMatchedBySourceDeleteClause, DeltaMergeIntoNotMatchedBySourceUpdateClause, + DeltaMergeIntoNotMatchedClause, DeltaMergeIntoNotMatchedInsertClause, LogicalPlan, Project} +import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap +import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} +import org.apache.spark.sql.execution.command.LeafRunnableCommand +import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.nvidia.DFUDFShims +import org.apache.spark.sql.types.{BooleanType, LongType, StringType, StructField, StructType} + +/** + * GPU version of Delta Lake's low shuffle merge implementation. + * + * Performs a merge of a source query/table into a Delta table. + * + * Issues an error message when the ON search_condition of the MERGE statement can match + * a single row from the target table with multiple rows of the source table-reference. + * Different from the original implementation, it optimized writing touched unmodified target files. + * + * Algorithm: + * + * Phase 1: Find the input files in target that are touched by the rows that satisfy + * the condition and verify that no two source rows match with the same target row. + * This is implemented as an inner-join using the given condition. See [[findTouchedFiles]] + * for more details. + * + * Phase 2: Read the touched files again and write new files with updated and/or inserted rows + * without copying unmodified rows. + * + * Phase 3: Read the touched files again and write new files with unmodified rows in target table, + * trying to keep its original order and avoid shuffle as much as possible. + * + * Phase 4: Use the Delta protocol to atomically remove the touched files and add the new files. + * + * @param source Source data to merge from + * @param target Target table to merge into + * @param gpuDeltaLog Delta log to use + * @param condition Condition for a source row to match with a target row + * @param matchedClauses All info related to matched clauses. + * @param notMatchedClauses All info related to not matched clause. + * @param migratedSchema The final schema of the target - may be changed by schema evolution. + */ +case class GpuLowShuffleMergeCommand( + @transient source: LogicalPlan, + @transient target: LogicalPlan, + @transient catalogTable: Option[CatalogTable], + @transient targetFileIndex: TahoeFileIndex, + @transient gpuDeltaLog: GpuDeltaLog, + condition: Expression, + matchedClauses: Seq[DeltaMergeIntoMatchedClause], + notMatchedClauses: Seq[DeltaMergeIntoNotMatchedClause], + notMatchedBySourceClauses: Seq[DeltaMergeIntoNotMatchedBySourceClause], + migratedSchema: Option[StructType], + trackHighWaterMarks: Set[String] = Set.empty, + schemaEvolutionEnabled: Boolean = false, + snapshotAtAnalysis: Option[Snapshot] = None)( + @transient val rapidsConf: RapidsConf) + extends LeafRunnableCommand + with DeltaCommand + with PredicateHelper + with AnalysisHelper + with ImplicitMetadataOperation + with MergeIntoMaterializeSource { + + import SQLMetrics._ + + override val otherCopyArgs: Seq[AnyRef] = Seq(rapidsConf) + + override val canMergeSchema: Boolean = schemaEvolutionEnabled + override val canOverwriteSchema: Boolean = false + + override val output: Seq[Attribute] = Seq( + AttributeReference("num_affected_rows", LongType)(), + AttributeReference("num_updated_rows", LongType)(), + AttributeReference("num_deleted_rows", LongType)(), + AttributeReference("num_inserted_rows", LongType)()) + + @transient private lazy val sc: SparkContext = SparkContext.getOrCreate() + @transient lazy val targetDeltaLog: DeltaLog = gpuDeltaLog.deltaLog + + override lazy val metrics = Map[String, SQLMetric]( + "numSourceRows" -> createMetric(sc, "number of source rows"), + "numSourceRowsInSecondScan" -> + createMetric(sc, "number of source rows (during repeated scan)"), + "numTargetRowsCopied" -> createMetric(sc, "number of target rows rewritten unmodified"), + "numTargetRowsInserted" -> createMetric(sc, "number of inserted rows"), + "numTargetRowsUpdated" -> createMetric(sc, "number of updated rows"), + "numTargetRowsDeleted" -> createMetric(sc, "number of deleted rows"), + "numTargetRowsMatchedUpdated" -> createMetric(sc, "number of target rows updated when matched"), + "numTargetRowsMatchedDeleted" -> createMetric(sc, "number of target rows deleted when matched"), + "numTargetRowsNotMatchedBySourceUpdated" -> createMetric(sc, + "number of target rows updated when not matched by source"), + "numTargetRowsNotMatchedBySourceDeleted" -> createMetric(sc, + "number of target rows deleted when not matched by source"), + "numTargetFilesBeforeSkipping" -> createMetric(sc, "number of target files before skipping"), + "numTargetFilesAfterSkipping" -> createMetric(sc, "number of target files after skipping"), + "numTargetFilesRemoved" -> createMetric(sc, "number of files removed to target"), + "numTargetFilesAdded" -> createMetric(sc, "number of files added to target"), + "numTargetChangeFilesAdded" -> + createMetric(sc, "number of change data capture files generated"), + "numTargetChangeFileBytes" -> + createMetric(sc, "total size of change data capture files generated"), + "numTargetBytesBeforeSkipping" -> createMetric(sc, "number of target bytes before skipping"), + "numTargetBytesAfterSkipping" -> createMetric(sc, "number of target bytes after skipping"), + "numTargetBytesRemoved" -> createMetric(sc, "number of target bytes removed"), + "numTargetBytesAdded" -> createMetric(sc, "number of target bytes added"), + "numTargetPartitionsAfterSkipping" -> + createMetric(sc, "number of target partitions after skipping"), + "numTargetPartitionsRemovedFrom" -> + createMetric(sc, "number of target partitions from which files were removed"), + "numTargetPartitionsAddedTo" -> + createMetric(sc, "number of target partitions to which files were added"), + "executionTimeMs" -> + createMetric(sc, "time taken to execute the entire operation"), + "scanTimeMs" -> + createMetric(sc, "time taken to scan the files for matches"), + "rewriteTimeMs" -> + createMetric(sc, "time taken to rewrite the matched files")) + + /** Whether this merge statement has only a single insert (NOT MATCHED) clause. */ + protected def isSingleInsertOnly: Boolean = matchedClauses.isEmpty && + notMatchedClauses.length == 1 + + private[rapids] def mergeSourceDF: DataFrame = getMergeSource.df + + private def checkIdentityColumnHighWaterMarks(deltaTxn: OptimisticTransaction): Unit = { + notMatchedClauses.foreach { clause => + val schema = deltaTxn.metadata.schema + if (schema.length != clause.resolvedActions.length) { + throw new IllegalStateException() + } + schema.zip(clause.resolvedActions.map(_.expr)).foreach { + case (field, expr: GenerateIdentityValues) => + val highWaterMark = IdentityColumn.getIdentityInfo(field).highWaterMark + if (highWaterMark != expr.generator.highWaterMarkOpt) { + IdentityColumn.logTransactionAbort(deltaTxn.deltaLog) + throw DeltaErrors.metadataChangedException(None) + } + case (field, _) => + if (ColumnWithDefaultExprUtils.isIdentityColumn(field) && + !IdentityColumn.allowExplicitInsert(field)) { + throw new IllegalStateException() + } + } + } + } + + private def runMerge(spark: SparkSession): Seq[Row] = { + recordDeltaOperation(targetDeltaLog, "delta.dml.lowshufflemerge") { + val startTime = System.nanoTime() + val result = gpuDeltaLog.withNewTransaction(catalogTable, snapshotAtAnalysis) { deltaTxn => + if (hasBeenExecuted(deltaTxn, spark)) { + val executionId = spark.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) + SQLMetrics.postDriverMetricUpdates(spark.sparkContext, executionId, metrics.values.toSeq) + return Seq.empty + } + + if (target.schema.size != deltaTxn.metadata.schema.size) { + throw DeltaErrors.schemaChangedSinceAnalysis( + atAnalysis = target.schema, latestSchema = deltaTxn.metadata.schema) + } + + TypeWidening.ensureFeatureConsistentlyEnabled( + protocol = targetFileIndex.protocol, + metadata = targetFileIndex.metadata, + otherProtocol = deltaTxn.protocol, + otherMetadata = deltaTxn.metadata) + + if (canMergeSchema) { + updateMetadata( + spark, deltaTxn, migratedSchema.getOrElse(target.schema), + deltaTxn.metadata.partitionColumns, deltaTxn.metadata.configuration, + isOverwriteMode = false, rearrangeOnly = false) + } + + checkIdentityColumnHighWaterMarks(deltaTxn) + deltaTxn.setTrackHighWaterMarks(trackHighWaterMarks) + + prepareMergeSource( + spark, + source, + condition, + matchedClauses, + notMatchedClauses, + isSingleInsertOnly) + + val executor: MergeExecutor = { + val context = MergeExecutorContext(this, spark, deltaTxn, rapidsConf) + if (isSingleInsertOnly && spark.conf.get(DeltaSQLConf.MERGE_INSERT_ONLY_ENABLED)) { + new InsertOnlyMergeExecutor(context) + } else { + new LowShuffleMergeExecutor(context) + } + } + + try { + val fallback = executor match { + case lowShuffle: LowShuffleMergeExecutor => lowShuffle.shouldFallback() + case _ => false + } + if (fallback) { + None + } else { + Some(runLowShuffleMerge(spark, startTime, deltaTxn, executor)) + } + } finally { + executor.close() + } + } + + result match { + case Some(row) => row + case None => + // We should rollback to normal gpu + new GpuMergeIntoCommand(source, target, catalogTable, targetFileIndex, gpuDeltaLog, + condition, matchedClauses, notMatchedClauses, notMatchedBySourceClauses, + migratedSchema, trackHighWaterMarks, schemaEvolutionEnabled, + snapshotAtAnalysis)(rapidsConf) + .run(spark) + } + } + } + + override def run(spark: SparkSession): Seq[Row] = { + val (materializeSource, _) = shouldMaterializeSource(spark, source, isSingleInsertOnly) + if (materializeSource) { + runWithMaterializedSourceLostRetries(spark, targetDeltaLog, metrics, runMerge) + } else { + runMerge(spark) + } + } + + + private def runLowShuffleMerge( + spark: SparkSession, + startTime: Long, + deltaTxn: GpuOptimisticTransactionBase, + mergeExecutor: MergeExecutor): Seq[Row] = { + val deltaActions = mergeExecutor.execute() + // Metrics should be recorded before commit (where they are written to delta logs). + metrics("executionTimeMs").set((System.nanoTime() - startTime) / 1000 / 1000) + deltaTxn.registerSQLMetrics(spark, metrics) + + // This is a best-effort sanity check. + if (metrics("numSourceRowsInSecondScan").value >= 0 && + metrics("numSourceRows").value != metrics("numSourceRowsInSecondScan").value) { + log.warn(s"Merge source has ${metrics("numSourceRows").value} rows in initial scan but " + + s"${metrics("numSourceRowsInSecondScan").value} rows in second scan") + if (conf.getConf(DeltaSQLConf.MERGE_FAIL_IF_SOURCE_CHANGED)) { + throw DeltaErrors.sourceNotDeterministicInMergeException(spark) + } + } + + val finalActions = createSetTransaction(spark, targetDeltaLog).toSeq ++ deltaActions + deltaTxn.commitIfNeeded( + finalActions, + DeltaOperations.Merge( + Option(condition), + matchedClauses.map(DeltaOperations.MergePredicate(_)), + notMatchedClauses.map(DeltaOperations.MergePredicate(_)), + // We do not support notMatchedBySourcePredicates yet and fall back to CPU + // See https://github.com/NVIDIA/spark-rapids/issues/8415 + notMatchedBySourcePredicates = Seq.empty[MergePredicate] + ), + RowTracking.addPreservedRowTrackingTagIfNotSet(deltaTxn.snapshot)) + + // Record metrics + val stats = GpuMergeStats.fromMergeSQLMetrics( + metrics, + condition, + matchedClauses, + notMatchedClauses, + deltaTxn.metadata.partitionColumns.nonEmpty) + recordDeltaEvent(targetDeltaLog, "delta.dml.merge.stats", data = stats) + + + spark.sharedState.cacheManager.recacheByPlan(spark, target) + + // This is needed to make the SQL metrics visible in the Spark UI. Also this needs + // to be outside the recordMergeOperation because this method will update some metric. + val executionId = spark.sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) + SQLMetrics.postDriverMetricUpdates(spark.sparkContext, executionId, metrics.values.toSeq) + Seq(Row(metrics("numTargetRowsUpdated").value + metrics("numTargetRowsDeleted").value + + metrics("numTargetRowsInserted").value, metrics("numTargetRowsUpdated").value, + metrics("numTargetRowsDeleted").value, metrics("numTargetRowsInserted").value)) + } + + /** + * Execute the given `thunk` and return its result while recording the time taken to do it. + * + * @param sqlMetricName name of SQL metric to update with the time taken by the thunk + * @param thunk the code to execute + */ + def recordMergeOperation[A](sqlMetricName: String)(thunk: => A): A = { + val startTimeNs = System.nanoTime() + val r = thunk + val timeTakenMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTimeNs) + if (sqlMetricName != null && timeTakenMs > 0) { + metrics(sqlMetricName) += timeTakenMs + } + r + } + + /** Expressions to increment SQL metrics */ + def makeMetricUpdateUDF(name: String, deterministic: Boolean = false): Column = { + // only capture the needed metric in a local variable + val metric = metrics(name) + var u = DeltaUDF.boolean(new GpuDeltaMetricUpdateUDF(metric)) + if (!deterministic) { + u = u.asNondeterministic() + } + u() + } + + @nowarn("cat=deprecation") + def metricUpdateExpr(name: String, deterministic: Boolean): Expression = { + makeMetricUpdateUDF(name, deterministic).expr + } +} + +/** + * Context merge execution. + */ +case class MergeExecutorContext(cmd: GpuLowShuffleMergeCommand, + spark: SparkSession, + deltaTxn: OptimisticTransaction, + rapidsConf: RapidsConf) + +trait MergeExecutor extends AnalysisHelper with PredicateHelper with Logging with AutoCloseable { + + val context: MergeExecutorContext + + + /** + * Map to get target output attributes by name. + * The case sensitivity of the map is set accordingly to Spark configuration. + */ + @transient private lazy val targetOutputAttributesMap: Map[String, Attribute] = { + val attrMap: Map[String, Attribute] = context.cmd.target + .outputSet.view + .map(attr => attr.name -> attr).toMap + if (context.cmd.conf.caseSensitiveAnalysis) { + attrMap + } else { + CaseInsensitiveMap(attrMap) + } + } + + def execute(): Seq[FileAction] + + override def close(): Unit = {} + + protected def targetOutputCols: Seq[NamedExpression] = { + context.deltaTxn.metadata.schema.map { col => + targetOutputAttributesMap + .get(col.name) + .map { a => + AttributeReference(col.name, col.dataType, col.nullable)(a.exprId) + } + .getOrElse(Alias(Literal(null, col.dataType), col.name)()) + } + } + + /** + * Build a DataFrame using the given `files` that has the same output columns (exprIds) + * as the `target` logical plan, so that existing update/insert expressions can be applied + * on this new plan. + */ + protected def buildTargetDFWithFiles(files: Seq[AddFile]): DataFrame = { + val targetOutputColsMap = { + val colsMap: Map[String, NamedExpression] = targetOutputCols.view + .map(col => col.name -> col).toMap + if (context.cmd.conf.caseSensitiveAnalysis) { + colsMap + } else { + CaseInsensitiveMap(colsMap) + } + } + + val plan = { + // We have to do surgery to use the attributes from `targetOutputCols` to scan the table. + // In cases of schema evolution, they may not be the same type as the original attributes. + val original = + context.deltaTxn.deltaLog.createDataFrame(context.deltaTxn.snapshot, files) + .queryExecution + .analyzed + val transformed = original.transform { + case r: LogicalRelation => + r.copy( + // We can ignore the new columns which aren't yet AttributeReferences. + output = targetOutputCols.collect { case a: AttributeReference => a }) + } + + // In case of schema evolution & column mapping, we would also need to rebuild the file + // format because under column mapping, the reference schema within DeltaParquetFileFormat + // that is used to populate metadata needs to be updated + if (context.deltaTxn.metadata.columnMappingMode != NoMapping) { + val updatedFileFormat = context.deltaTxn.deltaLog.fileFormat( + context.deltaTxn.deltaLog.unsafeVolatileSnapshot.protocol, context.deltaTxn.metadata) + DeltaTableUtils.replaceFileFormat(transformed, updatedFileFormat) + } else { + transformed + } + } + + // For each plan output column, find the corresponding target output column (by name) and + // create an alias + val aliases = plan.output.map { + case newAttrib: AttributeReference => + val existingTargetAttrib = targetOutputColsMap.getOrElse(newAttrib.name, + throw new AnalysisException( + s"Could not find ${newAttrib.name} among the existing target output " + + targetOutputCols.mkString(","))).asInstanceOf[AttributeReference] + + if (existingTargetAttrib.exprId == newAttrib.exprId) { + // It's not valid to alias an expression to its own exprId (this is considered a + // non-unique exprId by the analyzer), so we just use the attribute directly. + newAttrib + } else { + Alias(newAttrib, existingTargetAttrib.name)(exprId = existingTargetAttrib.exprId) + } + } + + Dataset.ofRows(context.spark, Project(aliases, plan)) + } + + + /** + * Repartitions the output DataFrame by the partition columns if table is partitioned + * and `merge.repartitionBeforeWrite.enabled` is set to true. + */ + protected def repartitionIfNeeded(df: DataFrame): DataFrame = { + val partitionColumns = context.deltaTxn.metadata.partitionColumns + // TODO: We should remove this method and use optimized write instead, see + // https://github.com/NVIDIA/spark-rapids/issues/10417 + if (partitionColumns.nonEmpty && context.spark.conf.get(DeltaSQLConf + .MERGE_REPARTITION_BEFORE_WRITE)) { + df.repartition(partitionColumns.map(col): _*) + } else { + df + } + } + + protected def sourceDF: DataFrame = { + // UDF to increment metrics + val incrSourceRowCountCol = context.cmd.makeMetricUpdateUDF("numSourceRows") + context.cmd.mergeSourceDF.filter(incrSourceRowCountCol) + } + + /** Whether this merge statement has no insert (NOT MATCHED) clause. */ + protected def hasNoInserts: Boolean = context.cmd.notMatchedClauses.isEmpty + + +} + +/** + * This is an optimization of the case when there is no update clause for the merge. + * We perform an left anti join on the source data to find the rows to be inserted. + * + * This will currently only optimize for the case when there is a _single_ notMatchedClause. + */ +class InsertOnlyMergeExecutor(override val context: MergeExecutorContext) extends MergeExecutor { + override def execute(): Seq[FileAction] = { + context.cmd.recordMergeOperation(sqlMetricName = "rewriteTimeMs") { + + // UDFs to update metrics + val incrSourceRowCountCol = context.cmd.makeMetricUpdateUDF("numSourceRows") + val incrInsertedCountCol = context.cmd.makeMetricUpdateUDF("numTargetRowsInserted") + + val outputColNames = targetOutputCols.map(_.name) + // we use head here since we know there is only a single notMatchedClause + val outputExprs = context.cmd.notMatchedClauses.head.resolvedActions.map(_.expr) + val outputCols = outputExprs.zip(outputColNames).map { case (expr, name) => + DFUDFShims.exprToColumn(Alias(expr, name)()) + } + + // source DataFrame + val sourceDF = context.cmd.mergeSourceDF + .filter(incrSourceRowCountCol) + .filter(DFUDFShims.exprToColumn(context.cmd.notMatchedClauses.head.condition + .getOrElse(Literal.TrueLiteral))) + + // Skip data based on the merge condition + val conjunctivePredicates = splitConjunctivePredicates(context.cmd.condition) + val targetOnlyPredicates = + conjunctivePredicates.filter(_.references.subsetOf(context.cmd.target.outputSet)) + val dataSkippedFiles = context.deltaTxn.filterFiles(targetOnlyPredicates) + + // target DataFrame + val targetDF = buildTargetDFWithFiles(dataSkippedFiles) + + val insertDf = sourceDF.join( + targetDF, DFUDFShims.exprToColumn(context.cmd.condition), "leftanti") + .select(outputCols: _*) + .filter(incrInsertedCountCol) + + val newFiles = context.deltaTxn.writeFiles(repartitionIfNeeded(insertDf)) + + // Update metrics + context.cmd.metrics("numTargetFilesBeforeSkipping") += context.deltaTxn.snapshot.numOfFiles + context.cmd.metrics("numTargetBytesBeforeSkipping") += context.deltaTxn.snapshot.sizeInBytes + val (afterSkippingBytes, afterSkippingPartitions) = + totalBytesAndDistinctPartitionValues(dataSkippedFiles) + context.cmd.metrics("numTargetFilesAfterSkipping") += dataSkippedFiles.size + context.cmd.metrics("numTargetBytesAfterSkipping") += afterSkippingBytes + context.cmd.metrics("numTargetPartitionsAfterSkipping") += afterSkippingPartitions + context.cmd.metrics("numTargetFilesRemoved") += 0 + context.cmd.metrics("numTargetBytesRemoved") += 0 + context.cmd.metrics("numTargetPartitionsRemovedFrom") += 0 + val (addedBytes, addedPartitions) = totalBytesAndDistinctPartitionValues(newFiles) + context.cmd.metrics("numTargetFilesAdded") += newFiles.count(_.isInstanceOf[AddFile]) + context.cmd.metrics("numTargetBytesAdded") += addedBytes + context.cmd.metrics("numTargetPartitionsAddedTo") += addedPartitions + newFiles + } + } +} + + +/** + * This is an optimized algorithm for merge statement, where we avoid shuffling the unmodified + * target data. + * + * The algorithm is as follows: + * 1. Find touched target files in the target table by joining the source and target data, with + * collecting joined row identifiers as (`__metadata_file_path`, `__metadata_row_idx`) pairs. + * 2. Read the touched files again and write new files with updated and/or inserted rows + * without coping unmodified data from target table, but filtering target table with collected + * rows mentioned above. + * 3. Read the touched files again, filtering unmodified rows with collected row identifiers + * collected in first step, and saving them without shuffle. + */ +class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extends MergeExecutor { + + private val scanRegistrationIds = new mutable.ArrayBuffer[String]() + private var touchedRowsBroadcast = Option.empty[ + org.apache.spark.broadcast.Broadcast[Map[URI, Array[Byte]]]] + + override def close(): Unit = { + scanRegistrationIds.foreach(GpuLowShuffleMergeScanRegistry.remove) + touchedRowsBroadcast.foreach(_.destroy()) + } + + // We over-count numTargetRowsDeleted when there are multiple matches; + // this is the amount of the overcount, so we can subtract it to get a correct final metric. + private var multipleMatchDeleteOnlyOvercount: Option[Long] = None + + // UDFs to update metrics + private val incrSourceRowCountExpr: Expression = context.cmd + .metricUpdateExpr("numSourceRowsInSecondScan", deterministic = false) + private val incrUpdatedCountExpr: Expression = context.cmd + .metricUpdateExpr("numTargetRowsUpdated", deterministic = false) + private val incrUpdatedMatchedCountExpr: Expression = context.cmd + .metricUpdateExpr("numTargetRowsMatchedUpdated", deterministic = false) + private val incrUpdatedNotMatchedBySourceCountExpr: Expression = context.cmd + .metricUpdateExpr("numTargetRowsNotMatchedBySourceUpdated", deterministic = false) + private val incrInsertedCountExpr: Expression = context.cmd + .metricUpdateExpr("numTargetRowsInserted", deterministic = false) + private val incrDeletedCountExpr: Expression = context.cmd + .metricUpdateExpr("numTargetRowsDeleted", deterministic = false) + private val incrDeletedMatchedCountExpr: Expression = context.cmd + .metricUpdateExpr("numTargetRowsMatchedDeleted", deterministic = false) + private val incrDeletedNotMatchedBySourceCountExpr: Expression = context.cmd + .metricUpdateExpr("numTargetRowsNotMatchedBySourceDeleted", deterministic = false) + + private def updateOutput(resolvedActions: Seq[DeltaMergeAction], incrExpr: Expression) + : Seq[Expression] = { + resolvedActions.map(_.expr) :+ + Literal.FalseLiteral :+ + UnresolvedAttribute(TARGET_ROW_PRESENT_COL) :+ + UnresolvedAttribute(SOURCE_ROW_PRESENT_COL) :+ + incrExpr + } + + private def deleteOutput(incrExpr: Expression): Seq[Expression] = { + targetOutputCols :+ + TrueLiteral :+ + UnresolvedAttribute(TARGET_ROW_PRESENT_COL) :+ + UnresolvedAttribute(SOURCE_ROW_PRESENT_COL) :+ + incrExpr + } + + private def insertOutput(resolvedActions: Seq[DeltaMergeAction], incrExpr: Expression) + : Seq[Expression] = { + resolvedActions.map(_.expr) :+ + Literal.FalseLiteral :+ + UnresolvedAttribute(TARGET_ROW_PRESENT_COL) :+ + UnresolvedAttribute(SOURCE_ROW_PRESENT_COL) :+ + incrExpr + } + + private def clauseOutput(clause: DeltaMergeIntoClause): Seq[Expression] = clause match { + case u: DeltaMergeIntoMatchedUpdateClause => + updateOutput(u.resolvedActions, And(incrUpdatedCountExpr, incrUpdatedMatchedCountExpr)) + case _: DeltaMergeIntoMatchedDeleteClause => + deleteOutput(And(incrDeletedCountExpr, incrDeletedMatchedCountExpr)) + case i: DeltaMergeIntoNotMatchedInsertClause => + insertOutput(i.resolvedActions, incrInsertedCountExpr) + case u: DeltaMergeIntoNotMatchedBySourceUpdateClause => + updateOutput(u.resolvedActions, + And(incrUpdatedCountExpr, incrUpdatedNotMatchedBySourceCountExpr)) + case _: DeltaMergeIntoNotMatchedBySourceDeleteClause => + deleteOutput(And(incrDeletedCountExpr, incrDeletedNotMatchedBySourceCountExpr)) + } + + private def clauseCondition(clause: DeltaMergeIntoClause): Expression = { + // if condition is None, then expression always evaluates to true + clause.condition.getOrElse(TrueLiteral) + } + + /** + * Though low shuffle merge algorithm performs better than traditional merge algorithm in some + * cases, there are some case we should fallback to traditional merge executor: + * + * 1. Low shuffle merge algorithm requires generating metadata columns such as + * [[METADATA_ROW_IDX_COL]], [[METADATA_ROW_DEL_COL]], which only implemented on + * [[org.apache.spark.sql.rapids.GpuFileSourceScanExec]]. That means we need to fallback to + * this normal executor when [[org.apache.spark.sql.rapids.GpuFileSourceScanExec]] is disabled + * for some reason. + * 2. Low shuffle merge algorithm currently needs to broadcast deletion vector, which may + * introduce extra overhead. It maybe better to fallback to this algorithm when the changeset + * it too large. + */ + def shouldFallback(): Boolean = { + // Trying to detect if we can execute finding touched files. + val touchFilePlanOverrideSucceed = verifyGpuPlan(planForFindingTouchedFiles()) { planMeta => + def check(meta: SparkPlanMeta[SparkPlan]): Boolean = { + meta match { + case scan if scan.isInstanceOf[FileSourceScanExecMeta] => scan + .asInstanceOf[FileSourceScanExecMeta] + .wrapped + .schema + .fieldNames + .contains(METADATA_ROW_IDX_COL) && scan.canThisBeReplaced + case m => m.childPlans.exists(check) + } + } + + check(planMeta) + } + if (!touchFilePlanOverrideSucceed) { + logWarning("Unable to override file scan for low shuffle merge for finding touched files " + + "plan, fallback to tradition merge.") + return true + } + + // Trying to detect if we can execute the merge plan. + val mergePlanOverrideSucceed = verifyGpuPlan(planForMergeExecution(touchedFiles)) { planMeta => + var overrideCount = 0 + def count(meta: SparkPlanMeta[SparkPlan]): Unit = { + meta match { + case scan if scan.isInstanceOf[FileSourceScanExecMeta] => + if (scan.asInstanceOf[FileSourceScanExecMeta] + .wrapped.schema.fieldNames.contains(METADATA_ROW_DEL_COL) && scan.canThisBeReplaced) { + overrideCount += 1 + } + case m => m.childPlans.foreach(count) + } + } + + count(planMeta) + overrideCount == 2 + } + + if (!mergePlanOverrideSucceed) { + logWarning("Unable to override file scan for low shuffle merge for merge plan, fallback to " + + "tradition merge.") + return true + } + + val deletionVectorSize = touchedFiles.values.map(_._1.serializedSizeInBytes()).sum + val maxDelVectorSize = context.rapidsConf + .get(DELTA_LOW_SHUFFLE_MERGE_DEL_VECTOR_BROADCAST_THRESHOLD) + if (deletionVectorSize > maxDelVectorSize) { + logWarning( + s"""Low shuffle merge can't be executed because broadcast deletion vector count + |$deletionVectorSize is large than max value $maxDelVectorSize """.stripMargin) + return true + } + + false + } + + private def verifyGpuPlan(input: DataFrame)(checkPlanMeta: SparkPlanMeta[SparkPlan] => Boolean) + : Boolean = { + val overridePlan = GpuOverrides.wrapAndTagPlan(input.queryExecution.sparkPlan, + context.rapidsConf) + checkPlanMeta(overridePlan) + } + + override def execute(): Seq[FileAction] = { + val newFiles = context.cmd.withStatusCode("DELTA", + s"Rewriting ${touchedFiles.size} files and saving modified data") { + val df = planForMergeExecution(touchedFiles) + context.deltaTxn.writeFiles(df) + } + + // Update metrics + val (addedBytes, addedPartitions) = totalBytesAndDistinctPartitionValues(newFiles) + context.cmd.metrics("numTargetFilesAdded") += newFiles.count(_.isInstanceOf[AddFile]) + context.cmd.metrics("numTargetChangeFilesAdded") += newFiles.count(_.isInstanceOf[AddCDCFile]) + context.cmd.metrics("numTargetChangeFileBytes") += newFiles.collect { + case f: AddCDCFile => f.size + } + .sum + context.cmd.metrics("numTargetBytesAdded") += addedBytes + context.cmd.metrics("numTargetPartitionsAddedTo") += addedPartitions + + if (multipleMatchDeleteOnlyOvercount.isDefined) { + // Compensate for counting duplicates during the query. + val actualRowsDeleted = + context.cmd.metrics("numTargetRowsDeleted").value - multipleMatchDeleteOnlyOvercount.get + assert(actualRowsDeleted >= 0) + context.cmd.metrics("numTargetRowsDeleted").set(actualRowsDeleted) + } + + touchedFiles.values.map(_._2).map(_.remove).toSeq ++ newFiles + } + + private lazy val dataSkippedFiles: Seq[AddFile] = { + // Skip data based on the merge condition + val targetOnlyPredicates = splitConjunctivePredicates(context.cmd.condition) + .filter(_.references.subsetOf(context.cmd.target.outputSet)) + context.deltaTxn.filterFiles(targetOnlyPredicates) + } + + private lazy val dataSkippedTargetDF: DataFrame = { + addRowIndexMetaColumn(buildTargetDFWithFiles(dataSkippedFiles)) + } + + private lazy val touchedFiles: Map[String, (Roaring64Bitmap, AddFile)] = this.findTouchedFiles() + + private def planForFindingTouchedFiles(): DataFrame = { + + // Apply inner join to between source and target using the merge condition to find matches + // In addition, we attach two columns + // - METADATA_ROW_IDX column to identify target row in file + // - FILE_PATH_COL the target file name the row is from to later identify the files touched + // by matched rows + val targetDF = dataSkippedTargetDF.withColumn(FILE_PATH_COL, input_file_name()) + + sourceDF.join(targetDF, DFUDFShims.exprToColumn(context.cmd.condition), "inner") + } + + private def planForMergeExecution(touchedFiles: Map[String, (Roaring64Bitmap, AddFile)]) + : DataFrame = { + getModifiedDF(touchedFiles).unionAll(getUnmodifiedDF(touchedFiles)) + } + + /** + * Find the target table files that contain the rows that satisfy the merge condition. This is + * implemented as an inner-join between the source query/table and the target table using + * the merge condition. + */ + private def findTouchedFiles(): Map[String, (Roaring64Bitmap, AddFile)] = + context.cmd.recordMergeOperation(sqlMetricName = "scanTimeMs") { + context.spark.udf.register("row_index_set", udaf(RoaringBitmapUDAF)) + // Process the matches from the inner join to record touched files and find multiple matches + val collectTouchedFiles = planForFindingTouchedFiles() + .select(col(FILE_PATH_COL), col(METADATA_ROW_IDX_COL)) + .groupBy(FILE_PATH_COL) + .agg( + expr(s"row_index_set($METADATA_ROW_IDX_COL) as row_idxes"), + count("*").as("count")) + .collect().map(row => { + val filename = row.getAs[String](FILE_PATH_COL) + val rowIdxSet = row.getAs[RoaringBitmapWrapper]("row_idxes").inner + val count = row.getAs[Long]("count") + (filename, (rowIdxSet, count)) + }) + .toMap + + val duplicateCount = { + val distinctMatchedRowCounts = collectTouchedFiles.values + .map(_._1.getLongCardinality).sum + val allMatchedRowCounts = collectTouchedFiles.values.map(_._2).sum + allMatchedRowCounts - distinctMatchedRowCounts + } + + val hasMultipleMatches = duplicateCount > 0 + + // Throw error if multiple matches are ambiguous or cannot be computed correctly. + val canBeComputedUnambiguously = { + // Multiple matches are not ambiguous when there is only one unconditional delete as + // all the matched row pairs in the 2nd join in `writeAllChanges` will get deleted. + val isUnconditionalDelete = context.cmd.matchedClauses.headOption match { + case Some(DeltaMergeIntoMatchedDeleteClause(None)) => true + case _ => false + } + context.cmd.matchedClauses.size == 1 && isUnconditionalDelete + } + + if (hasMultipleMatches && !canBeComputedUnambiguously) { + throw DeltaErrors.multipleSourceRowMatchingTargetRowInMergeException(context.spark) + } + + if (hasMultipleMatches) { + // This is only allowed for delete-only queries. + // This query will count the duplicates for numTargetRowsDeleted in Job 2, + // because we count matches after the join and not just the target rows. + // We have to compensate for this by subtracting the duplicates later, + // so we need to record them here. + multipleMatchDeleteOnlyOvercount = Some(duplicateCount) + } + + // Get the AddFiles using the touched file names. + val touchedFileNames = collectTouchedFiles.keys.toSeq + + val nameToAddFileMap = context.cmd.generateCandidateFileMap( + context.cmd.targetDeltaLog.dataPath, + dataSkippedFiles) + + val touchedAddFiles = touchedFileNames.map(f => + context.cmd.getTouchedFile(context.cmd.targetDeltaLog.dataPath, f, nameToAddFileMap)) + .map(f => (DeltaFileOperations + .absolutePath(context.cmd.targetDeltaLog.dataPath.toString, f.path) + .toString, f)).toMap + + // When the target table is empty, and the optimizer optimized away the join entirely + // numSourceRows will be incorrectly 0. + // We need to scan the source table once to get the correct + // metric here. + if (context.cmd.metrics("numSourceRows").value == 0 && + (dataSkippedFiles.isEmpty || dataSkippedTargetDF.take(1).isEmpty)) { + val numSourceRows = sourceDF.count() + context.cmd.metrics("numSourceRows").set(numSourceRows) + } + + // Update metrics + context.cmd.metrics("numTargetFilesBeforeSkipping") += context.deltaTxn.snapshot.numOfFiles + context.cmd.metrics("numTargetBytesBeforeSkipping") += context.deltaTxn.snapshot.sizeInBytes + val (afterSkippingBytes, afterSkippingPartitions) = + totalBytesAndDistinctPartitionValues(dataSkippedFiles) + context.cmd.metrics("numTargetFilesAfterSkipping") += dataSkippedFiles.size + context.cmd.metrics("numTargetBytesAfterSkipping") += afterSkippingBytes + context.cmd.metrics("numTargetPartitionsAfterSkipping") += afterSkippingPartitions + val (removedBytes, removedPartitions) = + totalBytesAndDistinctPartitionValues(touchedAddFiles.values.toSeq) + context.cmd.metrics("numTargetFilesRemoved") += touchedAddFiles.size + context.cmd.metrics("numTargetBytesRemoved") += removedBytes + context.cmd.metrics("numTargetPartitionsRemovedFrom") += removedPartitions + + collectTouchedFiles.map(kv => (kv._1, (kv._2._1, touchedAddFiles(kv._1)))) + } + + + /** + * Modify original data frame to insert + * [[GpuDeltaParquetFileFormatUtils.METADATA_ROW_IDX_COL]]. + */ + private def addRowIndexMetaColumn(baseDF: DataFrame): DataFrame = { + val rowIdxAttr = AttributeReference( + METADATA_ROW_IDX_COL, + METADATA_ROW_IDX_FIELD.dataType, + METADATA_ROW_IDX_FIELD.nullable)() + + val newPlan = baseDF.queryExecution.analyzed.transformUp { + case r: LogicalRelation if r.relation.isInstanceOf[HadoopFsRelation] => + val fs = r.relation.asInstanceOf[HadoopFsRelation] + val newSchema = StructType(fs.dataSchema.fields).add(METADATA_ROW_IDX_FIELD) + val newFs = lowShuffleScanRelation( + fs, newSchema, GpuLowShuffleMergeScanInfo(rowIndexMaps = None)) + + val newOutput = r.output :+ rowIdxAttr + r.copy(relation = newFs, output = newOutput) + case p@Project(projectList, _) => + val newProjectList = projectList :+ rowIdxAttr + p.copy(projectList = newProjectList) + } + + Dataset.ofRows(context.spark, newPlan) + } + + /** + * The result is scanning target table with touched files, and added an extra + * [[METADATA_ROW_DEL_COL]] to indicate whether filtered by joining with source table in first + * step. + */ + private def getTouchedTargetDF(touchedFiles: Map[String, (Roaring64Bitmap, AddFile)]) + : DataFrame = { + // Generate a new target dataframe that has same output attributes exprIds as the target plan. + // This allows us to apply the existing resolved update/insert expressions. + val baseTargetDF = buildTargetDFWithFiles(touchedFiles.values.map(_._2).toSeq) + + val newPlan = { + val rowDelAttr = AttributeReference( + METADATA_ROW_DEL_COL, + METADATA_ROW_DEL_FIELD.dataType, + METADATA_ROW_DEL_FIELD.nullable)() + + baseTargetDF.queryExecution.analyzed.transformUp { + case r: LogicalRelation if r.relation.isInstanceOf[HadoopFsRelation] => + val fs = r.relation.asInstanceOf[HadoopFsRelation] + val newSchema = StructType(fs.dataSchema.fields).add(METADATA_ROW_DEL_FIELD) + val broadcastRows = touchedRowsBroadcast.getOrElse { + val rows = touchedFiles.map { case (path, (bitmap, _)) => + new URI(path) -> RoaringBitmapWrapper(bitmap).serializeToBytes() + } + val broadcast = context.spark.sparkContext.broadcast(rows) + touchedRowsBroadcast = Some(broadcast) + broadcast + } + val newFs = lowShuffleScanRelation( + fs, newSchema, GpuLowShuffleMergeScanInfo(rowIndexMaps = Some(broadcastRows))) + + val newOutput = r.output :+ rowDelAttr + r.copy(relation = newFs, output = newOutput) + case p@Project(projectList, _) => + val newProjectList = projectList :+ rowDelAttr + p.copy(projectList = newProjectList) + } + } + + val df = Dataset.ofRows(context.spark, newPlan) + .withColumn(TARGET_ROW_PRESENT_COL, lit(true)) + + df + } + + private def lowShuffleScanRelation( + relation: HadoopFsRelation, + dataSchema: StructType, + scanInfo: GpuLowShuffleMergeScanInfo): HadoopFsRelation = { + val scanId = GpuLowShuffleMergeScanRegistry.register(scanInfo) + scanRegistrationIds += scanId + val fileFormat = relation.fileFormat.asInstanceOf[DeltaParquetFileFormat] + .copy(optimizationsEnabled = false) + relation.copy( + dataSchema = dataSchema, + fileFormat = fileFormat, + options = relation.options + (GpuLowShuffleMergeScanRegistry.OPTION_KEY -> scanId))( + context.spark) + } + + /** + * Generate a plan by calculating modified rows. It's computed by joining source and target + * tables, where target table has been filtered by (`__metadata_file_name`, + * `__metadata_row_idx`) pairs collected in first step. + * + * Schema of `modifiedDF`: + * + * targetSchema + ROW_DROPPED_COL + TARGET_ROW_PRESENT_COL + + * SOURCE_ROW_PRESENT_COL + INCR_METRICS_COL + * INCR_METRICS_COL + * + * It consists of several parts: + * + * 1. Unmatched source rows which are inserted + * 2. Unmatched source rows which are deleted + * 3. Target rows which are updated + * 4. Target rows which are deleted + */ + private def getModifiedDF(touchedFiles: Map[String, (Roaring64Bitmap, AddFile)]): DataFrame = { + val sourceDF = this.sourceDF + .withColumn(SOURCE_ROW_PRESENT_COL, DFUDFShims.exprToColumn(incrSourceRowCountExpr)) + + val targetDF = getTouchedTargetDF(touchedFiles) + + val joinedDF = { + val joinType = if (hasNoInserts && + context.spark.conf.get(DeltaSQLConf.MERGE_MATCHED_ONLY_ENABLED)) { + "inner" + } else { + "leftOuter" + } + val matchedTargetDF = targetDF.filter(METADATA_ROW_DEL_COL) + .drop(METADATA_ROW_DEL_COL) + + sourceDF.join(matchedTargetDF, DFUDFShims.exprToColumn(context.cmd.condition), joinType) + } + + val modifiedRowsSchema = context.deltaTxn.metadata.schema + .add(ROW_DROPPED_FIELD) + .add(TARGET_ROW_PRESENT_FIELD.copy(nullable = true)) + .add(SOURCE_ROW_PRESENT_FIELD.copy(nullable = true)) + .add(INCR_METRICS_FIELD) + + // Here we generate a case when statement to handle all cases: + // CASE + // WHEN + // CASE WHEN + // + // WHEN + // + // ELSE + // + // WHEN + // CASE WHEN + // + // WHEN + // + // ELSE + // + // END + + val notMatchedConditions = context.cmd.notMatchedClauses.map(clauseCondition) + val notMatchedExpr = { + val deletedNotMatchedRow = { + targetOutputCols :+ + Literal.TrueLiteral :+ + Literal.FalseLiteral :+ + Literal(null) :+ + Literal.TrueLiteral + } + if (context.cmd.notMatchedClauses.isEmpty) { + // If there no `WHEN NOT MATCHED` clause, we should just delete not matched row + deletedNotMatchedRow + } else { + val notMatchedOutputs = context.cmd.notMatchedClauses.map(clauseOutput) + modifiedRowsSchema.zipWithIndex.map { + case (_, idx) => + CaseWhen(notMatchedConditions.zip(notMatchedOutputs.map(_(idx))), + deletedNotMatchedRow(idx)) + } + } + } + + val matchedConditions = context.cmd.matchedClauses.map(clauseCondition) + val matchedOutputs = context.cmd.matchedClauses.map(clauseOutput) + val matchedExprs = { + val notMatchedRow = { + targetOutputCols :+ + Literal.FalseLiteral :+ + Literal.TrueLiteral :+ + Literal(null) :+ + Literal.TrueLiteral + } + if (context.cmd.matchedClauses.isEmpty) { + // If there is not matched clause, this is insert only, we should delete this row. + notMatchedRow + } else { + modifiedRowsSchema.zipWithIndex.map { + case (_, idx) => + CaseWhen(matchedConditions.zip(matchedOutputs.map(_(idx))), + notMatchedRow(idx)) + } + } + } + + val sourceRowHasNoMatch = IsNull(UnresolvedAttribute(TARGET_ROW_PRESENT_COL)) + + val modifiedCols = modifiedRowsSchema.zipWithIndex.map { case (col, idx) => + val caseWhen = CaseWhen( + Seq(sourceRowHasNoMatch -> notMatchedExpr(idx)), + matchedExprs(idx)) + DFUDFShims.exprToColumn(Alias(caseWhen, col.name)()) + } + + val modifiedDF = { + + // Make this a udf to avoid catalyst to be too aggressive to even remove the join! + val noopRowDroppedCol = udf(new GpuDeltaNoopUDF()).apply(!col(ROW_DROPPED_COL)) + + val modifiedDF = joinedDF.select(modifiedCols: _*) + // This will not filter anything since they always return true, but we need to avoid + // catalyst from optimizing these udf + .filter(noopRowDroppedCol && col(INCR_METRICS_COL)) + .drop(ROW_DROPPED_COL, INCR_METRICS_COL, TARGET_ROW_PRESENT_COL, SOURCE_ROW_PRESENT_COL) + + repartitionIfNeeded(modifiedDF) + } + + modifiedDF + } + + private def getUnmodifiedDF(touchedFiles: Map[String, (Roaring64Bitmap, AddFile)]): DataFrame = { + getTouchedTargetDF(touchedFiles) + .filter(!col(METADATA_ROW_DEL_COL)) + .drop(TARGET_ROW_PRESENT_COL, METADATA_ROW_DEL_COL) + } +} + + +object MergeExecutor { + + /** + * Spark UI will track all normal accumulators along with Spark tasks to show them on Web UI. + * However, the accumulator used by `MergeIntoCommand` can store a very large value since it + * tracks all files that need to be rewritten. We should ask Spark UI to not remember it, + * otherwise, the UI data may consume lots of memory. Hence, we use the prefix `internal.metrics.` + * to make this accumulator become an internal accumulator, so that it will not be tracked by + * Spark UI. + */ + val TOUCHED_FILES_ACCUM_NAME = "internal.metrics.MergeIntoDelta.touchedFiles" + + val ROW_ID_COL = "_row_id_" + val FILE_PATH_COL: String = GpuDeltaParquetFileFormatUtils.FILE_PATH_COL + val SOURCE_ROW_PRESENT_COL: String = "_source_row_present_" + val SOURCE_ROW_PRESENT_FIELD: StructField = StructField(SOURCE_ROW_PRESENT_COL, BooleanType, + nullable = false) + val TARGET_ROW_PRESENT_COL: String = "_target_row_present_" + val TARGET_ROW_PRESENT_FIELD: StructField = StructField(TARGET_ROW_PRESENT_COL, BooleanType, + nullable = false) + val ROW_DROPPED_COL: String = GpuDeltaMergeConstants.ROW_DROPPED_COL + val ROW_DROPPED_FIELD: StructField = StructField(ROW_DROPPED_COL, BooleanType, nullable = false) + val INCR_METRICS_COL: String = "_incr_metrics_" + val INCR_METRICS_FIELD: StructField = StructField(INCR_METRICS_COL, BooleanType, nullable = false) + val INCR_ROW_COUNT_COL: String = "_incr_row_count_" + + // Some Delta versions use Literal(null) which translates to a literal of NullType instead + // of the Literal(null, StringType) which is needed, so using a fixed version here + // rather than the version from Delta Lake. + val CDC_TYPE_NOT_CDC_LITERAL: Literal = Literal(null, StringType) + + /** Count the number of distinct partition values among the AddFiles in the given set. */ + def totalBytesAndDistinctPartitionValues(files: Seq[FileAction]): (Long, Int) = { + val distinctValues = new mutable.HashSet[Map[String, String]]() + var bytes = 0L + val iter = files.collect { case a: AddFile => a }.iterator + while (iter.hasNext) { + val file = iter.next() + distinctValues += file.partitionValues + bytes += file.size + } + // If the only distinct value map is an empty map, then it must be an unpartitioned table. + // Return 0 in that case. + val numDistinctValues = + if (distinctValues.size == 1 && distinctValues.head.isEmpty) 0 else distinctValues.size + (bytes, numDistinctValues) + } +} diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala index 3a0372dfeaa..1dda69328c8 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala @@ -122,7 +122,9 @@ object DeltaSpark400DB173Provider extends DatabricksDeltaProviderBase { override def getReadFileFormat( relation: HadoopFsRelation, rapidsConf: RapidsConf): FileFormat = { val fmt = relation.fileFormat.asInstanceOf[DeltaParquetFileFormat] - if (isPushDVPredicateDownEnabled(rapidsConf)) { + if (GpuLowShuffleMergeScanRegistry.lookup(relation.options).isDefined) { + GpuDeltaParquetFileFormat.convertToGpu(relation) + } else if (isPushDVPredicateDownEnabled(rapidsConf)) { GpuDeltaParquetFileFormatNativeDV( relation = relation, protocol = fmt.protocol, 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..55d21bba02d 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 @@ -16,6 +16,8 @@ package com.nvidia.spark.rapids.delta +import java.net.URI + import com.databricks.sql.io.RowIndexFilterType import com.databricks.sql.transaction.tahoe.{ DeltaColumnMapping, @@ -28,19 +30,26 @@ import com.databricks.sql.transaction.tahoe.{ import com.databricks.sql.transaction.tahoe.actions.{Metadata, Protocol} import com.databricks.sql.transaction.tahoe.files.TahoeFileIndex import com.databricks.sql.transaction.tahoe.schema.SchemaMergingUtils -import com.nvidia.spark.rapids.{GpuMetric, SparkPlanMeta} +import com.nvidia.spark.rapids.{GpuMetric, RapidsConf, SparkPlanMeta} +import com.nvidia.spark.rapids.delta.GpuDeltaParquetFileFormatUtils.{addMetadataColumnsToBatch, + addMetadataColumnToIterator} import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path +import org.apache.spark.broadcast.Broadcast import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory} import org.apache.spark.sql.execution.FileSourceScanExec -import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, PartitionedFile} +import org.apache.spark.sql.execution.datasources.{FilePartition, HadoopFsRelation, PartitionedFile} import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, InputFileUtils} import org.apache.spark.sql.rapids.shims.TrampolineConnectShims import org.apache.spark.sql.sources.Filter import org.apache.spark.sql.types.{MetadataBuilder, StructType} +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.SerializableConfiguration /** * GPU Delta Parquet file format for Databricks 17.3. @@ -62,7 +71,8 @@ case class GpuDeltaParquetFileFormat( nullableRowTrackingGeneratedFields: Boolean = false, optimizationsEnabled: Boolean = true, tablePath: Option[String] = None, - isCDCRead: Boolean = false + isCDCRead: Boolean = false, + lowShuffleMergeScan: Option[GpuLowShuffleMergeScanInfo] = None ) extends GpuDeltaParquetFileFormatBase { override val columnMappingMode: DeltaColumnMappingMode = metadata.columnMappingMode @@ -113,7 +123,7 @@ case class GpuDeltaParquetFileFormat( * Translates pushed filters to physical column names when Delta column mapping is enabled. */ private def prepareFiltersForRead(filters: Seq[Filter]): Seq[Filter] = { - if (!effectiveOptimizationsEnabled) { + if (lowShuffleMergeScan.isDefined || !effectiveOptimizationsEnabled) { Seq.empty } else if (columnMappingMode != NoMapping) { val physicalNameMap = DeltaColumnMapping.getLogicalNameToPhysicalNameMap(referenceSchema) @@ -131,7 +141,7 @@ case class GpuDeltaParquetFileFormat( override def isSplitable( sparkSession: SparkSession, options: Map[String, String], - path: Path): Boolean = effectiveOptimizationsEnabled + path: Path): Boolean = lowShuffleMergeScan.isEmpty && effectiveOptimizationsEnabled private def hasDeletionVectorRead: Boolean = GpuDeltaParquetFileFormat.isDeletionVectorRead( @@ -164,7 +174,7 @@ case class GpuDeltaParquetFileFormat( hadoopConf: Configuration, metrics: Map[String, GpuMetric]) : PartitionedFile => Iterator[InternalRow] = { - super.buildReaderWithPartitionValuesAndMetrics( + val dataReader = super.buildReaderWithPartitionValuesAndMetrics( sparkSession, dataSchema, partitionSchema, @@ -173,6 +183,127 @@ case class GpuDeltaParquetFileFormat( options, hadoopConf, metrics) + + lowShuffleMergeScan.map { scanInfo => + val maxBatchSize = RapidsConf.DELTA_LOW_SHUFFLE_MERGE_SCATTER_DEL_VECTOR_BATCH_SIZE + .get(sparkSession.sessionState.conf) + val scatterTime = metrics(GpuMetric.DELETION_VECTOR_SCATTER_TIME) + val deletionVectorSize = metrics(GpuMetric.DELETION_VECTOR_SIZE) + (file: PartitionedFile) => { + val bitmap = lookupBitmap(scanInfo, file.filePath.toString, deletionVectorSize) + addMetadataColumnToIterator( + prepareSchema(requiredSchema), + bitmap, + dataReader(file).asInstanceOf[Iterator[ColumnarBatch]], + maxBatchSize, + scatterTime).asInstanceOf[Iterator[InternalRow]] + } + }.getOrElse(dataReader) + } + + override def createMultiFileReaderFactory( + broadcastedConf: Broadcast[SerializableConfiguration], + pushedFilters: Array[Filter], + fileScan: GpuFileSourceScanExec): PartitionReaderFactory = { + lowShuffleMergeScan.map { scanInfo => + // Low-shuffle metadata is file-relative. Prevent the coalescing reader from combining + // rows from multiple files into one batch; the multithreaded reader remains enabled. + val delegate = super.createMultiFileReaderFactory( + broadcastedConf, + pushedFilters, + fileScan.copy(queryUsesInputFile = true)(fileScan.rapidsConf)) + new LowShuffleMergePartitionReaderFactory( + delegate, + prepareSchema(fileScan.requiredSchema), + scanInfo, + fileScan.rapidsConf.get( + RapidsConf.DELTA_LOW_SHUFFLE_MERGE_SCATTER_DEL_VECTOR_BATCH_SIZE), + fileScan.allMetrics) + }.getOrElse(super.createMultiFileReaderFactory(broadcastedConf, pushedFilters, fileScan)) + } + + private def lookupBitmap( + scanInfo: GpuLowShuffleMergeScanInfo, + path: String, + deletionVectorSize: GpuMetric): Option[org.roaringbitmap.longlong.Roaring64Bitmap] = { + scanInfo.rowIndexMaps.flatMap { broadcast => + broadcast.value.get(new URI(path)).map { bytes => + deletionVectorSize += bytes.length + RoaringBitmapWrapper.deserializeFromBytes(bytes).inner + } + } + } + + private class LowShuffleMergePartitionReaderFactory( + delegate: PartitionReaderFactory, + readSchema: StructType, + scanInfo: GpuLowShuffleMergeScanInfo, + maxScatterBatchSize: Int, + metrics: Map[String, GpuMetric]) extends PartitionReaderFactory { + + override def createReader(partition: InputPartition): PartitionReader[InternalRow] = + delegate.createReader(partition) + + override def supportColumnarReads(partition: InputPartition): Boolean = true + + override def createColumnarReader( + partition: InputPartition): PartitionReader[ColumnarBatch] = { + val files = partition.asInstanceOf[FilePartition].filesWithAbsolutePaths + val byPath = files.flatMap { file => + Seq(file.filePath.toString -> file, file.urlEncodedPath -> file) + }.toMap + new LowShuffleMergePartitionReader( + delegate.createColumnarReader(partition), + byPath, + readSchema, + scanInfo, + maxScatterBatchSize, + metrics) + } + } + + private class LowShuffleMergePartitionReader( + delegate: PartitionReader[ColumnarBatch], + files: Map[String, PartitionedFile], + readSchema: StructType, + scanInfo: GpuLowShuffleMergeScanInfo, + maxScatterBatchSize: Int, + metrics: Map[String, GpuMetric]) extends PartitionReader[ColumnarBatch] { + + private var currentFile: PartitionedFile = _ + private var currentBitmap: Option[org.roaringbitmap.longlong.Roaring64Bitmap] = None + private var rowIndex = 0L + + override def next(): Boolean = delegate.next() + + override def get(): ColumnarBatch = { + val batch = delegate.get() + val inputPath = InputFileUtils.getCurInputFilePath() + val inputStart = InputFileUtils.getCurInputFileStartOffset + val inputLength = InputFileUtils.getCurInputFileLength + if (currentFile == null || + (currentFile.filePath.toString != inputPath && currentFile.urlEncodedPath != inputPath) || + currentFile.start != inputStart || + currentFile.length != inputLength) { + currentFile = files.getOrElse(inputPath, + throw new IllegalStateException(s"Unknown low-shuffle input file $inputPath")) + rowIndex = 0L + currentBitmap = lookupBitmap( + scanInfo, currentFile.filePath.toString, metrics(GpuMetric.DELETION_VECTOR_SIZE)) + } + val numRows = batch.numRows() + val result = addMetadataColumnsToBatch( + readSchema, + currentBitmap, + batch, + maxScatterBatchSize, + rowIndex, + metrics(GpuMetric.DELETION_VECTOR_SCATTER_TIME)) + rowIndex += numRows + result + } + + override def close(): Unit = delegate.close() } } @@ -268,7 +399,8 @@ object GpuDeltaParquetFileFormat { nullableRowTrackingGeneratedFields = fmt.nullableRowTrackingGeneratedFields, optimizationsEnabled = fmt.optimizationsEnabled, tablePath = fmt.tablePath, - isCDCRead = fmt.isCDCRead) + isCDCRead = fmt.isCDCRead, + lowShuffleMergeScan = GpuLowShuffleMergeScanRegistry.lookup(relation.options)) } private def hasRowIndexFiltersInTahoeFileIndex(relation: HadoopFsRelation): Boolean = { diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuLowShuffleMergeScan.scala b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuLowShuffleMergeScan.scala new file mode 100644 index 00000000000..9a3834df2c3 --- /dev/null +++ b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuLowShuffleMergeScan.scala @@ -0,0 +1,50 @@ +/* + * 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.nvidia.spark.rapids.delta + +import java.net.URI +import java.util.UUID +import java.util.concurrent.ConcurrentHashMap + +import org.apache.spark.broadcast.Broadcast + +/** Driver-side information needed to build a low-shuffle target scan. */ +case class GpuLowShuffleMergeScanInfo( + rowIndexMaps: Option[Broadcast[Map[URI, Array[Byte]]]]) + +/** + * Bridges information created by the low-shuffle command into Delta file-format conversion. + * Only a small opaque ID is placed in the logical relation. The broadcast itself is attached to + * the GPU file format when the scan is converted and is therefore sent to executors normally. + */ +object GpuLowShuffleMergeScanRegistry { + val OPTION_KEY: String = "spark.rapids.internal.delta.lowShuffleMerge.scanId" + + private val scans = new ConcurrentHashMap[String, GpuLowShuffleMergeScanInfo]() + + def register(info: GpuLowShuffleMergeScanInfo): String = { + val id = UUID.randomUUID().toString + scans.put(id, info) + id + } + + def lookup(options: Map[String, String]): Option[GpuLowShuffleMergeScanInfo] = { + options.get(OPTION_KEY).flatMap(id => Option(scans.get(id))) + } + + def remove(id: String): Unit = scans.remove(id) +} diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala index 0d53ef18f49..1dafc316b9d 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala @@ -19,7 +19,8 @@ package com.nvidia.spark.rapids.delta.shims import com.databricks.sql.transaction.tahoe.DeltaLog import com.databricks.sql.transaction.tahoe.commands.{DeletionVectorUtils, MergeIntoCommand, MergeIntoCommandBase, MergeIntoCommandEdge} -import com.databricks.sql.transaction.tahoe.rapids.{GpuDeltaLog, GpuMergeIntoCommand} +import com.databricks.sql.transaction.tahoe.rapids.{GpuDeltaLog, GpuLowShuffleMergeCommand, + GpuMergeIntoCommand} import com.databricks.sql.transaction.tahoe.sources.DeltaSQLConf import com.nvidia.spark.rapids.{RapidsConf, RapidsMeta} import com.nvidia.spark.rapids.delta.{MergeIntoCommandEdgeMeta, MergeIntoCommandMeta} @@ -59,38 +60,71 @@ object MergeIntoCommandMetaShim { } def convertToGpu(mergeCmd: MergeIntoCommand, conf: RapidsConf): RunnableCommand = { - GpuMergeIntoCommand( - mergeCmd.source, - mergeCmd.target, - mergeCmd.catalogTable, - mergeCmd.targetFileIndex, - new GpuDeltaLog(mergeCmd.targetFileIndex.deltaLog, conf), - mergeCmd.condition, - mergeCmd.matchedClauses, - mergeCmd.notMatchedClauses, - mergeCmd.notMatchedBySourceClauses, - mergeCmd.migratedSchema, - mergeCmd.trackHighWaterMarks, - mergeCmd.schemaEvolutionEnabled)(conf) + if (conf.isDeltaLowShuffleMergeEnabled) { + GpuLowShuffleMergeCommand( + mergeCmd.source, + mergeCmd.target, + mergeCmd.catalogTable, + mergeCmd.targetFileIndex, + new GpuDeltaLog(mergeCmd.targetFileIndex.deltaLog, conf), + mergeCmd.condition, + mergeCmd.matchedClauses, + mergeCmd.notMatchedClauses, + mergeCmd.notMatchedBySourceClauses, + mergeCmd.migratedSchema, + mergeCmd.trackHighWaterMarks, + mergeCmd.schemaEvolutionEnabled)(conf) + } else { + GpuMergeIntoCommand( + mergeCmd.source, + mergeCmd.target, + mergeCmd.catalogTable, + mergeCmd.targetFileIndex, + new GpuDeltaLog(mergeCmd.targetFileIndex.deltaLog, conf), + mergeCmd.condition, + mergeCmd.matchedClauses, + mergeCmd.notMatchedClauses, + mergeCmd.notMatchedBySourceClauses, + mergeCmd.migratedSchema, + mergeCmd.trackHighWaterMarks, + mergeCmd.schemaEvolutionEnabled)(conf) + } } def convertToGpu(mergeCmd: MergeIntoCommandEdge, conf: RapidsConf): RunnableCommand = { - GpuMergeIntoCommand( - mergeCmd.source, - mergeCmd.target, - mergeCmd.catalogTable, - mergeCmd.targetFileIndex, - new GpuDeltaLog(mergeCmd.targetFileIndex.deltaLog, conf), - mergeCmd.condition, - mergeCmd.matchedClauses, - mergeCmd.notMatchedClauses, - mergeCmd.notMatchedBySourceClauses, - mergeCmd.migratedSchema, - mergeCmd.trackHighWaterMarks, - mergeCmd.schemaEvolutionEnabled, - // This is safe to forward as-is because DBR analysis has already encoded snapshot reuse - // eligibility in this Option: Some(snapshot) means the Edge command may reuse the analyzed - // snapshot, while None makes GpuDeltaLog open the transaction on the latest snapshot. - mergeCmd.snapshotAtAnalysis)(conf) + if (conf.isDeltaLowShuffleMergeEnabled) { + GpuLowShuffleMergeCommand( + mergeCmd.source, + mergeCmd.target, + mergeCmd.catalogTable, + mergeCmd.targetFileIndex, + new GpuDeltaLog(mergeCmd.targetFileIndex.deltaLog, conf), + mergeCmd.condition, + mergeCmd.matchedClauses, + mergeCmd.notMatchedClauses, + mergeCmd.notMatchedBySourceClauses, + mergeCmd.migratedSchema, + mergeCmd.trackHighWaterMarks, + mergeCmd.schemaEvolutionEnabled, + mergeCmd.snapshotAtAnalysis)(conf) + } else { + GpuMergeIntoCommand( + mergeCmd.source, + mergeCmd.target, + mergeCmd.catalogTable, + mergeCmd.targetFileIndex, + new GpuDeltaLog(mergeCmd.targetFileIndex.deltaLog, conf), + mergeCmd.condition, + mergeCmd.matchedClauses, + mergeCmd.notMatchedClauses, + mergeCmd.notMatchedBySourceClauses, + mergeCmd.migratedSchema, + mergeCmd.trackHighWaterMarks, + mergeCmd.schemaEvolutionEnabled, + // This is safe to forward as-is because DBR analysis has already encoded snapshot reuse + // eligibility in this Option: Some(snapshot) means the Edge command may reuse the analyzed + // snapshot, while None makes GpuDeltaLog open the transaction on the latest snapshot. + mergeCmd.snapshotAtAnalysis)(conf) + } } } diff --git a/docs/additional-functionality/advanced_configs.md b/docs/additional-functionality/advanced_configs.md index 5b5b60be78a..0b964776335 100644 --- a/docs/additional-functionality/advanced_configs.md +++ b/docs/additional-functionality/advanced_configs.md @@ -87,7 +87,7 @@ Name | Description | Default Value | Applicable at spark.rapids.sql.csv.read.float.enabled|CSV reading is not 100% compatible when reading floats.|true|Runtime spark.rapids.sql.decimalOverflowGuarantees|FOR TESTING ONLY. DO NOT USE IN PRODUCTION. Please see the decimal section of the compatibility documents for more information on this config.|true|Runtime spark.rapids.sql.delta.lowShuffleMerge.deletionVector.broadcast.threshold|Currently we need to broadcast deletion vector to all executors to perform low shuffle merge. When we detect the deletion vector broadcast size is larger than this value, we will fallback to normal shuffle merge.|20971520|Runtime -spark.rapids.sql.delta.lowShuffleMerge.enabled|Option to turn on the low shuffle merge for Delta Lake. Currently there are some limitations for this feature: 1. We only support Delta Lake 2.4. 2. The file scan mode must be set to PERFILE 3. The deletion vector size must be smaller than spark.rapids.sql.delta.lowShuffleMerge.deletionVector.broadcast.threshold |false|Runtime +spark.rapids.sql.delta.lowShuffleMerge.enabled|Option to turn on the low shuffle merge for Delta Lake. Currently there are some limitations for this feature: 1. We support Delta Lake 2.4 and Databricks Runtime 17.3. 2. Delta Lake 2.4 requires the PERFILE file scan mode; Databricks Runtime 17.3 supports all Parquet reader modes. 3. The deletion vector size must be smaller than spark.rapids.sql.delta.lowShuffleMerge.deletionVector.broadcast.threshold |false|Runtime spark.rapids.sql.detectDeltaCheckpointQueries|Queries against Delta Lake _delta_log checkpoint Parquet files are not efficient on the GPU. When this option is enabled, the plugin will attempt to detect these queries and fall back to the CPU.|true|Runtime spark.rapids.sql.detectDeltaLogQueries|Queries against Delta Lake _delta_log JSON files are not efficient on the GPU. When this option is enabled, the plugin will attempt to detect these queries and fall back to the CPU.|true|Runtime spark.rapids.sql.exec.opTimeTrackingRDD.enabled|Enable OpTimeTrackingRDD for all GPU operations. When true, OpTimeTrackingRDD wrappers will be created to track operation time. When false, can improve performance by avoiding overhead of operation time tracking.|true|Runtime diff --git a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py index a71d570b6ed..09e0f20b6eb 100644 --- a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py +++ b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py @@ -19,19 +19,54 @@ from delta_lake_merge_common import * from marks import * from pyspark.sql.types import * -from spark_session import spark_version +from spark_session import is_databricks173_or_later, spark_version delta_merge_enabled_conf = copy_and_update(delta_writes_enabled_conf, {"spark.rapids.sql.command.MergeIntoCommand": "true", "spark.rapids.sql.command.MergeIntoCommandEdge": "true", - "spark.rapids.sql.delta.lowShuffleMerge.enabled": "true", - "spark.rapids.sql.format.parquet.reader.type": "PERFILE"}) + "spark.rapids.sql.delta.lowShuffleMerge.enabled": "true"}) + +# The Delta 2.4 implementation still uses a per-file iterator. The DBR 17.3 implementation +# tracks file boundaries for every reader mode and deliberately does not force PERFILE. +if not is_databricks173_or_later(): + delta_merge_enabled_conf = copy_and_update( + delta_merge_enabled_conf, + {"spark.rapids.sql.format.parquet.reader.type": "PERFILE"}) +else: + # Disable AQE temporarily until https://github.com/NVIDIA/spark-rapids/issues/14319 is resolved. + delta_merge_enabled_conf = copy_and_update( + delta_merge_enabled_conf, + {"spark.sql.adaptive.enabled": "false"}) + + +def supports_delta_low_shuffle_merge(): + return is_databricks173_or_later() or \ + (not is_databricks_runtime() and spark_version().startswith("3.4")) + + +def assert_low_shuffle_merge(do_merge, data_path, conf, expect_write=True): + assert expect_write + cpu_result = with_cpu_session(lambda spark: do_merge(spark, data_path + "/CPU"), conf=conf) + + callback = spark_jvm().org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback + callback.startCapture() + try: + gpu_result = with_gpu_session( + lambda spark: do_merge(spark, data_path + "/GPU"), conf=conf) + captured_plans = callback.getResultsWithTimeout(10000) + finally: + callback.endCapture() + + assert_equal(cpu_result, gpu_result) + command_name = "GpuLowShuffleMergeCommand" + assert any(command_name in str(plan) for plan in captured_plans), \ + f"{command_name} was not found in the captured MERGE plans" @allow_non_gpu("ColumnarToRowExec", *delta_meta_allow) @delta_lake @ignore_order -@pytest.mark.skipif(is_databricks_runtime() or not spark_version().startswith("3.4"), - reason="Delta Lake Low Shuffle Merge only supports OSS Delta Lake 2.4") +@pytest.mark.skipif(not supports_delta_low_shuffle_merge(), + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") @pytest.mark.parametrize("use_cdf", [True, False], ids=idfn) @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn) def test_delta_low_shuffle_merge_when_gpu_file_scan_override_failed(spark_tmp_path, @@ -58,8 +93,8 @@ def test_delta_low_shuffle_merge_when_gpu_file_scan_override_failed(spark_tmp_pa @allow_non_gpu(*delta_meta_allow) @delta_lake @ignore_order -@pytest.mark.skipif(is_databricks_runtime() or not spark_version().startswith("3.4"), - reason="Delta Lake Low Shuffle Merge only supports OSS Delta Lake 2.4") +@pytest.mark.skipif(not supports_delta_low_shuffle_merge(), + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") @pytest.mark.parametrize("table_ranges", [(range(20), range(10)), # partial insert of source (range(5), range(5)), # no-op insert (range(10), range(20, 30)) # full insert of source @@ -76,8 +111,8 @@ def test_delta_merge_not_match_insert_only(spark_tmp_path, spark_tmp_table_facto @allow_non_gpu(*delta_meta_allow) @delta_lake @ignore_order -@pytest.mark.skipif(is_databricks_runtime() or not spark_version().startswith("3.4"), - reason="Delta Lake Low Shuffle Merge only supports OSS Delta Lake 2.4") +@pytest.mark.skipif(not supports_delta_low_shuffle_merge(), + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") @pytest.mark.parametrize("table_ranges", [(range(10), range(20)), # partial delete of target (range(5), range(5)), # full delete of target (range(10), range(20, 30)) # no-op delete @@ -94,19 +129,42 @@ def test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, @allow_non_gpu(*delta_meta_allow) @delta_lake @ignore_order -@pytest.mark.skipif(is_databricks_runtime() or not spark_version().startswith("3.4"), - reason="Delta Lake Low Shuffle Merge only supports OSS Delta Lake 2.4") +@pytest.mark.skipif(not supports_delta_low_shuffle_merge(), + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") @pytest.mark.parametrize("use_cdf", [pytest.param(True, marks=pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552")), False], ids=idfn) @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn) def test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, use_cdf, num_slices): do_test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, use_cdf, False, num_slices, False, delta_merge_enabled_conf) + +@allow_non_gpu(*delta_meta_allow) +@delta_lake +@ignore_order +@pytest.mark.skipif(not is_databricks173_or_later(), + reason="All-reader low shuffle merge is supported on DBR 17.3+") +@pytest.mark.parametrize( + "reader_type", ["AUTO", "PERFILE", "MULTITHREADED", "COALESCING"], ids=idfn) +def test_databricks_delta_low_shuffle_merge_reader_type( + spark_tmp_path, spark_tmp_table_factory, reader_type): + conf = copy_and_update( + delta_merge_enabled_conf, + {"spark.rapids.sql.format.parquet.reader.type": reader_type}) + do_test_delta_merge_standard_upsert( + spark_tmp_path, + spark_tmp_table_factory, + use_cdf=False, + enable_deletion_vectors=False, + num_slices=10, + compare_logs=False, + conf=conf, + assert_func=assert_low_shuffle_merge) + @allow_non_gpu(*delta_meta_allow) @delta_lake @ignore_order -@pytest.mark.skipif(is_databricks_runtime() or not spark_version().startswith("3.4"), - reason="Delta Lake Low Shuffle Merge only supports OSS Delta Lake 2.4") +@pytest.mark.skipif(not supports_delta_low_shuffle_merge(), + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") @pytest.mark.parametrize("use_cdf", [pytest.param(True, marks=pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552")), False], ids=idfn) @pytest.mark.parametrize("merge_sql", [ "MERGE INTO {dest_table} d USING {src_table} s ON d.a == s.a" \ @@ -128,8 +186,8 @@ def test_delta_merge_upsert_with_condition(spark_tmp_path, spark_tmp_table_facto @allow_non_gpu(*delta_meta_allow) @delta_lake @ignore_order -@pytest.mark.skipif(is_databricks_runtime() or not spark_version().startswith("3.4"), - reason="Delta Lake Low Shuffle Merge only supports OSS Delta Lake 2.4") +@pytest.mark.skipif(not supports_delta_low_shuffle_merge(), + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") @pytest.mark.parametrize("use_cdf", [True, False], ids=idfn) @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn) def test_delta_merge_upsert_with_unmatchable_match_condition(spark_tmp_path, spark_tmp_table_factory, use_cdf, num_slices): @@ -144,8 +202,8 @@ def test_delta_merge_upsert_with_unmatchable_match_condition(spark_tmp_path, spa @allow_non_gpu(*delta_meta_allow) @delta_lake @ignore_order -@pytest.mark.skipif(is_databricks_runtime() or not spark_version().startswith("3.4"), - reason="Delta Lake Low Shuffle Merge only supports OSS Delta Lake 2.4") +@pytest.mark.skipif(not supports_delta_low_shuffle_merge(), + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") @pytest.mark.parametrize("use_cdf", [pytest.param(True, marks=pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552")), False], ids=idfn) def test_delta_merge_update_with_aggregation(spark_tmp_path, spark_tmp_table_factory, use_cdf): do_test_delta_merge_update_with_aggregation(spark_tmp_path, spark_tmp_table_factory, use_cdf, False, diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala index ff2533e7d4c..46121479af9 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala @@ -2886,8 +2886,9 @@ val SHUFFLE_COMPRESSION_LZ4_CHUNK_SIZE = conf("spark.rapids.shuffle.compression. conf("spark.rapids.sql.delta.lowShuffleMerge.enabled") .doc("Option to turn on the low shuffle merge for Delta Lake. Currently there are some " + "limitations for this feature: " + - "1. We only support Delta Lake 2.4. " + - s"2. The file scan mode must be set to ${RapidsReaderType.PERFILE} " + + "1. We support Delta Lake 2.4 and Databricks Runtime 17.3. " + + s"2. Delta Lake 2.4 requires the ${RapidsReaderType.PERFILE} file scan mode; " + + "Databricks Runtime 17.3 supports all Parquet reader modes. " + "3. The deletion vector size must be smaller than " + s"${DELTA_LOW_SHUFFLE_MERGE_DEL_VECTOR_BROADCAST_THRESHOLD.key} ") .booleanConf From bb6859a15f7378e8ba23f946bbb21819eff96e54 Mon Sep 17 00:00:00 2001 From: Ray Liu Date: Tue, 8 Sep 2026 02:56:54 +0000 Subject: [PATCH 02/10] [FEA] Restrict DBR 17.3 low shuffle merge to per-file reads Signed-off-by: Ray Liu --- .../GpuDeltaParquetFileFormatUtils.scala | 52 +++----- .../delta/GpuDeltaParquetFileFormat.scala | 121 ++---------------- .../shims/MergeIntoCommandMetaShim.scala | 4 +- .../advanced_configs.md | 2 +- .../delta_lake_low_shuffle_merge_test.py | 25 ++-- .../com/nvidia/spark/rapids/RapidsConf.scala | 3 +- 6 files changed, 39 insertions(+), 168 deletions(-) diff --git a/delta-lake/common/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatUtils.scala b/delta-lake/common/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatUtils.scala index 5a92b821fd2..1ade53b21b9 100644 --- a/delta-lake/common/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatUtils.scala +++ b/delta-lake/common/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatUtils.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024-2026, NVIDIA CORPORATION. + * Copyright (c) 2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -28,8 +28,8 @@ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} object GpuDeltaParquetFileFormatUtils { /** * Row number of the row in the file. When used with [[FILE_PATH_COL]] together, it can be used - * as unique id of a row in file. To calculate this correctly, the caller needs to make each file - * unsplittable and reset the row offset when a multi-file reader advances to the next file. + * as unique id of a row in file. Currently to correctly calculate this, the caller needs to + * set both [[isSplitable]] to false, and [[RapidsConf.PARQUET_READER_TYPE]] to "PERFILE". */ val METADATA_ROW_IDX_COL: String = "__metadata_row_index" val METADATA_ROW_IDX_FIELD: StructField = StructField(METADATA_ROW_IDX_COL, LongType, @@ -63,37 +63,23 @@ object GpuDeltaParquetFileFormatUtils { } var rowIndex = 0L input.map { batch => - val numRows = batch.numRows() - val newBatch = addMetadataColumnsToBatch(schema, delVector, batch, maxBatchSize, - rowIndex, delVectorScatterTimeMetric) - rowIndex += numRows - newBatch - } - } + withResource(batch) { _ => + val rowIdxCol = if (metadataRowIndexCol == -1) { + None + } else { + Some(metadataRowIndexCol) + } - /** - * Add low-shuffle metadata columns to one batch at the specified file-global row offset. - * This entry point is used by multi-file readers, which reset the offset when the input file - * changes. - */ - def addMetadataColumnsToBatch( - schema: StructType, - delVector: Option[Roaring64Bitmap], - batch: ColumnarBatch, - maxBatchSize: Int, - rowIndex: Long, - delVectorScatterTimeMetric: GpuMetric): ColumnarBatch = { - val metadataRowIndexCol = schema.fieldNames.indexOf(METADATA_ROW_IDX_COL) - val delRowIdx = schema.fieldNames.indexOf(METADATA_ROW_DEL_COL) - withResource(batch) { _ => - addMetadataColumns( - if (metadataRowIndexCol == -1) None else Some(metadataRowIndexCol), - if (delRowIdx == -1) None else Some(delRowIdx), - delVector, - maxBatchSize, - rowIndex, - batch, - delVectorScatterTimeMetric) + val delRowIdx2 = if (delRowIdx == -1) { + None + } else { + Some(delRowIdx) + } + val newBatch = addMetadataColumns(rowIdxCol, delRowIdx2, delVector,maxBatchSize, + rowIndex, batch, delVectorScatterTimeMetric) + rowIndex += batch.numRows() + newBatch + } } } 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 55d21bba02d..c436d43dd3d 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 @@ -31,25 +31,20 @@ import com.databricks.sql.transaction.tahoe.actions.{Metadata, Protocol} import com.databricks.sql.transaction.tahoe.files.TahoeFileIndex import com.databricks.sql.transaction.tahoe.schema.SchemaMergingUtils import com.nvidia.spark.rapids.{GpuMetric, RapidsConf, SparkPlanMeta} -import com.nvidia.spark.rapids.delta.GpuDeltaParquetFileFormatUtils.{addMetadataColumnsToBatch, - addMetadataColumnToIterator} +import com.nvidia.spark.rapids.delta.GpuDeltaParquetFileFormatUtils.addMetadataColumnToIterator import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path -import org.apache.spark.broadcast.Broadcast import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral -import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory} import org.apache.spark.sql.execution.FileSourceScanExec -import org.apache.spark.sql.execution.datasources.{FilePartition, HadoopFsRelation, PartitionedFile} +import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, PartitionedFile} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, InputFileUtils} import org.apache.spark.sql.rapids.shims.TrampolineConnectShims import org.apache.spark.sql.sources.Filter import org.apache.spark.sql.types.{MetadataBuilder, StructType} import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.spark.util.SerializableConfiguration /** * GPU Delta Parquet file format for Databricks 17.3. @@ -190,7 +185,12 @@ case class GpuDeltaParquetFileFormat( val scatterTime = metrics(GpuMetric.DELETION_VECTOR_SCATTER_TIME) val deletionVectorSize = metrics(GpuMetric.DELETION_VECTOR_SIZE) (file: PartitionedFile) => { - val bitmap = lookupBitmap(scanInfo, file.filePath.toString, deletionVectorSize) + val bitmap = scanInfo.rowIndexMaps.flatMap { broadcast => + broadcast.value.get(new URI(file.filePath.toString)).map { bytes => + deletionVectorSize += bytes.length + RoaringBitmapWrapper.deserializeFromBytes(bytes).inner + } + } addMetadataColumnToIterator( prepareSchema(requiredSchema), bitmap, @@ -200,111 +200,6 @@ case class GpuDeltaParquetFileFormat( } }.getOrElse(dataReader) } - - override def createMultiFileReaderFactory( - broadcastedConf: Broadcast[SerializableConfiguration], - pushedFilters: Array[Filter], - fileScan: GpuFileSourceScanExec): PartitionReaderFactory = { - lowShuffleMergeScan.map { scanInfo => - // Low-shuffle metadata is file-relative. Prevent the coalescing reader from combining - // rows from multiple files into one batch; the multithreaded reader remains enabled. - val delegate = super.createMultiFileReaderFactory( - broadcastedConf, - pushedFilters, - fileScan.copy(queryUsesInputFile = true)(fileScan.rapidsConf)) - new LowShuffleMergePartitionReaderFactory( - delegate, - prepareSchema(fileScan.requiredSchema), - scanInfo, - fileScan.rapidsConf.get( - RapidsConf.DELTA_LOW_SHUFFLE_MERGE_SCATTER_DEL_VECTOR_BATCH_SIZE), - fileScan.allMetrics) - }.getOrElse(super.createMultiFileReaderFactory(broadcastedConf, pushedFilters, fileScan)) - } - - private def lookupBitmap( - scanInfo: GpuLowShuffleMergeScanInfo, - path: String, - deletionVectorSize: GpuMetric): Option[org.roaringbitmap.longlong.Roaring64Bitmap] = { - scanInfo.rowIndexMaps.flatMap { broadcast => - broadcast.value.get(new URI(path)).map { bytes => - deletionVectorSize += bytes.length - RoaringBitmapWrapper.deserializeFromBytes(bytes).inner - } - } - } - - private class LowShuffleMergePartitionReaderFactory( - delegate: PartitionReaderFactory, - readSchema: StructType, - scanInfo: GpuLowShuffleMergeScanInfo, - maxScatterBatchSize: Int, - metrics: Map[String, GpuMetric]) extends PartitionReaderFactory { - - override def createReader(partition: InputPartition): PartitionReader[InternalRow] = - delegate.createReader(partition) - - override def supportColumnarReads(partition: InputPartition): Boolean = true - - override def createColumnarReader( - partition: InputPartition): PartitionReader[ColumnarBatch] = { - val files = partition.asInstanceOf[FilePartition].filesWithAbsolutePaths - val byPath = files.flatMap { file => - Seq(file.filePath.toString -> file, file.urlEncodedPath -> file) - }.toMap - new LowShuffleMergePartitionReader( - delegate.createColumnarReader(partition), - byPath, - readSchema, - scanInfo, - maxScatterBatchSize, - metrics) - } - } - - private class LowShuffleMergePartitionReader( - delegate: PartitionReader[ColumnarBatch], - files: Map[String, PartitionedFile], - readSchema: StructType, - scanInfo: GpuLowShuffleMergeScanInfo, - maxScatterBatchSize: Int, - metrics: Map[String, GpuMetric]) extends PartitionReader[ColumnarBatch] { - - private var currentFile: PartitionedFile = _ - private var currentBitmap: Option[org.roaringbitmap.longlong.Roaring64Bitmap] = None - private var rowIndex = 0L - - override def next(): Boolean = delegate.next() - - override def get(): ColumnarBatch = { - val batch = delegate.get() - val inputPath = InputFileUtils.getCurInputFilePath() - val inputStart = InputFileUtils.getCurInputFileStartOffset - val inputLength = InputFileUtils.getCurInputFileLength - if (currentFile == null || - (currentFile.filePath.toString != inputPath && currentFile.urlEncodedPath != inputPath) || - currentFile.start != inputStart || - currentFile.length != inputLength) { - currentFile = files.getOrElse(inputPath, - throw new IllegalStateException(s"Unknown low-shuffle input file $inputPath")) - rowIndex = 0L - currentBitmap = lookupBitmap( - scanInfo, currentFile.filePath.toString, metrics(GpuMetric.DELETION_VECTOR_SIZE)) - } - val numRows = batch.numRows() - val result = addMetadataColumnsToBatch( - readSchema, - currentBitmap, - batch, - maxScatterBatchSize, - rowIndex, - metrics(GpuMetric.DELETION_VECTOR_SCATTER_TIME)) - rowIndex += numRows - result - } - - override def close(): Unit = delegate.close() - } } object GpuDeltaParquetFileFormat { diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala index 1dafc316b9d..edf4f0ce025 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala @@ -60,7 +60,7 @@ object MergeIntoCommandMetaShim { } def convertToGpu(mergeCmd: MergeIntoCommand, conf: RapidsConf): RunnableCommand = { - if (conf.isDeltaLowShuffleMergeEnabled) { + if (conf.isDeltaLowShuffleMergeEnabled && conf.isParquetPerFileReadEnabled) { GpuLowShuffleMergeCommand( mergeCmd.source, mergeCmd.target, @@ -92,7 +92,7 @@ object MergeIntoCommandMetaShim { } def convertToGpu(mergeCmd: MergeIntoCommandEdge, conf: RapidsConf): RunnableCommand = { - if (conf.isDeltaLowShuffleMergeEnabled) { + if (conf.isDeltaLowShuffleMergeEnabled && conf.isParquetPerFileReadEnabled) { GpuLowShuffleMergeCommand( mergeCmd.source, mergeCmd.target, diff --git a/docs/additional-functionality/advanced_configs.md b/docs/additional-functionality/advanced_configs.md index 0b964776335..67312c9fea0 100644 --- a/docs/additional-functionality/advanced_configs.md +++ b/docs/additional-functionality/advanced_configs.md @@ -87,7 +87,7 @@ Name | Description | Default Value | Applicable at spark.rapids.sql.csv.read.float.enabled|CSV reading is not 100% compatible when reading floats.|true|Runtime spark.rapids.sql.decimalOverflowGuarantees|FOR TESTING ONLY. DO NOT USE IN PRODUCTION. Please see the decimal section of the compatibility documents for more information on this config.|true|Runtime spark.rapids.sql.delta.lowShuffleMerge.deletionVector.broadcast.threshold|Currently we need to broadcast deletion vector to all executors to perform low shuffle merge. When we detect the deletion vector broadcast size is larger than this value, we will fallback to normal shuffle merge.|20971520|Runtime -spark.rapids.sql.delta.lowShuffleMerge.enabled|Option to turn on the low shuffle merge for Delta Lake. Currently there are some limitations for this feature: 1. We support Delta Lake 2.4 and Databricks Runtime 17.3. 2. Delta Lake 2.4 requires the PERFILE file scan mode; Databricks Runtime 17.3 supports all Parquet reader modes. 3. The deletion vector size must be smaller than spark.rapids.sql.delta.lowShuffleMerge.deletionVector.broadcast.threshold |false|Runtime +spark.rapids.sql.delta.lowShuffleMerge.enabled|Option to turn on the low shuffle merge for Delta Lake. Currently there are some limitations for this feature: 1. We support Delta Lake 2.4 and Databricks Runtime 17.3. 2. The file scan mode must be set to PERFILE. 3. The deletion vector size must be smaller than spark.rapids.sql.delta.lowShuffleMerge.deletionVector.broadcast.threshold |false|Runtime spark.rapids.sql.detectDeltaCheckpointQueries|Queries against Delta Lake _delta_log checkpoint Parquet files are not efficient on the GPU. When this option is enabled, the plugin will attempt to detect these queries and fall back to the CPU.|true|Runtime spark.rapids.sql.detectDeltaLogQueries|Queries against Delta Lake _delta_log JSON files are not efficient on the GPU. When this option is enabled, the plugin will attempt to detect these queries and fall back to the CPU.|true|Runtime spark.rapids.sql.exec.opTimeTrackingRDD.enabled|Enable OpTimeTrackingRDD for all GPU operations. When true, OpTimeTrackingRDD wrappers will be created to track operation time. When false, can improve performance by avoiding overhead of operation time tracking.|true|Runtime diff --git a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py index 09e0f20b6eb..0b42e72a8f8 100644 --- a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py +++ b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py @@ -24,15 +24,10 @@ delta_merge_enabled_conf = copy_and_update(delta_writes_enabled_conf, {"spark.rapids.sql.command.MergeIntoCommand": "true", "spark.rapids.sql.command.MergeIntoCommandEdge": "true", - "spark.rapids.sql.delta.lowShuffleMerge.enabled": "true"}) + "spark.rapids.sql.delta.lowShuffleMerge.enabled": "true", + "spark.rapids.sql.format.parquet.reader.type": "PERFILE"}) -# The Delta 2.4 implementation still uses a per-file iterator. The DBR 17.3 implementation -# tracks file boundaries for every reader mode and deliberately does not force PERFILE. -if not is_databricks173_or_later(): - delta_merge_enabled_conf = copy_and_update( - delta_merge_enabled_conf, - {"spark.rapids.sql.format.parquet.reader.type": "PERFILE"}) -else: +if is_databricks173_or_later(): # Disable AQE temporarily until https://github.com/NVIDIA/spark-rapids/issues/14319 is resolved. delta_merge_enabled_conf = copy_and_update( delta_merge_enabled_conf, @@ -142,14 +137,9 @@ def test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, us @delta_lake @ignore_order @pytest.mark.skipif(not is_databricks173_or_later(), - reason="All-reader low shuffle merge is supported on DBR 17.3+") -@pytest.mark.parametrize( - "reader_type", ["AUTO", "PERFILE", "MULTITHREADED", "COALESCING"], ids=idfn) -def test_databricks_delta_low_shuffle_merge_reader_type( - spark_tmp_path, spark_tmp_table_factory, reader_type): - conf = copy_and_update( - delta_merge_enabled_conf, - {"spark.rapids.sql.format.parquet.reader.type": reader_type}) + reason="Databricks low shuffle merge requires DBR 17.3+") +def test_databricks_delta_low_shuffle_merge_perfile( + spark_tmp_path, spark_tmp_table_factory): do_test_delta_merge_standard_upsert( spark_tmp_path, spark_tmp_table_factory, @@ -157,9 +147,10 @@ def test_databricks_delta_low_shuffle_merge_reader_type( enable_deletion_vectors=False, num_slices=10, compare_logs=False, - conf=conf, + conf=delta_merge_enabled_conf, assert_func=assert_low_shuffle_merge) + @allow_non_gpu(*delta_meta_allow) @delta_lake @ignore_order diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala index 46121479af9..28d8d564272 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala @@ -2887,8 +2887,7 @@ val SHUFFLE_COMPRESSION_LZ4_CHUNK_SIZE = conf("spark.rapids.shuffle.compression. .doc("Option to turn on the low shuffle merge for Delta Lake. Currently there are some " + "limitations for this feature: " + "1. We support Delta Lake 2.4 and Databricks Runtime 17.3. " + - s"2. Delta Lake 2.4 requires the ${RapidsReaderType.PERFILE} file scan mode; " + - "Databricks Runtime 17.3 supports all Parquet reader modes. " + + s"2. The file scan mode must be set to ${RapidsReaderType.PERFILE}. " + "3. The deletion vector size must be smaller than " + s"${DELTA_LOW_SHUFFLE_MERGE_DEL_VECTOR_BROADCAST_THRESHOLD.key} ") .booleanConf From 5ce3e91779ee581b7b297ff063a04033c4c3609e Mon Sep 17 00:00:00 2001 From: Ray Liu Date: Tue, 8 Sep 2026 16:40:58 +0800 Subject: [PATCH 03/10] [FEA] Reuse DBR deletion-vector scans for low shuffle merge Replace the custom low-shuffle scan registry and file-format path with Databricks' native metadata-row-index and deletion-vector scans. Also address the review comments for copyrights and identity-column validation documentation. Signed-off-by: Ray Liu --- .../rapids/GpuLowShuffleMergeCommand.scala | 222 ++++++++---------- .../delta/DeltaSpark400DB173Provider.scala | 4 +- .../delta/GpuDeltaParquetFileFormat.scala | 39 +-- .../rapids/delta/GpuLowShuffleMergeScan.scala | 50 ---- 4 files changed, 99 insertions(+), 216 deletions(-) delete mode 100644 delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuLowShuffleMergeScan.scala diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala index 3b3678e3f9c..30e96c63a6a 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024-2026, NVIDIA CORPORATION. + * Copyright (c) 2026, NVIDIA CORPORATION. * * This file was derived from MergeIntoCommand.scala * in the Delta Lake project at https://github.com/delta-io/delta. @@ -21,18 +21,22 @@ package com.databricks.sql.transaction.tahoe.rapids -import java.net.URI import java.util.concurrent.TimeUnit import scala.annotation.nowarn import scala.collection.mutable +import com.databricks.sql.io.RowIndexFilterType import com.databricks.sql.transaction.tahoe._ import com.databricks.sql.transaction.tahoe.DeltaOperations.MergePredicate -import com.databricks.sql.transaction.tahoe.actions.{AddCDCFile, AddFile, FileAction} -import com.databricks.sql.transaction.tahoe.commands.DeltaCommand +import com.databricks.sql.transaction.tahoe.actions.{AddCDCFile, AddFile, + DeletionVectorDescriptor, FileAction} +import com.databricks.sql.transaction.tahoe.commands.{DeltaCommand, + DMLWithDeletionVectorsHelper} import com.databricks.sql.transaction.tahoe.commands.merge.MergeIntoMaterializeSource -import com.databricks.sql.transaction.tahoe.files.TahoeFileIndex +import com.databricks.sql.transaction.tahoe.deletionvectors.{RoaringBitmapArray, + RoaringBitmapArrayFormat} +import com.databricks.sql.transaction.tahoe.files.{TahoeBatchFileIndex, TahoeFileIndex} import com.databricks.sql.transaction.tahoe.rapids.MergeExecutor.{ totalBytesAndDistinctPartitionValues, FILE_PATH_COL, @@ -50,12 +54,9 @@ import com.databricks.sql.transaction.tahoe.util.{AnalysisHelper, DeltaFileOpera import com.nvidia.spark.rapids.{GpuOverrides, RapidsConf, SparkPlanMeta} import com.nvidia.spark.rapids.RapidsConf.DELTA_LOW_SHUFFLE_MERGE_DEL_VECTOR_BROADCAST_THRESHOLD import com.nvidia.spark.rapids.delta._ -import com.nvidia.spark.rapids.delta.GpuDeltaParquetFileFormatUtils.{ - METADATA_ROW_DEL_COL, - METADATA_ROW_DEL_FIELD, - METADATA_ROW_IDX_COL, - METADATA_ROW_IDX_FIELD} +import com.nvidia.spark.rapids.delta.GpuDeltaParquetFileFormatUtils.METADATA_ROW_IDX_COL import com.nvidia.spark.rapids.shims.FileSourceScanExecMeta +import org.apache.hadoop.conf.Configuration import org.roaringbitmap.longlong.Roaring64Bitmap import org.apache.spark.SparkContext @@ -195,6 +196,10 @@ case class GpuLowShuffleMergeCommand( private[rapids] def mergeSourceDF: DataFrame = getMergeSource.df + /** + * Validates that identity-column metadata has not changed since the merge was analyzed and that + * insert actions do not explicitly populate identity columns that disallow explicit values. + */ private def checkIdentityColumnHighWaterMarks(deltaTxn: OptimisticTransaction): Unit = { notMatchedClauses.foreach { clause => val schema = deltaTxn.metadata.schema @@ -604,15 +609,6 @@ class InsertOnlyMergeExecutor(override val context: MergeExecutorContext) extend */ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extends MergeExecutor { - private val scanRegistrationIds = new mutable.ArrayBuffer[String]() - private var touchedRowsBroadcast = Option.empty[ - org.apache.spark.broadcast.Broadcast[Map[URI, Array[Byte]]]] - - override def close(): Unit = { - scanRegistrationIds.foreach(GpuLowShuffleMergeScanRegistry.remove) - touchedRowsBroadcast.foreach(_.destroy()) - } - // We over-count numTargetRowsDeleted when there are multiple matches; // this is the amount of the overcount, so we can subtract it to get a correct final metric. private var multipleMatchDeleteOnlyOvercount: Option[Long] = None @@ -684,26 +680,19 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend * Though low shuffle merge algorithm performs better than traditional merge algorithm in some * cases, there are some case we should fallback to traditional merge executor: * - * 1. Low shuffle merge algorithm requires generating metadata columns such as - * [[METADATA_ROW_IDX_COL]], [[METADATA_ROW_DEL_COL]], which only implemented on - * [[org.apache.spark.sql.rapids.GpuFileSourceScanExec]]. That means we need to fallback to - * this normal executor when [[org.apache.spark.sql.rapids.GpuFileSourceScanExec]] is disabled - * for some reason. - * 2. Low shuffle merge algorithm currently needs to broadcast deletion vector, which may - * introduce extra overhead. It maybe better to fallback to this algorithm when the changeset - * it too large. + * 1. Low shuffle merge requires GPU file scans for both Databricks' metadata row-index scan and + * the temporary deletion-vector scan used to retain unmodified rows. + * 2. The temporary deletion vectors introduce extra overhead, so it may be better to fall back + * when the changeset is too large. */ def shouldFallback(): Boolean = { // Trying to detect if we can execute finding touched files. val touchFilePlanOverrideSucceed = verifyGpuPlan(planForFindingTouchedFiles()) { planMeta => def check(meta: SparkPlanMeta[SparkPlan]): Boolean = { meta match { - case scan if scan.isInstanceOf[FileSourceScanExecMeta] => scan - .asInstanceOf[FileSourceScanExecMeta] - .wrapped - .schema - .fieldNames - .contains(METADATA_ROW_IDX_COL) && scan.canThisBeReplaced + case scan if scan.isInstanceOf[FileSourceScanExecMeta] && + isLowShuffleTargetScan(scan.asInstanceOf[FileSourceScanExecMeta]) => + scan.asInstanceOf[FileSourceScanExecMeta].canThisBeReplaced case m => m.childPlans.exists(check) } } @@ -718,20 +707,23 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend // Trying to detect if we can execute the merge plan. val mergePlanOverrideSucceed = verifyGpuPlan(planForMergeExecution(touchedFiles)) { planMeta => - var overrideCount = 0 + var targetScanCount = 0 + var gpuTargetScanCount = 0 def count(meta: SparkPlanMeta[SparkPlan]): Unit = { meta match { - case scan if scan.isInstanceOf[FileSourceScanExecMeta] => - if (scan.asInstanceOf[FileSourceScanExecMeta] - .wrapped.schema.fieldNames.contains(METADATA_ROW_DEL_COL) && scan.canThisBeReplaced) { - overrideCount += 1 + case scan if scan.isInstanceOf[FileSourceScanExecMeta] && + isLowShuffleTargetScan(scan.asInstanceOf[FileSourceScanExecMeta]) => + val fileScan = scan.asInstanceOf[FileSourceScanExecMeta] + targetScanCount += 1 + if (fileScan.canThisBeReplaced) { + gpuTargetScanCount += 1 } case m => m.childPlans.foreach(count) } } count(planMeta) - overrideCount == 2 + targetScanCount == 2 && gpuTargetScanCount == targetScanCount } if (!mergePlanOverrideSucceed) { @@ -753,6 +745,13 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend false } + private def isLowShuffleTargetScan(scan: FileSourceScanExecMeta): Boolean = { + scan.wrapped.relation.location match { + case index: TahoeBatchFileIndex => index.deltaLog == context.deltaTxn.deltaLog + case _ => false + } + } + private def verifyGpuPlan(input: DataFrame)(checkPlanMeta: SparkPlanMeta[SparkPlan] => Boolean) : Boolean = { val overridePlan = GpuOverrides.wrapAndTagPlan(input.queryExecution.sparkPlan, @@ -797,7 +796,7 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend } private lazy val dataSkippedTargetDF: DataFrame = { - addRowIndexMetaColumn(buildTargetDFWithFiles(dataSkippedFiles)) + addRowIndexMetaColumn(dataSkippedFiles) } private lazy val touchedFiles: Map[String, (Roaring64Bitmap, AddFile)] = this.findTouchedFiles() @@ -917,91 +916,25 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend /** - * Modify original data frame to insert + * Uses Databricks' deletion-vector scan preparation to expose the file metadata column, then + * copies its file-relative row index into * [[GpuDeltaParquetFileFormatUtils.METADATA_ROW_IDX_COL]]. */ - private def addRowIndexMetaColumn(baseDF: DataFrame): DataFrame = { - val rowIdxAttr = AttributeReference( - METADATA_ROW_IDX_COL, - METADATA_ROW_IDX_FIELD.dataType, - METADATA_ROW_IDX_FIELD.nullable)() - - val newPlan = baseDF.queryExecution.analyzed.transformUp { - case r: LogicalRelation if r.relation.isInstanceOf[HadoopFsRelation] => - val fs = r.relation.asInstanceOf[HadoopFsRelation] - val newSchema = StructType(fs.dataSchema.fields).add(METADATA_ROW_IDX_FIELD) - val newFs = lowShuffleScanRelation( - fs, newSchema, GpuLowShuffleMergeScanInfo(rowIndexMaps = None)) - - val newOutput = r.output :+ rowIdxAttr - r.copy(relation = newFs, output = newOutput) - case p@Project(projectList, _) => - val newProjectList = projectList :+ rowIdxAttr - p.copy(projectList = newProjectList) - } - - Dataset.ofRows(context.spark, newPlan) - } - - /** - * The result is scanning target table with touched files, and added an extra - * [[METADATA_ROW_DEL_COL]] to indicate whether filtered by joining with source table in first - * step. - */ - private def getTouchedTargetDF(touchedFiles: Map[String, (Roaring64Bitmap, AddFile)]) - : DataFrame = { - // Generate a new target dataframe that has same output attributes exprIds as the target plan. - // This allows us to apply the existing resolved update/insert expressions. - val baseTargetDF = buildTargetDFWithFiles(touchedFiles.values.map(_._2).toSeq) - - val newPlan = { - val rowDelAttr = AttributeReference( - METADATA_ROW_DEL_COL, - METADATA_ROW_DEL_FIELD.dataType, - METADATA_ROW_DEL_FIELD.nullable)() - - baseTargetDF.queryExecution.analyzed.transformUp { - case r: LogicalRelation if r.relation.isInstanceOf[HadoopFsRelation] => - val fs = r.relation.asInstanceOf[HadoopFsRelation] - val newSchema = StructType(fs.dataSchema.fields).add(METADATA_ROW_DEL_FIELD) - val broadcastRows = touchedRowsBroadcast.getOrElse { - val rows = touchedFiles.map { case (path, (bitmap, _)) => - new URI(path) -> RoaringBitmapWrapper(bitmap).serializeToBytes() - } - val broadcast = context.spark.sparkContext.broadcast(rows) - touchedRowsBroadcast = Some(broadcast) - broadcast - } - val newFs = lowShuffleScanRelation( - fs, newSchema, GpuLowShuffleMergeScanInfo(rowIndexMaps = Some(broadcastRows))) - - val newOutput = r.output :+ rowDelAttr - r.copy(relation = newFs, output = newOutput) - case p@Project(projectList, _) => - val newProjectList = projectList :+ rowDelAttr - p.copy(projectList = newProjectList) + private def addRowIndexMetaColumn(files: Seq[AddFile]): DataFrame = { + val fileIndex = context.deltaTxn.deltaLog.createDataFrame(context.deltaTxn.snapshot, files) + .queryExecution.analyzed.collectFirst { + case relation: LogicalRelation + if relation.relation.isInstanceOf[HadoopFsRelation] && + relation.relation.asInstanceOf[HadoopFsRelation] + .location.isInstanceOf[TahoeFileIndex] => + relation.relation.asInstanceOf[HadoopFsRelation] + .location.asInstanceOf[TahoeFileIndex] + }.getOrElse { + throw new IllegalStateException("Unable to find the Delta file index for low shuffle merge") } - } - - val df = Dataset.ofRows(context.spark, newPlan) - .withColumn(TARGET_ROW_PRESENT_COL, lit(true)) - - df - } - - private def lowShuffleScanRelation( - relation: HadoopFsRelation, - dataSchema: StructType, - scanInfo: GpuLowShuffleMergeScanInfo): HadoopFsRelation = { - val scanId = GpuLowShuffleMergeScanRegistry.register(scanInfo) - scanRegistrationIds += scanId - val fileFormat = relation.fileFormat.asInstanceOf[DeltaParquetFileFormat] - .copy(optimizationsEnabled = false) - relation.copy( - dataSchema = dataSchema, - fileFormat = fileFormat, - options = relation.options + (GpuLowShuffleMergeScanRegistry.OPTION_KEY -> scanId))( - context.spark) + DMLWithDeletionVectorsHelper.createTargetDfForScanningForMatches( + context.spark, context.cmd.target, fileIndex) + .withColumn(METADATA_ROW_IDX_COL, col("_metadata.row_index")) } /** @@ -1026,7 +959,10 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend val sourceDF = this.sourceDF .withColumn(SOURCE_ROW_PRESENT_COL, DFUDFShims.exprToColumn(incrSourceRowCountExpr)) - val targetDF = getTouchedTargetDF(touchedFiles) + // The join itself selects touched target rows, so this pass can scan the touched files without + // applying the temporary deletion vectors used by the unmodified-row pass. + val targetDF = buildTargetDFWithFiles(touchedFiles.values.map(_._2).toSeq) + .withColumn(TARGET_ROW_PRESENT_COL, lit(true)) val joinedDF = { val joinType = if (hasNoInserts && @@ -1035,10 +971,7 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend } else { "leftOuter" } - val matchedTargetDF = targetDF.filter(METADATA_ROW_DEL_COL) - .drop(METADATA_ROW_DEL_COL) - - sourceDF.join(matchedTargetDF, DFUDFShims.exprToColumn(context.cmd.condition), joinType) + sourceDF.join(targetDF, DFUDFShims.exprToColumn(context.cmd.condition), joinType) } val modifiedRowsSchema = context.deltaTxn.metadata.schema @@ -1136,9 +1069,16 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend } private def getUnmodifiedDF(touchedFiles: Map[String, (Roaring64Bitmap, AddFile)]): DataFrame = { - getTouchedTargetDF(touchedFiles) - .filter(!col(METADATA_ROW_DEL_COL)) - .drop(TARGET_ROW_PRESENT_COL, METADATA_ROW_DEL_COL) + val hadoopConf = context.deltaTxn.deltaLog.newDeltaHadoopConf() + val tablePath = context.deltaTxn.deltaLog.dataPath.toString + val filesWithTemporaryDVs = touchedFiles.values.map { case (bitmap, addFile) => + addFile.copy(deletionVector = MergeExecutor.toDeletionVector( + bitmap, + Option(addFile.deletionVector), + hadoopConf, + tablePath)) + }.toSeq + buildTargetDFWithFiles(filesWithTemporaryDVs) } } @@ -1174,6 +1114,28 @@ object MergeExecutor { // rather than the version from Delta Lake. val CDC_TYPE_NOT_CDC_LITERAL: Literal = Literal(null, StringType) + private[rapids] def toDeletionVector( + bitmap: Roaring64Bitmap, + existing: Option[DeletionVectorDescriptor], + hadoopConf: Configuration, + tablePath: String): DeletionVectorDescriptor = { + val combined = existing.map { descriptor => + RapidsDeletionVectors.loadScalaBitmap( + hadoopConf, + Some(descriptor.serializeToBase64()), + Some(RowIndexFilterType.IF_CONTAINED), + None, + tablePath) + }.getOrElse(new RoaringBitmapArray()) + val touchedIndexes = bitmap.getLongIterator + while (touchedIndexes.hasNext) { + combined.add(touchedIndexes.next()) + } + combined.runOptimize() + DeletionVectorDescriptor.inlineInLog( + combined.serializeAsByteArray(RoaringBitmapArrayFormat.Portable), combined.cardinality) + } + /** Count the number of distinct partition values among the AddFiles in the given set. */ def totalBytesAndDistinctPartitionValues(files: Seq[FileAction]): (Long, Int) = { val distinctValues = new mutable.HashSet[Map[String, String]]() diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala index 1dda69328c8..3a0372dfeaa 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala @@ -122,9 +122,7 @@ object DeltaSpark400DB173Provider extends DatabricksDeltaProviderBase { override def getReadFileFormat( relation: HadoopFsRelation, rapidsConf: RapidsConf): FileFormat = { val fmt = relation.fileFormat.asInstanceOf[DeltaParquetFileFormat] - if (GpuLowShuffleMergeScanRegistry.lookup(relation.options).isDefined) { - GpuDeltaParquetFileFormat.convertToGpu(relation) - } else if (isPushDVPredicateDownEnabled(rapidsConf)) { + if (isPushDVPredicateDownEnabled(rapidsConf)) { GpuDeltaParquetFileFormatNativeDV( relation = relation, protocol = fmt.protocol, 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 c436d43dd3d..1cd4d12cf86 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 @@ -16,8 +16,6 @@ package com.nvidia.spark.rapids.delta -import java.net.URI - import com.databricks.sql.io.RowIndexFilterType import com.databricks.sql.transaction.tahoe.{ DeltaColumnMapping, @@ -30,8 +28,7 @@ import com.databricks.sql.transaction.tahoe.{ import com.databricks.sql.transaction.tahoe.actions.{Metadata, Protocol} import com.databricks.sql.transaction.tahoe.files.TahoeFileIndex import com.databricks.sql.transaction.tahoe.schema.SchemaMergingUtils -import com.nvidia.spark.rapids.{GpuMetric, RapidsConf, SparkPlanMeta} -import com.nvidia.spark.rapids.delta.GpuDeltaParquetFileFormatUtils.addMetadataColumnToIterator +import com.nvidia.spark.rapids.{GpuMetric, SparkPlanMeta} import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path @@ -44,7 +41,6 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.rapids.shims.TrampolineConnectShims import org.apache.spark.sql.sources.Filter import org.apache.spark.sql.types.{MetadataBuilder, StructType} -import org.apache.spark.sql.vectorized.ColumnarBatch /** * GPU Delta Parquet file format for Databricks 17.3. @@ -66,8 +62,7 @@ case class GpuDeltaParquetFileFormat( nullableRowTrackingGeneratedFields: Boolean = false, optimizationsEnabled: Boolean = true, tablePath: Option[String] = None, - isCDCRead: Boolean = false, - lowShuffleMergeScan: Option[GpuLowShuffleMergeScanInfo] = None + isCDCRead: Boolean = false ) extends GpuDeltaParquetFileFormatBase { override val columnMappingMode: DeltaColumnMappingMode = metadata.columnMappingMode @@ -118,7 +113,7 @@ case class GpuDeltaParquetFileFormat( * Translates pushed filters to physical column names when Delta column mapping is enabled. */ private def prepareFiltersForRead(filters: Seq[Filter]): Seq[Filter] = { - if (lowShuffleMergeScan.isDefined || !effectiveOptimizationsEnabled) { + if (!effectiveOptimizationsEnabled) { Seq.empty } else if (columnMappingMode != NoMapping) { val physicalNameMap = DeltaColumnMapping.getLogicalNameToPhysicalNameMap(referenceSchema) @@ -136,7 +131,7 @@ case class GpuDeltaParquetFileFormat( override def isSplitable( sparkSession: SparkSession, options: Map[String, String], - path: Path): Boolean = lowShuffleMergeScan.isEmpty && effectiveOptimizationsEnabled + path: Path): Boolean = effectiveOptimizationsEnabled private def hasDeletionVectorRead: Boolean = GpuDeltaParquetFileFormat.isDeletionVectorRead( @@ -169,7 +164,7 @@ case class GpuDeltaParquetFileFormat( hadoopConf: Configuration, metrics: Map[String, GpuMetric]) : PartitionedFile => Iterator[InternalRow] = { - val dataReader = super.buildReaderWithPartitionValuesAndMetrics( + super.buildReaderWithPartitionValuesAndMetrics( sparkSession, dataSchema, partitionSchema, @@ -178,27 +173,6 @@ case class GpuDeltaParquetFileFormat( options, hadoopConf, metrics) - - lowShuffleMergeScan.map { scanInfo => - val maxBatchSize = RapidsConf.DELTA_LOW_SHUFFLE_MERGE_SCATTER_DEL_VECTOR_BATCH_SIZE - .get(sparkSession.sessionState.conf) - val scatterTime = metrics(GpuMetric.DELETION_VECTOR_SCATTER_TIME) - val deletionVectorSize = metrics(GpuMetric.DELETION_VECTOR_SIZE) - (file: PartitionedFile) => { - val bitmap = scanInfo.rowIndexMaps.flatMap { broadcast => - broadcast.value.get(new URI(file.filePath.toString)).map { bytes => - deletionVectorSize += bytes.length - RoaringBitmapWrapper.deserializeFromBytes(bytes).inner - } - } - addMetadataColumnToIterator( - prepareSchema(requiredSchema), - bitmap, - dataReader(file).asInstanceOf[Iterator[ColumnarBatch]], - maxBatchSize, - scatterTime).asInstanceOf[Iterator[InternalRow]] - } - }.getOrElse(dataReader) } } @@ -294,8 +268,7 @@ object GpuDeltaParquetFileFormat { nullableRowTrackingGeneratedFields = fmt.nullableRowTrackingGeneratedFields, optimizationsEnabled = fmt.optimizationsEnabled, tablePath = fmt.tablePath, - isCDCRead = fmt.isCDCRead, - lowShuffleMergeScan = GpuLowShuffleMergeScanRegistry.lookup(relation.options)) + isCDCRead = fmt.isCDCRead) } private def hasRowIndexFiltersInTahoeFileIndex(relation: HadoopFsRelation): Boolean = { diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuLowShuffleMergeScan.scala b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuLowShuffleMergeScan.scala deleted file mode 100644 index 9a3834df2c3..00000000000 --- a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuLowShuffleMergeScan.scala +++ /dev/null @@ -1,50 +0,0 @@ -/* - * 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.nvidia.spark.rapids.delta - -import java.net.URI -import java.util.UUID -import java.util.concurrent.ConcurrentHashMap - -import org.apache.spark.broadcast.Broadcast - -/** Driver-side information needed to build a low-shuffle target scan. */ -case class GpuLowShuffleMergeScanInfo( - rowIndexMaps: Option[Broadcast[Map[URI, Array[Byte]]]]) - -/** - * Bridges information created by the low-shuffle command into Delta file-format conversion. - * Only a small opaque ID is placed in the logical relation. The broadcast itself is attached to - * the GPU file format when the scan is converted and is therefore sent to executors normally. - */ -object GpuLowShuffleMergeScanRegistry { - val OPTION_KEY: String = "spark.rapids.internal.delta.lowShuffleMerge.scanId" - - private val scans = new ConcurrentHashMap[String, GpuLowShuffleMergeScanInfo]() - - def register(info: GpuLowShuffleMergeScanInfo): String = { - val id = UUID.randomUUID().toString - scans.put(id, info) - id - } - - def lookup(options: Map[String, String]): Option[GpuLowShuffleMergeScanInfo] = { - options.get(OPTION_KEY).flatMap(id => Option(scans.get(id))) - } - - def remove(id: String): Unit = scans.remove(id) -} From 125c82762aaa9d745327f0cc3cb68aa585f70b43 Mon Sep 17 00:00:00 2001 From: Ray Liu Date: Wed, 9 Sep 2026 14:49:14 +0800 Subject: [PATCH 04/10] [FEA] Restrict low shuffle merge tests to DBR 17.3 Remove the duplicate Databricks-only upsert test and its plan-capture helper, drop the AQE override, and use the exact DBR 17.3 runtime predicate for the shared low-shuffle-merge coverage. Signed-off-by: Ray Liu --- .../delta_lake_low_shuffle_merge_test.py | 62 +++---------------- 1 file changed, 9 insertions(+), 53 deletions(-) diff --git a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py index 0b42e72a8f8..4cb9ad3acc3 100644 --- a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py +++ b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py @@ -19,7 +19,7 @@ from delta_lake_merge_common import * from marks import * from pyspark.sql.types import * -from spark_session import is_databricks173_or_later, spark_version +from spark_session import is_databricks_version, spark_version delta_merge_enabled_conf = copy_and_update(delta_writes_enabled_conf, {"spark.rapids.sql.command.MergeIntoCommand": "true", @@ -27,41 +27,16 @@ "spark.rapids.sql.delta.lowShuffleMerge.enabled": "true", "spark.rapids.sql.format.parquet.reader.type": "PERFILE"}) -if is_databricks173_or_later(): - # Disable AQE temporarily until https://github.com/NVIDIA/spark-rapids/issues/14319 is resolved. - delta_merge_enabled_conf = copy_and_update( - delta_merge_enabled_conf, - {"spark.sql.adaptive.enabled": "false"}) - - def supports_delta_low_shuffle_merge(): - return is_databricks173_or_later() or \ + return is_databricks_version(17, 3) or \ (not is_databricks_runtime() and spark_version().startswith("3.4")) -def assert_low_shuffle_merge(do_merge, data_path, conf, expect_write=True): - assert expect_write - cpu_result = with_cpu_session(lambda spark: do_merge(spark, data_path + "/CPU"), conf=conf) - - callback = spark_jvm().org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback - callback.startCapture() - try: - gpu_result = with_gpu_session( - lambda spark: do_merge(spark, data_path + "/GPU"), conf=conf) - captured_plans = callback.getResultsWithTimeout(10000) - finally: - callback.endCapture() - - assert_equal(cpu_result, gpu_result) - command_name = "GpuLowShuffleMergeCommand" - assert any(command_name in str(plan) for plan in captured_plans), \ - f"{command_name} was not found in the captured MERGE plans" - @allow_non_gpu("ColumnarToRowExec", *delta_meta_allow) @delta_lake @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), - reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") @pytest.mark.parametrize("use_cdf", [True, False], ids=idfn) @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn) def test_delta_low_shuffle_merge_when_gpu_file_scan_override_failed(spark_tmp_path, @@ -89,7 +64,7 @@ def test_delta_low_shuffle_merge_when_gpu_file_scan_override_failed(spark_tmp_pa @delta_lake @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), - reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") @pytest.mark.parametrize("table_ranges", [(range(20), range(10)), # partial insert of source (range(5), range(5)), # no-op insert (range(10), range(20, 30)) # full insert of source @@ -107,7 +82,7 @@ def test_delta_merge_not_match_insert_only(spark_tmp_path, spark_tmp_table_facto @delta_lake @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), - reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") @pytest.mark.parametrize("table_ranges", [(range(10), range(20)), # partial delete of target (range(5), range(5)), # full delete of target (range(10), range(20, 30)) # no-op delete @@ -125,37 +100,18 @@ def test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, @delta_lake @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), - reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") @pytest.mark.parametrize("use_cdf", [pytest.param(True, marks=pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552")), False], ids=idfn) @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn) def test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, use_cdf, num_slices): do_test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, use_cdf, False, num_slices, False, delta_merge_enabled_conf) - -@allow_non_gpu(*delta_meta_allow) -@delta_lake -@ignore_order -@pytest.mark.skipif(not is_databricks173_or_later(), - reason="Databricks low shuffle merge requires DBR 17.3+") -def test_databricks_delta_low_shuffle_merge_perfile( - spark_tmp_path, spark_tmp_table_factory): - do_test_delta_merge_standard_upsert( - spark_tmp_path, - spark_tmp_table_factory, - use_cdf=False, - enable_deletion_vectors=False, - num_slices=10, - compare_logs=False, - conf=delta_merge_enabled_conf, - assert_func=assert_low_shuffle_merge) - - @allow_non_gpu(*delta_meta_allow) @delta_lake @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), - reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") @pytest.mark.parametrize("use_cdf", [pytest.param(True, marks=pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552")), False], ids=idfn) @pytest.mark.parametrize("merge_sql", [ "MERGE INTO {dest_table} d USING {src_table} s ON d.a == s.a" \ @@ -178,7 +134,7 @@ def test_delta_merge_upsert_with_condition(spark_tmp_path, spark_tmp_table_facto @delta_lake @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), - reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") @pytest.mark.parametrize("use_cdf", [True, False], ids=idfn) @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn) def test_delta_merge_upsert_with_unmatchable_match_condition(spark_tmp_path, spark_tmp_table_factory, use_cdf, num_slices): @@ -194,7 +150,7 @@ def test_delta_merge_upsert_with_unmatchable_match_condition(spark_tmp_path, spa @delta_lake @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), - reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3+") + reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") @pytest.mark.parametrize("use_cdf", [pytest.param(True, marks=pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552")), False], ids=idfn) def test_delta_merge_update_with_aggregation(spark_tmp_path, spark_tmp_table_factory, use_cdf): do_test_delta_merge_update_with_aggregation(spark_tmp_path, spark_tmp_table_factory, use_cdf, False, From 2fb0190e5834438ee5395cc76a650f141cc2829a Mon Sep 17 00:00:00 2001 From: Ray Liu Date: Wed, 9 Sep 2026 17:00:58 +0800 Subject: [PATCH 05/10] [BUG] Fall back from low shuffle merge for CDF Signed-off-by: Ray Liu --- .../tahoe/rapids/GpuLowShuffleMergeCommand.scala | 6 ++++++ .../python/delta_lake_low_shuffle_merge_test.py | 15 +++++++++++---- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala index 30e96c63a6a..34a57e5e606 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala @@ -684,8 +684,14 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend * the temporary deletion-vector scan used to retain unmodified rows. * 2. The temporary deletion vectors introduce extra overhead, so it may be better to fall back * when the changeset is too large. + * 3. Low shuffle merge does not generate change data feed rows. */ def shouldFallback(): Boolean = { + if (DeltaConfigs.CHANGE_DATA_FEED.fromMetaData(context.deltaTxn.metadata)) { + logWarning("Change data feed is enabled, falling back to traditional merge.") + return true + } + // Trying to detect if we can execute finding touched files. val touchFilePlanOverrideSucceed = verifyGpuPlan(planForFindingTouchedFiles()) { planMeta => def check(meta: SparkPlanMeta[SparkPlan]): Boolean = { diff --git a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py index 4cb9ad3acc3..91c5b5bde7a 100644 --- a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py +++ b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py @@ -32,6 +32,13 @@ def supports_delta_low_shuffle_merge(): (not is_databricks_runtime() and spark_version().startswith("3.4")) +low_shuffle_cdf_param = pytest.param( + True, + marks=pytest.mark.xfail( + condition=not is_databricks_version(17, 3), + reason="https://github.com/NVIDIA/spark-rapids/issues/13552")) + + @allow_non_gpu("ColumnarToRowExec", *delta_meta_allow) @delta_lake @ignore_order @@ -87,7 +94,7 @@ def test_delta_merge_not_match_insert_only(spark_tmp_path, spark_tmp_table_facto (range(5), range(5)), # full delete of target (range(10), range(20, 30)) # no-op delete ], ids=idfn) -@pytest.mark.parametrize("use_cdf", [pytest.param(True, marks=pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552")), False], ids=idfn) +@pytest.mark.parametrize("use_cdf", [low_shuffle_cdf_param, False], ids=idfn) @pytest.mark.parametrize("partition_columns", [None, ["a"], ["b"], ["a", "b"]], ids=idfn) @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn) def test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, table_ranges, @@ -101,7 +108,7 @@ def test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") -@pytest.mark.parametrize("use_cdf", [pytest.param(True, marks=pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552")), False], ids=idfn) +@pytest.mark.parametrize("use_cdf", [low_shuffle_cdf_param, False], ids=idfn) @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn) def test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, use_cdf, num_slices): do_test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, use_cdf, False, @@ -112,7 +119,7 @@ def test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, us @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") -@pytest.mark.parametrize("use_cdf", [pytest.param(True, marks=pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552")), False], ids=idfn) +@pytest.mark.parametrize("use_cdf", [low_shuffle_cdf_param, False], ids=idfn) @pytest.mark.parametrize("merge_sql", [ "MERGE INTO {dest_table} d USING {src_table} s ON d.a == s.a" \ " WHEN MATCHED AND s.b > 'q' THEN UPDATE SET d.a = s.a / 2, d.b = s.b" \ @@ -151,7 +158,7 @@ def test_delta_merge_upsert_with_unmatchable_match_condition(spark_tmp_path, spa @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") -@pytest.mark.parametrize("use_cdf", [pytest.param(True, marks=pytest.mark.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552")), False], ids=idfn) +@pytest.mark.parametrize("use_cdf", [low_shuffle_cdf_param, False], ids=idfn) def test_delta_merge_update_with_aggregation(spark_tmp_path, spark_tmp_table_factory, use_cdf): do_test_delta_merge_update_with_aggregation(spark_tmp_path, spark_tmp_table_factory, use_cdf, False, delta_merge_enabled_conf) From 15ddbd15cd96e715a8a88c554cc44fbd959c9438 Mon Sep 17 00:00:00 2001 From: Ray Liu Date: Wed, 9 Sep 2026 18:40:53 +0800 Subject: [PATCH 06/10] [TEST] Narrow low shuffle CDF xfail on DBR 17.3 Signed-off-by: Ray Liu --- .../src/main/python/delta_lake_low_shuffle_merge_test.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py index 91c5b5bde7a..8e3b954873c 100644 --- a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py +++ b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py @@ -99,6 +99,9 @@ def test_delta_merge_not_match_insert_only(spark_tmp_path, spark_tmp_table_facto @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn) def test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, table_ranges, use_cdf, partition_columns, num_slices): + if (use_cdf and is_databricks_version(17, 3) + and table_ranges == (range(10), range(20, 30))): + pytest.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552") do_test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, table_ranges, use_cdf, False, partition_columns, num_slices, False, delta_merge_enabled_conf) From 9c3acceb1aab61a0fce9f8c897ec9c4a56c5f07c Mon Sep 17 00:00:00 2001 From: Ray Liu Date: Thu, 10 Sep 2026 10:08:51 +0800 Subject: [PATCH 07/10] [BUG] Reconcile low shuffle merge with upstream clauses Signed-off-by: Ray Liu --- .../tahoe/rapids/GpuLowShuffleMergeCommand.scala | 4 ++-- .../spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala index 34a57e5e606..4f304221111 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala @@ -335,8 +335,7 @@ case class GpuLowShuffleMergeCommand( Option(condition), matchedClauses.map(DeltaOperations.MergePredicate(_)), notMatchedClauses.map(DeltaOperations.MergePredicate(_)), - // We do not support notMatchedBySourcePredicates yet and fall back to CPU - // See https://github.com/NVIDIA/spark-rapids/issues/8415 + // The command shim selects traditional GPU merge when these clauses are present. notMatchedBySourcePredicates = Seq.empty[MergePredicate] ), RowTracking.addPreservedRowTrackingTagIfNotSet(deltaTxn.snapshot)) @@ -347,6 +346,7 @@ case class GpuLowShuffleMergeCommand( condition, matchedClauses, notMatchedClauses, + notMatchedBySourceClauses, deltaTxn.metadata.partitionColumns.nonEmpty) recordDeltaEvent(targetDeltaLog, "delta.dml.merge.stats", data = stats) diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala index c090ed242ce..fff12993803 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/shims/MergeIntoCommandMetaShim.scala @@ -56,7 +56,8 @@ object MergeIntoCommandMetaShim { } def convertToGpu(mergeCmd: MergeIntoCommand, conf: RapidsConf): RunnableCommand = { - if (conf.isDeltaLowShuffleMergeEnabled && conf.isParquetPerFileReadEnabled) { + if (conf.isDeltaLowShuffleMergeEnabled && conf.isParquetPerFileReadEnabled && + mergeCmd.notMatchedBySourceClauses.isEmpty) { GpuLowShuffleMergeCommand( mergeCmd.source, mergeCmd.target, @@ -88,7 +89,8 @@ object MergeIntoCommandMetaShim { } def convertToGpu(mergeCmd: MergeIntoCommandEdge, conf: RapidsConf): RunnableCommand = { - if (conf.isDeltaLowShuffleMergeEnabled && conf.isParquetPerFileReadEnabled) { + if (conf.isDeltaLowShuffleMergeEnabled && conf.isParquetPerFileReadEnabled && + mergeCmd.notMatchedBySourceClauses.isEmpty) { GpuLowShuffleMergeCommand( mergeCmd.source, mergeCmd.target, From b873db07c2e181754c27e6fa12787d1714c5894e Mon Sep 17 00:00:00 2001 From: Ray Liu Date: Thu, 10 Sep 2026 13:12:19 +0800 Subject: [PATCH 08/10] [FEA] Support CDF with low shuffle merge Signed-off-by: Ray Liu --- .../delta24x/GpuLowShuffleMergeCommand.scala | 278 ++++++++++++++- .../rapids/GpuLowShuffleMergeCommand.scala | 328 +++++++++++++++++- .../delta_lake_low_shuffle_merge_test.py | 18 +- 3 files changed, 595 insertions(+), 29 deletions(-) diff --git a/delta-lake/delta-24x/src/main/scala/org/apache/spark/sql/delta/rapids/delta24x/GpuLowShuffleMergeCommand.scala b/delta-lake/delta-24x/src/main/scala/org/apache/spark/sql/delta/rapids/delta24x/GpuLowShuffleMergeCommand.scala index 9c27d28ebd3..da57c28bebb 100644 --- a/delta-lake/delta-24x/src/main/scala/org/apache/spark/sql/delta/rapids/delta24x/GpuLowShuffleMergeCommand.scala +++ b/delta-lake/delta-24x/src/main/scala/org/apache/spark/sql/delta/rapids/delta24x/GpuLowShuffleMergeCommand.scala @@ -26,7 +26,7 @@ import java.util.concurrent.TimeUnit import scala.collection.mutable -import com.nvidia.spark.rapids.{GpuOverrides, RapidsConf, SparkPlanMeta} +import com.nvidia.spark.rapids.{BaseExprMeta, GpuOverrides, RapidsConf, SparkPlanMeta} import com.nvidia.spark.rapids.RapidsConf.DELTA_LOW_SHUFFLE_MERGE_DEL_VECTOR_BROADCAST_THRESHOLD import com.nvidia.spark.rapids.delta._ import com.nvidia.spark.rapids.delta.GpuDeltaParquetFileFormatUtils._ @@ -37,17 +37,26 @@ import org.apache.spark.SparkContext import org.apache.spark.internal.Logging import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute -import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeReference, CaseWhen, Expression, Literal, NamedExpression, PredicateHelper} +import org.apache.spark.sql.catalyst.encoders.RowEncoder +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeReference, + CaseWhen, Expression, Literal, NamedExpression, PredicateHelper} import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.plans.logical.{DeltaMergeAction, DeltaMergeIntoClause, DeltaMergeIntoMatchedClause, DeltaMergeIntoMatchedDeleteClause, DeltaMergeIntoMatchedUpdateClause, DeltaMergeIntoNotMatchedBySourceClause, DeltaMergeIntoNotMatchedBySourceDeleteClause, DeltaMergeIntoNotMatchedBySourceUpdateClause, DeltaMergeIntoNotMatchedClause, DeltaMergeIntoNotMatchedInsertClause, LogicalPlan, Project} import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap -import org.apache.spark.sql.delta.{DeltaErrors, DeltaLog, DeltaOperations, DeltaParquetFileFormat, DeltaTableUtils, DeltaUDF, NoMapping, OptimisticTransaction, RowIndexFilterType} +import org.apache.spark.sql.delta.{DeltaConfigs, DeltaErrors, DeltaLog, DeltaOperations, + DeltaParquetFileFormat, DeltaTableUtils, DeltaUDF, NoMapping, OptimisticTransaction, + RowIndexFilterType} import org.apache.spark.sql.delta.DeltaOperations.MergePredicate import org.apache.spark.sql.delta.DeltaParquetFileFormat.DeletionVectorDescriptorWithFilterType import org.apache.spark.sql.delta.actions.{AddCDCFile, AddFile, DeletionVectorDescriptor, FileAction} import org.apache.spark.sql.delta.commands.DeltaCommand +import org.apache.spark.sql.delta.commands.cdc.CDCReader._ import org.apache.spark.sql.delta.rapids.{GpuDeltaLog, GpuOptimisticTransactionBase} -import org.apache.spark.sql.delta.rapids.delta24x.MergeExecutor.{toDeletionVector, totalBytesAndDistinctPartitionValues, INCR_METRICS_COL, INCR_METRICS_FIELD, ROW_DROPPED_COL, ROW_DROPPED_FIELD, SOURCE_ROW_PRESENT_COL, SOURCE_ROW_PRESENT_FIELD, TARGET_ROW_PRESENT_COL, TARGET_ROW_PRESENT_FIELD} +import org.apache.spark.sql.delta.rapids.delta24x.MergeExecutor.{toDeletionVector, + totalBytesAndDistinctPartitionValues, CDC_TYPE_NOT_CDC_LITERAL, INCR_METRICS_COL, + INCR_METRICS_FIELD, INCR_ROW_COUNT_COL, + ROW_DROPPED_COL, ROW_DROPPED_FIELD, SOURCE_ROW_PRESENT_COL, SOURCE_ROW_PRESENT_FIELD, + TARGET_ROW_PRESENT_COL, TARGET_ROW_PRESENT_FIELD} import org.apache.spark.sql.delta.schema.ImplicitMetadataOperation import org.apache.spark.sql.delta.sources.DeltaSQLConf import org.apache.spark.sql.delta.util.{AnalysisHelper, DeltaFileOperations} @@ -890,6 +899,254 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend df } + private def addMergeJoinProcessor( + joinedPlan: LogicalPlan, + outputRowSchema: StructType, + targetRowHasNoMatch: Expression, + sourceRowHasNoMatch: Expression, + matchedConditions: Seq[Expression], + matchedOutputs: Seq[Seq[Seq[Expression]]], + notMatchedConditions: Seq[Expression], + notMatchedOutputs: Seq[Seq[Seq[Expression]]], + notMatchedBySourceConditions: Seq[Expression], + notMatchedBySourceOutputs: Seq[Seq[Seq[Expression]]], + noopCopyOutput: Seq[Expression], + deleteRowOutput: Seq[Expression]): Dataset[Row] = { + def wrap(e: Expression): BaseExprMeta[Expression] = { + GpuOverrides.wrapExpr(e, context.rapidsConf, None) + } + + val targetRowHasNoMatchMeta = wrap(targetRowHasNoMatch) + val sourceRowHasNoMatchMeta = wrap(sourceRowHasNoMatch) + val matchedConditionsMetas = matchedConditions.map(wrap) + val matchedOutputsMetas = matchedOutputs.map(_.map(_.map(wrap))) + val notMatchedConditionsMetas = notMatchedConditions.map(wrap) + val notMatchedOutputsMetas = notMatchedOutputs.map(_.map(_.map(wrap))) + val notMatchedBySourceConditionsMetas = notMatchedBySourceConditions.map(wrap) + val notMatchedBySourceOutputsMetas = notMatchedBySourceOutputs.map(_.map(_.map(wrap))) + val noopCopyOutputMetas = noopCopyOutput.map(wrap) + val deleteRowOutputMetas = deleteRowOutput.map(wrap) + val allMetas = Seq(targetRowHasNoMatchMeta, sourceRowHasNoMatchMeta) ++ + matchedConditionsMetas ++ matchedOutputsMetas.flatten.flatten ++ + notMatchedConditionsMetas ++ notMatchedOutputsMetas.flatten.flatten ++ + notMatchedBySourceConditionsMetas ++ notMatchedBySourceOutputsMetas.flatten.flatten ++ + noopCopyOutputMetas ++ deleteRowOutputMetas + allMetas.foreach(_.tagForGpu()) + val canReplace = allMetas.forall(_.canExprTreeBeReplaced) && + context.rapidsConf.isOperatorEnabled( + "spark.rapids.sql.exec.RapidsProcessDeltaMergeJoinExec", false, false) + if (context.rapidsConf.shouldExplainAll || (context.rapidsConf.shouldExplain && !canReplace)) { + val exprExplains = allMetas.map(_.explain(context.rapidsConf.shouldExplainAll)) + val execWorkInfo = if (canReplace) { + "will run on GPU" + } else { + "cannot run on GPU because not all merge processing expressions can be replaced" + } + logWarning(s" $execWorkInfo:\n" + + s" ${exprExplains.mkString(" ")}") + } + + if (canReplace) { + val processedJoinPlan = RapidsProcessDeltaMergeJoin( + joinedPlan, + outputRowSchema.toAttributes, + targetRowHasNoMatch = targetRowHasNoMatch, + sourceRowHasNoMatch = sourceRowHasNoMatch, + matchedConditions = matchedConditions, + matchedOutputs = matchedOutputs, + notMatchedConditions = notMatchedConditions, + notMatchedOutputs = notMatchedOutputs, + notMatchedBySourceConditions = notMatchedBySourceConditions, + notMatchedBySourceOutputs = notMatchedBySourceOutputs, + noopCopyOutput = noopCopyOutput, + deleteRowOutput = deleteRowOutput) + Dataset.ofRows(context.spark, processedJoinPlan) + } else { + val joinedRowEncoder = RowEncoder(joinedPlan.schema) + val outputRowEncoder = RowEncoder(outputRowSchema).resolveAndBind() + val processor = new GpuMergeIntoCommand.JoinedRowProcessor( + targetRowHasNoMatch = targetRowHasNoMatch, + sourceRowHasNoMatch = sourceRowHasNoMatch, + matchedConditions = matchedConditions, + matchedOutputs = matchedOutputs, + notMatchedConditions = notMatchedConditions, + notMatchedOutputs = notMatchedOutputs, + notMatchedBySourceConditions = notMatchedBySourceConditions, + notMatchedBySourceOutputs = notMatchedBySourceOutputs, + noopCopyOutput = noopCopyOutput, + deleteRowOutput = deleteRowOutput, + joinedAttributes = joinedPlan.output, + joinedRowEncoder = joinedRowEncoder, + outputRowEncoder = outputRowEncoder) + Dataset.ofRows(context.spark, joinedPlan) + .mapPartitions(processor.processPartition)(outputRowEncoder) + } + } + + /** Generate both rewritten table rows and explicit change-data-feed rows. */ + private def getModifiedDFWithCdf( + touchedFiles: Map[String, (Roaring64Bitmap, AddFile)]): DataFrame = { + import org.apache.spark.sql.catalyst.expressions.Literal.{FalseLiteral, TrueLiteral} + + val isDeleteWithDuplicateMatches = multipleMatchDeleteOnlyOvercount.nonEmpty + var sourceDF = this.sourceDF + .withColumn(SOURCE_ROW_PRESENT_COL, new Column(incrSourceRowCountExpr)) + var targetDF = getTouchedTargetDF(touchedFiles) + .filter(METADATA_ROW_DEL_COL) + .drop(METADATA_ROW_DEL_COL) + if (isDeleteWithDuplicateMatches) { + targetDF = targetDF.withColumn( + GpuMergeIntoCommand.TARGET_ROW_ID_COL, monotonically_increasing_id()) + if (context.cmd.notMatchedClauses.nonEmpty) { + sourceDF = sourceDF.withColumn( + GpuMergeIntoCommand.SOURCE_ROW_ID_COL, monotonically_increasing_id()) + } + } + + val joinType = if (hasNoInserts && + context.spark.conf.get(DeltaSQLConf.MERGE_MATCHED_ONLY_ENABLED)) { + "inner" + } else { + "leftOuter" + } + val joinedPlan = sourceDF.join( + targetDF, new Column(context.cmd.condition), joinType).queryExecution.analyzed + + def resolveOnJoinedPlan(exprs: Seq[Expression]): Seq[Expression] = { + tryResolveReferencesForExpressions(context.spark, exprs, joinedPlan) + } + + val incrUpdatedCount = context.cmd.makeMetricUpdateUDF( + "numTargetRowsUpdated", deterministic = true) + val incrUpdatedMatchedCount = context.cmd.makeMetricUpdateUDF( + "numTargetRowsMatchedUpdated", deterministic = true) + val incrUpdatedNotMatchedBySourceCount = context.cmd.makeMetricUpdateUDF( + "numTargetRowsNotMatchedBySourceUpdated", deterministic = true) + val incrInsertedCount = context.cmd.makeMetricUpdateUDF( + "numTargetRowsInserted", deterministic = true) + val incrDeletedCount = context.cmd.makeMetricUpdateUDF( + "numTargetRowsDeleted", deterministic = true) + val incrDeletedMatchedCount = context.cmd.makeMetricUpdateUDF( + "numTargetRowsMatchedDeleted", deterministic = true) + val incrDeletedNotMatchedBySourceCount = context.cmd.makeMetricUpdateUDF( + "numTargetRowsNotMatchedBySourceDeleted", deterministic = true) + + var cdfTargetOutputCols: Seq[Expression] = targetOutputCols + var outputRowSchema = context.deltaTxn.metadata.schema + if (isDeleteWithDuplicateMatches) { + cdfTargetOutputCols = cdfTargetOutputCols :+ + UnresolvedAttribute(GpuMergeIntoCommand.TARGET_ROW_ID_COL) + outputRowSchema = outputRowSchema.add(GpuMergeIntoCommand.TARGET_ROW_ID_COL, LongType) + if (context.cmd.notMatchedClauses.nonEmpty) { + cdfTargetOutputCols = cdfTargetOutputCols :+ + Alias(Literal(null, LongType), GpuMergeIntoCommand.SOURCE_ROW_ID_COL)() + outputRowSchema = outputRowSchema.add(GpuMergeIntoCommand.SOURCE_ROW_ID_COL, LongType) + } + } + outputRowSchema = outputRowSchema + .add(ROW_DROPPED_COL, BooleanType) + .add(INCR_ROW_COUNT_COL, BooleanType) + .add(CDC_TYPE_COLUMN_NAME, StringType) + + def updateOutput( + actions: Seq[DeltaMergeAction], + incrMetricExpr: Expression): Seq[Seq[Expression]] = { + val mainDataOutput = actions.map(_.expr) :+ FalseLiteral :+ incrMetricExpr :+ + CDC_TYPE_NOT_CDC_LITERAL + val preImageOutput = cdfTargetOutputCols :+ FalseLiteral :+ TrueLiteral :+ + Literal(CDC_TYPE_UPDATE_PREIMAGE) + val postImageOutput = mainDataOutput.dropRight(2) :+ TrueLiteral :+ + Literal(CDC_TYPE_UPDATE_POSTIMAGE) + Seq(mainDataOutput, preImageOutput, postImageOutput).map(resolveOnJoinedPlan) + } + + def deleteOutput(incrMetricExpr: Expression): Seq[Seq[Expression]] = { + val mainDataOutput = cdfTargetOutputCols :+ TrueLiteral :+ incrMetricExpr :+ + CDC_TYPE_NOT_CDC_LITERAL + val deleteCdfOutput = cdfTargetOutputCols :+ FalseLiteral :+ TrueLiteral :+ + Literal(CDC_TYPE_DELETE) + Seq(mainDataOutput, deleteCdfOutput).map(resolveOnJoinedPlan) + } + + def insertOutput( + actions: Seq[DeltaMergeAction], + incrMetricExpr: Expression): Seq[Seq[Expression]] = { + val insertExprs = actions.map(_.expr) + val outputExprs = if (isDeleteWithDuplicateMatches) { + insertExprs :+ + Alias(Literal(null, LongType), GpuMergeIntoCommand.TARGET_ROW_ID_COL)() :+ + UnresolvedAttribute(GpuMergeIntoCommand.SOURCE_ROW_ID_COL) + } else { + insertExprs + } + val mainDataOutput = resolveOnJoinedPlan( + outputExprs :+ FalseLiteral :+ incrMetricExpr :+ CDC_TYPE_NOT_CDC_LITERAL) + val insertCdfOutput = mainDataOutput.dropRight(2) :+ TrueLiteral :+ + Literal(CDC_TYPE_INSERT) + Seq(mainDataOutput, insertCdfOutput) + } + + def clauseOutput(clause: DeltaMergeIntoClause): Seq[Seq[Expression]] = clause match { + case u: DeltaMergeIntoMatchedUpdateClause => + updateOutput(u.resolvedActions, And(incrUpdatedCount, incrUpdatedMatchedCount)) + case _: DeltaMergeIntoMatchedDeleteClause => + deleteOutput(And(incrDeletedCount, incrDeletedMatchedCount)) + case i: DeltaMergeIntoNotMatchedInsertClause => + insertOutput(i.resolvedActions, incrInsertedCount) + case u: DeltaMergeIntoNotMatchedBySourceUpdateClause => + updateOutput(u.resolvedActions, + And(incrUpdatedCount, incrUpdatedNotMatchedBySourceCount)) + case _: DeltaMergeIntoNotMatchedBySourceDeleteClause => + deleteOutput(And(incrDeletedCount, incrDeletedNotMatchedBySourceCount)) + } + + def clauseCondition(clause: DeltaMergeIntoClause): Expression = { + resolveOnJoinedPlan(Seq(clause.condition.getOrElse(TrueLiteral))).head + } + + val targetRowHasNoMatch = resolveOnJoinedPlan( + Seq(col(SOURCE_ROW_PRESENT_COL).isNull.expr)).head + val sourceRowHasNoMatch = resolveOnJoinedPlan( + Seq(col(TARGET_ROW_PRESENT_COL).isNull.expr)).head + val matchedConditions = context.cmd.matchedClauses.map(clauseCondition) + val matchedOutputs = context.cmd.matchedClauses.map(clauseOutput) + val notMatchedConditions = context.cmd.notMatchedClauses.map(clauseCondition) + val notMatchedOutputs = context.cmd.notMatchedClauses.map(clauseOutput) + val notMatchedBySourceConditions = + context.cmd.notMatchedBySourceClauses.map(clauseCondition) + val notMatchedBySourceOutputs = context.cmd.notMatchedBySourceClauses.map(clauseOutput) + val noopCopyOutput = resolveOnJoinedPlan( + cdfTargetOutputCols :+ FalseLiteral :+ TrueLiteral :+ CDC_TYPE_NOT_CDC_LITERAL) + val deleteRowOutput = resolveOnJoinedPlan( + cdfTargetOutputCols :+ TrueLiteral :+ TrueLiteral :+ CDC_TYPE_NOT_CDC_LITERAL) + + var outputDF = addMergeJoinProcessor( + joinedPlan, + outputRowSchema, + targetRowHasNoMatch, + sourceRowHasNoMatch, + matchedConditions, + matchedOutputs, + notMatchedConditions, + notMatchedOutputs, + notMatchedBySourceConditions, + notMatchedBySourceOutputs, + noopCopyOutput, + deleteRowOutput) + + if (isDeleteWithDuplicateMatches) { + val columnsToDedupeBy = if (context.cmd.notMatchedClauses.nonEmpty) { + Seq(GpuMergeIntoCommand.TARGET_ROW_ID_COL, + GpuMergeIntoCommand.SOURCE_ROW_ID_COL, CDC_TYPE_COLUMN_NAME) + } else { + Seq(GpuMergeIntoCommand.TARGET_ROW_ID_COL) + } + outputDF = outputDF.dropDuplicates(columnsToDedupeBy) + .drop(GpuMergeIntoCommand.TARGET_ROW_ID_COL, GpuMergeIntoCommand.SOURCE_ROW_ID_COL) + } + repartitionIfNeeded(outputDF.drop(ROW_DROPPED_COL, INCR_ROW_COUNT_COL)) + } + /** * Generate a plan by calculating modified rows. It's computed by joining source and target * tables, where target table has been filtered by (`__metadata_file_name`, @@ -909,6 +1166,10 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend * 4. Target rows which are deleted */ private def getModifiedDF(touchedFiles: Map[String, (Roaring64Bitmap, AddFile)]): DataFrame = { + if (DeltaConfigs.CHANGE_DATA_FEED.fromMetaData(context.deltaTxn.metadata)) { + return getModifiedDFWithCdf(touchedFiles) + } + val sourceDF = this.sourceDF .withColumn(SOURCE_ROW_PRESENT_COL, new Column(incrSourceRowCountExpr)) @@ -1022,9 +1283,14 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend } private def getUnmodifiedDF(touchedFiles: Map[String, (Roaring64Bitmap, AddFile)]): DataFrame = { - getTouchedTargetDF(touchedFiles) + val unmodifiedDF = getTouchedTargetDF(touchedFiles) .filter(!col(METADATA_ROW_DEL_COL)) .drop(TARGET_ROW_PRESENT_COL, METADATA_ROW_DEL_COL) + if (DeltaConfigs.CHANGE_DATA_FEED.fromMetaData(context.deltaTxn.metadata)) { + unmodifiedDF.withColumn(CDC_TYPE_COLUMN_NAME, new Column(CDC_TYPE_NOT_CDC_LITERAL)) + } else { + unmodifiedDF + } } } @@ -1081,4 +1347,4 @@ object MergeExecutor { if (distinctValues.size == 1 && distinctValues.head.isEmpty) 0 else distinctValues.size (bytes, numDistinctValues) } -} \ No newline at end of file +} diff --git a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala index 4f304221111..fbd1edf5094 100644 --- a/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala +++ b/delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuLowShuffleMergeCommand.scala @@ -33,15 +33,18 @@ import com.databricks.sql.transaction.tahoe.actions.{AddCDCFile, AddFile, DeletionVectorDescriptor, FileAction} import com.databricks.sql.transaction.tahoe.commands.{DeltaCommand, DMLWithDeletionVectorsHelper} +import com.databricks.sql.transaction.tahoe.commands.cdc.CDCReader._ import com.databricks.sql.transaction.tahoe.commands.merge.MergeIntoMaterializeSource import com.databricks.sql.transaction.tahoe.deletionvectors.{RoaringBitmapArray, RoaringBitmapArrayFormat} import com.databricks.sql.transaction.tahoe.files.{TahoeBatchFileIndex, TahoeFileIndex} import com.databricks.sql.transaction.tahoe.rapids.MergeExecutor.{ totalBytesAndDistinctPartitionValues, + CDC_TYPE_NOT_CDC_LITERAL, FILE_PATH_COL, INCR_METRICS_COL, INCR_METRICS_FIELD, + INCR_ROW_COUNT_COL, ROW_DROPPED_COL, ROW_DROPPED_FIELD, SOURCE_ROW_PRESENT_COL, @@ -51,7 +54,7 @@ import com.databricks.sql.transaction.tahoe.rapids.MergeExecutor.{ import com.databricks.sql.transaction.tahoe.schema.ImplicitMetadataOperation import com.databricks.sql.transaction.tahoe.sources.DeltaSQLConf import com.databricks.sql.transaction.tahoe.util.{AnalysisHelper, DeltaFileOperations} -import com.nvidia.spark.rapids.{GpuOverrides, RapidsConf, SparkPlanMeta} +import com.nvidia.spark.rapids.{BaseExprMeta, GpuOverrides, RapidsConf, SparkPlanMeta} import com.nvidia.spark.rapids.RapidsConf.DELTA_LOW_SHUFFLE_MERGE_DEL_VECTOR_BROADCAST_THRESHOLD import com.nvidia.spark.rapids.delta._ import com.nvidia.spark.rapids.delta.GpuDeltaParquetFileFormatUtils.METADATA_ROW_IDX_COL @@ -64,14 +67,16 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute import org.apache.spark.sql.catalyst.catalog.CatalogTable +import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, RowEncoder} import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeReference, - CaseWhen, Expression, IsNull, Literal, NamedExpression, PredicateHelper} + CaseWhen, EqualNullSafe, Expression, If, IsNull, Literal, NamedExpression, Not, PredicateHelper} import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.plans.logical.{DeltaMergeAction, DeltaMergeIntoClause, DeltaMergeIntoMatchedClause, DeltaMergeIntoMatchedDeleteClause, DeltaMergeIntoMatchedUpdateClause, DeltaMergeIntoNotMatchedBySourceClause, DeltaMergeIntoNotMatchedBySourceDeleteClause, DeltaMergeIntoNotMatchedBySourceUpdateClause, DeltaMergeIntoNotMatchedClause, DeltaMergeIntoNotMatchedInsertClause, LogicalPlan, Project} +import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} import org.apache.spark.sql.execution.command.LeafRunnableCommand @@ -684,14 +689,8 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend * the temporary deletion-vector scan used to retain unmodified rows. * 2. The temporary deletion vectors introduce extra overhead, so it may be better to fall back * when the changeset is too large. - * 3. Low shuffle merge does not generate change data feed rows. */ def shouldFallback(): Boolean = { - if (DeltaConfigs.CHANGE_DATA_FEED.fromMetaData(context.deltaTxn.metadata)) { - logWarning("Change data feed is enabled, falling back to traditional merge.") - return true - } - // Trying to detect if we can execute finding touched files. val touchFilePlanOverrideSucceed = verifyGpuPlan(planForFindingTouchedFiles()) { planMeta => def check(meta: SparkPlanMeta[SparkPlan]): Boolean = { @@ -943,6 +942,307 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend .withColumn(METADATA_ROW_IDX_COL, col("_metadata.row_index")) } + private def uniqueColumnName(base: String, existing: Seq[String]): String = { + val resolver = context.cmd.conf.resolver + Iterator.from(0) + .map(i => if (i == 0) base else s"$base$i") + .find(candidate => !existing.exists(name => resolver(name, candidate))) + .get + } + + private def addMergeJoinProcessor( + joinedPlan: LogicalPlan, + outputRowSchema: StructType, + targetRowHasNoMatch: Expression, + sourceRowHasNoMatch: Expression, + matchedConditions: Seq[Expression], + matchedOutputs: Seq[Seq[Seq[Expression]]], + notMatchedConditions: Seq[Expression], + notMatchedOutputs: Seq[Seq[Seq[Expression]]], + notMatchedBySourceConditions: Seq[Expression], + notMatchedBySourceOutputs: Seq[Seq[Seq[Expression]]], + noopCopyOutput: Seq[Expression], + deleteRowOutput: Seq[Expression], + rowDroppedColumnIndex: Int): Dataset[Row] = { + def wrap(e: Expression): BaseExprMeta[Expression] = { + GpuOverrides.wrapExpr(e, context.rapidsConf, None) + } + + val targetRowHasNoMatchMeta = wrap(targetRowHasNoMatch) + val sourceRowHasNoMatchMeta = wrap(sourceRowHasNoMatch) + val matchedConditionsMetas = matchedConditions.map(wrap) + val matchedOutputsMetas = matchedOutputs.map(_.map(_.map(wrap))) + val notMatchedConditionsMetas = notMatchedConditions.map(wrap) + val notMatchedOutputsMetas = notMatchedOutputs.map(_.map(_.map(wrap))) + val notMatchedBySourceConditionsMetas = notMatchedBySourceConditions.map(wrap) + val notMatchedBySourceOutputsMetas = notMatchedBySourceOutputs.map(_.map(_.map(wrap))) + val noopCopyOutputMetas = noopCopyOutput.map(wrap) + val deleteRowOutputMetas = deleteRowOutput.map(wrap) + val allMetas = Seq(targetRowHasNoMatchMeta, sourceRowHasNoMatchMeta) ++ + matchedConditionsMetas ++ matchedOutputsMetas.flatten.flatten ++ + notMatchedConditionsMetas ++ notMatchedOutputsMetas.flatten.flatten ++ + notMatchedBySourceConditionsMetas ++ notMatchedBySourceOutputsMetas.flatten.flatten ++ + noopCopyOutputMetas ++ deleteRowOutputMetas + allMetas.foreach(_.tagForGpu()) + val canReplace = allMetas.forall(_.canExprTreeBeReplaced) && + context.rapidsConf.isOperatorEnabled( + "spark.rapids.sql.exec.RapidsProcessDeltaMergeJoinExec", false, false) + if (context.rapidsConf.shouldExplainAll || (context.rapidsConf.shouldExplain && !canReplace)) { + val exprExplains = allMetas.map(_.explain(context.rapidsConf.shouldExplainAll)) + val execWorkInfo = if (canReplace) { + "will run on GPU" + } else { + "cannot run on GPU because not all merge processing expressions can be replaced" + } + logWarning(s" $execWorkInfo:\n" + + s" ${exprExplains.mkString(" ")}") + } + + if (canReplace) { + val processedJoinPlan = RapidsProcessDeltaMergeJoin( + joinedPlan, + toAttributes(outputRowSchema), + targetRowHasNoMatch = targetRowHasNoMatch, + sourceRowHasNoMatch = sourceRowHasNoMatch, + matchedConditions = matchedConditions, + matchedOutputs = matchedOutputs, + notMatchedConditions = notMatchedConditions, + notMatchedOutputs = notMatchedOutputs, + notMatchedBySourceConditions = notMatchedBySourceConditions, + notMatchedBySourceOutputs = notMatchedBySourceOutputs, + noopCopyOutput = noopCopyOutput, + deleteRowOutput = deleteRowOutput, + rowDroppedColumnIndex = Some(rowDroppedColumnIndex)) + Dataset.ofRows(context.spark, processedJoinPlan) + } else { + val joinedRowEncoder = ExpressionEncoder(RowEncoder.encoderFor(joinedPlan.schema)) + val outputRowEncoder = ExpressionEncoder(RowEncoder.encoderFor(outputRowSchema)) + .resolveAndBind() + val processor = new GpuMergeIntoCommand.JoinedRowProcessor( + targetRowHasNoMatch = targetRowHasNoMatch, + sourceRowHasNoMatch = sourceRowHasNoMatch, + matchedConditions = matchedConditions, + matchedOutputs = matchedOutputs, + notMatchedConditions = notMatchedConditions, + notMatchedOutputs = notMatchedOutputs, + notMatchedBySourceConditions = notMatchedBySourceConditions, + notMatchedBySourceOutputs = notMatchedBySourceOutputs, + noopCopyOutput = noopCopyOutput, + deleteRowOutput = deleteRowOutput, + joinedAttributes = joinedPlan.output, + joinedRowEncoder = joinedRowEncoder, + outputRowEncoder = outputRowEncoder, + rowDroppedColumnIndex = rowDroppedColumnIndex) + Dataset.ofRows(context.spark, joinedPlan) + .mapPartitions(processor.processPartition)(outputRowEncoder) + } + } + + /** Generate both rewritten table rows and explicit change-data-feed rows. */ + private def getModifiedDFWithCdf( + touchedFiles: Map[String, (Roaring64Bitmap, AddFile)]): DataFrame = { + import org.apache.spark.sql.catalyst.expressions.Literal.{FalseLiteral, TrueLiteral} + + val isDeleteWithDuplicateMatches = multipleMatchDeleteOnlyOvercount.nonEmpty + val sourcePlanDF = this.sourceDF + val targetPlanDF = buildTargetDFWithFiles(touchedFiles.values.map(_._2).toSeq) + val userColumns = sourcePlanDF.columns.toSeq ++ targetPlanDF.columns.toSeq + val sourceRowPresentCol = uniqueColumnName(SOURCE_ROW_PRESENT_COL, userColumns) + val targetRowPresentCol = uniqueColumnName( + TARGET_ROW_PRESENT_COL, userColumns :+ sourceRowPresentCol) + val taken = userColumns ++ Seq(sourceRowPresentCol, targetRowPresentCol) + val targetRowIdCol = uniqueColumnName(GpuMergeIntoCommand.TARGET_ROW_ID_COL, taken) + val sourceRowIdCol = uniqueColumnName( + GpuMergeIntoCommand.SOURCE_ROW_ID_COL, taken :+ targetRowIdCol) + + var sourceDF = sourcePlanDF.withColumn( + sourceRowPresentCol, DFUDFShims.exprToColumn(incrSourceRowCountExpr)) + var targetDF = targetPlanDF.withColumn(targetRowPresentCol, lit(true)) + if (isDeleteWithDuplicateMatches) { + targetDF = targetDF.withColumn(targetRowIdCol, monotonically_increasing_id()) + if (context.cmd.notMatchedClauses.nonEmpty) { + sourceDF = sourceDF.withColumn(sourceRowIdCol, monotonically_increasing_id()) + } + } + + val joinType = if (hasNoInserts && + context.spark.conf.get(DeltaSQLConf.MERGE_MATCHED_ONLY_ENABLED)) { + "inner" + } else { + "leftOuter" + } + val joinedPlan = sourceDF + .join(targetDF, DFUDFShims.exprToColumn(context.cmd.condition), joinType) + .queryExecution.analyzed + + def resolveOnJoinedPlan(exprs: Seq[Expression]): Seq[Expression] = { + tryResolveReferencesForExpressions(context.spark, exprs, joinedPlan) + } + + val incrUpdatedCount = context.cmd.metricUpdateExpr( + "numTargetRowsUpdated", deterministic = true) + val incrUpdatedMatchedCount = context.cmd.metricUpdateExpr( + "numTargetRowsMatchedUpdated", deterministic = true) + val incrInsertedCount = context.cmd.metricUpdateExpr( + "numTargetRowsInserted", deterministic = true) + val incrDeletedCount = context.cmd.metricUpdateExpr( + "numTargetRowsDeleted", deterministic = true) + val incrDeletedMatchedCount = context.cmd.metricUpdateExpr( + "numTargetRowsMatchedDeleted", deterministic = true) + + var cdfTargetOutputCols: Seq[Expression] = targetOutputCols + var outputRowSchema = context.deltaTxn.metadata.schema + if (isDeleteWithDuplicateMatches) { + cdfTargetOutputCols = cdfTargetOutputCols :+ UnresolvedAttribute(targetRowIdCol) + outputRowSchema = outputRowSchema.add(targetRowIdCol, LongType) + if (context.cmd.notMatchedClauses.nonEmpty) { + cdfTargetOutputCols = cdfTargetOutputCols :+ + Alias(Literal(null, LongType), sourceRowIdCol)() + outputRowSchema = outputRowSchema.add(sourceRowIdCol, LongType) + } + } + val rowDroppedColumnIndex = cdfTargetOutputCols.size + outputRowSchema = outputRowSchema + .add(ROW_DROPPED_COL, BooleanType) + .add(INCR_ROW_COUNT_COL, BooleanType) + .add(CDC_TYPE_COLUMN_NAME, StringType) + + val materializedValues = mutable.ArrayBuffer[NamedExpression]() + def materializeNonDeterministic( + exprs: Seq[Expression], + takesClause: Expression): Seq[Expression] = exprs.map { + case e if !e.deterministic => + val resolved = resolveOnJoinedPlan(Seq(e)).head + val existing = joinedPlan.output.map(_.name) ++ materializedValues.map(_.name) + val alias = Alias(If(takesClause, resolved, Literal(null, resolved.dataType)), + uniqueColumnName(GpuMergeIntoCommand.NON_DETERMINISTIC_VALUE_COL, existing))() + materializedValues += alias + alias.toAttribute + case e => e + } + + def clauseRouting( + rowKind: Expression, + conditions: Seq[Expression], + index: Int): Expression = { + val earlierNotTaken = conditions.take(index) + .map(condition => Not(EqualNullSafe(condition, TrueLiteral))) + (rowKind +: earlierNotTaken :+ EqualNullSafe(conditions(index), TrueLiteral)).reduce(And) + } + + def updateOutput( + updateExprs: Seq[Expression], + incrMetricExpr: Expression): Seq[Seq[Expression]] = { + val mainDataOutput = updateExprs :+ FalseLiteral :+ incrMetricExpr :+ + CDC_TYPE_NOT_CDC_LITERAL + val preImageOutput = cdfTargetOutputCols :+ FalseLiteral :+ TrueLiteral :+ + Literal(CDC_TYPE_UPDATE_PREIMAGE) + val postImageOutput = mainDataOutput.dropRight(2) :+ TrueLiteral :+ + Literal(CDC_TYPE_UPDATE_POSTIMAGE) + Seq(mainDataOutput, preImageOutput, postImageOutput).map(resolveOnJoinedPlan) + } + + def deleteOutput(incrMetricExpr: Expression): Seq[Seq[Expression]] = { + val mainDataOutput = cdfTargetOutputCols :+ TrueLiteral :+ incrMetricExpr :+ + CDC_TYPE_NOT_CDC_LITERAL + val deleteCdfOutput = cdfTargetOutputCols :+ FalseLiteral :+ TrueLiteral :+ + Literal(CDC_TYPE_DELETE) + Seq(mainDataOutput, deleteCdfOutput).map(resolveOnJoinedPlan) + } + + def insertOutput( + insertExprs: Seq[Expression], + incrMetricExpr: Expression): Seq[Seq[Expression]] = { + val outputExprs = if (isDeleteWithDuplicateMatches) { + insertExprs :+ Alias(Literal(null, LongType), targetRowIdCol)() :+ + UnresolvedAttribute(sourceRowIdCol) + } else { + insertExprs + } + val mainDataOutput = resolveOnJoinedPlan( + outputExprs :+ FalseLiteral :+ incrMetricExpr :+ CDC_TYPE_NOT_CDC_LITERAL) + val insertCdfOutput = mainDataOutput.dropRight(2) :+ TrueLiteral :+ + Literal(CDC_TYPE_INSERT) + Seq(mainDataOutput, insertCdfOutput) + } + + def clauseOutput(clause: DeltaMergeIntoClause, routing: Expression) + : Seq[Seq[Expression]] = clause match { + case u: DeltaMergeIntoMatchedUpdateClause => + updateOutput(materializeNonDeterministic(u.resolvedActions.map(_.expr), routing), + And(incrUpdatedCount, incrUpdatedMatchedCount)) + case _: DeltaMergeIntoMatchedDeleteClause => + deleteOutput(And(incrDeletedCount, incrDeletedMatchedCount)) + case i: DeltaMergeIntoNotMatchedInsertClause => + insertOutput(materializeNonDeterministic(i.resolvedActions.map(_.expr), routing), + incrInsertedCount) + case other => + throw new IllegalArgumentException(s"Unsupported low-shuffle merge clause: " + + other.getClass.getName) + } + + def clauseCondition(clause: DeltaMergeIntoClause): Expression = { + resolveOnJoinedPlan(Seq(clause.condition.getOrElse(TrueLiteral))).head + } + + val targetRowHasNoMatch = resolveOnJoinedPlan( + Seq(IsNull(UnresolvedAttribute(sourceRowPresentCol)))).head + val sourceRowHasNoMatch = resolveOnJoinedPlan( + Seq(IsNull(UnresolvedAttribute(targetRowPresentCol)))).head + val matchedRow = And(Not(targetRowHasNoMatch), Not(sourceRowHasNoMatch)) + val matchedConditions = context.cmd.matchedClauses.map(clauseCondition) + val matchedOutputs = context.cmd.matchedClauses.zipWithIndex.map { case (clause, index) => + clauseOutput(clause, clauseRouting(matchedRow, matchedConditions, index)) + } + val notMatchedConditions = context.cmd.notMatchedClauses.map(clauseCondition) + val notMatchedOutputs = context.cmd.notMatchedClauses.zipWithIndex.map { + case (clause, index) => + clauseOutput(clause, clauseRouting(sourceRowHasNoMatch, notMatchedConditions, index)) + } + val noopCopyOutput = resolveOnJoinedPlan( + cdfTargetOutputCols :+ FalseLiteral :+ TrueLiteral :+ CDC_TYPE_NOT_CDC_LITERAL) + val deleteRowOutput = resolveOnJoinedPlan( + cdfTargetOutputCols :+ TrueLiteral :+ TrueLiteral :+ CDC_TYPE_NOT_CDC_LITERAL) + val processorInputPlan = if (materializedValues.isEmpty) { + joinedPlan + } else { + Project(joinedPlan.output ++ materializedValues, joinedPlan) + } + + var outputDF = addMergeJoinProcessor( + processorInputPlan, + outputRowSchema, + targetRowHasNoMatch, + sourceRowHasNoMatch, + matchedConditions, + matchedOutputs, + notMatchedConditions, + notMatchedOutputs, + Seq.empty, + Seq.empty, + noopCopyOutput, + deleteRowOutput, + rowDroppedColumnIndex) + + if (isDeleteWithDuplicateMatches) { + val columnsToDedupeBy = if (context.cmd.notMatchedClauses.nonEmpty) { + Seq(targetRowIdCol, sourceRowIdCol, CDC_TYPE_COLUMN_NAME) + } else { + Seq(targetRowIdCol) + } + outputDF = outputDF.dropDuplicates(columnsToDedupeBy) + } + + val outputAttributes = outputDF.queryExecution.analyzed.output + outputDF = Seq(ROW_DROPPED_COL, INCR_ROW_COUNT_COL) + .flatMap(name => outputAttributes.reverse.find(_.name == name)) + .foldLeft(outputDF)((df, attr) => df.drop(DFUDFShims.exprToColumn(attr))) + if (isDeleteWithDuplicateMatches) { + outputDF = outputDF.drop(targetRowIdCol, sourceRowIdCol) + } + repartitionIfNeeded(outputDF) + } + /** * Generate a plan by calculating modified rows. It's computed by joining source and target * tables, where target table has been filtered by (`__metadata_file_name`, @@ -962,6 +1262,10 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend * 4. Target rows which are deleted */ private def getModifiedDF(touchedFiles: Map[String, (Roaring64Bitmap, AddFile)]): DataFrame = { + if (DeltaConfigs.CHANGE_DATA_FEED.fromMetaData(context.deltaTxn.metadata)) { + return getModifiedDFWithCdf(touchedFiles) + } + val sourceDF = this.sourceDF .withColumn(SOURCE_ROW_PRESENT_COL, DFUDFShims.exprToColumn(incrSourceRowCountExpr)) @@ -1084,7 +1388,13 @@ class LowShuffleMergeExecutor(override val context: MergeExecutorContext) extend hadoopConf, tablePath)) }.toSeq - buildTargetDFWithFiles(filesWithTemporaryDVs) + val unmodifiedDF = buildTargetDFWithFiles(filesWithTemporaryDVs) + if (DeltaConfigs.CHANGE_DATA_FEED.fromMetaData(context.deltaTxn.metadata)) { + unmodifiedDF.withColumn( + CDC_TYPE_COLUMN_NAME, DFUDFShims.exprToColumn(CDC_TYPE_NOT_CDC_LITERAL)) + } else { + unmodifiedDF + } } } diff --git a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py index 8e3b954873c..ab578424874 100644 --- a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py +++ b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py @@ -32,13 +32,6 @@ def supports_delta_low_shuffle_merge(): (not is_databricks_runtime() and spark_version().startswith("3.4")) -low_shuffle_cdf_param = pytest.param( - True, - marks=pytest.mark.xfail( - condition=not is_databricks_version(17, 3), - reason="https://github.com/NVIDIA/spark-rapids/issues/13552")) - - @allow_non_gpu("ColumnarToRowExec", *delta_meta_allow) @delta_lake @ignore_order @@ -94,14 +87,11 @@ def test_delta_merge_not_match_insert_only(spark_tmp_path, spark_tmp_table_facto (range(5), range(5)), # full delete of target (range(10), range(20, 30)) # no-op delete ], ids=idfn) -@pytest.mark.parametrize("use_cdf", [low_shuffle_cdf_param, False], ids=idfn) +@pytest.mark.parametrize("use_cdf", [True, False], ids=idfn) @pytest.mark.parametrize("partition_columns", [None, ["a"], ["b"], ["a", "b"]], ids=idfn) @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn) def test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, table_ranges, use_cdf, partition_columns, num_slices): - if (use_cdf and is_databricks_version(17, 3) - and table_ranges == (range(10), range(20, 30))): - pytest.xfail(reason="https://github.com/NVIDIA/spark-rapids/issues/13552") do_test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, table_ranges, use_cdf, False, partition_columns, num_slices, False, delta_merge_enabled_conf) @@ -111,7 +101,7 @@ def test_delta_merge_match_delete_only(spark_tmp_path, spark_tmp_table_factory, @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") -@pytest.mark.parametrize("use_cdf", [low_shuffle_cdf_param, False], ids=idfn) +@pytest.mark.parametrize("use_cdf", [True, False], ids=idfn) @pytest.mark.parametrize("num_slices", num_slices_to_test, ids=idfn) def test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, use_cdf, num_slices): do_test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, use_cdf, False, @@ -122,7 +112,7 @@ def test_delta_merge_standard_upsert(spark_tmp_path, spark_tmp_table_factory, us @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") -@pytest.mark.parametrize("use_cdf", [low_shuffle_cdf_param, False], ids=idfn) +@pytest.mark.parametrize("use_cdf", [True, False], ids=idfn) @pytest.mark.parametrize("merge_sql", [ "MERGE INTO {dest_table} d USING {src_table} s ON d.a == s.a" \ " WHEN MATCHED AND s.b > 'q' THEN UPDATE SET d.a = s.a / 2, d.b = s.b" \ @@ -161,7 +151,7 @@ def test_delta_merge_upsert_with_unmatchable_match_condition(spark_tmp_path, spa @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), reason="Low Shuffle Merge requires Delta Lake 2.4 or DBR 17.3") -@pytest.mark.parametrize("use_cdf", [low_shuffle_cdf_param, False], ids=idfn) +@pytest.mark.parametrize("use_cdf", [True, False], ids=idfn) def test_delta_merge_update_with_aggregation(spark_tmp_path, spark_tmp_table_factory, use_cdf): do_test_delta_merge_update_with_aggregation(spark_tmp_path, spark_tmp_table_factory, use_cdf, False, delta_merge_enabled_conf) From a18416c1394ccc88c8c5fd4dbbf4b64034afc4d3 Mon Sep 17 00:00:00 2001 From: Ray Liu Date: Thu, 10 Sep 2026 14:14:10 +0800 Subject: [PATCH 09/10] [TEST] Allow DBR AQE empty relation in low shuffle merge Signed-off-by: Ray Liu --- .../src/main/python/delta_lake_low_shuffle_merge_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py index ab578424874..9177cfcd3f1 100644 --- a/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py +++ b/integration_tests/src/main/python/delta_lake_low_shuffle_merge_test.py @@ -78,7 +78,8 @@ def test_delta_merge_not_match_insert_only(spark_tmp_path, spark_tmp_table_facto table_ranges, use_cdf, False, partition_columns, num_slices, False, delta_merge_enabled_conf) -@allow_non_gpu(*delta_meta_allow) +# DBR 17.3 AQE can replace a no-match join with its row-based EmptyRelationExec. +@allow_non_gpu("EmptyRelationExec", *delta_meta_allow) @delta_lake @ignore_order @pytest.mark.skipif(not supports_delta_low_shuffle_merge(), From 45a7bc35a1e40b0e96fcfd4d4649386c660ed178 Mon Sep 17 00:00:00 2001 From: Ray Liu Date: Thu, 10 Sep 2026 16:38:05 +0800 Subject: [PATCH 10/10] [DOC] Update low shuffle merge copyright Signed-off-by: Ray Liu --- .../sql/delta/rapids/delta24x/GpuLowShuffleMergeCommand.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/delta-lake/delta-24x/src/main/scala/org/apache/spark/sql/delta/rapids/delta24x/GpuLowShuffleMergeCommand.scala b/delta-lake/delta-24x/src/main/scala/org/apache/spark/sql/delta/rapids/delta24x/GpuLowShuffleMergeCommand.scala index da57c28bebb..612f69c2a01 100644 --- a/delta-lake/delta-24x/src/main/scala/org/apache/spark/sql/delta/rapids/delta24x/GpuLowShuffleMergeCommand.scala +++ b/delta-lake/delta-24x/src/main/scala/org/apache/spark/sql/delta/rapids/delta24x/GpuLowShuffleMergeCommand.scala @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-2026, NVIDIA CORPORATION. * * This file was derived from MergeIntoCommand.scala * in the Delta Lake project at https://github.com/delta-io/delta.