Skip to content

Commit 6d858cb

Browse files
malinjawimalinjawi
committed
[VL][Delta] Narrow DML row-index scan fallback
1 parent f313556 commit 6d858cb

6 files changed

Lines changed: 112 additions & 23 deletions

File tree

‎backends-velox/src-delta/main/scala/org/apache/gluten/component/VeloxDeltaComponent.scala‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ package org.apache.gluten.component
1818

1919
import org.apache.gluten.backendsapi.velox.VeloxBackend
2020
import org.apache.gluten.config.GlutenConfig
21-
import org.apache.gluten.extension.{DeltaPostTransformRules, OffloadDeltaFilter, OffloadDeltaProject, OffloadDeltaScan}
21+
import org.apache.gluten.extension.{DeltaDeletionVectorDmlUtils, DeltaPostTransformRules, OffloadDeltaFilter, OffloadDeltaProject, OffloadDeltaScan}
2222
import org.apache.gluten.extension.columnar.heuristic.HeuristicTransform
2323
import org.apache.gluten.extension.columnar.validator.Validators
2424
import org.apache.gluten.extension.injector.Injector
@@ -43,6 +43,7 @@ class VeloxDeltaComponent extends Component {
4343
override def injectRules(injector: Injector): Unit = {
4444
val legacy = injector.gluten.legacy
4545
injector.spark.injectOptimizerRule(deltaDvPreprocessRule)
46+
legacy.injectPreTransform(_ => DeltaDeletionVectorDmlUtils.tagDmlRowIndexScans)
4647
legacy.injectTransform {
4748
c =>
4849
val offload = Seq(OffloadDeltaScan(), OffloadDeltaProject(), OffloadDeltaFilter())

‎backends-velox/src-delta33/test/scala/org/apache/spark/sql/delta/DeltaDeletionVectorHandoffSuite.scala‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,22 @@ class DeltaDeletionVectorHandoffSuite
9393
assert(!executedPlans.exists(hasNativeParentOverDmlFallbackScan), planText)
9494
}
9595

96+
private def assertReadPlanAfterDmlFallback(path: String, useMetadataRowIndex: Boolean): Unit = {
97+
withSQLConf(
98+
DeltaSQLConf.DELETION_VECTORS_USE_METADATA_ROW_INDEX.key -> useMetadataRowIndex.toString) {
99+
val df = spark.read.format("delta").load(path)
100+
val executedPlan = df.queryExecution.executedPlan
101+
val planText = executedPlan.treeString
102+
if (useMetadataRowIndex) {
103+
assert(executedPlan.collect { case _: DeltaScanTransformer => true }.nonEmpty, planText)
104+
assert(!planText.contains(DmlFallbackReason), planText)
105+
} else {
106+
assert(executedPlan.collect { case _: DeltaScanTransformer => true }.isEmpty, planText)
107+
}
108+
checkAnswer(df, Seq((1, "a"), (2, "b")).toDF())
109+
}
110+
}
111+
96112
private def activeDvCardinality(path: String): Long = {
97113
val log = DeltaLog.forTable(spark, new Path(path))
98114
log.update().allFiles.collect().flatMap(
@@ -169,7 +185,7 @@ class DeltaDeletionVectorHandoffSuite
169185

170186
val log = DeltaLog.forTable(spark, new Path(path))
171187
assert(log.update().allFiles.collect().exists(_.deletionVector != null))
172-
checkAnswer(spark.read.format("delta").load(path), Seq((1, "a"), (2, "b")).toDF())
188+
assertReadPlanAfterDmlFallback(path, useMetadataRowIndex)
173189
}
174190
}
175191
}
@@ -196,7 +212,7 @@ class DeltaDeletionVectorHandoffSuite
196212
assertSparkDmlFallback(captureDeletePlans(path, "id IN (3, 4)", useMetadataRowIndex = true))
197213
assert(activeDvCardinality(path) === 4L)
198214

199-
checkAnswer(spark.read.format("delta").load(path), Seq((1, "a"), (2, "b")).toDF())
215+
assertReadPlanAfterDmlFallback(path, useMetadataRowIndex = true)
200216
}
201217
}
202218
}

‎backends-velox/src-delta40/test/scala/org/apache/spark/sql/delta/DeltaDeletionVectorHandoffSuite.scala‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,22 @@ class DeltaDeletionVectorHandoffSuite
9393
assert(!executedPlans.exists(hasNativeParentOverDmlFallbackScan), planText)
9494
}
9595

