Use Shims v2 for ShuffledBatchRDD (#3459)

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>
NVIDIA · Sep 14, 2021 · 126c7d3 · 126c7d3
1 parent 5fcecae
commit 126c7d3
Show file tree

Hide file tree

Showing 30 changed files with 494 additions and 583 deletions.
diff --git a/pom.xml b/pom.xml
@@ -102,6 +102,7 @@
                                 <phase>generate-sources</phase>
                                 <configuration>
                                     <sources>
+                                        <source>${project.basedir}/src/main/301until310-nondb/scala</source>
                                         <source>${project.basedir}/src/main/301until320-nondb/scala</source>
                                         <source>${project.basedir}/src/main/301until320-all/scala</source>
                                     </sources>
@@ -136,6 +137,7 @@
                                 <phase>generate-sources</phase>
                                 <configuration>
                                     <sources>
+                                        <source>${project.basedir}/src/main/301until310-nondb/scala</source>
                                         <source>${project.basedir}/src/main/301until320-nondb/scala</source>
                                         <source>${project.basedir}/src/main/301until320-all/scala</source>
                                     </sources>
@@ -179,6 +181,7 @@
                                 <phase>generate-sources</phase>
                                 <configuration>
                                     <sources>
+                                        <source>${project.basedir}/src/main/301until310-nondb/scala</source>
                                         <source>${project.basedir}/src/main/301until320-nondb/scala</source>
                                         <source>${project.basedir}/src/main/301until320-all/scala</source>
                                     </sources>
@@ -217,6 +220,7 @@
                                 <phase>generate-sources</phase>
                                 <configuration>
                                     <sources>
+                                        <source>${project.basedir}/src/main/301until310-nondb/scala</source>
                                         <source>${project.basedir}/src/main/301until320-nondb/scala</source>
                                         <source>${project.basedir}/src/main/301until320-all/scala</source>
                                     </sources>
@@ -257,6 +261,7 @@
                                     <sources>
                                         <source>${project.basedir}/src/main/301until320-nondb/scala</source>
                                         <source>${project.basedir}/src/main/301until320-all/scala</source>
+                                        <source>${project.basedir}/src/main/311until320-all/scala</source>
                                         <source>${project.basedir}/src/main/311+-all/scala</source>
                                         <source>${project.basedir}/src/main/311+-apache/scala</source>
                                     </sources>
@@ -363,6 +368,7 @@
                                     <sources>
                                         <source>${project.basedir}/src/main/311db/scala</source>
                                         <source>${project.basedir}/src/main/301until320-all/scala</source>
+                                        <source>${project.basedir}/src/main/311until320-all/scala</source>
                                         <source>${project.basedir}/src/main/311+-all/scala</source>
                                         <source>${project.basedir}/src/main/311+-apache/scala</source>
                                     </sources>
@@ -402,6 +408,7 @@
                                     <sources>
                                         <source>${project.basedir}/src/main/301until320-nondb/scala</source>
                                         <source>${project.basedir}/src/main/301until320-all/scala</source>
+                                        <source>${project.basedir}/src/main/311until320-all/scala</source>
                                         <source>${project.basedir}/src/main/311+-all/scala</source>
                                         <source>${project.basedir}/src/main/311+-apache/scala</source>
                                     </sources>
@@ -442,6 +449,7 @@
                                     <sources>
                                         <source>${project.basedir}/src/main/301until320-nondb/scala</source>
                                         <source>${project.basedir}/src/main/301until320-all/scala</source>
+                                        <source>${project.basedir}/src/main/311until320-all/scala</source>
                                         <source>${project.basedir}/src/main/311+-all/scala</source>
                                         <source>${project.basedir}/src/main/311+-apache/scala</source>
                                     </sources>
@@ -520,6 +528,7 @@
                                         <source>${project.basedir}/src/main/301until320-nondb/scala</source>
                                         <source>${project.basedir}/src/main/311cdh/scala</source>
                                         <source>${project.basedir}/src/main/301until320-all/scala</source>
+                                        <source>${project.basedir}/src/main/311until320-all/scala</source>
                                         <source>${project.basedir}/src/main/311+-all/scala</source>
                                     </sources>
                                 </configuration>

diff --git a/shims/spark301/src/main/scala/com/nvidia/spark/rapids/shims/spark301/SparkBaseShims.scala b/shims/spark301/src/main/scala/com/nvidia/spark/rapids/shims/spark301/SparkBaseShims.scala
@@ -52,11 +52,11 @@ import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNes
 import org.apache.spark.sql.execution.python.{AggregateInPandasExec, ArrowEvalPythonExec, FlatMapGroupsInPandasExec, MapInPandasExec, WindowInPandasExec}
 import org.apache.spark.sql.execution.window.WindowExecBase
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, GpuStringReplace, GpuTimeSub, ShuffleManagerShimBase}
+import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, GpuStringReplace, GpuTimeSub}
 import org.apache.spark.sql.rapids.execution.{GpuBroadcastExchangeExecBase, GpuBroadcastNestedLoopJoinExecBase, GpuShuffleExchangeExecBase, JoinTypeChecks}
 import org.apache.spark.sql.rapids.execution.python.GpuPythonUDF
 import org.apache.spark.sql.rapids.execution.python.shims.spark301._
