Merge pull request #562 from databrickslabs/checkpoint_0.4.2

Global Checkpoint

.github/workflows/build_main.yml

@@ -19,7 +19,7 @@ jobs:
         python: [ 3.10.12 ]
         numpy: [ 1.22.4 ]
         gdal: [ 3.4.1 ]
-        spark: [ 3.4.0 ]
+        spark: [ 3.4.1 ]
         R: [ 4.2.2 ]
     steps:
       - name: checkout code

.github/workflows/build_python.yml

@@ -15,7 +15,7 @@ jobs:
         python: [ 3.10.12 ]
         numpy: [ 1.22.4 ]
         gdal: [ 3.4.1 ]
-        spark: [ 3.4.0 ]
+        spark: [ 3.4.1 ]
         R: [ 4.2.2 ]
     steps:
       - name: checkout code

.github/workflows/build_r.yml

@@ -16,7 +16,7 @@ jobs:
         python: [ 3.10.12 ]
         numpy: [ 1.22.4 ]
         gdal: [ 3.4.1 ]
-        spark: [ 3.4.0 ]
+        spark: [ 3.4.1 ]
         R: [ 4.2.2 ]
     steps:
       - name: checkout code

.github/workflows/build_scala.yml

@@ -14,7 +14,7 @@ jobs:
         python: [ 3.10.12 ]
         numpy: [ 1.22.4 ]
         gdal: [ 3.4.1 ]
-        spark: [ 3.4.0 ]
+        spark: [ 3.4.1 ]
         R: [ 4.2.2 ]
     steps:
       - name: checkout code

.github/workflows/pypi-release.yml

@@ -12,7 +12,7 @@ jobs:
         python: [ 3.10.12 ]
         numpy: [ 1.22.4 ]
         gdal: [ 3.4.1 ]
-        spark: [ 3.4.0 ]
+        spark: [ 3.4.1 ]
         R: [ 4.2.2 ]
     steps:
       - name: checkout code

.gitignore

@@ -159,3 +159,9 @@ spark-warehouse
 .DS_Store
 .Rproj.user
 docker/.m2/
+/python/notebooks/
+/scripts/m2/
+/python/mosaic_test/
+/python/checkpoint/
+/python/checkpoint-new/
+/scripts/docker/docker-build/ubuntu-22-spark-3.4/Dockerfile

CHANGELOG.md

@@ -1,3 +1,12 @@
+## v0.4.3 [DBR 13.3 LTS]
+- Pyspark requirement removed from python setup.cfg as it is supplied by DBR
+- Python version limited to "<3.11,>=3.10" for DBR
+- iPython dependency limited to "<8.11,>=7.4.2" for both DBR and keplergl-jupyter 
+- Expanded support for fuse-based checkpointing (persisted raster storage), managed through:
+  - spark config 'spark.databricks.labs.mosaic.raster.use.checkpoint' in addition to 'spark.databricks.labs.mosaic.raster.checkpoint'.
+  - python: `mos.enable_gdal(spark, with_checkpoint_path=path)`.
+  - scala: `MosaicGDAL.enableGDALWithCheckpoint(spark, path)`.
+
 ## v0.4.2 [DBR 13.3 LTS]
 - Geopandas now fixed to "<0.14.4,>=0.14" due to conflict with minimum numpy version in geopandas 0.14.4.
 - H3 python changed from "==3.7.0" to "<4.0,>=3.7" to pick up patches.

CONTRIBUTING.md

@@ -84,7 +84,7 @@ The repository is structured as follows:
 ## Test & build Mosaic
 
 Given that DBR 13.3 is Ubuntu 22.04, we recommend using docker, 
