Python/fix/raster tests

.github/workflows/build_main.yml

@@ -2,10 +2,10 @@ name: build main
 on:
   push:
     branches-ignore:
-      - "R/*"
-      - "r/*"
-      - "python/*"
-      - "scala/*"
+      - "R/**"
+      - "r/**"
+      - "python/**"
+      - "scala/**"
   pull_request:
     branches:
       - "**"

.github/workflows/build_python.yml

@@ -3,7 +3,7 @@ name: build_python
 on:
   push:
     branches:
-      - "python/*"
+      - "python/**"
 
 jobs:
   build:

.github/workflows/build_r.yml

@@ -3,8 +3,8 @@ name: build_R
 on:
   push:
     branches:
-      - 'r/*'
-      - 'R/*'
+      - 'r/**'
+      - 'R/**'
 
 jobs:
   build:

.github/workflows/build_scala.yml

@@ -2,7 +2,7 @@ name: build_scala
 on:
   push:
     branches:
-      - "scala/"
+      - "scala/**"
 
 jobs:
   build:

python/mosaic/api/aggregators.py

@@ -48,9 +48,7 @@ def st_intersection_aggregate(
     )
 
 
-def st_intersection_agg(
-        leftIndex: ColumnOrName, rightIndex: ColumnOrName
-) -> Column:
+def st_intersection_agg(leftIndex: ColumnOrName, rightIndex: ColumnOrName) -> Column:
     """
     Computes the intersection of all `leftIndex` : `rightIndex` pairs
     and unions these to produce a single geometry.
@@ -100,9 +98,7 @@ def st_intersects_aggregate(
     )
 
 
-def st_intersects_agg(
-        leftIndex: ColumnOrName, rightIndex: ColumnOrName
-) -> Column:
+def st_intersects_agg(leftIndex: ColumnOrName, rightIndex: ColumnOrName) -> Column:
     """
     Tests if any `leftIndex` : `rightIndex` pairs intersect.
 

python/mosaic/api/functions.py

@@ -208,7 +208,9 @@ def st_bufferloop(
     )
 
 
-def st_buffer_cap_style(geom: ColumnOrName, radius: ColumnOrName, cap_style: ColumnOrName) -> Column:
+def st_buffer_cap_style(
+    geom: ColumnOrName, radius: ColumnOrName, cap_style: ColumnOrName
+) -> Column:
     """
     Compute the buffered geometry based on geom and radius.
 
@@ -231,7 +233,7 @@ def st_buffer_cap_style(geom: ColumnOrName, radius: ColumnOrName, cap_style: Col
         "st_buffer_cap_style",
         pyspark_to_java_column(geom),
         pyspark_to_java_column(radius),
-        pyspark_to_java_column(cap_style)
+        pyspark_to_java_column(cap_style),
     )
 
 

python/mosaic/api/gdal.py

@@ -25,14 +25,11 @@ def setup_gdal(
     -------
 
     """
-    sc = spark.sparkContext
-    mosaicContextClass = getattr(
-        sc._jvm.com.databricks.labs.mosaic.functions, "MosaicContext"
+    mosaicGDALObject = getattr(
+        spark.sparkContext._jvm.com.databricks.labs.mosaic.gdal, "MosaicGDAL"
     )
-    mosaicGDALObject = getattr(sc._jvm.com.databricks.labs.mosaic.gdal, "MosaicGDAL")
     mosaicGDALObject.prepareEnvironment(spark._jsparkSession, init_script_path)
     print("GDAL setup complete.\n")
-    print(f"Shared objects (*.so) stored in: {shared_objects_path}.\n")
     print(f"Init script stored in: {init_script_path}.\n")
     print(
         "Please restart the cluster with the generated init script to complete the setup.\n"

python/mosaic/api/raster.py

@@ -13,29 +13,29 @@
     "rst_boundingbox",
     "rst_clip",
     "rst_combineavg",
-    "rst_fromfile",
     "rst_frombands",
+    "rst_fromfile",
     "rst_georeference",
-    "ret_getnodata",
+    "rst_getnodata",
     "rst_getsubdataset",
     "rst_height",
-    "rst_isempty",
     "rst_initnodata",
+    "rst_isempty",
     "rst_memsize",
-    "rst_metadata",
     "rst_merge",
-    "rst_numbands",
+    "rst_metadata",
     "rst_ndvi",
+    "rst_numbands",
     "rst_pixelheight",
     "rst_pixelwidth",
     "rst_rastertogridavg",
     "rst_rastertogridcount",
     "rst_rastertogridmax",
-    "rst_rastertogridmin",
     "rst_rastertogridmedian",
-    "rst_rastertoworldcoord",
+    "rst_rastertogridmin",
     "rst_rastertoworldcoordx",
     "rst_rastertoworldcoordy",
+    "rst_rastertoworldcoord",
     "rst_retile",
     "rst_rotation",
     "rst_scalex",
@@ -45,17 +45,17 @@
     "rst_skewy",
     "rst_srid",
     "rst_subdatasets",
-    "rst_summary",
     "rst_subdivide",
+    "rst_summary",
     "rst_tessellate",
     "rst_to_overlapping_tiles",
     "rst_tryopen",
     "rst_upperleftx",
     "rst_upperlefty",
     "rst_width",
-    "rst_worldtorastercoord",
     "rst_worldtorastercoordx",
     "rst_worldtorastercoordy",
+    "rst_worldtorastercoord",
 ]
 
 
@@ -172,7 +172,7 @@ def rst_georeference(raster: ColumnOrName) -> Column:
     )
 
 
