Skip to content

Commit 5891b3e

Browse files
nartal1greptile-apps[bot]
nartal1
and
greptile-apps[bot]
authored
Fallback Delta file scans to CPU on DB-17.3 [databricks] (#14615)
Fixes #14607 ### Description - On DB-17.3, UC managed tables use `DeltaParquetFileFormat` which requires Delta-specific columns (e.g. skip_row). Since there is no Delta provider for DB-17.3, the GPU reader cannot fill these columns. - Forked `FileSourceScanExecMeta` for the 400db173 shim to detect unsupported `ParquetFileFormat` subclasses and fall back the scan to CPU. - Created `FileSourceScanExecMetaBase` to share the common logic across DB shims. - Plain Parquet reads are unaffected. ### Tested this manually on DB cluster - [x] UC managed Delta table on DB-17.3 — scan on CPU, rest on GPU - [x] Raw Parquet read on DB-17.3 — scan stays on GPU ### Checklists Documentation - [ ] Updated for new or modified user-facing features or behaviors - [x] No user-facing change Testing - [ ] Added or modified tests to cover new code paths - [ ] Covered by existing tests (Please provide the names of the existing tests in the PR description.) - [ ] Not required Performance - [ ] Tests ran and results are added in the PR description - [ ] Issue filed with a link in the PR description - [x] Not required --------- Signed-off-by: Niranjan Artal <nartal@nvidia.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
1 parent 74203a3 commit 5891b3e

3 files changed

Lines changed: 209 additions & 128 deletions

File tree

‎sql-plugin/src/main/spark330db/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala‎

Lines changed: 2 additions & 128 deletions
Original file line numberDiff line numberDiff line change
@@ -19,141 +19,15 @@
1919
{"spark": "332db"}
2020
{"spark": "341db"}
2121
{"spark": "350db143"}
22-
{"spark": "400db173"}
2322
spark-rapids-shim-json-lines ***/
2423
package com.nvidia.spark.rapids.shims
2524

2625
import com.nvidia.spark.rapids._
2726

28-
import org.apache.spark.internal.Logging
29-
import org.apache.spark.sql.catalyst.expressions.Expression
30-
import org.apache.spark.sql.execution._
31-
import org.apache.spark.sql.execution.datasources.HadoopFsRelation
32-
import org.apache.spark.sql.execution.datasources.json.JsonFileFormat
33-
import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
34-
import org.apache.spark.sql.rapids.GpuFileSourceScanExec
35-
import org.apache.spark.sql.rapids.execution.{GpuBroadcastExchangeExec, GpuSubqueryBroadcastExec}
27+
import org.apache.spark.sql.execution.FileSourceScanExec
3628

3729
class FileSourceScanExecMeta(plan: FileSourceScanExec,
3830
conf: RapidsConf,
3931
parent: Option[RapidsMeta[_, _, _]],
4032
rule: DataFromReplacementRule)
41-
extends SparkPlanMeta[FileSourceScanExec](plan, conf, parent, rule) with Logging {
42-
43-
// Replaces SubqueryBroadcastExec inside dynamic pruning filters with GPU counterpart
44-
// if possible. Instead regarding filters as childExprs of current Meta, we create
45-
// a new meta for SubqueryBroadcastExec. The reason is that the GPU replacement of
46-
// FileSourceScan is independent from the replacement of the partitionFilters. It is
47-
// possible that the FileSourceScan is on the CPU, while the dynamic partitionFilters
48-
// are on the GPU. And vice versa. The same applies for dataFilters in the case of
49-
// Dynamic File Pruning
50-
private def convertBroadcast(bc: SubqueryBroadcastExec): BaseSubqueryExec = {
51-
val meta = GpuOverrides.wrapAndTagPlan(bc, conf)
52-
meta.tagForExplain()
53-
if (conf.shouldExplain) {
54-
val explain = meta.explain(conf.shouldExplainAll)
55-
if (explain.nonEmpty) {
56-
logWarning(s"\n$explain")
57-
}
58-
}
59-
val converted = meta.convertIfNeeded()
60-
// Because the PlanSubqueries rule is not called (and does not work as expected),
61-
// we might actually have to fully convert the subquery plan as the plugin would
62-
// intend (in this case calling GpuTransitionOverrides to insert GpuCoalesceBatches,
63-
// etc.) to match the other side of the join to reuse the BroadcastExchange.
64-
// This happens when SubqueryBroadcast has the original (Gpu)BroadcastExchangeExec
65-
converted match {
66-
case e: GpuSubqueryBroadcastExec => e.child match {
67-
// If the GpuBroadcastExchange is here, then we will need to run the transition
68-
// overrides here
69-
case _: GpuBroadcastExchangeExec =>
70-
var updated = ApplyColumnarRulesAndInsertTransitions(Seq(), true)
71-
.apply(converted)
72-
updated = (new GpuTransitionOverrides()).apply(updated)
73-
updated match {
74-
case h: GpuBringBackToHost =>
75-
h.child.asInstanceOf[BaseSubqueryExec]
76-
case c2r: GpuColumnarToRowExec =>
77-
c2r.child.asInstanceOf[BaseSubqueryExec]
78-
case _: GpuSubqueryBroadcastExec =>
79-
updated.asInstanceOf[BaseSubqueryExec]
80-
}
81-
// Otherwise, if this SubqueryBroadcast is using a ReusedExchange, then we don't
82-
// do anything further
83-
case _: ReusedExchangeExec =>
84-
converted.asInstanceOf[BaseSubqueryExec]
85-
}
86-
case _ =>
87-
converted.asInstanceOf[BaseSubqueryExec]
88-
}
89-
}
90-
91-
private def convertDynamicPruningFilters(filters: Seq[Expression]): Seq[Expression] = {
92-
filters.map { filter =>
93-
filter.transformDown {
94-
case dpe @ DynamicPruningShims(inSub: InSubqueryExec) =>
95-
inSub.plan match {
96-
case bc: SubqueryBroadcastExec =>
97-
DynamicPruningShims(inSub.copy(plan = convertBroadcast(bc)))
98-
case reuse @ ReusedSubqueryExec(bc: SubqueryBroadcastExec) =>
99-
DynamicPruningShims(inSub.copy(plan = reuse.copy(convertBroadcast(bc))))
100-
case _ =>
101-
dpe
102-
}
103-
}
104-
}
105-
}
106-
107-
// Support partitionFilters in Dynamic Partition Pruning
108-
private lazy val partitionFilters = convertDynamicPruningFilters(wrapped.partitionFilters)
109-
110-
// Support dataFilters in Dynamic File Pruning
111-
private lazy val dataFilters = convertDynamicPruningFilters(wrapped.dataFilters)
112-
113-
// partition filters and data filters are not run on the GPU
114-
override val childExprs: Seq[ExprMeta[_]] = Seq.empty
115-
116-
override def tagPlanForGpu(): Unit = {
117-
// this is very specific check to have any of the Delta log metadata queries
118-
// fallback and run on the CPU since there is some incompatibilities in
119-
// Databricks Spark and Apache Spark.
120-
if (wrapped.relation.fileFormat.isInstanceOf[JsonFileFormat] &&
121-
wrapped.relation.location.getClass.getCanonicalName() ==
122-
"com.databricks.sql.transaction.tahoe.DeltaLogFileIndex") {
123-
this.entirePlanWillNotWork("Plans that read Delta Index JSON files can not run " +
124-
"any part of the plan on the GPU!")
125-
}
126-
ScanExecShims.tagGpuFileSourceScanExecSupport(this)
127-
}
128-
129-
override def convertToCpu(): SparkPlan = {
130-
val cpu = wrapped.copy(partitionFilters = partitionFilters, dataFilters = dataFilters)
131-
cpu.copyTagsFrom(wrapped)
132-
cpu
133-
}
134-
135-
override def convertToGpu(): GpuExec = {
136-
val sparkSession = wrapped.relation.sparkSession
137-
val options = wrapped.relation.options
138-
val newRelation = HadoopFsRelation(
139-
wrapped.relation.location,
140-
wrapped.relation.partitionSchema,
141-
wrapped.relation.dataSchema,
142-
wrapped.relation.bucketSpec,
143-
GpuFileSourceScanExec.convertFileFormat(wrapped.relation, conf),
144-
options)(sparkSession)
145-
146-
GpuFileSourceScanExec(
147-
newRelation,
148-
wrapped.output,
149-
wrapped.requiredSchema,
150-
partitionFilters,
151-
wrapped.optionalBucketSet,
152-
// TODO: Does Databricks have coalesced bucketing implemented?
153-
None,
154-
dataFilters,
155-
wrapped.tableIdentifier,
156-
wrapped.disableBucketedScan,
157-
queryUsesInputFile = false)(conf)
158-
}
159-
}
33+
extends FileSourceScanExecMetaBase(plan, conf, parent, rule)

‎sql-plugin/src/main/spark330db/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMetaBase.scala‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
/*
2+
* Copyright (c) 2022-2026, NVIDIA CORPORATION.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
/*** spark-rapids-shim-json-lines
18+
{"spark": "330db"}
19+
{"spark": "332db"}
20+
{"spark": "341db"}
21+
{"spark": "350db143"}
22+
{"spark": "400db173"}
23+
spark-rapids-shim-json-lines ***/
24+
package com.nvidia.spark.rapids.shims
25+
26+
import com.nvidia.spark.rapids._
27+
28+
import org.apache.spark.internal.Logging
29+
import org.apache.spark.sql.catalyst.expressions.Expression
30+
import org.apache.spark.sql.execution._
31+
import org.apache.spark.sql.execution.datasources.HadoopFsRelation
32+
import org.apache.spark.sql.execution.datasources.json.JsonFileFormat
33+
import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
34+
import org.apache.spark.sql.rapids.GpuFileSourceScanExec
35+
import org.apache.spark.sql.rapids.execution.{GpuBroadcastExchangeExec, GpuSubqueryBroadcastExec}
36+
37+
abstract class FileSourceScanExecMetaBase(plan: FileSourceScanExec,
38+
conf: RapidsConf,
39+
parent: Option[RapidsMeta[_, _, _]],
40+
rule: DataFromReplacementRule)
41+
extends SparkPlanMeta[FileSourceScanExec](plan, conf, parent, rule) with Logging {
42+
43+
// Replaces SubqueryBroadcastExec inside dynamic pruning filters with GPU counterpart
44+
// if possible. Instead regarding filters as childExprs of current Meta, we create
45+
// a new meta for SubqueryBroadcastExec. The reason is that the GPU replacement of
46+
// FileSourceScan is independent from the replacement of the partitionFilters. It is
47+
// possible that the FileSourceScan is on the CPU, while the dynamic partitionFilters
48+
// are on the GPU. And vice versa. The same applies for dataFilters in the case of
49+
// Dynamic File Pruning
50+
private def convertBroadcast(bc: SubqueryBroadcastExec): BaseSubqueryExec = {
51+
val meta = GpuOverrides.wrapAndTagPlan(bc, conf)
52+
meta.tagForExplain()
53+
if (conf.shouldExplain) {
54+
val explain = meta.explain(conf.shouldExplainAll)
55+
if (explain.nonEmpty) {
56+
logWarning(s"\n$explain")
57+
}
58+
}
59+
val converted = meta.convertIfNeeded()
60+
// Because the PlanSubqueries rule is not called (and does not work as expected),
61+
// we might actually have to fully convert the subquery plan as the plugin would
62+
// intend (in this case calling GpuTransitionOverrides to insert GpuCoalesceBatches,
63+
// etc.) to match the other side of the join to reuse the BroadcastExchange.
64+
// This happens when SubqueryBroadcast has the original (Gpu)BroadcastExchangeExec
65+
converted match {
66+
case e: GpuSubqueryBroadcastExec => e.child match {
67+
// If the GpuBroadcastExchange is here, then we will need to run the transition
68+
// overrides here
69+
case _: GpuBroadcastExchangeExec =>
70+
var updated = ApplyColumnarRulesAndInsertTransitions(Seq(), true)
71+
.apply(converted)
72+
updated = (new GpuTransitionOverrides()).apply(updated)
73+
updated match {
74+
case h: GpuBringBackToHost =>
75+
h.child.asInstanceOf[BaseSubqueryExec]
76+
case c2r: GpuColumnarToRowExec =>
77+
c2r.child.asInstanceOf[BaseSubqueryExec]
78+
case _: GpuSubqueryBroadcastExec =>
79+
updated.asInstanceOf[BaseSubqueryExec]
80+
}
81+
// Otherwise, if this SubqueryBroadcast is using a ReusedExchange, then we don't
82+
// do anything further
83+
case _: ReusedExchangeExec =>
84+
converted.asInstanceOf[BaseSubqueryExec]
85+
}
86+
case _ =>
87+
converted.asInstanceOf[BaseSubqueryExec]
88+
}
89+
}
90+
91+
private def convertDynamicPruningFilters(filters: Seq[Expression]): Seq[Expression] = {
92+
filters.map { filter =>
93+
filter.transformDown {
94+
case dpe @ DynamicPruningShims(inSub: InSubqueryExec) =>
95+
inSub.plan match {
96+
case bc: SubqueryBroadcastExec =>
97+
DynamicPruningShims(inSub.copy(plan = convertBroadcast(bc)))
98+
case reuse @ ReusedSubqueryExec(bc: SubqueryBroadcastExec) =>
99+
DynamicPruningShims(inSub.copy(plan = reuse.copy(convertBroadcast(bc))))
100+
case _ =>
101+
dpe
102+
}
103+
}
104+
}
105+
}
106+
107+
// Support partitionFilters in Dynamic Partition Pruning
108+
protected lazy val partitionFilters =
109+
convertDynamicPruningFilters(wrapped.partitionFilters)
110+
111+
// Support dataFilters in Dynamic File Pruning
112+
protected lazy val dataFilters = convertDynamicPruningFilters(wrapped.dataFilters)
113+
114+
// partition filters and data filters are not run on the GPU
115+
override val childExprs: Seq[ExprMeta[_]] = Seq.empty
116+
117+
override def tagPlanForGpu(): Unit = {
118+
// this is very specific check to have any of the Delta log metadata queries
119+
// fallback and run on the CPU since there is some incompatibilities in
120+
// Databricks Spark and Apache Spark.
121+
if (wrapped.relation.fileFormat.isInstanceOf[JsonFileFormat] &&
122+
wrapped.relation.location.getClass.getCanonicalName() ==
123+
"com.databricks.sql.transaction.tahoe.DeltaLogFileIndex") {
124+
this.entirePlanWillNotWork("Plans that read Delta Index JSON files can not run " +
125+
"any part of the plan on the GPU!")
126+
}
127+
ScanExecShims.tagGpuFileSourceScanExecSupport(this)
128+
}
129+
130+
override def convertToCpu(): SparkPlan = {
131+
val cpu = wrapped.copy(partitionFilters = partitionFilters, dataFilters = dataFilters)
132+
cpu.copyTagsFrom(wrapped)
133+
cpu
134+
}
135+
136+
override def convertToGpu(): GpuExec = {
137+
val sparkSession = wrapped.relation.sparkSession
138+
val options = wrapped.relation.options
139+
val newRelation = HadoopFsRelation(
140+
wrapped.relation.location,
141+
wrapped.relation.partitionSchema,
142+
wrapped.relation.dataSchema,
143+
wrapped.relation.bucketSpec,
144+
GpuFileSourceScanExec.convertFileFormat(wrapped.relation, conf),
145+
options)(sparkSession)
146+
147+
GpuFileSourceScanExec(
148+
newRelation,
149+
wrapped.output,
150+
wrapped.requiredSchema,
151+
partitionFilters,
152+
wrapped.optionalBucketSet,
153+
// TODO: Does Databricks have coalesced bucketing implemented?
154+
None,
155+
dataFilters,
156+
wrapped.tableIdentifier,
157+
wrapped.disableBucketedScan,
158+
queryUsesInputFile = false)(conf)
159+
}
160+
}

‎sql-plugin/src/main/spark400db173/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
/*
2+
* Copyright (c) 2022-2026, NVIDIA CORPORATION.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
/*** spark-rapids-shim-json-lines
18+
{"spark": "400db173"}
19+
spark-rapids-shim-json-lines ***/
20+
package com.nvidia.spark.rapids.shims
21+
22+
import com.nvidia.spark.rapids._
23+
24+
import org.apache.spark.sql.execution.FileSourceScanExec
25+
import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
26+
import org.apache.spark.sql.rapids.ExternalSource
27+
28+
class FileSourceScanExecMeta(plan: FileSourceScanExec,
29+
conf: RapidsConf,
30+
parent: Option[RapidsMeta[_, _, _]],
31+
rule: DataFromReplacementRule)
32+
extends FileSourceScanExecMetaBase(plan, conf, parent, rule) {
33+
34+
override def tagPlanForGpu(): Unit = {
35+
super.tagPlanForGpu()
36+
// DB-17.3 has no Delta provider. If the file format is a ParquetFileFormat subclass
37+
// (e.g. DeltaParquetFileFormat) not handled by ExternalSource, fall back the scan
38+
// to CPU so the subclass can handle format-specific features like deletion vectors
39+
// and skip_row columns.
40+
val fmtCls = wrapped.relation.fileFormat.getClass
41+
if (classOf[ParquetFileFormat].isAssignableFrom(fmtCls) &&
42+
fmtCls != classOf[ParquetFileFormat] &&
43+
!ExternalSource.isSupportedFormat(fmtCls)) {
44+
willNotWorkOnGpu(s"unsupported file format: ${fmtCls.getCanonicalName}")
45+
}
46+
}
47+
}

0 commit comments

Comments
 (0)