-import org.apache.spark.sql.rapids.shims.spark301.{GpuSchemaUtils, ShuffleManagerShim}
+import org.apache.spark.sql.rapids.shims.spark301.GpuSchemaUtils
 import org.apache.spark.sql.sources.BaseRelation
 import org.apache.spark.sql.types._
 import org.apache.spark.storage.{BlockId, BlockManagerId}
@@ -380,10 +380,6 @@ abstract class SparkBaseShims extends Spark30XShims {
     GpuJoinUtils.getGpuBuildSide(join.buildSide)
   }
 
-  override def getShuffleManagerShims(): ShuffleManagerShimBase = {
-    new ShuffleManagerShim
-  }
-
   override def getPartitionFileNames(
       partitions: Seq[PartitionDirectory]): Seq[String] = {
     val files = partitions.flatMap(partition => partition.files)

diff --git a/...ark301/src/main/scala/org/apache/spark/sql/rapids/shims/spark301/ShuffleManagerShim.scala b/...ark301/src/main/scala/org/apache/spark/sql/rapids/shims/spark301/ShuffleManagerShim.scala
diff --git a/...s/spark301db/src/main/scala/com/nvidia/spark/rapids/shims/spark301db/SparkBaseShims.scala b/...s/spark301db/src/main/scala/com/nvidia/spark/rapids/shims/spark301db/SparkBaseShims.scala
@@ -54,7 +54,7 @@ import org.apache.spark.sql.execution.joins._
 import org.apache.spark.sql.execution.python._
 import org.apache.spark.sql.execution.window.WindowExecBase
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, GpuStringReplace, GpuTimeSub, ShuffleManagerShimBase}
+import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, GpuStringReplace, GpuTimeSub}
 import org.apache.spark.sql.rapids.execution.{GpuBroadcastExchangeExecBase, GpuBroadcastNestedLoopJoinExecBase, GpuShuffleExchangeExecBase, JoinTypeChecks, TrampolineUtil}
 import org.apache.spark.sql.rapids.execution.python.{GpuFlatMapGroupsInPandasExecMeta, GpuPythonUDF}
 import org.apache.spark.sql.rapids.execution.python.shims.spark301db._
@@ -403,10 +403,6 @@ abstract class SparkBaseShims extends Spark30XShims {
     GpuJoinUtils.getGpuBuildSide(join.buildSide)
   }
 
