Fix the incomplete capture of SubqueryBroadcast

integration_tests/src/main/python/dpp_test.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2021, NVIDIA CORPORATION.
+# Copyright (c) 2021-2022, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -100,6 +100,25 @@ def fn(spark):
     WHERE dim.filter = {2}
     GROUP BY fact.key, fact.skey, fact.ex_key
     ''',
+    # This query checks the pattern of reused broadcast subquery: ReusedSubquery(SubqueryBroadcast(...))
+    # https://github.com/NVIDIA/spark-rapids/issues/4625
+    """
+    SELECT key, max(value)
+    FROM (
+        SELECT fact.key as key, fact.value as value
+        FROM {0} fact
+        JOIN {1} dim
+        ON fact.key = dim.key
+        WHERE dim.filter = {2}
+    UNION ALL
+        SELECT fact.key as key, fact.value as value
+        FROM {0} fact
+        JOIN {1} dim
+        ON fact.key = dim.key
+        WHERE dim.filter = {2}
+    )
+    GROUP BY key
+    """
 ]
 
 

sql-plugin/src/main/301db/scala/com/nvidia/spark/rapids/shims/v2/Spark30XdbShims.scala

@@ -151,15 +151,24 @@ abstract class Spark30XdbShims extends Spark30XdbShimsBase with Logging {
           // FileSourceScan is independent from the replacement of the partitionFilters. It is
           // possible that the FileSourceScan is on the CPU, while the dynamic partitionFilters
           // are on the GPU. And vice versa.
-          private lazy val partitionFilters = wrapped.partitionFilters.map { filter =>
-            filter.transformDown {
-              case dpe @ DynamicPruningExpression(inSub: InSubqueryExec)
-                if inSub.plan.isInstanceOf[SubqueryBroadcastExec] =>
-
-                val subBcMeta = GpuOverrides.wrapAndTagPlan(inSub.plan, conf)
-                subBcMeta.tagForExplain()
-                val gpuSubBroadcast = subBcMeta.convertIfNeeded().asInstanceOf[BaseSubqueryExec]
-                dpe.copy(inSub.copy(plan = gpuSubBroadcast))
+          private lazy val partitionFilters = {
+            val convertBroadcast = (bc: SubqueryBroadcastExec) => {
+              val meta = GpuOverrides.wrapAndTagPlan(bc, conf)
+              meta.tagForExplain()
+              meta.convertIfNeeded().asInstanceOf[BaseSubqueryExec]
+            }
+            wrapped.partitionFilters.map { filter =>
+              filter.transformDown {
+                case dpe @ DynamicPruningExpression(inSub: InSubqueryExec) =>
+                  inSub.plan match {
+                    case bc: SubqueryBroadcastExec =>
+                      dpe.copy(inSub.copy(plan = convertBroadcast(bc)))
+                    case reuse @ ReusedSubqueryExec(bc: SubqueryBroadcastExec) =>
+                      dpe.copy(inSub.copy(plan = reuse.copy(convertBroadcast(bc))))
+                    case _ =>
+                      dpe
+                  }
+              }
             }
           }
 

sql-plugin/src/main/301until310-nondb/scala/com/nvidia/spark/rapids/shims/v2/Spark30XShims.scala

@@ -130,15 +130,24 @@ abstract class Spark30XShims extends Spark301until320Shims with Logging {
           // FileSourceScan is independent from the replacement of the partitionFilters. It is
           // possible that the FileSourceScan is on the CPU, while the dynamic partitionFilters
           // are on the GPU. And vice versa.
-          private lazy val partitionFilters = wrapped.partitionFilters.map { filter =>
-            filter.transformDown {
-              case dpe @ DynamicPruningExpression(inSub: InSubqueryExec)
-                if inSub.plan.isInstanceOf[SubqueryBroadcastExec] =>
-
-                val subBcMeta = GpuOverrides.wrapAndTagPlan(inSub.plan, conf)
-                subBcMeta.tagForExplain()
-                val gpuSubBroadcast = subBcMeta.convertIfNeeded().asInstanceOf[BaseSubqueryExec]
-                dpe.copy(inSub.copy(plan = gpuSubBroadcast))
+          private lazy val partitionFilters = {
+            val convertBroadcast = (bc: SubqueryBroadcastExec) => {
+              val meta = GpuOverrides.wrapAndTagPlan(bc, conf)
+              meta.tagForExplain()
+              meta.convertIfNeeded().asInstanceOf[BaseSubqueryExec]
+            }
+            wrapped.partitionFilters.map { filter =>
+              filter.transformDown {
+                case dpe @ DynamicPruningExpression(inSub: InSubqueryExec) =>
+                  inSub.plan match {
+                    case bc: SubqueryBroadcastExec =>
+                      dpe.copy(inSub.copy(plan = convertBroadcast(bc)))
+                    case reuse @ ReusedSubqueryExec(bc: SubqueryBroadcastExec) =>
+                      dpe.copy(inSub.copy(plan = reuse.copy(convertBroadcast(bc))))
+                    case _ =>
+                      dpe
+                  }
+              }
             }
           }
 

sql-plugin/src/main/311until320-nondb/scala/com/nvidia/spark/rapids/shims/v2/Spark31XShims.scala

@@ -370,15 +370,24 @@ abstract class Spark31XShims extends Spark301until320Shims with Logging {
           // FileSourceScan is independent from the replacement of the partitionFilters. It is
           // possible that the FileSourceScan is on the CPU, while the dynamic partitionFilters
           // are on the GPU. And vice versa.
-          private lazy val partitionFilters = wrapped.partitionFilters.map { filter =>
-            filter.transformDown {
-              case dpe @ DynamicPruningExpression(inSub: InSubqueryExec)
-                if inSub.plan.isInstanceOf[SubqueryBroadcastExec] =>
-
-                val subBcMeta = GpuOverrides.wrapAndTagPlan(inSub.plan, conf)
-                subBcMeta.tagForExplain()
-                val gpuSubBroadcast = subBcMeta.convertIfNeeded().asInstanceOf[BaseSubqueryExec]
-                dpe.copy(inSub.copy(plan = gpuSubBroadcast))
+          private lazy val partitionFilters = {
+            val convertBroadcast = (bc: SubqueryBroadcastExec) => {
+              val meta = GpuOverrides.wrapAndTagPlan(bc, conf)
+              meta.tagForExplain()
+              meta.convertIfNeeded().asInstanceOf[BaseSubqueryExec]
+            }
+            wrapped.partitionFilters.map { filter =>
+              filter.transformDown {
+                case dpe @ DynamicPruningExpression(inSub: InSubqueryExec) =>
+                  inSub.plan match {
+                    case bc: SubqueryBroadcastExec =>
+                      dpe.copy(inSub.copy(plan = convertBroadcast(bc)))
+                    case reuse @ ReusedSubqueryExec(bc: SubqueryBroadcastExec) =>
+                      dpe.copy(inSub.copy(plan = reuse.copy(convertBroadcast(bc))))
+                    case _ =>
+                      dpe
+                  }
+              }
             }
           }
 

sql-plugin/src/main/31xdb/scala/com/nvidia/spark/rapids/shims/v2/Spark31XdbShims.scala

@@ -382,15 +382,24 @@ abstract class Spark31XdbShims extends Spark31XdbShimsBase with Logging {
           // FileSourceScan is independent from the replacement of the partitionFilters. It is
           // possible that the FileSourceScan is on the CPU, while the dynamic partitionFilters
           // are on the GPU. And vice versa.
-          private lazy val partitionFilters = wrapped.partitionFilters.map { filter =>
-            filter.transformDown {
-              case dpe @ DynamicPruningExpression(inSub: InSubqueryExec)
-                if inSub.plan.isInstanceOf[SubqueryBroadcastExec] =>
-
-                val subBcMeta = GpuOverrides.wrapAndTagPlan(inSub.plan, conf)
-                subBcMeta.tagForExplain()
-                val gpuSubBroadcast = subBcMeta.convertIfNeeded().asInstanceOf[BaseSubqueryExec]
-                dpe.copy(inSub.copy(plan = gpuSubBroadcast))
+          private lazy val partitionFilters = {
+            val convertBroadcast = (bc: SubqueryBroadcastExec) => {
+              val meta = GpuOverrides.wrapAndTagPlan(bc, conf)
+              meta.tagForExplain()
+              meta.convertIfNeeded().asInstanceOf[BaseSubqueryExec]
+            }
+            wrapped.partitionFilters.map { filter =>
+              filter.transformDown {
+                case dpe @ DynamicPruningExpression(inSub: InSubqueryExec) =>
+                  inSub.plan match {
+                    case bc: SubqueryBroadcastExec =>
+                      dpe.copy(inSub.copy(plan = convertBroadcast(bc)))
+                    case reuse @ ReusedSubqueryExec(bc: SubqueryBroadcastExec) =>
+                      dpe.copy(inSub.copy(plan = reuse.copy(convertBroadcast(bc))))
+                    case _ =>
+                      dpe
+                  }
+              }
             }
           }
 

sql-plugin/src/main/320+/scala/com/nvidia/spark/rapids/shims/v2/Spark320PlusShims.scala

@@ -44,7 +44,7 @@ import org.apache.spark.sql.catalyst.plans.physical.{BroadcastMode, Partitioning
 import org.apache.spark.sql.catalyst.trees.TreeNode
 import org.apache.spark.sql.catalyst.util.DateFormatter
 import org.apache.spark.sql.connector.read.{Scan, SupportsRuntimeFiltering}
-import org.apache.spark.sql.execution.{BaseSubqueryExec, CommandResultExec, FileSourceScanExec, InSubqueryExec, PartitionedFileUtil, SparkPlan, SubqueryBroadcastExec}
+import org.apache.spark.sql.execution.{BaseSubqueryExec, CommandResultExec, FileSourceScanExec, InSubqueryExec, PartitionedFileUtil, ReusedSubqueryExec, SparkPlan, SubqueryBroadcastExec}
 import org.apache.spark.sql.execution.adaptive._
 import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
 import org.apache.spark.sql.execution.command._
@@ -534,15 +534,24 @@ trait Spark320PlusShims extends SparkShims with RebaseShims with Logging {
           // FileSourceScan is independent from the replacement of the partitionFilters. It is
           // possible that the FileSourceScan is on the CPU, while the dynamic partitionFilters
           // are on the GPU. And vice versa.
-          private lazy val partitionFilters = wrapped.partitionFilters.map { filter =>
-            filter.transformDown {
-              case dpe@DynamicPruningExpression(inSub: InSubqueryExec)
-                if inSub.plan.isInstanceOf[SubqueryBroadcastExec] =>
-
-                val subBcMeta = GpuOverrides.wrapAndTagPlan(inSub.plan, conf)
-                subBcMeta.tagForExplain()
-                val gpuSubBroadcast = subBcMeta.convertIfNeeded().asInstanceOf[BaseSubqueryExec]
-                dpe.copy(inSub.copy(plan = gpuSubBroadcast))
+          private lazy val partitionFilters = {
+            val convertBroadcast = (bc: SubqueryBroadcastExec) => {
+              val meta = GpuOverrides.wrapAndTagPlan(bc, conf)
+              meta.tagForExplain()
+              meta.convertIfNeeded().asInstanceOf[BaseSubqueryExec]
+            }
+            wrapped.partitionFilters.map { filter =>
+              filter.transformDown {
+                case dpe @ DynamicPruningExpression(inSub: InSubqueryExec) =>
+                  inSub.plan match {
+                    case bc: SubqueryBroadcastExec =>
+                      dpe.copy(inSub.copy(plan = convertBroadcast(bc)))
+                    case reuse @ ReusedSubqueryExec(bc: SubqueryBroadcastExec) =>
+                      dpe.copy(inSub.copy(plan = reuse.copy(convertBroadcast(bc))))
+                    case _ =>
+                      dpe
+                  }
+              }
             }
           }