Add ability to ignore tests depending on spark shim version (NVIDIA#504)

* Add ability to ignore tests depending on spark shim version Signed-off-by: Andy Grove <andygrove@nvidia.com> * better version check Signed-off-by: Andy Grove <andygrove@nvidia.com> * remove duplicate VERSION strings and use case classes instead Signed-off-by: Andy Grove <andygrove@nvidia.com> * fix formatting Signed-off-by: Andy Grove <andygrove@nvidia.com>
nartal1 · Aug 4, 2020 · 296da89 · 296da89
1 parent cb75393
commit 296da89
Show file tree

Hide file tree

Showing 10 changed files with 56 additions and 17 deletions.
diff --git a/shims/spark300/src/main/scala/com/nvidia/spark/rapids/shims/spark300/Spark300Shims.scala b/shims/spark300/src/main/scala/com/nvidia/spark/rapids/shims/spark300/Spark300Shims.scala
@@ -41,6 +41,8 @@ import org.apache.spark.unsafe.types.CalendarInterval
 
 class Spark300Shims extends SparkShims {
 
+  override def getSparkShimVersion: ShimVersion = SparkShimServiceProvider.VERSION
+
   override def getScalaUDFAsExpression(
       function: AnyRef,
       dataType: DataType,

diff --git a/...k300/src/main/scala/com/nvidia/spark/rapids/shims/spark300/SparkShimServiceProvider.scala b/...k300/src/main/scala/com/nvidia/spark/rapids/shims/spark300/SparkShimServiceProvider.scala
@@ -16,14 +16,16 @@
 
 package com.nvidia.spark.rapids.shims.spark300
 
-import com.nvidia.spark.rapids.SparkShims
+import com.nvidia.spark.rapids.{SparkShims, SparkShimVersion}
 
-class SparkShimServiceProvider extends com.nvidia.spark.rapids.SparkShimServiceProvider {
+object SparkShimServiceProvider {
+  val VERSION = SparkShimVersion(3, 0, 0)
+}
 
-  val VERSIONNAME = "3.0.0"
+class SparkShimServiceProvider extends com.nvidia.spark.rapids.SparkShimServiceProvider {
 
   def matchesVersion(version: String): Boolean = {
-    version == VERSIONNAME
+    version == SparkShimServiceProvider.VERSION.toString()
   }
 
   def buildShim: SparkShims = {

diff --git a/.../spark300db/src/main/scala/com/nvidia/spark/rapids/shims/spark300db/Spark300dbShims.scala b/.../spark300db/src/main/scala/com/nvidia/spark/rapids/shims/spark300db/Spark300dbShims.scala
@@ -37,6 +37,8 @@ import org.apache.spark.storage.{BlockId, BlockManagerId}
 
 class Spark300dbShims extends Spark300Shims {
 
+  override def getSparkShimVersion: ShimVersion = SparkShimVersionProvider.VERSION
+
   override def getGpuBroadcastNestedLoopJoinShim(
       left: SparkPlan,
       right: SparkPlan,

diff --git a/...db/src/main/scala/com/nvidia/spark/rapids/shims/spark300db/SparkShimServiceProvider.scala b/...db/src/main/scala/com/nvidia/spark/rapids/shims/spark300db/SparkShimServiceProvider.scala
@@ -16,14 +16,17 @@
 
 package com.nvidia.spark.rapids.shims.spark300db
 
-import com.nvidia.spark.rapids.SparkShims
+import com.nvidia.spark.rapids.{DatabricksShimVersion, SparkShims}
 
-class SparkShimServiceProvider extends com.nvidia.spark.rapids.SparkShimServiceProvider {
+object SparkShimServiceProvider {
+  val VERSION = DatabricksShimVersion(3, 0, 0)
+  val VERSIONNAMES = Seq(s"$VERSION")
+}
 
-  val VERSIONNAMES = Seq("3.0.0-databricks")
+class SparkShimServiceProvider extends com.nvidia.spark.rapids.SparkShimServiceProvider {
 
   def matchesVersion(version: String): Boolean = {
-    VERSIONNAMES.contains(version)
+    SparkShimServiceProvider.VERSIONNAMES.contains(version)
   }
 
   def buildShim: SparkShims = {

diff --git a/shims/spark301/src/main/scala/com/nvidia/spark/rapids/shims/spark301/Spark301Shims.scala b/shims/spark301/src/main/scala/com/nvidia/spark/rapids/shims/spark301/Spark301Shims.scala
@@ -28,6 +28,8 @@ import org.apache.spark.sql.execution.SparkPlan
 
 class Spark301Shims extends Spark300Shims {
 
+  override def getSparkShimVersion: ShimVersion = SparkShimServiceProvider.VERSION
+
   def exprs301: Map[Class[_ <: Expression], ExprRule[_ <: Expression]] = Seq(
     GpuOverrides.expr[First](
       "first aggregate operator",

diff --git a/...k301/src/main/scala/com/nvidia/spark/rapids/shims/spark301/SparkShimServiceProvider.scala b/...k301/src/main/scala/com/nvidia/spark/rapids/shims/spark301/SparkShimServiceProvider.scala
@@ -16,14 +16,16 @@
 
 package com.nvidia.spark.rapids.shims.spark301
 
-import com.nvidia.spark.rapids.SparkShims
+import com.nvidia.spark.rapids.{SparkShims, SparkShimVersion}
 
+object SparkShimServiceProvider {
+  val VERSION = SparkShimVersion(3, 0, 1)
+  val VERSIONNAMES = Seq(s"$VERSION", s"$VERSION-SNAPSHOT")
+}
 class SparkShimServiceProvider extends com.nvidia.spark.rapids.SparkShimServiceProvider {
 
-  val VERSIONNAMES = Seq("3.0.1-SNAPSHOT", "3.0.1")
-
   def matchesVersion(version: String): Boolean = {
-    VERSIONNAMES.contains(version)
+    SparkShimServiceProvider.VERSIONNAMES.contains(version)
   }
 
   def buildShim: SparkShims = {

diff --git a/shims/spark310/src/main/scala/com/nvidia/spark/rapids/shims/spark310/Spark310Shims.scala b/shims/spark310/src/main/scala/com/nvidia/spark/rapids/shims/spark310/Spark310Shims.scala
@@ -39,6 +39,8 @@ import org.apache.spark.unsafe.types.CalendarInterval
 
 class Spark310Shims extends Spark301Shims {
 
+  override def getSparkShimVersion: ShimVersion = SparkShimServiceProvider.VERSION
+
   override def getScalaUDFAsExpression(
       function: AnyRef,
       dataType: DataType,

diff --git a/...k310/src/main/scala/com/nvidia/spark/rapids/shims/spark310/SparkShimServiceProvider.scala b/...k310/src/main/scala/com/nvidia/spark/rapids/shims/spark310/SparkShimServiceProvider.scala
@@ -16,14 +16,17 @@
 
 package com.nvidia.spark.rapids.shims.spark310
 
-import com.nvidia.spark.rapids.SparkShims
+import com.nvidia.spark.rapids.{SparkShims, SparkShimVersion}
 
-class SparkShimServiceProvider extends com.nvidia.spark.rapids.SparkShimServiceProvider {
+object SparkShimServiceProvider {
+  val VERSION = SparkShimVersion(3, 1, 0)
+  val VERSIONNAMES = Seq(s"$VERSION", s"$VERSION-SNAPSHOT")
+}
 
-  val VERSIONNAMES = Seq("3.1.0-SNAPSHOT", "3.1.0")
+class SparkShimServiceProvider extends com.nvidia.spark.rapids.SparkShimServiceProvider {
 
   def matchesVersion(version: String): Boolean = {
-    VERSIONNAMES.contains(version)
+    SparkShimServiceProvider.VERSIONNAMES.contains(version)
   }
 
   def buildShim: SparkShims = {

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/SparkShims.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/SparkShims.scala
@@ -37,7 +37,18 @@ case object GpuBuildRight extends GpuBuildSide
 
 case object GpuBuildLeft extends GpuBuildSide
 
+sealed abstract class ShimVersion
+
+case class SparkShimVersion(major: Int, minor: Int, patch: Int) extends ShimVersion {
+  override def toString(): String = s"$major.$minor.$patch"
+}
+
+case class DatabricksShimVersion(major: Int, minor: Int, patch: Int) extends ShimVersion {
+  override def toString(): String = s"$major.$minor.$patch-databricks"
+}
+
 trait SparkShims {
+  def getSparkShimVersion: ShimVersion
   def isGpuHashJoin(plan: SparkPlan): Boolean
   def isGpuBroadcastHashJoin(plan: SparkPlan): Boolean
   def isGpuShuffledHashJoin(plan: SparkPlan): Boolean

diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/StringFallbackSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/StringFallbackSuite.scala
@@ -29,7 +29,17 @@ class StringFallbackSuite extends SparkQueryCompareTestSuite {
     "RegExpReplace",
     nullableStringsFromCsv, execsAllowedNonGpu = Seq("ProjectExec", "Alias",
       "RegExpReplace", "AttributeReference", "Literal")) {
-    frame => frame.selectExpr("regexp_replace(strings,null,'D')")
+    frame => {
+      // this test is only valid in Spark 3.0.x because the expression is NullIntolerant
+      // since Spark 3.1.0 and gets replaced with a null literal instead
+      val isValidTestForSparkVersion = ShimLoader.getSparkShims.getSparkShimVersion match {
+        case SparkShimVersion(major, minor, _) => major == 3 && minor == 0
+        case DatabricksShimVersion(major, minor, _) => major == 3 && minor == 0
+        case _ => true
+      }
+      assume(isValidTestForSparkVersion)
+      frame.selectExpr("regexp_replace(strings,null,'D')")
+    }
   }
 
   testGpuFallback("String regexp_replace input empty cpu fall back",