96+
private def assertReadPlanAfterDmlFallback(path: String, useMetadataRowIndex: Boolean): Unit = {
97+
withSQLConf(
98+
DeltaSQLConf.DELETION_VECTORS_USE_METADATA_ROW_INDEX.key -> useMetadataRowIndex.toString) {
99+
val df = spark.read.format("delta").load(path)
100+
val executedPlan = df.queryExecution.executedPlan
101+
val planText = executedPlan.treeString
102+
if (useMetadataRowIndex) {
103+
assert(executedPlan.collect { case _: DeltaScanTransformer => true }.nonEmpty, planText)
104+
assert(!planText.contains(DmlFallbackReason), planText)
105+
} else {
106+
assert(executedPlan.collect { case _: DeltaScanTransformer => true }.isEmpty, planText)
107+
}
108+
checkAnswer(df, Seq((1, "a"), (2, "b")).toDF())
109+
}
110+
}
111+
96112
private def activeDvCardinality(path: String): Long = {
97113
val log = DeltaLog.forTable(spark, new Path(path))
98114
log.update().allFiles.collect().flatMap(
@@ -166,7 +182,7 @@ class DeltaDeletionVectorHandoffSuite
166182

167183
val log = DeltaLog.forTable(spark, new Path(path))
168184
assert(log.update().allFiles.collect().exists(_.deletionVector != null))
169-
checkAnswer(spark.read.format("delta").load(path), Seq((1, "a"), (2, "b")).toDF())
185+
assertReadPlanAfterDmlFallback(path, useMetadataRowIndex)
170186
}
171187
}
172188
}
@@ -225,7 +241,7 @@ class DeltaDeletionVectorHandoffSuite
225241
assertSparkDmlFallback(captureDeletePlans(path, "id IN (3, 4)", useMetadataRowIndex = true))
226242
assert(activeDvCardinality(path) === 4L)
227243

228-
checkAnswer(spark.read.format("delta").load(path), Seq((1, "a"), (2, "b")).toDF())
244+
assertReadPlanAfterDmlFallback(path, useMetadataRowIndex = true)
229245
}
230246
}
231247
}

‎gluten-delta/src/main/scala/org/apache/gluten/extension/DeltaDeletionVectorDmlUtils.scala‎

Lines changed: 70 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,21 +16,60 @@
1616
*/
1717
package org.apache.gluten.extension
1818

19+
import org.apache.spark.sql.catalyst.expressions.Expression
20+
import org.apache.spark.sql.catalyst.rules.Rule
21+
import org.apache.spark.sql.catalyst.trees.TreeNodeTag
1922
import org.apache.spark.sql.delta.DeltaParquetFileFormat
2023
import org.apache.spark.sql.delta.files.TahoeFileIndex
2124
import org.apache.spark.sql.delta.stats.PreparedDeltaFileIndex
22-
import org.apache.spark.sql.execution.FileSourceScanExec
25+
import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan}
2326

2427
object DeltaDeletionVectorDmlUtils {
28+
private val DmlRowIndexScanTag: TreeNodeTag[Boolean] =
29+
TreeNodeTag[Boolean]("org.apache.gluten.delta.dml.row.index.scan")
30+
2531
// Spark 3.5+ exposes this as ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME.
2632
private val parquetTemporaryRowIndexColumnName = "_tmp_metadata_row_index"
2733
private val deletionVectorRowIndexColumnNames =
2834
Set(
2935
"__delta_internal_row_index",
3036
DeltaParquetFileFormat.ROW_INDEX_COLUMN_NAME,
31-
parquetTemporaryRowIndexColumnName)
37+
parquetTemporaryRowIndexColumnName,
38+
"row_index",
39+
"rowIndexCol")
3240
private val filePathColumnNames = Set("file_path", "filePath")
3341

42+
val tagDmlRowIndexScans: Rule[SparkPlan] = (plan: SparkPlan) => {
43+
def visit(
44+
node: SparkPlan,
45+
hasRowIndexReference: Boolean,
46+
hasFilePathReference: Boolean): Unit = {
47+
val nextHasRowIndexReference =
48+
hasRowIndexReference || node.expressions.exists(referencesRowIndexColumn)
49+
val nextHasFilePathReference =
50+
hasFilePathReference || node.expressions.exists(referencesFilePathColumn)
51+
52+
node.children.foreach {
53+
case scan: FileSourceScanExec
54+
if nextHasRowIndexReference &&
55+
nextHasFilePathReference &&
56+
isDeletionVectorDmlRowIndexScanCandidate(scan) =>
57+
scan.setTagValue(DmlRowIndexScanTag, true)
58+
case child =>
59+
visit(child, nextHasRowIndexReference, nextHasFilePathReference)
60+
}
61+
}
62+
63+
visit(plan, hasRowIndexReference = false, hasFilePathReference = false)
64+
plan
65+
}
66+
67+
def copyDmlRowIndexScanTag(from: SparkPlan, to: SparkPlan): Unit = {
68+
if (from.getTagValue(DmlRowIndexScanTag).contains(true)) {
69+
to.setTagValue(DmlRowIndexScanTag, true)
70+
}
71+
}
72+
3473
def isDeltaScan(scan: FileSourceScanExec): Boolean = {
3574
isDeltaFileIndex(scan) || isDeltaParquetScan(scan)
3675
}
@@ -47,12 +86,39 @@ object DeltaDeletionVectorDmlUtils {
4786
}
4887

