Skip to content

Commit

Permalink
Fix bug where AQE does not respect entirePlanWillNotWork (#6250)
Browse files Browse the repository at this point in the history
* Fix bug where AQE does not respect entirePlanWillNotWork

Signed-off-by: Andy Grove <andygrove@nvidia.com>

* use Set instead of Seq for getReasonsNotToReplaceEntirePlan
  • Loading branch information
andygrove authored Aug 9, 2022
1 parent d2b6a96 commit 8a50bc5
Showing 1 changed file with 7 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,10 @@ abstract class RapidsMeta[INPUT <: BASE, BASE, OUTPUT <: BASE](
*/
final def entirePlanWillNotWork(because: String): Unit = {
cannotReplaceAnyOfPlanReasons.get.add(because)
// recursively tag the plan so that AQE does not attempt
// to run any of the child query stages on the GPU
willNotWorkOnGpu(because)
childPlans.foreach(_.entirePlanWillNotWork(because))
}

final def shouldBeRemoved(because: String): Unit =
Expand Down Expand Up @@ -229,8 +233,8 @@ abstract class RapidsMeta[INPUT <: BASE, BASE, OUTPUT <: BASE](
* set means the entire plan is ok to be replaced, do the normal checking
* per exec and children.
*/
final def entirePlanExcludedReasons: Seq[String] = {
cannotReplaceAnyOfPlanReasons.getOrElse(mutable.Set.empty).toSeq
final def entirePlanExcludedReasons: Set[String] = {
cannotReplaceAnyOfPlanReasons.getOrElse(mutable.Set.empty).toSet
}

/**
Expand Down Expand Up @@ -609,7 +613,7 @@ abstract class SparkPlanMeta[INPUT <: SparkPlan](plan: INPUT,
wrapped.withNewChildren(childPlans.map(_.convertIfNeeded()))
}

def getReasonsNotToReplaceEntirePlan: Seq[String] = {
def getReasonsNotToReplaceEntirePlan: Set[String] = {
val childReasons = childPlans.flatMap(_.getReasonsNotToReplaceEntirePlan)
entirePlanExcludedReasons ++ childReasons
}
Expand Down

0 comments on commit 8a50bc5

Please sign in to comment.