-def ret_getnodata(raster: ColumnOrName) -> Column:
+def rst_getnodata(raster: ColumnOrName) -> Column:
     """
     Returns the nodata value of the band.
 
@@ -190,7 +190,7 @@ def ret_getnodata(raster: ColumnOrName) -> Column:
 
     """
     return config.mosaic_context.invoke_function(
-        "ret_getnodata", pyspark_to_java_column(raster)
+        "rst_getnodata", pyspark_to_java_column(raster)
     )
 
 
@@ -253,8 +253,7 @@ def rst_initnodata(raster: ColumnOrName) -> Column:
 
     """
     return config.mosaic_context.invoke_function(
-        "rst_initnodata",
-        pyspark_to_java_column(raster)
+        "rst_initnodata", pyspark_to_java_column(raster)
     )
 
 
@@ -897,13 +896,16 @@ def rst_fromfile(raster: ColumnOrName, sizeInMB: ColumnOrName) -> Column:
     """
 
     return config.mosaic_context.invoke_function(
-        "rst_fromfile",
-        pyspark_to_java_column(raster),
-        pyspark_to_java_column(sizeInMB)
+        "rst_fromfile", pyspark_to_java_column(raster), pyspark_to_java_column(sizeInMB)
     )
 
 
-def rst_to_overlapping_tiles(raster: ColumnOrName, width: ColumnOrName, height: ColumnOrName, overlap: ColumnOrName) -> Column:
+def rst_to_overlapping_tiles(
+    raster: ColumnOrName,
+    width: ColumnOrName,
+    height: ColumnOrName,
+    overlap: ColumnOrName,
+) -> Column:
     """
     Tiles the raster into tiles of the given size.
     :param raster:
@@ -916,7 +918,7 @@ def rst_to_overlapping_tiles(raster: ColumnOrName, width: ColumnOrName, height:
         pyspark_to_java_column(raster),
         pyspark_to_java_column(width),
         pyspark_to_java_column(height),
-        pyspark_to_java_column(overlap)
+        pyspark_to_java_column(overlap),
     )
 
 
@@ -1048,7 +1050,10 @@ def rst_worldtorastercoord(
 
     """
     return config.mosaic_context.invoke_function(
-        "rst_worldtorastercoord", pyspark_to_java_column(raster)
+        "rst_worldtorastercoord",
+        pyspark_to_java_column(raster),
+        pyspark_to_java_column(x),
+        pyspark_to_java_column(y),
     )
 
 
@@ -1074,7 +1079,10 @@ def rst_worldtorastercoordx(
 
     """
     return config.mosaic_context.invoke_function(
-        "rst_worldtorastercoordx", pyspark_to_java_column(raster)
+        "rst_worldtorastercoordx",
+        pyspark_to_java_column(raster),
+        pyspark_to_java_column(x),
+        pyspark_to_java_column(y),
     )
 
 
@@ -1100,5 +1108,8 @@ def rst_worldtorastercoordy(
 
     """
     return config.mosaic_context.invoke_function(
-        "rst_worldtorastercoordy", pyspark_to_java_column(raster)
+        "rst_worldtorastercoordy",
+        pyspark_to_java_column(raster),
+        pyspark_to_java_column(x),
+        pyspark_to_java_column(y),
     )

python/mosaic/config/config.py

@@ -10,3 +10,4 @@
 display_handler: DisplayHandler
 ipython_hook: InteractiveShell
 notebook_utils = None
+default_gdal_init_script_path: str = "/dbfs/FileStore/geospatial/mosaic/gdal/"

python/mosaic/core/mosaic_context.py

@@ -51,9 +51,7 @@ def __init__(self, spark: SparkSession):
         IndexSystem = self._indexSystemFactory.getIndexSystem(self._index_system)
         GeometryAPIClass = getattr(self._mosaicPackageObject, self._geometry_api)
 
-        self._context = self._mosaicContextClass.build(
-            IndexSystem, GeometryAPIClass()
-        )
+        self._context = self._mosaicContextClass.build(IndexSystem, GeometryAPIClass())
 
     def invoke_function(self, name: str, *args: Any) -> MosaicColumn:
         func = getattr(self._context.functions(), name)

python/setup.cfg

@@ -17,10 +17,14 @@ classifiers =
 [options]
 packages = find:
 python_requires = >=3.7.0
+setup_requires =
+    pyspark==3.3.2
+    ipython>=7.22.0
+
 install_requires =
     keplergl==0.3.2
     h3==3.7.3
-    ipython>=7.22.0
+    gdal[numpy]==3.4.3
 
 [options.package_data]
 mosaic =

python/test/test_gdal_install.py

@@ -0,0 +1,20 @@
+from .utils import SparkTestCase, GDALInstaller
+
+
+class TestGDALInstall(SparkTestCase):
+    def test_setup_gdal(self):
+        installer = GDALInstaller(self.spark)
+        try:
+            installer.copy_objects()
+        except Exception:
+            self.fail("Copying objects with `setup_gdal()` raised an exception.")
+
+        try:
+            installer_result = installer.run_init_script()
+        except Exception:
+            self.fail("Execution of GDAL init script raised an exception.")
+
+        self.assertEqual(installer_result, 0)
+
+        gdalinfo_result = installer.test_gdalinfo()
+        self.assertEqual(gdalinfo_result, "GDAL 3.4.3, released 2022/04/22\n")