4988
def isDeletionVectorDmlRowIndexScan(scan: FileSourceScanExec): Boolean = {
89+
scan.getTagValue(DmlRowIndexScanTag).contains(true) &&
90+
isDeletionVectorDmlRowIndexScanCandidate(scan)
91+
}
92+
93+
def isDeletionVectorDmlRowIndexScan(plan: SparkPlan): Boolean = {
94+
plan.getTagValue(DmlRowIndexScanTag).contains(true)
95+
}
96+
97+
private def isDeletionVectorDmlRowIndexScanCandidate(scan: FileSourceScanExec): Boolean = {
5098
if (!isDeltaScan(scan)) {
5199
return false
52100
}
53101

102+
scanContainsColumnName(scan, deletionVectorRowIndexColumnNames) &&
103+
scanContainsColumnName(scan, filePathColumnNames)
104+
}
105+
106+
private def scanContainsColumnName(
107+
scan: FileSourceScanExec,
108+
columnNames: Set[String]): Boolean = {
54109
val scanColumnNames = (scan.output.map(_.name) ++ scan.requiredSchema.fieldNames).toSet
55-
scanColumnNames.exists(deletionVectorRowIndexColumnNames.contains) &&
56-
scanColumnNames.exists(filePathColumnNames.contains)
110+
scanColumnNames.exists(columnNames.contains) || columnNames.exists(scan.treeString.contains)
111+
}
112+
113+
private def referencesRowIndexColumn(expr: Expression): Boolean = {
114+
val expressionText = expr.toString()
115+
expr.references.exists(attr => deletionVectorRowIndexColumnNames.contains(attr.name)) ||
116+
deletionVectorRowIndexColumnNames.exists(expressionText.contains)
117+
}
118+
119+
private def referencesFilePathColumn(expr: Expression): Boolean = {
120+
val expressionText = expr.toString()
121+
expr.references.exists(attr => filePathColumnNames.contains(attr.name)) ||
122+
filePathColumnNames.exists(expressionText.contains)
57123
}
58124
}

‎gluten-delta/src/main/scala/org/apache/gluten/extension/DeltaPostTransformRules.scala‎

Lines changed: 1 addition & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,6 @@ object DeltaPostTransformRules {
5353
Set(deletionVectorDeletedRowColumnName, deletionVectorRowIndexColumnName)
5454
private val deletionVectorPredicateColumnNames =
5555
deletionVectorInternalColumnNames ++ deletionVectorRowIndexColumnNames
56-
private val deletionVectorFilePathColumnNames = Set("file_path", "filePath")
5756

5857
private val COLUMN_MAPPING_RULE_TAG: TreeNodeTag[String] =
5958
TreeNodeTag[String]("org.apache.gluten.delta.column.mapping")
@@ -239,18 +238,7 @@ object DeltaPostTransformRules {
239238
}
240239

241240
private def isDeletionVectorDmlRowIndexScan(plan: SparkPlan): Boolean = {
242-
val scanColumnNames = plan match {
243-
case scan: DeltaScanTransformer =>
244-
(scan.output.map(_.name) ++ scan.requiredSchema.fieldNames).toSet
245-
case scan: FileSourceScanExec =>
246-
(scan.output.map(_.name) ++ scan.requiredSchema.fieldNames).toSet
247-
case _ =>
248-
Set.empty[String]
249-
}
250-
val hasRowIndex = scanColumnNames.exists(deletionVectorRowIndexColumnNames.contains)
251-
val hasFilePath = scanColumnNames.exists(deletionVectorFilePathColumnNames.contains)
252-
val hasDeletedRowMarker = scanColumnNames.contains(deletionVectorDeletedRowColumnName)
253-
hasRowIndex && (hasFilePath || !hasDeletedRowMarker)
241+
DeltaDeletionVectorDmlUtils.isDeletionVectorDmlRowIndexScan(plan)
254242
}
255243

256244
private def shouldStripDeletionVectorInternalColumn(

‎gluten-delta/src/main/scala/org/apache/gluten/extension/OffloadDeltaScan.scala‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,9 @@ case class OffloadDeltaScan() extends OffloadSingleNode {
5050
FallbackTags.add(scan, "fallback Delta DV scan without metadata row index")
5151
scan
5252
case scan: FileSourceScanExec if isDeltaScan(scan) =>
53-
DeltaScanTransformer(scan)
53+
val transformer = DeltaScanTransformer(scan)
54+
DeltaDeletionVectorDmlUtils.copyDmlRowIndexScanTag(scan, transformer)
55+
transformer
5456
case other => other
5557
}
5658

0 commit comments

Comments
 (0)