Draft: Added custom grid (#265)

* Added custom grid

* Fixed missing IndexID

* Changed index system parameter type

* Fixed snapshot version in python

* Passing IndexSystem instead of IndexSystemName

* Updated python path loader

* Fixed CustomIndexSystem polyfill for empty polygon

* Fixed mosaic context h3

* Fixed mosaic unite tests for custom grid

* Added K-ring and K-loop for custom grid

* Fixed cell center offset for custom grid

* Fixed lon lat order in custom grid

* Fixed grid offset for cell geometries

* Fixed KNN tests for custom grid

* Fixed buffer radius custom grid

* Reduced KNN iterations to speed up tests

* Fixed python test

* Fixed python import

* Changed custom grid conf structure

* Fixed unit tests for custom grid

* Refactored IndexSystemFactory

* Fixed unit test for KNN

* Added custom grid for python

* Fixed R bindings for custom grid

* Updated docs for custom grid

* Updated docs for custom grid

* Fix KNN tests.

* Fixed index Y splits

* Fixed index KNN resolution

* Fixed double build for feature branch

* Fixed grid root cell size calculation

* Fixed Ring Neighbours tests

* Merge main and fix conflicts.

---------

Co-authored-by: milos.colic <[email protected]>

.github/workflows/build_main.yml

@@ -8,7 +8,10 @@ on:
       - "scala/*"
   pull_request:
     branches:
-      - '**'
+      - "R/*"
+      - "r/*"
+      - "python/*"
+      - "scala/*"
 jobs:
   build:
     runs-on: ubuntu-20.04

.github/workflows/build_python.yml

@@ -4,9 +4,7 @@ on:
   push:
     branches:
       - "python/*"
-  pull_request:
-    branches:
-      - "python/*"
+
 jobs:
   build:
     runs-on: ubuntu-20.04

.github/workflows/build_r.yml

@@ -5,10 +5,7 @@ on:
     branches:
       - 'r/*'
       - 'R/*'
-  pull_request:
-    branches:
-      - 'r/*'
-      - 'R/*'
+
 jobs:
   build:
     runs-on: ubuntu-20.04

.github/workflows/build_scala.yml

@@ -3,9 +3,7 @@ on:
   push:
     branches:
       - "scala/"
-  pull_request:
-    branches:
-      - "scala/"
+
 jobs:
   build:
     runs-on: ubuntu-20.04

R/sparkR-mosaic/enableMosaic.R

@@ -20,8 +20,7 @@ enableMosaic <- function(
   ,rasterAPI="GDAL"
 ){
   geometry_api <- sparkR.callJStatic(x="com.databricks.labs.mosaic.core.geometry.api.GeometryAPI", methodName="apply", geometryAPI)
-  index_system_id <- sparkR.callJStatic(x="com.databricks.labs.mosaic.core.index.IndexSystemID", methodName="apply", indexSystem)
-  indexing_system <- sparkR.callJStatic(x="com.databricks.labs.mosaic.core.index.IndexSystemID", methodName="getIndexSystem", index_system_id)
+  indexing_system <- sparkR.callJStatic(x="com.databricks.labs.mosaic.core.index.IndexSystemFactory", methodName="getIndexSystem", indexSystem)
 
   raster_api <- sparkR.callJStatic(x="com.databricks.labs.mosaic.core.raster.api.RasterAPI", methodName="apply", rasterAPI)
 

docs/source/api/spatial-indexing.rst

@@ -5,6 +5,30 @@ Spatial grid indexing
 Spatial grid indexing is the process of mapping a geometry (or a point) to one or more cells (or cell ID)
 from the selected spatial grid.
 
+The grid system can be specified by using the spark configuration `spark.databricks.labs.mosaic.index.system`
+before enabling Mosaic.
+
+The valid values are
+* `H3` - Good all-rounder for any location on earth
+* `BNG` - Local grid system Great Britain (EPSG:27700)
+* `CUSTOM(minX,maxX,minY,maxY,splits,rootCellSizeX,rootCellSizeY)` - Can be used with any local or global CRS
+    * `minX`,`maxX`,`minY`,`maxY` can be positive or negative integers defining the grid bounds
+    * `splits` defines how many splits are applied to each cell for an increase in resolution step (usually 2 or 10)
+    * `rootCellSizeX`,`rootCellSizeY` define the size of the cells on resolution 0
+
+Example
+
+.. tabs::
+   .. code-tab:: py
+
+    spark.conf.set("spark.databricks.labs.mosaic.index.system", "H3") # Default
+    # spark.conf.set("spark.databricks.labs.mosaic.index.system", "BNG")
+    # spark.conf.set("spark.databricks.labs.mosaic.index.system", "CUSTOM(-180,180,-90,90,2,30,30)")
+
+    import mosaic as mos
+    mos.enable_mosaic(spark, dbutils)
+
+
 grid_longlatascellid
 ********************
 

python/mosaic/core/library_handler.py

@@ -9,7 +9,7 @@ class MosaicLibraryHandler:
     spark = None
     sc = None
     _jar_path = None
-    _jar_filename = f"mosaic-{importlib.metadata.version('databricks-mosaic')}-jar-with-dependencies.jar"
+    _jar_filename = None
     _auto_attached_enabled = None
 
     def __init__(self, spark):
@@ -50,8 +50,14 @@ def mosaic_library_location(self):
                 )
                 self._jar_filename = self._jar_path.split("/")[-1]
             except Py4JJavaError as e:
-                with importlib.resources.path("mosaic.lib", self._jar_filename) as p:
-                    self._jar_path = p.as_posix()
+                self._jar_filename = f"mosaic-{importlib.metadata.version('databricks-mosaic')}-jar-with-dependencies.jar"
+                try:
+                    with importlib.resources.path("mosaic.lib", self._jar_filename) as p:
+                        self._jar_path = p.as_posix()
+                except FileNotFoundError as fnf:
+                    self._jar_filename = f"mosaic-{importlib.metadata.version('databricks-mosaic')}-SNAPSHOT-jar-with-dependencies.jar"
+                    with importlib.resources.path("mosaic.lib", self._jar_filename) as p:
+                        self._jar_path = p.as_posix()
         return self._jar_path
 
     def auto_attach(self):

python/mosaic/core/mosaic_context.py

@@ -24,6 +24,7 @@ def __init__(self, spark: SparkSession):
         self._mosaicPackageRef = getattr(sc._jvm.com.databricks.labs.mosaic, "package$")
         self._mosaicPackageObject = getattr(self._mosaicPackageRef, "MODULE$")
         self._mosaicGDALObject = getattr(sc._jvm.com.databricks.labs.mosaic.gdal, "MosaicGDAL")
+        self._indexSystemFactory = getattr(sc._jvm.com.databricks.labs.mosaic.core.index, "IndexSystemFactory")
 
         try:
             self._geometry_api = spark.conf.get(
@@ -46,12 +47,12 @@ def __init__(self, spark: SparkSession):
         except Py4JJavaError as e:
             self._raster_api = "GDAL"
 
-        IndexSystemClass = getattr(self._mosaicPackageObject, self._index_system)
+        IndexSystem = self._indexSystemFactory.getIndexSystem(self._index_system)
         GeometryAPIClass = getattr(self._mosaicPackageObject, self._geometry_api)
         RasterAPIClass   = getattr(self._mosaicPackageObject, self._raster_api)
 
         self._context = self._mosaicContextClass.build(
-            IndexSystemClass(), GeometryAPIClass(), RasterAPIClass()
+            IndexSystem, GeometryAPIClass(), RasterAPIClass()
         )
 
     def invoke_function(self, name: str, *args: Any) -> MosaicColumn:

python/test/utils/spark_test_case.py

@@ -1,4 +1,5 @@
 import unittest
+import os
 from importlib.metadata import version
 
 from pyspark.sql import SparkSession
@@ -14,6 +15,9 @@ class SparkTestCase(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None:
         cls.library_location = f"{mosaic.__path__[0]}/lib/mosaic-{version('databricks-mosaic')}-jar-with-dependencies.jar"
+        if not os.path.exists(cls.library_location):
+            cls.library_location = f"{mosaic.__path__[0]}/lib/mosaic-{version('databricks-mosaic')}-SNAPSHOT-jar-with-dependencies.jar"
+
         cls.spark = (
             SparkSession.builder.master("local")
             .config("spark.jars", cls.library_location)

src/main/scala/com/databricks/labs/mosaic/core/index/BNGIndexSystem.scala

@@ -27,6 +27,8 @@ import scala.util.{Success, Try}
   */
 object BNGIndexSystem extends IndexSystem(StringType) with Serializable {
 
+    val name = "BNG"
+
     /**
       * Quadrant encodings. The order is determined in a way that preserves
       * similarity to space filling curves.
@@ -201,15 +203,6 @@ object BNGIndexSystem extends IndexSystem(StringType) with Serializable {
         }
     }
 
-    /**
-      * Returns the index system ID instance that uniquely identifies an index
-      * system. This instance is used to select appropriate Mosaic expressions.
-      *
-      * @return
-      *   An instance of [[IndexSystemID]]
-      */
-    override def getIndexSystemID: IndexSystemID = BNG
-
     /**
       * Get the k ring of indices around the provided index id.
       *