merged in main

databrickslabs · Jan 30, 2024 · deab4ae · deab4ae
1 parent 7cfcfce
commit deab4ae
Show file tree

Hide file tree

Showing 16 changed files with 152 additions and 117 deletions.
diff --git a/python/mosaic/api/aggregators.py b/python/mosaic/api/aggregators.py
@@ -45,9 +45,7 @@ def st_intersection_agg(leftIndex: ColumnOrName, rightIndex: ColumnOrName) -> Co
     )
 
 
-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.
 
@@ -186,7 +184,9 @@ def rst_combineavg_agg(raster_tile: ColumnOrName) -> Column:
     )
 
 
-def rst_derivedband_agg(raster_tile: ColumnOrName, python_func: ColumnOrName, func_name: ColumnOrName) -> Column:
+def rst_derivedband_agg(
+    raster_tile: ColumnOrName, python_func: ColumnOrName, func_name: ColumnOrName
+) -> Column:
     """
     Returns the raster tile representing the aggregation of rasters using provided python function.
 

diff --git a/python/mosaic/api/enable.py b/python/mosaic/api/enable.py
@@ -11,8 +11,11 @@
 
 
 def enable_mosaic(
-        spark: SparkSession, dbutils = None, log_info: bool = False,
-        jar_path: str = None, jar_autoattach: bool = True
+    spark: SparkSession,
+    dbutils=None,
+    log_info: bool = False,
+    jar_path: str = None,
+    jar_autoattach: bool = True,
 ) -> None:
     """
     Enable Mosaic functions.
