[FEA] Add GPU low shuffle merge for Databricks 17.3 [databricks]

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")
           }
         })
 }

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,

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,7 +30,8 @@ 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.addMetadataColumnToIterator
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 
@@ -41,6 +44,7 @@ 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.
@@ -62,7 +66,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 +118,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 +136,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 +169,7 @@ case class GpuDeltaParquetFileFormat(
       hadoopConf: Configuration,
       metrics: Map[String, GpuMetric])
   : PartitionedFile => Iterator[InternalRow] = {
-    super.buildReaderWithPartitionValuesAndMetrics(
+    val dataReader = super.buildReaderWithPartitionValuesAndMetrics(
       sparkSession,
       dataSchema,
       partitionSchema,
@@ -173,6 +178,27 @@ 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)
   }
 }
 
@@ -268,7 +294,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 = {

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)
+}

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 && conf.isParquetPerFileReadEnabled) {
+      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 && conf.isParquetPerFileReadEnabled) {
+      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)
+    }
   }
 }

docs/additional-functionality/advanced_configs.md

@@ -87,7 +87,7 @@ Name | Description | Default Value | Applicable at
 <a name="sql.csv.read.float.enabled"></a>spark.rapids.sql.csv.read.float.enabled|CSV reading is not 100% compatible when reading floats.|true|Runtime
 <a name="sql.decimalOverflowGuarantees"></a>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
 <a name="sql.delta.lowShuffleMerge.deletionVector.broadcast.threshold"></a>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
-<a name="sql.delta.lowShuffleMerge.enabled"></a>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
+<a name="sql.delta.lowShuffleMerge.enabled"></a>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
 <a name="sql.detectDeltaCheckpointQueries"></a>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
 <a name="sql.detectDeltaLogQueries"></a>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
 <a name="sql.exec.opTimeTrackingRDD.enabled"></a>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