felipepessoto commented on code in PR #12215:
URL: https://github.com/apache/gluten/pull/12215#discussion_r3899041447
##########
gluten-delta/src/main/scala/org/apache/gluten/extension/OffloadDeltaScan.scala:
##########
@@ -20,47 +20,80 @@ import org.apache.gluten.execution.DeltaScanTransformer
import org.apache.gluten.extension.columnar.FallbackTags
import org.apache.gluten.extension.columnar.offload.OffloadSingleNode
-import org.apache.spark.sql.delta.DeltaParquetFileFormat
-import org.apache.spark.sql.delta.SnapshotDescriptor
+import org.apache.spark.sql.delta.{DeltaParquetFileFormat, SnapshotDescriptor}
import
org.apache.spark.sql.delta.commands.DeletionVectorUtils.deletionVectorsReadable
-import org.apache.spark.sql.delta.files.{CdcAddFileIndex, TahoeFileIndex,
TahoeRemoveFileIndex}
+import org.apache.spark.sql.delta.files.{CdcAddFileIndex, TahoeBatchFileIndex,
TahoeFileIndex, TahoeRemoveFileIndex}
import org.apache.spark.sql.delta.stats.PreparedDeltaFileIndex
import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan}
+import org.apache.spark.sql.types.{DataType, StructType}
import org.apache.spark.util.SparkVersionUtil
-case class OffloadDeltaScan() extends OffloadSingleNode {
+case class OffloadDeltaScan(enableNativeDmlRowIndexScan: Boolean) extends
OffloadSingleNode {
private val DeletionVectorsUseMetadataRowIndexKey =
"spark.databricks.delta.deletionVectors.useMetadataRowIndex"
+ // Spark 3.5+ exposes this as
ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME.
+ private val parquetTemporaryRowIndexColumnName = "_tmp_metadata_row_index"
+ private val rowIndexColumnNames =
+ Set(
+ "__delta_internal_row_index",
+ DeltaParquetFileFormat.ROW_INDEX_COLUMN_NAME,
+ parquetTemporaryRowIndexColumnName,
+ "row_index")
+ // TahoeBatchFileIndex.actionType as set by Delta's DELETE, UPDATE and MERGE
commands.
+ private val dmlActionTypes = Set("delete", "update", "merge")
+
override def offload(plan: SparkPlan): SparkPlan = plan match {
case scan: FileSourceScanExec if isDeltaLogScan(scan) =>
FallbackTags.add(scan, "fallback Delta _delta_log scan")
scan
+ case scan: FileSourceScanExec if shouldFallbackDeletionVectorDmlScan(scan)
=>
+ FallbackTags.add(scan, "fallback Delta DV DML row-index scan by
configuration")
+ scan
case scan: FileSourceScanExec if
shouldFallbackSpark34DeletionVectorScan(scan) =>
FallbackTags.add(scan, "fallback Spark 3.4 Delta DV scan")
scan
case scan: FileSourceScanExec
if shouldFallbackDeletionVectorScanWithoutMetadataRowIndex(scan) =>
FallbackTags.add(scan, "fallback Delta DV scan without metadata row
index")
scan
- case scan: FileSourceScanExec if isDeltaScan(scan) =>
+ case scan: FileSourceScanExec if DeltaScanUtils.isDeltaScan(scan) =>
DeltaScanTransformer(scan)
case other => other
}
- private def isDeltaScan(scan: FileSourceScanExec): Boolean = {
- isDeltaFileIndex(scan) || isDeltaParquetScan(scan)
+ /**
+ * The scoped escape hatch: with the config off, the DELETE/UPDATE/MERGE
target scan that produces
+ * file paths and row indexes for deletion-vector writes stays on Spark,
while every other scan
+ * keeps offloading. The whole check lives here, on the scan alone: Delta
builds every DML target
+ * relation over a [[TahoeBatchFileIndex]] carrying the command name, which
survives AQE stage
+ * splits and arbitrary join placement, and only DV-writing DML reads a
row-index column from that
+ * relation; DML that rewrites whole files does not, and remains eligible
for native execution.
+ */
+ private def shouldFallbackDeletionVectorDmlScan(scan: FileSourceScanExec):
Boolean = {
+ !enableNativeDmlRowIndexScan && isDmlTargetScan(scan) &&
scanReadsRowIndexColumn(scan)
}
- private def isDeltaParquetScan(scan: FileSourceScanExec): Boolean = {
- val fileFormatClass = scan.relation.fileFormat.getClass
- fileFormatClass == classOf[DeltaParquetFileFormat] ||
- fileFormatClass.getSimpleName == "GlutenDeltaParquetFileFormat"
+ private def isDmlTargetScan(scan: FileSourceScanExec): Boolean = {
+ scan.relation.location match {
+ case index: TahoeBatchFileIndex =>
dmlActionTypes.contains(index.actionType)
+ case _ => false
+ }
}
- private def isDeltaFileIndex(scan: FileSourceScanExec): Boolean = {
- scan.relation.location.isInstanceOf[TahoeFileIndex] ||
- scan.relation.location.isInstanceOf[PreparedDeltaFileIndex]
+ private def scanReadsRowIndexColumn(scan: FileSourceScanExec): Boolean = {
+ def nestedFieldNames(dataType: DataType): Seq[String] = dataType match {
+ case struct: StructType =>
+ struct.fields.flatMap(field => field.name +:
nestedFieldNames(field.dataType)).toSeq
+ case _ => Seq.empty
+ }
+
+ val outputColumnNames =
+ scan.output.flatMap(attribute => attribute.name +:
nestedFieldNames(attribute.dataType))
+ val requiredColumnNames = scan.requiredSchema.fields.flatMap {
+ field => field.name +: nestedFieldNames(field.dataType)
+ }
+ (outputColumnNames ++
requiredColumnNames).exists(rowIndexColumnNames.contains)
Review Comment:
Non-blocking: `scanReadsRowIndexColumn` recursively flattens every struct
and treats any field named `row_index` as Delta metadata. Could we narrow that
match to the actual `_metadata.row_index` field? A regular Delta table may have
a top-level or nested user column with that name, and non-DV DELETE/UPDATE
paths still use a `TahoeBatchFileIndex` with action type `delete`/`update`.
With this config disabled, those ordinary scans would therefore fall back to
Spark even though the config promises that other scans are unaffected.
For example, the check could retain exact matching for generated top-level
columns while inspecting `row_index` only under `_metadata`:
```scala
private val generatedRowIndexColumnNames = Set(
"__delta_internal_row_index",
DeltaParquetFileFormat.ROW_INDEX_COLUMN_NAME,
parquetTemporaryRowIndexColumnName)
private def isRowIndexColumn(name: String, dataType: DataType): Boolean = {
generatedRowIndexColumnNames.contains(name) ||
(name == "_metadata" && (dataType match {
case struct: StructType => struct.fieldNames.contains("row_index")
case _ => false
}))
}
private def scanReadsRowIndexColumn(scan: FileSourceScanExec): Boolean = {
val outputFields = scan.output.iterator.map(attr => attr.name ->
attr.dataType)
val requiredFields = scan.requiredSchema.fields.iterator.map(field =>
field.name -> field.dataType)
(outputFields ++ requiredFields).exists { case (name, dataType) =>
isRowIndexColumn(name, dataType)
}
}
```
A non-DV DML regression test using a user-defined `row_index` column would
also lock in the intended scope.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
