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

delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/GpuDeltaParquetFileFormatBase2.scala

@@ -1028,7 +1028,9 @@ class GpuDeltaParquetFileFormatBase2(
                 bitmap, isRetention, info.rowGroupOffsets, info.rowGroupNumRows)
             }
             RmmRapidsRetryIterator.withRetryNoSplit {
-              DeletionVector.computeNumDeletedRows(hostDvInfos, maxReadBatchSizeRows)
+              hostDvInfos.map { info =>
+                DeletionVector.computeNumDeletedRows(info, maxReadBatchSizeRows)
+              }.sum
             }
           }
         }.sum
@@ -1229,7 +1231,9 @@ class GpuDeltaParquetFileFormatBase2(
                   bitmap, isRetention, entry.rowGroupOffsets, entry.rowGroupNumRows)
               }.toArray
               RmmRapidsRetryIterator.withRetryNoSplit {
-                DeletionVector.computeNumDeletedRows(dvInfos, maxReadBatchSizeRows)
+                dvInfos.map { info =>
+                  DeletionVector.computeNumDeletedRows(info, maxReadBatchSizeRows)
+                }.sum
               }
             }
           }.sum

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

@@ -125,7 +125,9 @@ object DeltaSpark400DB173Provider extends DatabricksDeltaProviderBase {
   override def getReadFileFormat(
       relation: HadoopFsRelation, rapidsConf: RapidsConf): FileFormat = {
     val fmt = relation.fileFormat.asInstanceOf[DeltaParquetFileFormat]
-    if (isPushDVPredicateDownEnabled(rapidsConf)) {
+    if (GpuDeltaParquetFileFormat.isLowShuffleMergeScan(relation.options)) {
+      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

@@ -28,7 +28,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 +42,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 +64,8 @@ case class GpuDeltaParquetFileFormat(
     nullableRowTrackingGeneratedFields: Boolean = false,
     optimizationsEnabled: Boolean = true,
     tablePath: Option[String] = None,
-    isCDCRead: Boolean = false
+    isCDCRead: Boolean = false,
+    lowShuffleMergeScan: Boolean = false
   ) extends GpuDeltaParquetFileFormatBase {
 
   override val columnMappingMode: DeltaColumnMappingMode = metadata.columnMappingMode
@@ -113,7 +116,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 || !effectiveOptimizationsEnabled) {
       Seq.empty
     } else if (columnMappingMode != NoMapping) {
       val physicalNameMap = DeltaColumnMapping.getLogicalNameToPhysicalNameMap(referenceSchema)
@@ -131,7 +134,7 @@ case class GpuDeltaParquetFileFormat(
   override def isSplitable(
       sparkSession: SparkSession,
       options: Map[String, String],
-      path: Path): Boolean = effectiveOptimizationsEnabled
+      path: Path): Boolean = !lowShuffleMergeScan && effectiveOptimizationsEnabled
 
   private def hasDeletionVectorRead: Boolean =
     GpuDeltaParquetFileFormat.isDeletionVectorRead(
@@ -164,7 +167,7 @@ case class GpuDeltaParquetFileFormat(
       hadoopConf: Configuration,
       metrics: Map[String, GpuMetric])
   : PartitionedFile => Iterator[InternalRow] = {
-    super.buildReaderWithPartitionValuesAndMetrics(
+    val dataReader = super.buildReaderWithPartitionValuesAndMetrics(
       sparkSession,
       dataSchema,
       partitionSchema,
@@ -173,13 +176,35 @@ case class GpuDeltaParquetFileFormat(
       options,
       hadoopConf,
       metrics)
+
+    if (lowShuffleMergeScan) {
+      val maxBatchSize = RapidsConf.DELTA_LOW_SHUFFLE_MERGE_SCATTER_DEL_VECTOR_BATCH_SIZE
+        .get(sparkSession.sessionState.conf)
+      val scatterTime = metrics(GpuMetric.DELETION_VECTOR_SCATTER_TIME)
+      (file: PartitionedFile) => {
+        addMetadataColumnToIterator(
+          prepareSchema(requiredSchema),
+          None,
+          dataReader(file).asInstanceOf[Iterator[ColumnarBatch]],
+          maxBatchSize,
+          scatterTime).asInstanceOf[Iterator[InternalRow]]
+      }
+    } else {
+      dataReader
+    }
   }
 }
 
 object GpuDeltaParquetFileFormat {
   private[delta] val EDGE_COMPUTED_COLUMN_SKIP_ROW =
     "_databricks_internal_edge_computed_column_skip_row"
 
+  val LOW_SHUFFLE_MERGE_SCAN_OPTION =
+    "spark.rapids.internal.delta.lowShuffleMerge.scan"
+
+  def isLowShuffleMergeScan(options: Map[String, String]): Boolean =
+    options.get(LOW_SHUFFLE_MERGE_SCAN_OPTION).contains("true")
+
   def isDeletionVectorRead(format: DeltaParquetFileFormat): Boolean =
     isDeletionVectorRead(
       format.generateRowIndexFilterId,
@@ -268,7 +293,8 @@ object GpuDeltaParquetFileFormat {
       nullableRowTrackingGeneratedFields = fmt.nullableRowTrackingGeneratedFields,
       optimizationsEnabled = fmt.optimizationsEnabled,
       tablePath = fmt.tablePath,
-      isCDCRead = fmt.isCDCRead)
+      isCDCRead = fmt.isCDCRead,
+      lowShuffleMergeScan = isLowShuffleMergeScan(relation.options))
   }
 
   private def hasRowIndexFiltersInTahoeFileIndex(relation: HadoopFsRelation): Boolean = {

delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatNativeDV.scala

@@ -31,7 +31,7 @@ import com.databricks.sql.transaction.tahoe.{
   NameMapping,
   NoMapping
 }
-import com.databricks.sql.transaction.tahoe.actions.{Metadata, Protocol}
+import com.databricks.sql.transaction.tahoe.actions.{DeletionVectorDescriptor, Metadata, Protocol}
 import com.databricks.sql.transaction.tahoe.schema.SchemaMergingUtils
 import com.databricks.sql.transaction.tahoe.sources.DeltaSQLConf
 import com.nvidia.spark.rapids._
@@ -477,9 +477,9 @@ case class GpuDeltaParquetFileFormatNativeDV(
       dateRebaseMode: DateTimeRebaseMode,
       timestampRebaseMode: DateTimeRebaseMode,
       hasInt96Timestamps: Boolean,
-      // Base64-encoded DV descriptor string for this block's source file. None if no DV.
+      // DV descriptor for this block's source file. None if no DV.
       // The filter type is always RowIndexFilterType.IF_CONTAINED.
-      val dvDescriptor: Option[String],
+      val dvDescriptor: Option[DeletionVectorDescriptor],
       val rowIndexFilterProvider: Option[RowIndexFilterProvider],
       // Within-file row-index ordinal of this row group's first row.
       // Captured from BlockMetaData before any merging; invariant to computeBlockMetaData().
@@ -490,14 +490,14 @@ case class GpuDeltaParquetFileFormatNativeDV(
   /**
    * Per-file DV entry assembled during [[augmentChunkMeta]].
    *
-   * @param dvDescriptor base64-encoded DV descriptor for this file; None if no DV
+   * @param dvDescriptor DV descriptor for this file; None if no DV
    * @param rowIndexFilterProvider serialized row-index filter provider if no descriptor exists
    * @param rowGroupOffsets within-file row-index ordinals of each row group's first row
    * @param rowGroupNumRows number of rows in each row group
    * @param partitionIndex index into rowsPerPartition / allPartValues this file contributes to
    */
   case class PerFileDVEntry(
-      dvDescriptor: Option[String],
+      dvDescriptor: Option[DeletionVectorDescriptor],
       rowIndexFilterProvider: Option[RowIndexFilterProvider],
       rowGroupOffsets: Array[Long],
       rowGroupNumRows: Array[Int],
@@ -1107,7 +1107,9 @@ case class GpuDeltaParquetFileFormatNativeDV(
           }
           GpuSemaphore.acquireIfNecessary(TaskContext.get())
           RmmRapidsRetryIterator.withRetryNoSplit {
-            DeletionVector.computeNumDeletedRows(hostDvInfos, maxReadBatchSizeRows)
+            hostDvInfos.map { info =>
+              DeletionVector.computeNumDeletedRows(info, maxReadBatchSizeRows)
+            }.sum
           }
         }
       }
@@ -1297,7 +1299,9 @@ case class GpuDeltaParquetFileFormatNativeDV(
                 bitmap, false, entry.rowGroupOffsets, entry.rowGroupNumRows)
             }.toArray
             RmmRapidsRetryIterator.withRetryNoSplit {
-              DeletionVector.computeNumDeletedRows(dvInfos, maxReadBatchSizeRows)
+              dvInfos.map { info =>
+                DeletionVector.computeNumDeletedRows(info, maxReadBatchSizeRows)
+              }.sum
             }
           }
         }
@@ -1328,8 +1332,9 @@ case class GpuDeltaParquetFileFormatNativeDV(
       val loadFutures = batchExtra.perFileEntries.map { entry =>
         val loadTask = new FutureTask[SpillableHostBuffer](new Callable[SpillableHostBuffer] {
           override def call(): SpillableHostBuffer = {
-            val rawBitmap = RapidsDeletionVectors.loadDeletionVector(
-              conf, entry.dvDescriptor, entry.rowIndexFilterProvider, tp)
+            val filterTypeOpt = entry.dvDescriptor.map(_ => RowIndexFilterType.IF_CONTAINED)
+            val rawBitmap = RapidsDeletionVectors.loadDeletionVectorDescriptor(
+              conf, entry.dvDescriptor, filterTypeOpt, entry.rowIndexFilterProvider, tp)
             // DeltaBatchExtraInfo.close() releases the SpillableHostBuffer when the decode
             // phase completes (via withRetryNoSplit in readBatchData).
             closeOnExcept(rawBitmap) { raw =>

delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/RapidsDeletionVectors.scala

@@ -50,7 +50,7 @@ object RapidsDeletionVectors extends Logging {
   private val MISSING_ROW_INDEX_FILTER_MESSAGE = "Row index filter not found. file="
 
   case class DeletionVectorLookupResult(
-      dvDescriptor: Option[String],
+      dvDescriptor: Option[DeletionVectorDescriptor],
       filterType: Option[RowIndexFilterType],
       rowIndexFilterProvider: Option[RowIndexFilterProvider])
 
@@ -174,37 +174,43 @@ object RapidsDeletionVectors extends Logging {
       partitionedFile: PartitionedFile,
       deletionVectorReadInfo: Option[RapidsDeletionVectorReadInfo])
   : DeletionVectorLookupResult = {
-    val (dvDescriptorOpt, filterTypeOpt) = deletionVectorDescriptorAndFilter(partitionedFile)
+    val (encodedDescriptorOpt, filterTypeOpt) =
+      deletionVectorDescriptorAndFilter(partitionedFile)
     val rowIndexFilterProviderOpt = partitionedFile.getRowIndexFilter.toSeq.headOption
 
-    if (dvDescriptorOpt.isDefined && filterTypeOpt.isDefined) {
+    if (encodedDescriptorOpt.isDefined && filterTypeOpt.isDefined) {
       return DeletionVectorLookupResult(
-        dvDescriptorOpt, filterTypeOpt, rowIndexFilterProviderOpt)
-    } else if (dvDescriptorOpt.isDefined || filterTypeOpt.isDefined) {
+        encodedDescriptorOpt.map(DeletionVectorDescriptor.deserializeFromBase64),
+        filterTypeOpt,
+        rowIndexFilterProviderOpt)
+    } else if (encodedDescriptorOpt.isDefined || filterTypeOpt.isDefined) {
       throw new IllegalStateException(
         s"Both $FILE_ROW_INDEX_FILTER_ID_ENCODED and $FILE_ROW_INDEX_FILTER_TYPE " +
           "should either both have values or no values at all.")
     }
 
     rowIndexFilterProviderOpt
       .map(provider => DeletionVectorLookupResult(None, None, Some(provider)))
+      .orElse {
+        val lookupKeys = fileLookupKeys(partitionedFile)
+        lookupKeys.flatMap(key =>
+          deletionVectorReadInfo.flatMap(_.filePathToFilterProvider.get(key))).headOption
+          .map(provider => DeletionVectorLookupResult(None, None, Some(provider)))
+      }
       .orElse {
         val lookupKeys = fileLookupKeys(partitionedFile)
         lookupKeys.flatMap(key =>
           deletionVectorReadInfo.flatMap(_.filePathToDVMap.get(key))).headOption
           .map { descriptorWithFilterType =>
+            // Keep the descriptor object intact. Temporary low-shuffle-merge DVs can exceed
+            // 64 KiB, while DBR's serializeToBase64 uses DataOutput.writeUTF and cannot encode
+            // strings above that limit.
             DeletionVectorLookupResult(
-              Some(descriptorWithFilterType.descriptor.serializeToBase64()),
+              Some(descriptorWithFilterType.descriptor),
               Some(descriptorWithFilterType.filterType),
               None)
           }
       }
-      .orElse {
-        val lookupKeys = fileLookupKeys(partitionedFile)
-        lookupKeys.flatMap(key =>
-          deletionVectorReadInfo.flatMap(_.filePathToFilterProvider.get(key))).headOption
-          .map(provider => DeletionVectorLookupResult(None, None, Some(provider)))
-      }
       .getOrElse(DeletionVectorLookupResult(None, None, None))
   }
 
@@ -225,7 +231,7 @@ object RapidsDeletionVectors extends Logging {
       tablePath: String,
       deletionVectorReadInfo: Option[RapidsDeletionVectorReadInfo]): HostMemoryBuffer = {
     val dv = lookupDeletionVector(partitionedFile, deletionVectorReadInfo)
-    loadDeletionVector(
+    loadDeletionVectorDescriptor(
       conf,
       dv.dvDescriptor,
       dv.filterType,
@@ -262,8 +268,22 @@ object RapidsDeletionVectors extends Logging {
       filterTypeOpt: Option[RowIndexFilterType],
       rowIndexFilterProviderOpt: Option[RowIndexFilterProvider],
       tablePath: String): HostMemoryBuffer = {
+    loadDeletionVectorDescriptor(
+      conf,
+      dvDescriptorOpt.map(DeletionVectorDescriptor.deserializeFromBase64),
+      filterTypeOpt,
+      rowIndexFilterProviderOpt,
+      tablePath)
+  }
+
+  def loadDeletionVectorDescriptor(
+      conf: Configuration,
+      dvDescriptorOpt: Option[DeletionVectorDescriptor],
+      filterTypeOpt: Option[RowIndexFilterType],
+      rowIndexFilterProviderOpt: Option[RowIndexFilterProvider],
+      tablePath: String): HostMemoryBuffer = {
     if (dvDescriptorOpt.isDefined && filterTypeOpt.isDefined) {
-      val dvDesc = DeletionVectorDescriptor.deserializeFromBase64(dvDescriptorOpt.get)
+      val dvDesc = dvDescriptorOpt.get
       filterTypeOpt.get match {
         case RowIndexFilterType.IF_CONTAINED =>
           if (dvDesc.cardinality == 0) {
@@ -313,7 +333,7 @@ object RapidsDeletionVectors extends Logging {
       tablePath: String,
       deletionVectorReadInfo: Option[RapidsDeletionVectorReadInfo]): RoaringBitmapArray = {
     val dv = lookupDeletionVector(partitionedFile, deletionVectorReadInfo)
-    loadScalaBitmap(
+    loadScalaBitmapDescriptor(
       conf,
       dv.dvDescriptor,
       dv.filterType,
@@ -327,8 +347,22 @@ object RapidsDeletionVectors extends Logging {
       filterTypeOpt: Option[RowIndexFilterType],
       rowIndexFilterProviderOpt: Option[RowIndexFilterProvider],
       tablePath: String): RoaringBitmapArray = {
+    loadScalaBitmapDescriptor(
+      conf,
+      dvDescriptorOpt.map(DeletionVectorDescriptor.deserializeFromBase64),
+      filterTypeOpt,
+      rowIndexFilterProviderOpt,
+      tablePath)
+  }
+
+  def loadScalaBitmapDescriptor(
+      conf: Configuration,
+      dvDescriptorOpt: Option[DeletionVectorDescriptor],
+      filterTypeOpt: Option[RowIndexFilterType],
+      rowIndexFilterProviderOpt: Option[RowIndexFilterProvider],
+      tablePath: String): RoaringBitmapArray = {
     if (dvDescriptorOpt.isDefined && filterTypeOpt.isDefined) {
-      val dvDesc = DeletionVectorDescriptor.deserializeFromBase64(dvDescriptorOpt.get)
+      val dvDesc = dvDescriptorOpt.get
       filterTypeOpt.get match {
         case RowIndexFilterType.IF_CONTAINED =>
           val dvStore = new com.databricks.sql.transaction.tahoe.storage.dv.HadoopFileSystemDVStore(