-  override def getShuffleManagerShims(): ShuffleManagerShimBase = {
-    new ShuffleManagerShim
-  }
-
   override def getPartitionFileNames(
       partitions: Seq[PartitionDirectory]): Seq[String] = {
     val files = partitions.flatMap(partition => partition.files)

diff --git a/...01db/src/main/scala/org/apache/spark/sql/rapids/shims/spark301db/ShuffleManagerShim.scala b/...01db/src/main/scala/org/apache/spark/sql/rapids/shims/spark301db/ShuffleManagerShim.scala
diff --git a/shims/spark302/src/main/scala/com/nvidia/spark/rapids/shims/spark302/SparkBaseShims.scala b/shims/spark302/src/main/scala/com/nvidia/spark/rapids/shims/spark302/SparkBaseShims.scala
@@ -52,11 +52,11 @@ import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNes
 import org.apache.spark.sql.execution.python.{AggregateInPandasExec, ArrowEvalPythonExec, FlatMapGroupsInPandasExec, MapInPandasExec, WindowInPandasExec}
 import org.apache.spark.sql.execution.window.WindowExecBase
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, GpuStringReplace, GpuTimeSub, ShuffleManagerShimBase}
+import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, GpuStringReplace, GpuTimeSub}
 import org.apache.spark.sql.rapids.execution.{GpuBroadcastExchangeExecBase, GpuBroadcastNestedLoopJoinExecBase, GpuShuffleExchangeExecBase, JoinTypeChecks}
 import org.apache.spark.sql.rapids.execution.python.GpuPythonUDF
 import org.apache.spark.sql.rapids.execution.python.shims.spark302._
-import org.apache.spark.sql.rapids.shims.spark302.{GpuSchemaUtils, ShuffleManagerShim}
+import org.apache.spark.sql.rapids.shims.spark302.GpuSchemaUtils
 import org.apache.spark.sql.sources.BaseRelation
 import org.apache.spark.sql.types._
 import org.apache.spark.storage.{BlockId, BlockManagerId}
@@ -380,10 +380,6 @@ abstract class SparkBaseShims extends Spark30XShims {
     GpuJoinUtils.getGpuBuildSide(join.buildSide)
   }
 
-  override def getShuffleManagerShims(): ShuffleManagerShimBase = {
-    new ShuffleManagerShim
-  }
-
   override def getPartitionFileNames(
       partitions: Seq[PartitionDirectory]): Seq[String] = {
     val files = partitions.flatMap(partition => partition.files)

diff --git a/...ark302/src/main/scala/org/apache/spark/sql/rapids/shims/spark302/ShuffleManagerShim.scala b/...ark302/src/main/scala/org/apache/spark/sql/rapids/shims/spark302/ShuffleManagerShim.scala
diff --git a/shims/spark303/src/main/scala/com/nvidia/spark/rapids/shims/spark303/SparkBaseShims.scala b/shims/spark303/src/main/scala/com/nvidia/spark/rapids/shims/spark303/SparkBaseShims.scala
@@ -52,11 +52,11 @@ import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNes
 import org.apache.spark.sql.execution.python.{AggregateInPandasExec, ArrowEvalPythonExec, FlatMapGroupsInPandasExec, MapInPandasExec, WindowInPandasExec}
 import org.apache.spark.sql.execution.window.WindowExecBase
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, GpuStringReplace, GpuTimeSub, ShuffleManagerShimBase}
+import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, GpuStringReplace, GpuTimeSub}
 import org.apache.spark.sql.rapids.execution.{GpuBroadcastExchangeExecBase, GpuBroadcastNestedLoopJoinExecBase, GpuShuffleExchangeExecBase, JoinTypeChecks}
 import org.apache.spark.sql.rapids.execution.python.GpuPythonUDF
 import org.apache.spark.sql.rapids.execution.python.shims.spark303._
-import org.apache.spark.sql.rapids.shims.spark303.{GpuSchemaUtils, ShuffleManagerShim}
+import org.apache.spark.sql.rapids.shims.spark303.GpuSchemaUtils
 import org.apache.spark.sql.sources.BaseRelation
 import org.apache.spark.sql.types._
 import org.apache.spark.storage.{BlockId, BlockManagerId}
