Cleans up redundant code in proxy/internal RAPIDS Shuffle Managers (#…

…6030)

Signed-off-by: Alessandro Bellina <abellina@nvidia.com>

dist/unshimmed-common-from-spark311.txt

@@ -28,6 +28,6 @@ com/nvidia/spark/rapids/SparkShimVersion*
 com/nvidia/spark/rapids/SparkShims*
 com/nvidia/spark/udf/Plugin*
 org/apache/spark/sql/rapids/ProxyRapidsShuffleInternalManagerBase*
-org/apache/spark/sql/rapids/VisibleShuffleManager*
+org/apache/spark/sql/rapids/RapidsShuffleManagerLike*
 rapids/*.py
 rapids4spark-version-info.properties

sql-plugin/src/main/311-nondb/scala/org/apache/spark/sql/rapids/shims/spark311/RapidsShuffleInternalManager.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-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.
@@ -16,7 +16,7 @@
 
 package org.apache.spark.sql.rapids.shims.spark311
 
-import org.apache.spark.{SparkConf, TaskContext}
+import org.apache.spark.SparkConf
 import org.apache.spark.shuffle._
 import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, RapidsShuffleInternalManagerBase}
 
@@ -26,34 +26,8 @@ import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, Rapid
  *       `ShuffleManager` and `SortShuffleManager` classes.
  */
 class RapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends RapidsShuffleInternalManagerBase(conf, isDriver) {
-
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
-    getReaderInternal(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
+    extends RapidsShuffleInternalManagerBase(conf, isDriver)
 
 class ProxyRapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver) with ShuffleManager {
-
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter
-  ): ShuffleReader[K, C] = {
-    self.getReader(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
+    extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver)
+      with ShuffleManager

sql-plugin/src/main/312-nondb/scala/org/apache/spark/sql/rapids/shims/spark312/RapidsShuffleInternalManager.scala

@@ -1,5 +1,5 @@
 /*
- * 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.
@@ -16,7 +16,7 @@
 
 package org.apache.spark.sql.rapids.shims.spark312
 
-import org.apache.spark.{SparkConf, TaskContext}
+import org.apache.spark.SparkConf
 import org.apache.spark.shuffle._
 import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, RapidsShuffleInternalManagerBase}
 
@@ -26,35 +26,8 @@ import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, Rapid
  *       `ShuffleManager` and `SortShuffleManager` classes.
  */
 class RapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends RapidsShuffleInternalManagerBase(conf, isDriver) {
-
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
-    getReaderInternal(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
-
+    extends RapidsShuffleInternalManagerBase(conf, isDriver)
 
 class ProxyRapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver) with ShuffleManager {
-
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter
-  ): ShuffleReader[K, C] = {
-    self.getReader(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
+  extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver)
+    with ShuffleManager

sql-plugin/src/main/312db/scala/org/apache/spark/sql/rapids/shims/spark312db/RapidsShuffleInternalManager.scala

@@ -1,5 +1,5 @@
 /*
- * 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.
@@ -16,7 +16,7 @@
 
 package org.apache.spark.sql.rapids.shims.spark312db
 
-import org.apache.spark.{SparkConf, TaskContext}
+import org.apache.spark.SparkConf
 import org.apache.spark.shuffle._
 import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, RapidsShuffleInternalManagerBase}
 
@@ -26,34 +26,8 @@ import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, Rapid
  *       `ShuffleManager` and `SortShuffleManager` classes.
  */
 class RapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends RapidsShuffleInternalManagerBase(conf, isDriver) {
-
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
-    getReaderInternal(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
+    extends RapidsShuffleInternalManagerBase(conf, isDriver)
 
 class ProxyRapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver) with ShuffleManager {
-
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter
-  ): ShuffleReader[K, C] = {
-    self.getReader(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
+    extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver)
+      with ShuffleManager

sql-plugin/src/main/313/scala/org/apache/spark/sql/rapids/shims/spark313/RapidsShuffleInternalManager.scala

@@ -1,5 +1,5 @@
 /*
- * 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.
@@ -16,7 +16,7 @@
 
 package org.apache.spark.sql.rapids.shims.spark313
 
-import org.apache.spark.{SparkConf, TaskContext}
+import org.apache.spark.SparkConf
 import org.apache.spark.shuffle._
 import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, RapidsShuffleInternalManagerBase}
 
@@ -26,34 +26,8 @@ import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, Rapid
  *       `ShuffleManager` and `SortShuffleManager` classes.
  */
 class RapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends RapidsShuffleInternalManagerBase(conf, isDriver) {
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
-    getReaderInternal(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
-
+    extends RapidsShuffleInternalManagerBase(conf, isDriver)
 
 class ProxyRapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver) with ShuffleManager {
-
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter
-  ): ShuffleReader[K, C] = {
-    self.getReader(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
+    extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver)
+      with ShuffleManager

sql-plugin/src/main/314/scala/org/apache/spark/sql/rapids/shims/spark314/RapidsShuffleInternalManager.scala

@@ -16,7 +16,7 @@
 
 package org.apache.spark.sql.rapids.shims.spark314
 
-import org.apache.spark.{SparkConf, TaskContext}
+import org.apache.spark.SparkConf
 import org.apache.spark.shuffle._
 import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, RapidsShuffleInternalManagerBase}
 
@@ -26,34 +26,9 @@ import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, Rapid
  *       `ShuffleManager` and `SortShuffleManager` classes.
  */
 class RapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends RapidsShuffleInternalManagerBase(conf, isDriver) {
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
-    getReaderInternal(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
+    extends RapidsShuffleInternalManagerBase(conf, isDriver)
 
 
 class ProxyRapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver) with ShuffleManager {
-
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter
-  ): ShuffleReader[K, C] = {
-    self.getReader(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
+    extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver)
+      with ShuffleManager

sql-plugin/src/main/320/scala/org/apache/spark/sql/rapids/shims/spark320/RapidsShuffleInternalManager.scala

@@ -1,5 +1,5 @@
 /*
- * 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.
@@ -16,7 +16,7 @@
 
 package org.apache.spark.sql.rapids.shims.spark320
 
-import org.apache.spark.{SparkConf, TaskContext}
+import org.apache.spark.SparkConf
 import org.apache.spark.shuffle._
 import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, RapidsShuffleInternalManagerBase}
 
@@ -26,35 +26,8 @@ import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, Rapid
  *       `ShuffleManager` and `SortShuffleManager` classes.
  */
 class RapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends RapidsShuffleInternalManagerBase(conf, isDriver) {
-
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
-    getReaderInternal(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
-
+    extends RapidsShuffleInternalManagerBase(conf, isDriver)
 
 class ProxyRapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-  extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver) with ShuffleManager {
-
-  def getReader[K, C](
-    handle: ShuffleHandle,
-    startMapIndex: Int,
-    endMapIndex: Int,
-    startPartition: Int,
-    endPartition: Int,
-    context: TaskContext,
-    metrics: ShuffleReadMetricsReporter
-  ): ShuffleReader[K, C] = {
-    self.getReader(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
+  extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver)
+    with ShuffleManager

sql-plugin/src/main/321/scala/org/apache/spark/sql/rapids/shims/spark321/RapidsShuffleInternalManager.scala

@@ -1,5 +1,5 @@
 /*
- * 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.
@@ -16,7 +16,7 @@
 
 package org.apache.spark.sql.rapids.shims.spark321
 
-import org.apache.spark.{SparkConf, TaskContext}
+import org.apache.spark.SparkConf
 import org.apache.spark.shuffle._
 import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, RapidsShuffleInternalManagerBase}
 
@@ -26,35 +26,8 @@ import org.apache.spark.sql.rapids.{ProxyRapidsShuffleInternalManagerBase, Rapid
  *       `ShuffleManager` and `SortShuffleManager` classes.
  */
 class RapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-    extends RapidsShuffleInternalManagerBase(conf, isDriver) {
-
-  def getReader[K, C](
-      handle: ShuffleHandle,
-      startMapIndex: Int,
-      endMapIndex: Int,
-      startPartition: Int,
-      endPartition: Int,
-      context: TaskContext,
-      metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
-    getReaderInternal(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
-
+    extends RapidsShuffleInternalManagerBase(conf, isDriver)
 
 class ProxyRapidsShuffleInternalManager(conf: SparkConf, isDriver: Boolean)
-  extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver) with ShuffleManager {
-
-  def getReader[K, C](
-    handle: ShuffleHandle,
-    startMapIndex: Int,
-    endMapIndex: Int,
-    startPartition: Int,
-    endPartition: Int,
-    context: TaskContext,
-    metrics: ShuffleReadMetricsReporter
-  ): ShuffleReader[K, C] = {
-    self.getReader(handle, startMapIndex, endMapIndex, startPartition, endPartition, context,
-      metrics)
-  }
-}
+  extends ProxyRapidsShuffleInternalManagerBase(conf, isDriver)
+    with ShuffleManager