fix cudf version checker

NVIDIA · May 21, 2021 · 7a5002b · 7a5002b
1 parent eddb8d6
commit 7a5002b
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/SQLPlugin.scala b/sql-plugin/src/main/scala/com/nvidia/spark/SQLPlugin.scala
@@ -17,7 +17,6 @@
 package com.nvidia.spark
 
 import com.nvidia.spark.rapids.{RapidsDriverPlugin, RapidsExecutorPlugin}
-
 import org.apache.spark.api.plugin.{DriverPlugin, ExecutorPlugin, SparkPlugin}
 import org.apache.spark.internal.Logging
 
@@ -27,5 +26,6 @@ import org.apache.spark.internal.Logging
  */
 class SQLPlugin extends SparkPlugin with Logging {
   override def driverPlugin(): DriverPlugin = new RapidsDriverPlugin
+
   override def executorPlugin(): ExecutorPlugin = new RapidsExecutorPlugin
 }
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/Plugin.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/Plugin.scala
@@ -254,8 +254,12 @@ object RapidsExecutorPlugin {
    * version 7.1.1.
    */
   def cudfVersionSatisfied(expected: String, actual: String): Boolean = {
-    val (expMajorMinor, expPatch) = expected.split('.').splitAt(2)
-    val (actMajorMinor, actPatch) = actual.split('.').splitAt(2)
+    val expHyphen = if (expected.indexOf('-') >= 0) expected.indexOf('-') else expected.length
+    val actHyphen = if (actual.indexOf('-') >= 0) actual.indexOf('-') else actual.length
+    if (actual.substring(actHyphen) != expected.substring(expHyphen)) return false
+
+    val (expMajorMinor, expPatch) = expected.substring(0, expHyphen).split('.').splitAt(2)
+    val (actMajorMinor, actPatch) = actual.substring(0, actHyphen).split('.').splitAt(2)
     actMajorMinor.startsWith(expMajorMinor) && {
       val expPatchInts = expPatch.map(_.toInt)
       val actPatchInts = actPatch.map(v => Try(v.toInt).getOrElse(Int.MinValue))

diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/RapidsExecutorPluginSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/RapidsExecutorPluginSuite.scala
@@ -40,6 +40,10 @@ class RapidsExecutorPluginSuite extends FunSuite {
     assert(!RapidsExecutorPlugin.cudfVersionSatisfied("7.0.1", "7.1"))
     assert(!RapidsExecutorPlugin.cudfVersionSatisfied("7.0.1", "7.1.1"))
     assert(!RapidsExecutorPlugin.cudfVersionSatisfied("7.0.1", "7.0.1-special"))
+    assert(!RapidsExecutorPlugin.cudfVersionSatisfied("7.0.1-special", "7.0.1"))
+    assert(RapidsExecutorPlugin.cudfVersionSatisfied("7.0.1-special", "7.0.1-special"))
+    assert(!RapidsExecutorPlugin.cudfVersionSatisfied("7.0.2-special", "7.0.1-special"))
+    assert(RapidsExecutorPlugin.cudfVersionSatisfied("7.0.1-special", "7.0.2-special"))
     assert(!RapidsExecutorPlugin.cudfVersionSatisfied("7.0.2.2.2", "7.0.2.2"))
     assert(RapidsExecutorPlugin.cudfVersionSatisfied("7.0.2.2.2", "7.0.2.2.2"))
   }