-see [mosaic-docker.sh](https://github.com/databrickslabs/mosaic/blob/main/scripts/mosaic-docker.sh).
+see [mosaic-docker.sh](https://github.com/databrickslabs/mosaic/blob/main/scripts/docker/mosaic-docker.sh).
 
 ### Scala JAR
 

R/sparkR-mosaic/sparkrMosaic/DESCRIPTION

@@ -1,6 +1,6 @@
 Package: sparkrMosaic
 Title: SparkR bindings for Databricks Mosaic
-Version: 0.4.2
+Version: 0.4.3
 Authors@R:
     person("Robert", "Whiffin", , "[email protected]", role = c("aut", "cre")
            )

R/sparklyr-mosaic/sparklyrMosaic/DESCRIPTION

@@ -1,6 +1,6 @@
 Package: sparklyrMosaic
 Title: sparklyr bindings for Databricks Mosaic
-Version: 0.4.2
+Version: 0.4.3
 Authors@R:
     person("Robert", "Whiffin", , "[email protected]", role = c("aut", "cre")
            )

R/sparklyr-mosaic/tests.R

@@ -9,7 +9,7 @@ library(sparklyr.nested)
 spark_home <- Sys.getenv("SPARK_HOME")
 spark_home_set(spark_home)
 
-install.packages("sparklyrMosaic_0.4.2.tar.gz", repos = NULL)
+install.packages("sparklyrMosaic_0.4.3.tar.gz", repos = NULL)
 library(sparklyrMosaic)
 
 # find the mosaic jar in staging

docs/source/api/rasterio-udfs.rst

@@ -248,7 +248,6 @@ depending on your needs.
     def write_raster(raster, driver, file_id, fuse_dir):
         from io import BytesIO
         from pathlib import Path
-        from pyspark.sql.functions import udf
         from rasterio.io import MemoryFile
         import numpy as np
         import rasterio

docs/source/usage/install-gdal.rst

@@ -112,7 +112,7 @@ Here are spark session configs available for raster, e.g. :code:`spark.conf.set(
      - Checkpoint location, e.g. :ref:`rst_maketiles`
    * - spark.databricks.labs.mosaic.raster.use.checkpoint
      - "false"
-     - Checkpoint for session, in 0.4.2+
+     - Checkpoint for session, in 0.4.3+
    * - spark.databricks.labs.mosaic.raster.tmp.prefix
      - "" (will use "/tmp")
      - Local directory for workers

pom.xml

@@ -146,27 +146,6 @@
                     </execution>
                 </executions>
             </plugin>
-            <plugin>
-                <groupId>org.scoverage</groupId>
-                <artifactId>scoverage-maven-plugin</artifactId>
-                <version>2.0.2</version>
-                <executions>
-                    <execution>
-                        <id>scoverage-report</id>
-                        <phase>package</phase>
-                        <goals>
-                            <goal>check</goal>
-                            <goal>report-only</goal>
-                        </goals>
-                    </execution>
-                </executions>
-                <configuration>
-                    <minimumCoverage>${minimum.coverage}</minimumCoverage>
-                    <failOnMinimumCoverage>true</failOnMinimumCoverage>
-                    <scalaVersion>${scala.version}</scalaVersion>
-                    <additionalForkedProjectProperties>skipTests=false</additionalForkedProjectProperties>
-                </configuration>
-            </plugin>
             <plugin>
                 <!-- see http://davidb.github.com/scala-maven-plugin -->
                 <groupId>net.alchim31.maven</groupId>
@@ -277,8 +256,45 @@
             <properties>
                 <scala.version>2.12.10</scala.version>
                 <scala.compat.version>2.12</scala.compat.version>
-                <spark.version>3.4.0</spark.version>
-                <mosaic.version>0.4.2</mosaic.version>
+                <spark.version>3.4.1</spark.version>
+                <mosaic.version>0.4.3</mosaic.version>
+            </properties>
+            <build>
+                <plugins>
+                    <plugin>
+                        <groupId>org.scoverage</groupId>
+                        <artifactId>scoverage-maven-plugin</artifactId>
+                        <version>2.0.2</version>
+                        <executions>
+                            <execution>
+                                <id>scoverage-report</id>
+                                <phase>package</phase>
+                                <goals>
+                                    <goal>check</goal>
+                                    <goal>report-only</goal>
+                                </goals>
+                            </execution>
+                        </executions>
+                        <configuration>
+                            <minimumCoverage>${minimum.coverage}</minimumCoverage>
+                            <failOnMinimumCoverage>true</failOnMinimumCoverage>
+                            <scalaVersion>${scala.version}</scalaVersion>
+                            <additionalForkedProjectProperties>skipTests=false</additionalForkedProjectProperties>
+                        </configuration>
+                    </plugin>
+                </plugins>
+            </build>
+        </profile>
+        <profile>
+            <!-- local testing `mvn test -PskipScoverage -DskipTests=false -Dsuite=...` -->
+            <id>skipScoverage</id>
+            <properties>
+                <scala.version>2.12.10</scala.version>
+                <scala.compat.version>2.12</scala.compat.version>
+                <spark.version>3.4.1</spark.version>
+                <mosaic.version>0.4.3</mosaic.version>
+                <maven.compiler.source>1.8</maven.compiler.source>
+                <maven.compiler.target>1.8</maven.compiler.target>
             </properties>
         </profile>
     </profiles>

python/mosaic/__init__.py

@@ -4,4 +4,4 @@
 from .models import SpatialKNN
 from .readers import read
 
-__version__ = "0.4.2"
+__version__ = "0.4.3"

python/mosaic/api/__init__.py

@@ -1,7 +1,7 @@
 from .accessors import *
 from .aggregators import *
 from .constructors import *
-from .enable import enable_mosaic
+from .enable import enable_mosaic, get_install_version, get_install_lib_dir
 from .functions import *
 from .fuse import *
 from .predicates import *

python/mosaic/api/enable.py

@@ -1,3 +1,5 @@
+import importlib.metadata
+import importlib.resources
 import warnings
 
 from IPython.core.getipython import get_ipython
@@ -72,24 +74,25 @@ def enable_mosaic(
     if not jar_autoattach:
         spark.conf.set("spark.databricks.labs.mosaic.jar.autoattach", "false")
         print("...set 'spark.databricks.labs.mosaic.jar.autoattach' to false")
+        config.jar_autoattach=False
     if jar_path is not None:
         spark.conf.set("spark.databricks.labs.mosaic.jar.path", jar_path)
         print(f"...set 'spark.databricks.labs.mosaic.jar.path' to '{jar_path}'")
+        config.jar_path=jar_path
     if log_info:
         spark.sparkContext.setLogLevel("info")
+        config.log_info=True
+
+    # Config global objects
+    # - add MosaicContext after MosaicLibraryHandler
     config.mosaic_spark = spark
-    _ = MosaicLibraryHandler(config.mosaic_spark, log_info=log_info)
-    config.mosaic_context = MosaicContext(config.mosaic_spark)
-
-    # Register SQL functions
-    optionClass = getattr(spark._sc._jvm.scala, "Option$")
-    optionModule = getattr(optionClass, "MODULE$")
-    config.mosaic_context._context.register(
-        spark._jsparkSession, optionModule.apply(None)
-    )
-
-    isSupported = config.mosaic_context._context.checkDBR(spark._jsparkSession)
-    if not isSupported:
+    _ = MosaicLibraryHandler(spark, log_info=log_info)
+    config.mosaic_context = MosaicContext(spark)
+    config.mosaic_context.jRegister(spark)
+
+    _jcontext = config.mosaic_context.jContext()
+    is_supported = _jcontext.checkDBR(spark._jsparkSession)
+    if not is_supported:
         # unexpected - checkDBR returns true or throws exception
         print("""WARNING: checkDBR returned False.""")
 
@@ -104,3 +107,31 @@ def enable_mosaic(
         from mosaic.utils.kepler_magic import MosaicKepler
 
         config.ipython_hook.register_magics(MosaicKepler)
+
+
+def get_install_version() -> str:
+    """
+    :return: mosaic version installed
+    """
+    return importlib.metadata.version("databricks-mosaic")
+
+
+def get_install_lib_dir(override_jar_filename=None) -> str:
+    """
+    This is looking for the library dir under site packages using the jar name.
+    :return: located library dir.
+    """
+    v = get_install_version()
+    jar_filename = f"mosaic-{v}-jar-with-dependencies.jar"
+    if override_jar_filename:
+        jar_filename = override_jar_filename
+    with importlib.resources.path("mosaic.lib", jar_filename) as p:
+        return p.parent.as_posix()
+
+
+def refresh_context():
+    """
+    Refresh mosaic context, using previously configured information.
+    - This is needed when spark configs change, such as for checkpointing.
+    """
+    config.mosaic_context.jContextReset()