[AURON #1875] Minor refactor: Move legacy Spark version compatibility methods to Shims.scala

Author: yew1eb

spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala

@@ -18,9 +18,7 @@ package org.apache.spark.sql.auron
 
 import java.io.File
 import java.util.UUID
-
 import scala.collection.mutable
-
 import org.apache.commons.lang3.reflect.FieldUtils
 import org.apache.spark.{OneToOneDependency, ShuffleDependency, SparkContext, SparkEnv, SparkException, TaskContext}
 import org.apache.spark.internal.Logging
@@ -95,7 +93,7 @@ import org.apache.spark.sql.execution.auron.plan.NativeWindowBase
 import org.apache.spark.sql.execution.auron.plan.NativeWindowExec
 import org.apache.spark.sql.execution.auron.shuffle.{AuronBlockStoreShuffleReaderBase, AuronRssShuffleManagerBase, RssPartitionWriterBase}
 import org.apache.spark.sql.execution.datasources.PartitionedFile
-import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec}
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec, ShuffleExchangeExec}
 import org.apache.spark.sql.execution.joins.auron.plan.NativeBroadcastJoinExec
 import org.apache.spark.sql.execution.joins.auron.plan.NativeShuffledHashJoinExecProvider
 import org.apache.spark.sql.execution.joins.auron.plan.NativeSortMergeJoinExecProvider
@@ -108,11 +106,11 @@ import org.apache.spark.sql.types.StringType
 import org.apache.spark.status.ElementTrackingStore
 import org.apache.spark.storage.BlockManagerId
 import org.apache.spark.storage.FileSegment
-
-import org.apache.auron.{protobuf => pb, sparkver}
+import org.apache.auron.{sparkver, protobuf => pb}
 import org.apache.auron.common.AuronBuildInfo
 import org.apache.auron.metric.SparkMetricNode
 import org.apache.auron.spark.ui.AuronBuildInfoEvent
+import org.apache.spark.sql.execution.joins.ShuffledHashJoinExec
 
 class ShimsImpl extends Shims with Logging {
 
@@ -1035,6 +1033,18 @@ class ShimsImpl extends Shims with Logging {
   override def getAdaptiveInputPlan(exec: AdaptiveSparkPlanExec): SparkPlan = {
     exec.inputPlan
   }
+
+  @sparkver(" 3.2 / 3.3 / 3.4 / 3.5")
+  override def getIsSkewJoinFromSHJ(exec: ShuffledHashJoinExec): Boolean = exec.isSkewJoin
+
+  @sparkver("3.0 / 3.1")
+  override def getIsSkewJoinFromSHJ(exec: ShuffledHashJoinExec): Boolean = false
+
+  @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
+  override def getShuffleOrigin(exec: ShuffleExchangeExec): Option[Any] = Some(exec.shuffleOrigin)
+
+  @sparkver("3.0")
+  override def getShuffleOrigin(exec: ShuffleExchangeExec): Option[Any] = None
 }
 
 case class ForceNativeExecutionWrapper(override val child: SparkPlan)

spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala

@@ -79,7 +79,6 @@ import org.apache.auron.metric.SparkMetricNode
 import org.apache.auron.protobuf.EmptyPartitionsExecNode
 import org.apache.auron.protobuf.PhysicalPlanNode
 import org.apache.auron.spark.configuration.SparkAuronConfiguration
-import org.apache.auron.sparkver
 
 object AuronConverters extends Logging {
   def enableScan: Boolean =
@@ -421,21 +420,9 @@ object AuronConverters extends Logging {
     Shims.get.createNativeShuffleExchangeExec(
       outputPartitioning,
       addRenameColumnsExec(convertedChild),
-      getShuffleOrigin(exec))
+      Shims.get.getShuffleOrigin(exec))
   }
 
-  @sparkver(" 3.2 / 3.3 / 3.4 / 3.5")
-  def getIsSkewJoinFromSHJ(exec: ShuffledHashJoinExec): Boolean = exec.isSkewJoin
-
-  @sparkver("3.0 / 3.1")
-  def getIsSkewJoinFromSHJ(exec: ShuffledHashJoinExec): Boolean = false
-
-  @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5")
-  def getShuffleOrigin(exec: ShuffleExchangeExec): Option[Any] = Some(exec.shuffleOrigin)
-
-  @sparkver("3.0")
-  def getShuffleOrigin(exec: ShuffleExchangeExec): Option[Any] = None
-
   def convertFileSourceScanExec(exec: FileSourceScanExec): SparkPlan = {
     val (
       relation,
@@ -624,7 +611,7 @@ object AuronConverters extends Logging {
           case BuildLeft => org.apache.spark.sql.execution.auron.plan.BuildLeft
           case BuildRight => org.apache.spark.sql.execution.auron.plan.BuildRight
         },
-        getIsSkewJoinFromSHJ(exec))
+        Shims.get.getIsSkewJoinFromSHJ(exec))
 
     } catch {
       case _ if sparkAuronConfig.getBoolean(SparkAuronConfiguration.FORCE_SHUFFLED_HASH_JOIN) =>

spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala

@@ -47,7 +47,8 @@ import org.apache.spark.sql.execution.auron.plan.NativeBroadcastJoinBase
 import org.apache.spark.sql.execution.auron.plan.NativeSortMergeJoinBase
 import org.apache.spark.sql.execution.auron.shuffle.RssPartitionWriterBase
 import org.apache.spark.sql.execution.datasources.PartitionedFile
-import org.apache.spark.sql.execution.exchange.BroadcastExchangeLike
+import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.joins.ShuffledHashJoinExec
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.hive.execution.InsertIntoHiveTable
 import org.apache.spark.sql.types.DataType
@@ -260,6 +261,10 @@ abstract class Shims {
   def postTransform(plan: SparkPlan, sc: SparkContext): Unit = {}
 
   def getAdaptiveInputPlan(exec: AdaptiveSparkPlanExec): SparkPlan
+
+  def getIsSkewJoinFromSHJ(exec: ShuffledHashJoinExec): Boolean
+
+  def getShuffleOrigin(exec: ShuffleExchangeExec): Option[Any]
 }
 
 object Shims {