Don't enable facilities requiring rasterio, run basic tests without r…

…asterio installed
apache · Nov 23, 2024 · b627b90 · b627b90
1 parent 4909fe1
commit b627b90
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 8 deletions.
diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml
@@ -143,14 +143,25 @@ jobs:
       - env:
           PYTHON_VERSION: ${{ matrix.python }}
         run: find spark-shaded/target -name sedona-*.jar -exec cp {} ${VENV_PATH}/lib/python${PYTHON_VERSION}/site-packages/pyspark/jars/ \;
-      - env:
+      - name: Run tests
+        env:
           PYTHON_VERSION: ${{ matrix.python }}
         run: |
           export SPARK_HOME=${VENV_PATH}/lib/python${PYTHON_VERSION}/site-packages/pyspark
           cd python
           source ${VENV_PATH}/bin/activate
-          pytest tests
-      - env:
+          pytest -sv tests
+      - name: Run basic tests without rasterio
+        env:
+          PYTHON_VERSION: ${{ matrix.python }}
+        run: |
+          export SPARK_HOME=${VENV_PATH}/lib/python${PYTHON_VERSION}/site-packages/pyspark
+          cd python
+          source ${VENV_PATH}/bin/activate
+          pip uninstall -y rasterio
+          pytest -sv tests/core/test_rdd.py tests/sql/test_dataframe_api.py
+      - name: Run Spark Connect tests
+        env:
           PYTHON_VERSION: ${{ matrix.python }}
         run: |
           if [ ! -f "${VENV_PATH}/lib/python${PYTHON_VERSION}/site-packages/pyspark/sbin/start-connect-server.sh" ]

diff --git a/python/sedona/sql/types.py b/python/sedona/sql/types.py
@@ -17,8 +17,21 @@
 
 from pyspark.sql.types import BinaryType, UserDefinedType
 
-from ..raster import raster_serde
-from ..raster.sedona_raster import SedonaRaster
+# Only support RasterType when rasterio is installed
+try:
+    import rasterio
+except ImportError:
+    rasterio = None
+
+if rasterio is not None:
+    from ..raster import raster_serde
+    from ..raster.sedona_raster import SedonaRaster
+else:
+    # We'll skip RasterType UDT registration and raise error when deserializing
+    # RasterUDT objects if rasterio is not installed
+    raster_serde = None
+    SedonaRaster = None
+
 from ..utils import geometry_serde
 
 
@@ -57,7 +70,12 @@ def serialize(self, obj):
         raise NotImplementedError("RasterType.serialize is not implemented yet")
 
     def deserialize(self, datum):
-        return raster_serde.deserialize(datum)
+        if raster_serde is not None:
+            return raster_serde.deserialize(datum)
+        else:
+            raise NotImplementedError(
+                "rasterio is not installed. Please install it to support RasterType deserialization"
+            )
 
     @classmethod
     def module(cls):
@@ -71,4 +89,5 @@ def scalaUDT(cls):
         return "org.apache.spark.sql.sedona_sql.UDT.RasterUDT"
 
 
-SedonaRaster.__UDT__ = RasterType()
+if SedonaRaster is not None:
+    SedonaRaster.__UDT__ = RasterType()
diff --git a/python/setup.py b/python/setup.py
@@ -63,7 +63,13 @@
         "spark": ["pyspark>=2.3.0"],
         "pydeck-map": ["geopandas", "pydeck==0.8.0"],
         "kepler-map": ["geopandas", "keplergl==0.3.2"],
-        "all": ["pyspark>=2.3.0", "geopandas", "pydeck==0.8.0", "keplergl==0.3.2", "rasterio>=1.2.10"],
+        "all": [
+            "pyspark>=2.3.0",
+            "geopandas",
+            "pydeck==0.8.0",
+            "keplergl==0.3.2",
+            "rasterio>=1.2.10",
+        ],
     },
     project_urls={
         "Documentation": "https://sedona.apache.org",