Merge pull request #479 from databrickslabs/feature/dbr_warning

Feature/dbr warning
databrickslabs · Dec 12, 2023 · 232c955 · 232c955
2 parents bd424e5 + 6f1ac14
commit 232c955
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 18 deletions.
diff --git a/pom.xml b/pom.xml
@@ -93,7 +93,7 @@
             <artifactId>h3</artifactId>
             <!--H3 fixed at 3.7.0 until Photon updates the version-->
             <!--suppress MavenPackageUpdate-->
-            <version>3.7.3</version>
+            <version>3.7.0</version>
         </dependency>
         <dependency>
             <groupId>org.locationtech.jts</groupId>

diff --git a/src/main/scala/com/databricks/labs/mosaic/functions/MosaicContext.scala b/src/main/scala/com/databricks/labs/mosaic/functions/MosaicContext.scala
@@ -965,6 +965,8 @@ object MosaicContext extends Logging {
 
     val tmpDir: String = Files.createTempDirectory("mosaic").toAbsolutePath.toString
 
+    val mosaicVersion: String = "0.3.12"
+
     private var instance: Option[MosaicContext] = None
 
     def build(indexSystem: IndexSystem, geometryAPI: GeometryAPI): MosaicContext = {
@@ -988,23 +990,35 @@ object MosaicContext extends Logging {
     def reset(): Unit = instance = None
 
     // noinspection ScalaStyle,ScalaWeakerAccess
-    def checkDBR(spark: SparkSession): Boolean = {
-        val sparkVersion = spark.conf.get("spark.databricks.clusterUsageTags.sparkVersion", "")
+    def checkDBR(spark: SparkSession): Unit = {
+        val sparkVersion = spark.conf.get("spark.databricks.clusterUsageTags.sparkVersion", "0")
         val isML = sparkVersion.contains("-ml-")
         val isPhoton = spark.conf.getOption("spark.databricks.photon.enabled").getOrElse("false").toBoolean
         val isTest = spark.conf.getOption("spark.databricks.clusterUsageTags.clusterType").isEmpty
 
+        val dbrMajor = sparkVersion.split("-").head.split("\\.").head.toInt
+        if (
+          (dbrMajor < 13 && mosaicVersion >= "0.4.0") ||
+          (dbrMajor > 12 && mosaicVersion < "0.4.0")
+        ) {
+            val msg = """|DEPRECATION ERROR:
+                         |    Mosaic v0.3.x series only supports Databricks Runtime 12 and below.
+                         |    You can specify `%pip install 'databricks-mosaic<0.4,>=0.3'` for DBR < 13.""".stripMargin
+
+            logError(msg)
+            println(msg)
+            throw new Exception(msg)
+        }
+
         if (!isML && !isPhoton && !isTest) {
-            // Print out the warnings both to the log and to the console
-            logWarning("DEPRECATION WARNING: Mosaic is not supported on the selected Databricks Runtime")
-            logWarning("DEPRECATION WARNING: Mosaic will stop working on this cluster after v0.3.x.")
-            logWarning("Please use a Databricks Photon-enabled Runtime (for performance benefits) or Runtime ML (for spatial AI benefits).")
-            println("DEPRECATION WARNING: Mosaic is not supported on the selected Databricks Runtime")
-            println("DEPRECATION WARNING: Mosaic will stop working on this cluster after v0.3.x.")
-            println("Please use a Databricks Photon-enabled Runtime (for performance benefits) or Runtime ML (for spatial AI benefits).")
-            false
-        } else {
-            true
+            val msg = """|DEPRECATION WARNING: 
+                         |  Please use a Databricks:
+                         |      - Photon-enabled Runtime for performance benefits
+                         |      - Runtime ML for spatial AI benefits
+                         |  Mosaic will stop working on this cluster after v0.3.x.""".stripMargin
+            logWarning(msg)
+            println(msg)
+
         }
     }
 

diff --git a/src/test/scala/com/databricks/labs/mosaic/functions/MosaicContextBehaviors.scala b/src/test/scala/com/databricks/labs/mosaic/functions/MosaicContextBehaviors.scala
@@ -1,15 +1,15 @@
 package com.databricks.labs.mosaic.functions
 
-import com.databricks.labs.mosaic.{H3, _}
 import com.databricks.labs.mosaic.core.index._
 import com.databricks.labs.mosaic.test._
+import com.databricks.labs.mosaic._
 import org.apache.spark.sql.adapters.Column
 import org.apache.spark.sql.catalyst.FunctionIdentifier
 import org.apache.spark.sql.catalyst.analysis.FunctionRegistry.FunctionBuilder
 import org.apache.spark.sql.catalyst.catalog.CatalogDatabase
 import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo, Literal}
 import org.apache.spark.sql.functions._
-import org.apache.spark.sql.types.{BinaryType, LongType, StringType}
+import org.apache.spark.sql.types.{LongType, StringType}
 import org.scalamock.scalatest.MockFactory
 import org.scalatest.matchers.must.Matchers.{be, noException}
 import org.scalatest.matchers.should.Matchers.{an, convertToAnyShouldWrapper}
@@ -249,10 +249,16 @@ trait MosaicContextBehaviors extends MosaicSpatialQueryTest {
     }
 
     def printWarnings(): Unit = {
-        spark.conf.set("spark.databricks.clusterUsageTags.sparkVersion", "x")
+        spark.conf.set("spark.databricks.clusterUsageTags.sparkVersion", "1-x")
         spark.conf.set("spark.databricks.photon.enabled", "false")
-        spark.conf.set("spark.databricks.clusterUsageTags.clusterType", "x")
-        MosaicContext.checkDBR(spark) should be(false)
+        spark.conf.set("spark.databricks.clusterUsageTags.clusterType", "1-x")
+        noException should be thrownBy MosaicContext.checkDBR(spark)
+    }
+
+    def throwError(): Unit = {
+        spark.conf.set("spark.databricks.clusterUsageTags.sparkVersion", "14-x")
+        spark.conf.set("spark.databricks.clusterUsageTags.clusterType", "14-x")
+        an[Exception] should be thrownBy MosaicContext.checkDBR(spark)
     }
 
 }

diff --git a/src/test/scala/com/databricks/labs/mosaic/functions/MosaicContextTest.scala b/src/test/scala/com/databricks/labs/mosaic/functions/MosaicContextTest.scala
@@ -18,5 +18,6 @@ class MosaicContextTest extends MosaicSpatialQueryTest with SharedSparkSession w
     test("MosaicContext should use databricks h3") { callDatabricksH3() }
     test("MosaicContext should correctly reflect functions") { reflectedMethods() }
     test("MosaicContext should printWarning") { printWarnings() }
+    test("MosaicContext should throwError") { throwError() }
 
 }