feat: support pyspark 3 (via a databricks.koalas stub)

Author: percevalw

.github/workflows/testing.yml

@@ -49,6 +49,20 @@ jobs:
     needs: check_skip
     if: ${{ needs.check_skip.outputs.skip == 'false' }}
     runs-on: "ubuntu-latest"
+    strategy:
+      fail-fast: true
+      matrix:
+        include:
+          - python-version: "3.7"
+            spark: "spark2"
+          - python-version: "3.7"
+            spark: "spark3"
+          - python-version: "3.8"
+            spark: "spark3"
+          - python-version: "3.9"
+            spark: "spark3"
+          - python-version: "3.10"
+            spark: "spark3"
     name: 'Testing on ubuntu'
     defaults:
       run:
@@ -61,13 +75,15 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
-          python-version: '3.7'
-      - name: Install eds-scikit
-        shell: bash {0}
-        run: ./build_tools/github/install.sh
-      - name: Run tests
-        shell: bash {0}
-        run: ./build_tools/github/test.sh
+          python-version: ${{ matrix.python-version }}
+          cache: 'pip'
+      - name: Install dependencies
+        run: |
+          pip install -U "pip<23"
+          pip install --progress-bar off ".[${{ matrix.spark }}, dev, doc]"
+      - name: Run pytest
+        run: |
+          python -m pytest --pyargs tests -m "" --cov=eds_scikit
       - name: Upload coverage to CodeCov
         uses: codecov/codecov-action@v3
         if: success()

build_tools/github/install.sh

@@ -2,9 +2,11 @@
 
 ## Unreleased
 
-### Changed
+### Added
 
 - Support for pyarrow > 0.17.0
+- Support for Python 3.7 to 3.10 (3.11 or higher is not tested)
+- Support for pyspark 3 (to force pyspark 2, use `pip install eds-scikit[spark2]`)
 
 ### Fixed
 - Caching in spark instead of koalas to improve speed

build_tools/github/test.sh

@@ -124,6 +124,7 @@ The goal of **Koalas** is precisely to avoid this issue. It aims at allowing cod
 
 ```python
 from databricks import koalas as ks
+# or from pyspark import pandas as ks, if you have spark 3
 
 # Converting the Spark DataFrame into a Koalas DataFrame
 visit_occurrence_koalas = visit_occurrence_spark.to_koalas()

changelog.md

@@ -26,15 +26,15 @@
 from pyspark import SparkContext
 from pyspark.sql import SparkSession
 
-import eds_scikit.biology  # noqa: F401 --> To register functions
-
 pyarrow.open_stream = pyarrow.ipc.open_stream
 
 sys.path.insert(
     0, (pathlib.Path(__file__).parent / "package-override").absolute().as_posix()
 )
 os.environ["PYTHONPATH"] = ":".join(sys.path)
 
+import eds_scikit.biology  # noqa: F401 --> To register functions
+
 # Remove SettingWithCopyWarning
 pd.options.mode.chained_assignment = None
 

docs/project_description.md

@@ -0,0 +1,17 @@
+# This file is used to override the databricks.koalas package with the pyspark.pandas
+# package, if the databricks.koalas package is not available (python >= 3.8)
+import sys
+import pyarrow  # noqa: E402, F401
+
+old_sys_path = sys.path.copy()
+sys.path.remove(next((p for p in sys.path if "package-override" in p), None))
+databricks = sys.modules.pop("databricks")
+sys.modules.pop("databricks.koalas")
+try:
+    from databricks.koalas import *  # noqa: E402, F401, F403
+except ImportError:
+    from pyspark.pandas import *  # noqa: E402, F401, F403
+
+    sys.modules["databricks"] = databricks
+    sys.modules["databricks.koalas"] = sys.modules["pyspark.pandas"]
+sys.path[:] = old_sys_path

eds_scikit/__init__.py

@@ -21,18 +21,17 @@
    is the only one that resolves to this very module, still gets what it asked for:
    the pyarrow module's content.
 """
-
 import sys
 
+old_sys_path = sys.path.copy()
 sys.path.remove(next((p for p in sys.path if "package-override" in p), None))
 del sys.modules["pyarrow"]
-import pyarrow  # noqa: E402, F401
 
-try:
-    import pyarrow.ipc
+import pyarrow  # noqa: E402, F401
+from pyarrow.ipc import open_stream  # noqa: E402, F401
 
-    pyarrow.open_stream = pyarrow.ipc.open_stream
-except ImportError:
-    pass
+pyarrow.open_stream = open_stream
 
 from pyarrow import *  # noqa: F401, F403, E402
+
+sys.path[:] = old_sys_path

eds_scikit/package-override/databricks/__init__.py

@@ -35,19 +35,18 @@ dependencies = [
     "pgpasslib>=1.1.0, <2.0.0",
     "psycopg2-binary>=2.9.0, <3.0.0",
     "pandas>=1.3.0, <2.0.0",
-    "numpy>=1.0.0, <1.20",
-    "koalas>=1.8.1, <2.0.0",
+    "numpy>=1.0.0",
     "altair>=5.0.0, <6.0.0",
     "loguru==0.7.0",
     "pypandoc==1.7.5",
-    "pyspark==2.4.3",
+    "pyspark",
     "pyarrow>=0.10.0",
     "pretty-html-table>=0.9.15, <0.10.0",
     "catalogue",
     "schemdraw>=0.15.0, <1.0.0",
-    "ipython>=7.32.0, <8.0.0",
-    "packaging==21.3",
-    "tomli==2.0.1",
+    "ipython>=7.32.0",
+    "packaging>=21.3",
+    "tomli>=2.0.1",
 ]
 dynamic = ['version']
 
@@ -66,6 +65,10 @@ Documentation = "https://aphp.github.io/eds-scikit"
 "Bug Tracker" = "https://github.com/aphp/eds-scikit/issues"
 
 [project.optional-dependencies]
+spark2 = [
+    "pyspark==2.4.3",
+    "koalas>=1.8.1,<2.0.0",
+]
 dev = [
     "black>=22.3.0, <23.0.0",
     "flake8==3.9.2",