@@ -380,10 +380,6 @@ abstract class SparkBaseShims extends Spark30XShims {
     GpuJoinUtils.getGpuBuildSide(join.buildSide)
   }
 
-  override def getShuffleManagerShims(): ShuffleManagerShimBase = {
-    new ShuffleManagerShim
-  }
-
   override def getPartitionFileNames(
       partitions: Seq[PartitionDirectory]): Seq[String] = {
     val files = partitions.flatMap(partition => partition.files)

diff --git a/...ark303/src/main/scala/org/apache/spark/sql/rapids/shims/spark303/ShuffleManagerShim.scala b/...ark303/src/main/scala/org/apache/spark/sql/rapids/shims/spark303/ShuffleManagerShim.scala
diff --git a/shims/spark304/src/main/scala/com/nvidia/spark/rapids/shims/spark304/SparkBaseShims.scala b/shims/spark304/src/main/scala/com/nvidia/spark/rapids/shims/spark304/SparkBaseShims.scala
@@ -51,11 +51,11 @@ import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNes
 import org.apache.spark.sql.execution.python.{AggregateInPandasExec, ArrowEvalPythonExec, FlatMapGroupsInPandasExec, MapInPandasExec, WindowInPandasExec}
 import org.apache.spark.sql.execution.window.WindowExecBase
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, GpuStringReplace, GpuTimeSub, ShuffleManagerShimBase}
+import org.apache.spark.sql.rapids.{GpuFileSourceScanExec, GpuStringReplace, GpuTimeSub}
 import org.apache.spark.sql.rapids.execution.{GpuBroadcastExchangeExecBase, GpuBroadcastNestedLoopJoinExecBase, GpuShuffleExchangeExecBase, JoinTypeChecks}
 import org.apache.spark.sql.rapids.execution.python.GpuPythonUDF
 import org.apache.spark.sql.rapids.execution.python.shims.spark304._
-import org.apache.spark.sql.rapids.shims.spark304.{GpuSchemaUtils, ShuffleManagerShim}
+import org.apache.spark.sql.rapids.shims.spark304.GpuSchemaUtils
 import org.apache.spark.sql.sources.BaseRelation
 import org.apache.spark.sql.types._
 import org.apache.spark.storage.{BlockId, BlockManagerId}
@@ -380,10 +380,6 @@ abstract class SparkBaseShims extends Spark30XShims {
     GpuJoinUtils.getGpuBuildSide(join.buildSide)
   }
 
-  override def getShuffleManagerShims(): ShuffleManagerShimBase = {
-    new ShuffleManagerShim
-  }
-
   override def getPartitionFileNames(
       partitions: Seq[PartitionDirectory]): Seq[String] = {
     val files = partitions.flatMap(partition => partition.files)

diff --git a/...ark304/src/main/scala/org/apache/spark/sql/rapids/shims/spark304/ShuffleManagerShim.scala b/...ark304/src/main/scala/org/apache/spark/sql/rapids/shims/spark304/ShuffleManagerShim.scala
diff --git a/shims/spark311/src/main/scala/com/nvidia/spark/rapids/shims/spark311/SparkBaseShims.scala b/shims/spark311/src/main/scala/com/nvidia/spark/rapids/shims/spark311/SparkBaseShims.scala
@@ -550,10 +550,6 @@ abstract class SparkBaseShims extends Spark30XShims {
     GpuJoinUtils.getGpuBuildSide(join.buildSide)
   }
 
-  override def getShuffleManagerShims(): ShuffleManagerShimBase = {
-    new ShuffleManagerShim
-  }
-
   override def getPartitionFileNames(
       partitions: Seq[PartitionDirectory]): Seq[String] = {
     val files = partitions.flatMap(partition => partition.files)