@@ -30,20 +33,20 @@ def enable_mosaic(
             Logging cannot be adjusted with Unity Catalog Shared Access clusters;
             if you try to do so, will throw a Py4JSecurityException.
              - True will try to setLogLevel to 'info'
-             - False will not; Default is False 
+             - False will not; Default is False
     jar_path : str
             Convenience when you need to change the JAR path for Unity Catalog
             Volumes with Shared Access clusters
-              - Default is None; if provided, sets 
-                "spark.databricks.labs.mosaic.jar.path" 
+              - Default is None; if provided, sets
+                "spark.databricks.labs.mosaic.jar.path"
     jar_autoattach : bool
             Convenience when you need to turn off JAR auto-attach for Unity
-            Catalog Volumes with Shared Access clusters. 
+            Catalog Volumes with Shared Access clusters.
               - False will not registers the JAR; sets
                 "spark.databricks.labs.mosaic.jar.autoattach" to "false"
               - True will register the JAR; Default is True
 
-              
+
     Returns
     -------
 
@@ -62,7 +65,7 @@ def enable_mosaic(
        Explicitly specify the index system to use for optimized spatial joins. (Optional)
 
     """
-    # Set spark session, conditionally: 
+    # Set spark session, conditionally:
     # - set conf for jar autoattach
     # - set conf for jar path
     # - set log level to 'info'
@@ -73,9 +76,9 @@ def enable_mosaic(
         spark.conf.set("spark.databricks.labs.mosaic.jar.path", jar_path)
         print(f"...set 'spark.databricks.labs.mosaic.jar.path' to '{jar_path}'")
     if log_info:
-        spark.sparkContext.setLogLevel('info')
+        spark.sparkContext.setLogLevel("info")
     config.mosaic_spark = spark
-    _ = MosaicLibraryHandler(config.mosaic_spark, log_info = log_info)
+    _ = MosaicLibraryHandler(config.mosaic_spark, log_info=log_info)
     config.mosaic_context = MosaicContext(config.mosaic_spark)
 
     # Register SQL functions

diff --git a/python/mosaic/api/functions.py b/python/mosaic/api/functions.py
@@ -155,7 +155,9 @@ def st_convexhull(geom: ColumnOrName) -> Column:
     )
 
 
-def st_concavehull(geom: ColumnOrName, concavity: ColumnOrName, has_holes: Any = False) -> Column:
+def st_concavehull(
+    geom: ColumnOrName, concavity: ColumnOrName, has_holes: Any = False
+) -> Column:
     """
     Compute the concave hull of a geometry or multi-geometry object.
     It uses lengthRatio and
@@ -190,7 +192,7 @@ def st_concavehull(geom: ColumnOrName, concavity: ColumnOrName, has_holes: Any =
         "st_concavehull",
         pyspark_to_java_column(geom),
         pyspark_to_java_column(concavity),
-        pyspark_to_java_column(has_holes)
+        pyspark_to_java_column(has_holes),
     )
 
 
@@ -217,7 +219,7 @@ def st_buffer(geom: ColumnOrName, radius: ColumnOrName) -> Column:
 
 
 def st_bufferloop(
-        geom: ColumnOrName, inner_radius: ColumnOrName, outer_radius: ColumnOrName
+    geom: ColumnOrName, inner_radius: ColumnOrName, outer_radius: ColumnOrName
 ) -> Column:
     """
     Compute the buffered geometry loop (hollow ring) based on geom and provided radius-es.
@@ -363,7 +365,7 @@ def st_transform(geom: ColumnOrName, srid: ColumnOrName) -> Column:
 
 
 def st_hasvalidcoordinates(
-        geom: ColumnOrName, crs: ColumnOrName, which: ColumnOrName
+    geom: ColumnOrName, crs: ColumnOrName, which: ColumnOrName
 ) -> Column:
     """
     Checks if all points in geometry are valid with respect to crs bounds.
@@ -570,7 +572,7 @@ def st_distance(geom1: ColumnOrName, geom2: ColumnOrName) -> Column:
 
 
 def st_haversine(
-        lat1: ColumnOrName, lng1: ColumnOrName, lat2: ColumnOrName, lng2: ColumnOrName
+    lat1: ColumnOrName, lng1: ColumnOrName, lat2: ColumnOrName, lng2: ColumnOrName
 ) -> Column:
     """
     Compute the haversine distance in kilometers between two latitude/longitude pairs.
@@ -722,7 +724,7 @@ def st_unaryunion(geom: ColumnOrName) -> Column:
 
 
 def st_updatesrid(
-        geom: ColumnOrName, srcSRID: ColumnOrName, destSRID: ColumnOrName
+    geom: ColumnOrName, srcSRID: ColumnOrName, destSRID: ColumnOrName
 ) -> Column:
     """
     Updates the SRID of the input geometry `geom` from `srcSRID` to `destSRID`.
@@ -991,7 +993,7 @@ def grid_boundary(index_id: ColumnOrName, format_name: ColumnOrName) -> Column:
 
 
 def grid_longlatascellid(
-        lon: ColumnOrName, lat: ColumnOrName, resolution: ColumnOrName
+    lon: ColumnOrName, lat: ColumnOrName, resolution: ColumnOrName
 ) -> Column:
     """
     Returns the grid's cell ID associated with the input `lng` and `lat` coordinates at a given grid `resolution`.
@@ -1059,7 +1061,7 @@ def grid_polyfill(geom: ColumnOrName, resolution: ColumnOrName) -> Column:
 
 
 def grid_tessellate(
-        geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
+    geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
 ) -> Column:
     """
     Generates:
@@ -1094,7 +1096,7 @@ def grid_tessellate(
 
 
 def grid_tessellateexplode(
-        geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
+    geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
 ) -> Column:
     """
     Generates:
@@ -1254,7 +1256,7 @@ def grid_cellkloopexplode(cellid: ColumnOrName, k: ColumnOrName) -> Column:
 
 
 def grid_geometrykring(
-        geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
+    geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
 ) -> Column:
     """
     Returns the k-ring of cells around the input geometry.
@@ -1279,7 +1281,7 @@ def grid_geometrykring(
 
 
 def grid_geometrykloop(
-        geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
+    geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
 ) -> Column:
     """
     Returns the k loop (hollow ring) of cells around the input geometry.
@@ -1304,7 +1306,7 @@ def grid_geometrykloop(
 
 
 def grid_geometrykringexplode(
-        geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
+    geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
 ) -> Column:
     """
     Returns the exploded k-ring of cells around the input geometry.
@@ -1329,7 +1331,7 @@ def grid_geometrykringexplode(
 
 
 def grid_geometrykloopexplode(
-        geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
+    geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
 ) -> Column:
     """
     Returns the exploded k loop (hollow ring) of cells around the input geometry.
@@ -1376,7 +1378,7 @@ def point_index_geom(geom: ColumnOrName, resolution: ColumnOrName) -> Column:
 
 
 def point_index_lonlat(
-        lon: ColumnOrName, lat: ColumnOrName, resolution: ColumnOrName
+    lon: ColumnOrName, lat: ColumnOrName, resolution: ColumnOrName
 ) -> Column:
     """
     [Deprecated] alias for `grid_longlatascellid`
@@ -1433,7 +1435,7 @@ def polyfill(geom: ColumnOrName, resolution: ColumnOrName) -> Column:
 
 
 def mosaic_explode(
-        geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
+    geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
 ) -> Column:
     """
     [Deprecated] alias for `grid_tessellateexplode`
@@ -1468,7 +1470,7 @@ def mosaic_explode(
 
 
 def mosaicfill(
-        geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
+    geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
 ) -> Column:
     """
     [Deprecated] alias for `grid_tessellate`