Repository: G-Research/spark-extension (language: python)

Issue:
Get spark-extension work in Python dev environment
We probably need a pip installable package in order to get the spark-extension API accessible during development time.


Discussion / hints:

Fix patch:
diff --git a/.github/actions/build-whl/action.yml b/.github/actions/build-whl/action.yml
new file mode 100644
index 00000000..63232f83
--- /dev/null
+++ b/.github/actions/build-whl/action.yml
@@ -0,0 +1,78 @@
+name: 'Build Whl'
+author: 'EnricoMi'
+description: 'A GitHub Action that builds pyspark-extension package'
+
+inputs:
+  spark-version:
+    description: Spark version, e.g. 3.4.0 or 3.4.0-SNAPSHOT
+    required: true
+  scala-version:
+    description: Scala version, e.g. 2.12.15
+    required: true
+  spark-compat-version:
+    description: Spark compatibility version, e.g. 3.4
+    required: true
+  python-version:
+    description: Python version, e.g. 3.8
+    required: true
+
+runs:
+  using: 'composite'
+  steps:
+  - name: Set versions in pom.xml
+    run: |
+      ./set-version.sh ${{ inputs.spark-version }} ${{ inputs.scala-version }}
+      git diff
+    shell: bash
+
+  - name: Fetch Binaries Artifact
+    uses: actions/download-artifact@v3
+    with:
+      name: Binaries-${{ inputs.spark-compat-version }}-${{ inputs.scala-version }}
+      path: .
+
+  - name: Cache Maven packages
+    uses: actions/cache@v3
+    with:
+      path: ~/.m2/repository
+      key: ${{ runner.os }}-mvn-build-${{ inputs.spark-version }}-${{ inputs.scala-version }}-${{ hashFiles('pom.xml') }}
+      restore-keys: ${{ runner.os }}-mvn-build-${{ inputs.spark-version }}-${{ inputs.scala-version }}-
+
+  - name: Cache Pip packages
+    uses: actions/cache@v3
+    with:
+      path: ~/.cache/pip
+      key: ${{ runner.os }}-pip-whl-${{ inputs.python-version }}-${{ inputs.spark-version }}
+
+  - name: Setup Python
+    uses: actions/setup-python@v4
+    with:
+      python-version: ${{ inputs.python-version }}
+
+  - name: Install Python dependencies
+    run: |
+      python -m pip install --upgrade pip build twine
+    shell: bash
+
+  - name: Build whl
+    run: |
+      ./build-whl.sh
+    shell: bash
+
+  - name: Test whl
+    run: |
+      twine check python/dist/*
+      pip install python/dist/*.whl
+      python test-release.py
+    shell: bash
+
+  - name: Upload whl
+    uses: actions/upload-artifact@v3
+    with:
+      name: Whl (Spark ${{ inputs.spark-version }} Scala ${{ inputs.scala-version }})
+      path: |
+        python/dist/*.whl
+
+branding:
+  icon: 'check-circle'
+  color: 'green'
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index e168ee09..f3e35c3e 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -131,6 +131,52 @@ jobs:
         scala-version: ${{ matrix.scala-version }}
         spark-compat-version: ${{ matrix.spark-compat-version }}-SNAPSHOT
 
+  # pyspark is not available for snapshots or scala other than 2.12
+  package-python:
+    name: Build whl (Spark ${{ matrix.spark-version }} Scala ${{ matrix.scala-version }})
+    runs-on: ubuntu-latest
+    needs: build
+
+    strategy:
+      fail-fast: false
+      matrix:
+        spark-compat-version: ['3.0', '3.1', '3.2', '3.3']
+        scala-compat-version: ['2.12']
+        include:
+          - spark-compat-version: '3.0'
+            spark-version: '3.0.3'
+          - spark-compat-version: '3.1'
+            spark-version: '3.1.3'
+          - spark-compat-version: '3.2'
+            spark-version: '3.2.3'
+          - spark-compat-version: '3.3'
+            spark-version: '3.3.2'
+
+          - spark-compat-version: '3.0'
+            scala-compat-version: '2.12'
+            scala-version: '2.12.10'
+          - spark-compat-version: '3.1'
+            scala-compat-version: '2.12'
+            scala-version: '2.12.10'
+          - spark-compat-version: '3.2'
+            scala-compat-version: '2.12'
+            scala-version: '2.12.15'
+          - spark-compat-version: '3.3'
+            scala-compat-version: '2.12'
+            scala-version: '2.12.15'
+
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v3
+
+    - name: Build
+      uses: ./.github/actions/build-whl
+      with:
+        spark-version: ${{ matrix.spark-version }}
+        scala-version: ${{ matrix.scala-version }}
+        spark-compat-version: ${{ matrix.spark-compat-version }}
+        python-version: "3.8"
+
   test-jvm:
     name: Test (Spark ${{ matrix.spark-compat-version }}.${{ matrix.spark-patch-version }} Scala ${{ matrix.scala-version }})
     runs-on: ubuntu-latest
@@ -308,7 +354,7 @@ jobs:
   test_success:
     name: "Test success"
     runs-on: ubuntu-latest
-    needs: [test-jvm, test-python, test-snapshots-jvm]
+    needs: [package-python, test-jvm, test-python, test-snapshots-jvm]
     if: always()
     steps:
       - name: "Inspect test outcome"
diff --git a/README.md b/README.md
index aa9ad13e..0d66872d 100644
--- a/README.md
+++ b/README.md
@@ -135,6 +135,20 @@ on a filesystem where it is accessible by the notebook, and reference that jar f
 
 Check the documentation of your favorite notebook to learn how to add jars to your Spark environment.
 
+### PyPi package (local Spark cluster only)
+
+You may want to install the `pyspark-extension` python package from PyPi into your development environment.
+This provides you code completion, typing and test capabilities during your development phase.
+
+Running your Python application on a Spark cluster will still require one of the above ways
+to add the Scala package to the Spark environment.
+
+```shell script
+pip install pyspark-extension==2.4.0.3.3
+```
+
+Note: Pick the right Spark version (here 3.3) depending on your PySpark version.
+
 ## Build
 
 You can build this project against different versions of Spark and Scala.
@@ -171,3 +185,17 @@ pip install pytest
 ### Run Python tests
 
 Run the Python tests via `env PYTHONPATH=python:python/test python -m pytest python/test`.
+
+Note: you first have to [build the Scala sources](#build-the-scala-project).
+
+### Build Python package
+
+Run the following sequence of commands in the project root directory:
+
+```shell script
+mkdir -p python/pyspark/jars/
+cp -v target/spark-extension_*-*.jar python/pyspark/jars/
+pip install build
+```
+
+Then execute `python -m build python/` to create a whl from the sources. It can be found in `python/dist/`.
diff --git a/RELEASE.md b/RELEASE.md
index b2c44fe8..8d49c2fd 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -20,17 +20,17 @@ The following steps release a snapshot and test it. Test all versions listed [fu
 Follow this procedure to release a new version:
 
 - Add a new entry to `CHANGELOG.md` listing all notable changes of this release.
-  Use the heading `## [VERSION] - YYYY-MM-dd`, e.g. `## [1.0.0] - 2020-03-12`.
-- Remove the `-SNAPSHOT` suffix from `<version>` in the [`pom.xml`](pom.xml) file, e.g. `1.1.0-SNAPSHOT` → `1.1.0`.
-- Update the versions in the `README.md` file to the version of your `pom.xml` to reflect the latest version,
-  e.g. replace all `1.0.0-3.1` with `1.1.0-3.1`.
+  Use the heading `## [VERSION] - YYYY-MM-dd`, e.g. `## [1.1.0] - 2020-03-12`.
+- Remove the `-SNAPSHOT` suffix from the version, e.g. `./set-version 1.1.0`.
+- Update the versions in the `README.md` and `python/README.md` file to the version of your `pom.xml` to reflect the latest version,
+  e.g. replace all `1.0.0-3.1` with `1.1.0-3.1` and `1.0.0.3.1` with `1.1.0.3.1`, respectively.
 - Commit the change to your local git repository, use a commit message like `Releasing 1.1.0`. Do not push to github yet.
 - Tag that commit with a version tag like `v1.1.0` and message like `Release v1.1.0`. Do not push to github yet.
 - Release the version with `mvn clean deploy`. This will be put into a staging repository and not automatically released (due to `<autoReleaseAfterClose>false</autoReleaseAfterClose>` in your [`pom.xml`](pom.xml) file).
 - Inspect and test the staged version. Use `./test-release.sh` or the `spark-examples` project for that. If you are happy with everything:
   - Push the commit and tag to origin.
   - Release the package with `mvn nexus-staging:release`.
-  - Bump the version to the next [minor version](https://semver.org/) in `pom.xml` and append the `-SNAPSHOT` suffix again, e.g. `1.1.0` → `1.2.0-SNAPSHOT`.
+  - Bump the version to the next [minor version](https://semver.org/) and append the `-SNAPSHOT` suffix again: `./set-version 1.2.0-SNAPSHOT`.
   - Commit this change to your local git repository, use a commit message like `Post-release version bump to 1.2.0`.
   - Push all local commits to origin.
 - Otherwise drop it with `mvn nexus-staging:drop`. Remove the last two commits from your local history.
@@ -78,17 +78,17 @@ This is very similar to [releasing from master](#releasing-from-master),
 but the version increment occurs on [patch level](https://semver.org/):
 
 - Add a new entry to `CHANGELOG.md` listing all notable changes of this release.
-  Use the heading `## [VERSION] - YYYY-MM-dd`, e.g. `## [1.0.0] - 2020-03-12`.
-- Remove the `-SNAPSHOT` suffix from `<version>` in the [`pom.xml`](pom.xml) file, e.g. `1.1.1-SNAPSHOT` → `1.1.1`.
-- Update the versions in the `README.md` file to the version of your `pom.xml` to reflect the latest version,
-  e.g. replace all `1.0.0-3.1` with `1.1.0-3.1`.
+  Use the heading `## [VERSION] - YYYY-MM-dd`, e.g. `## [1.1.1] - 2020-03-12`.
+- Remove the `-SNAPSHOT` suffix from the version, e.g. `./set-version 1.1.1`.
+- Update the versions in the `README.md` and `python/README.md` file to the version of your `pom.xml` to reflect the latest version,
+  e.g. replace all `1.1.0-3.1` with `1.1.1-3.1` and `1.1.0.3.1` with `1.1.1.3.1`, respectively.
 - Commit the change to your local git repository, use a commit message like `Releasing 1.1.1`. Do not push to github yet.
-- Tag that commit with a version tag like `v1.1.0` and message like `Release v1.1.1`. Do not push to github yet.
+- Tag that commit with a version tag like `v1.1.1` and message like `Release v1.1.1`. Do not push to github yet.
 - Release the version with `mvn clean deploy`. This will be put into a staging repository and not automatically released (due to `<autoReleaseAfterClose>false</autoReleaseAfterClose>` in your [`pom.xml`](pom.xml) file).
 - Inspect and test the staged version. Use `spark-examples` for that. If you are happy with everything:
   - Push the commit and tag to origin.
   - Release the package with `mvn nexus-staging:release`.
-  - Bump the version to the next [patch version](https://semver.org/) in `pom.xml` and append the `-SNAPSHOT` suffix again, e.g. `1.1.1` → `1.1.2-SNAPSHOT`.
+  - Bump the version to the next [patch version](https://semver.org/) and append the `-SNAPSHOT` suffix again: `./set-version 1.1.2-SNAPSHOT`.
   - Commit this change to your local git repository, use a commit message like `Post-release version bump to 1.1.2`.
   - Push all local commits to origin.
 - Otherwise drop it with `mvn nexus-staging:drop`. Remove the last two commits from your local history.
diff --git a/build-whl.sh b/build-whl.sh
new file mode 100755
index 00000000..21b5265b
--- /dev/null
+++ b/build-whl.sh
@@ -0,0 +1,5 @@
+mkdir -p python/pyspark/jars/
+cp -v target/spark-extension_*-*.jar python/pyspark/jars/
+pip install build
+python -m build python/
+
diff --git a/pom.xml b/pom.xml
index 3722ac99..34097673 100644
--- a/pom.xml
+++ b/pom.xml
@@ -1,7 +1,7 @@
 <project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
   <modelVersion>4.0.0</modelVersion>
   <groupId>uk.co.gresearch.spark</groupId>
-  <artifactId>spark-extension_2.12</artifactId>
+  <artifactId>spark-extension_2.13</artifactId>
   <version>2.5.0-3.3-SNAPSHOT</version>
   <name>Spark Extension</name>
   <description>A library that provides useful extensions to Apache Spark.</description>
@@ -37,11 +37,11 @@
     <maven.compiler.source>${java.version}</maven.compiler.source>
     <maven.compiler.target>${java.version}</maven.compiler.target>
     <encoding>UTF-8</encoding>
-    <scala.compat.version>2.12</scala.compat.version>
-    <scala.version>${scala.compat.version}.15</scala.version>
+    <scala.compat.version>2.13</scala.compat.version>
+    <scala.version>${scala.compat.version}.8</scala.version>
     <!-- keep in-sync with python/requirements-3.3_2.13.txt -->
     <spark.compat.version>3.3</spark.compat.version>
-    <spark.version>${spark.compat.version}.0</spark.version>
+    <spark.version>${spark.compat.version}.1</spark.version>
   </properties>
 
   <dependencies>
diff --git a/python/README.md b/python/README.md
new file mode 100644
index 00000000..a34eb200
--- /dev/null
+++ b/python/README.md
@@ -0,0 +1,78 @@
+# Spark Extension
+
+This project provides extensions to the [Apache Spark project](https://spark.apache.org/) in Scala and Python:
+
+**Diff:** A `diff` transformation for `Dataset`s that computes the differences between
+two datasets, i.e. which rows to _add_, _delete_ or _change_ to get from one dataset to the other.
+
+**Global Row Number:** A `withRowNumbers` transformation that provides the global row number w.r.t.
+the current order of the Dataset, or any given order. In contrast to the existing SQL function `row_number`, which
+requires a window spec, this transformation provides the row number across the entire Dataset without scaling problems.
+
+For details, see the [README.md](https://github.com/G-Research/spark-extension#spark-extension) at the project homepage.
+
+## Using Spark Extension
+
+### Python
+
+### PyPi package (local Spark cluster only)
+
+You may want to install the `pyspark-extension` python package from PyPi into your development environment.
+This provides you code completion, typing and test capabilities during your development phase.
+
+Running your Python application on a Spark cluster will still require one of the ways below
+to add the Scala package to the Spark environment.
+
+```shell script
+pip install pyspark-extension==2.4.0.3.3
+```
+
+Note: Pick the right Spark version (here 3.3) depending on your PySpark version.
+
+#### PySpark API
+
+Start a PySpark session with the Spark Extension dependency (version ≥1.1.0) as follows:
+
+```python
+from pyspark.sql import SparkSession
+
+spark = SparkSession \
+    .builder \
+    .config("spark.jars.packages", "uk.co.gresearch.spark:spark-extension_2.12:2.4.0-3.3") \
+    .getOrCreate()
+```
+
+Note: Pick the right Scala version (here 2.12) and Spark version (here 3.3) depending on your PySpark version.
+
+#### PySpark REPL
+
+Launch the Python Spark REPL with the Spark Extension dependency (version ≥1.1.0) as follows:
+
+```shell script
+pyspark --packages uk.co.gresearch.spark:spark-extension_2.12:2.4.0-3.3
+```
+
+Note: Pick the right Scala version (here 2.12) and Spark version (here 3.3) depending on your PySpark version.
+
+#### PySpark `spark-submit`
+
+Run your Python scripts that use PySpark via `spark-submit`:
+
+```shell script
+spark-submit --packages uk.co.gresearch.spark:spark-extension_2.12:2.4.0-3.3 [script.py]
+```
+
+Note: Pick the right Scala version (here 2.12) and Spark version (here 3.3) depending on your Spark version.
+
+### Your favorite Data Science notebook
+
+There are plenty of [Data Science notebooks](https://datasciencenotebook.org/) around. To use this library,
+add **a jar dependency** to your notebook using these **Maven coordinates**:
+
+    uk.co.gresearch.spark:spark-extension_2.12:2.4.0-3.3
+
+Or [download the jar](https://mvnrepository.com/artifact/uk.co.gresearch.spark/spark-extension) and place it
+on a filesystem where it is accessible by the notebook, and reference that jar file directly.
+
+Check the documentation of your favorite notebook to learn how to add jars to your Spark environment.
+
diff --git a/python/pyproject.toml b/python/pyproject.toml
new file mode 100644
index 00000000..fed528d4
--- /dev/null
+++ b/python/pyproject.toml
@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
diff --git a/python/setup.py b/python/setup.py
new file mode 100755
index 00000000..d21129e9
--- /dev/null
+++ b/python/setup.py
@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+
+#  Copyright 2023 G-Research
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from pathlib import Path
+from setuptools import setup
+
+jar_version = '2.5.0-3.3-SNAPSHOT'
+scala_version = '2.13.8'
+scala_compat_version = '.'.join(scala_version.split('.')[:2])
+spark_compat_version = jar_version.split('-')[1]
+version = jar_version.replace('SNAPSHOT', 'dev0').replace('-', '.')
+
+# read the contents of the README.md file
+long_description = (Path(__file__).parent / "README.md").read_text()
+
+setup(
+    name="pyspark-extension",
+    version=version,
+    description="A library that provides useful extensions to Apache Spark.",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    author="Enrico Minack",
+    author_email="github@enrico.minack.dev",
+    url="https://github.com/G-Research/spark-extension",
+    install_requires=[f"pyspark>={spark_compat_version},<4", "py4j"],
+    packages=[
+        "gresearch",
+        "gresearch.spark",
+        "gresearch.spark.diff",
+        "gresearch.spark.diff.comparator",
+        "pyspark.jars",
+    ],
+    include_package_data=False,
+    package_data={
+        "pyspark.jars": [f"*_{scala_compat_version}-{jar_version}.jar"],
+    },
+    license="http://www.apache.org/licenses/LICENSE-2.0.html",
+    python_requires=">=3.7",
+    classifiers=[
+        "Development Status :: 5 - Production/Stable",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: Implementation :: CPython",
+        "Programming Language :: Python :: Implementation :: PyPy",
+        "Typing :: Typed",
+    ],
+)
diff --git a/release.sh b/release.sh
index 3100e53f..3e97ac29 100755
--- a/release.sh
+++ b/release.sh
@@ -54,18 +54,19 @@ then
 fi
 
 # testing all versions
-./set-version.sh 3.0.3 2.12.10 && mvn clean deploy && ./test-release.sh || exit 1
-./set-version.sh 3.1.3 2.12.10 && mvn clean deploy && ./test-release.sh || exit 1
-./set-version.sh 3.2.3 2.12.15 && mvn clean deploy && ./test-release.sh || exit 1
-./set-version.sh 3.3.1 2.12.16 && mvn clean deploy && ./test-release.sh || exit 1
+./set-version.sh 3.0.3 2.12.10 && mvn clean deploy && ./build-whl.sh && ./test-release.sh || exit 1
+./set-version.sh 3.1.3 2.12.10 && mvn clean deploy && ./build-whl.sh && ./test-release.sh || exit 1
+./set-version.sh 3.2.3 2.12.15 && mvn clean deploy && ./build-whl.sh && ./test-release.sh || exit 1
+./set-version.sh 3.3.2 2.12.16 && mvn clean deploy && ./build-whl.sh && ./test-release.sh || exit 1
+rm -rf python/dist
 
 ./set-version.sh 3.2.3 2.13.5 && mvn clean deploy && ./test-release.sh || exit 1
-./set-version.sh 3.3.1 2.13.8 && mvn clean deploy && ./test-release.sh || exit 1
+./set-version.sh 3.3.2 2.13.8 && mvn clean deploy && ./test-release.sh || exit 1
 
 # all SNAPSHOT versions build, test and complete the example, releasing
 
-# revert pom.xml changes
-git checkout pom.xml
+# revert pom.xml and python/setup.py changes
+git checkout pom.xml python/setup.py
 
 # get latest and release version
 latest=$(grep --max-count=1 "<version>.*</version>" README.md | sed -E -e "s/\s*<[^>]+>//g" -e "s/-[0-9.]+//g")
@@ -75,8 +76,8 @@ echo "Releasing ${#changes[@]} changes as version $version:"
 for (( i=0; i<${#changes[@]}; i++ )); do echo "${changes[$i]}" ; done
 
 sed -i "s/## \[UNRELEASED\] - YYYY-MM-DD/## [$version] - $(date +%Y-%m-%d)/" CHANGELOG.md
-sed -i "s/-SNAPSHOT//g" pom.xml
-sed -i "s/$latest-/$version-/g" README.md
+sed -i -e "s/$latest-/$version-/g" -e "s/$latest\./$version./g" README.md python/README.md
+./set-version.sh $version
 
 # commit changes to local repo
 echo
@@ -96,13 +97,25 @@ echo
 
 # create release
 echo "Creating release packages"
-./set-version.sh 3.0.3 2.12.10 && mvn clean deploy -Dsign && mvn nexus-staging:release
-./set-version.sh 3.1.3 2.12.10 && mvn clean deploy -Dsign && mvn nexus-staging:release
-./set-version.sh 3.2.3 2.12.15 && mvn clean deploy -Dsign && mvn nexus-staging:release
-./set-version.sh 3.3.1 2.12.16 && mvn clean deploy -Dsign && mvn nexus-staging:release
+mkdir -p python/pyspark/jars/
+./set-version.sh 3.0.3 2.12.10 && mvn clean deploy -Dsign && mvn nexus-staging:release && ./build-whl.sh
+./set-version.sh 3.1.3 2.12.10 && mvn clean deploy -Dsign && mvn nexus-staging:release && ./build-whl.sh
+./set-version.sh 3.2.3 2.12.15 && mvn clean deploy -Dsign && mvn nexus-staging:release && ./build-whl.sh
+./set-version.sh 3.3.2 2.12.16 && mvn clean deploy -Dsign && mvn nexus-staging:release && ./build-whl.sh
 
 ./set-version.sh 3.2.3 2.13.5 && mvn clean deploy -Dsign && mvn nexus-staging:release
-./set-version.sh 3.3.1 2.13.8 && mvn clean deploy -Dsign && mvn nexus-staging:release
+./set-version.sh 3.3.2 2.13.8 && mvn clean deploy -Dsign && mvn nexus-staging:release
+
+# upload to test PyPi
+pip install twine
+twine check python/dist/*
+python3 -m twine upload --repository testpypi python/dist/*
+
+echo "Press <ENTER> to upload to PyPi"
+read
+
+# upload to PyPi
+python3 -m twine upload python/dist/*
 
 echo
 
diff --git a/set-version.sh b/set-version.sh
index e9815a8e..204e6b35 100755
--- a/set-version.sh
+++ b/set-version.sh
@@ -1,27 +1,50 @@
 #!/bin/bash
 
-if [ $# -ne 2 ]
+if [ $# -eq 1 ]
 then
-	echo "Provide the Spark and Scala version"
-	exit 1
-fi
+    IFS=-
+    read version flavour <<< "$1"
+
+    echo "setting version=$version${flavour:+ with }$flavour"
+
+    sed -i -E \
+        -e "s%^(  <version>)[^-]+-([^-]+).*(</version>)$%\1$version-\2${flavour:+-}$flavour\3%" \
+        pom.xml
+
+    version=$(grep -m 1 version pom.xml | sed "s/\s*<[^>]*>\s*//g")
 
-spark=$1
-scala=$2
+    sed -i -E \
+        -e "s/(jar_version *= *).*/\1'$version'/" \
+        python/setup.py
+elif [ $# -eq 2 ]
+then
+    spark=$1
+    scala=$2
+
+    spark_compat=${spark%.*}
+    scala_compat=${scala%.*}
 
-spark_compat=${spark%.*}
-scala_compat=${scala%.*}
+    scala_patch=${scala/*./}
+    spark_patch=${spark/*./}
 
-scala_patch=${scala/*./}
-spark_patch=${spark/*./}
+    echo "setting spark=$spark and scala=$scala"
+    sed -i -E \
+        -e "s%^(  <artifactId>)([^_]+)[_0-9.]+(</artifactId>)$%\1\2_${scala_compat}\3%" \
+        -e "s%^(  <version>)([^-]+)-[^-]+(.*</version>)$%\1\2-$spark_compat\3%" \
+        -e "s%^(    <scala.compat.version>).+(</scala.compat.version>)$%\1${scala_compat}\2%" \
+        -e "s%^(    <scala.version>\\\$\{scala.compat.version\}.).+(</scala.version>)$%\1$scala_patch\2%" \
+        -e "s%^(    <spark.compat.version>).+(</spark.compat.version>)$%\1${spark_compat}\2%" \
+        -e "s%^(    <spark.version>\\\$\{spark.compat.version\}.).+(</spark.version>)$%\1$spark_patch\2%" \
+        pom.xml
 
-echo "setting spark=$spark and scala=$scala"
-sed -i -E \
- -e "s%^(  <artifactId>)([^_]+)[_0-9.]+(</artifactId>)$%\1\2_${scala_compat}\3%" \
- -e "s%^(  <version>)([^-]+)-[^-]+(.*</version>)$%\1\2-$spark_compat\3%" \
- -e "s%^(    <scala.compat.version>).+(</scala.compat.version>)$%\1${scala_compat}\2%" \
- -e "s%^(    <scala.version>\\\$\{scala.compat.version\}.).+(</scala.version>)$%\1$scala_patch\2%" \
- -e "s%^(    <spark.compat.version>).+(</spark.compat.version>)$%\1${spark_compat}\2%" \
- -e "s%^(    <spark.version>\\\$\{spark.compat.version\}.).+(</spark.version>)$%\1$spark_patch\2%" \
- pom.xml
+    version=$(grep -m 1 version pom.xml | sed "s/\s*<[^>]*>\s*//g")
+
+    sed -i -E \
+        -e "s/(jar_version *= *).*/\1'$version'/" \
+        -e "s/(scala_version *= *).*/\1'$scala'/" \
+        python/setup.py
+else
+    echo "Provide the Spark-Extension version (e.g. 2.5.0 or 2.5.0-SNAPSHOT), or the Spark and Scala version"
+    exit 1
+fi
 


Test patch:
diff --git a/python/test/spark_common.py b/python/test/spark_common.py
index ebe623c0..44df2eb1 100644
--- a/python/test/spark_common.py
+++ b/python/test/spark_common.py
@@ -62,18 +62,7 @@ def get_pom_path() -> str:
         raise RuntimeError('Could not find path to pom.xml, looked here: {}'.format(', '.join(paths)))
 
     @staticmethod
-    def get_dependencies_from_mvn(path) -> str:
-        logging.info('running mvn to get JVM dependencies')
-        dependencies_process = subprocess.run(
-            ['/bin/bash', '-c', 'mvn dependency:build-classpath 2>/dev/null | grep -A 1 "Dependencies classpath:$" | tail -n 1'],
-            cwd=path, stdout=subprocess.PIPE
-        )
-        if dependencies_process.returncode != 0:
-            raise RuntimeError("failed to run mvn to get classpath for JVM")
-        return str(dependencies_process.stdout.strip())
-
-    @staticmethod
-    def get_spark_config(path, dependencies) -> SparkConf:
+    def get_spark_config(path) -> SparkConf:
         master = 'local[2]'
         conf = SparkConf().setAppName('unit test').setMaster(master)
         return conf.setAll([
@@ -83,7 +72,6 @@ def get_spark_config(path, dependencies) -> SparkConf:
             ('spark.driver.extraClassPath', '{}'.format(':'.join([
                 os.path.join(os.getcwd(), path, 'target', 'classes'),
                 os.path.join(os.getcwd(), path, 'target', 'test-classes'),
-                dependencies
             ]))),
         ])
 
@@ -96,9 +84,7 @@ def get_spark_session(cls) -> SparkSession:
         else:
             logging.info('Setting up Spark environment')
             path = cls.get_pom_path()
-            dependencies = cls.get_dependencies_from_mvn(path)
-            logging.info('found {} JVM dependencies'.format(len(dependencies.split(':'))))
-            conf = cls.get_spark_config(path, dependencies)
+            conf = cls.get_spark_config(path)
             builder.config(conf=conf)
 
         return builder.getOrCreate()
diff --git a/test-release.sh b/test-release.sh
index 3c274019..92e8f700 100755
--- a/test-release.sh
+++ b/test-release.sh
@@ -13,41 +13,44 @@ echo
 
 if [ ! -e "spark-$spark-$scala_compat" ]
 then
-	if [ "$scala_compat" == "2.12" ]
-	then
-		if [ "$spark_compat" == "3.3" || "$spark_compat" > "3.3"]
-		then
-			hadoop="hadoop2"
-		else
-			hadoop="hadoop2.7"
-		fi
-	elif [ "$scala_compat" == "2.13" ]
-	then
-		if [ "$spark_compat" == "3.3" || "$spark_compat" > "3.3"]
-		then
-			hadoop="hadoop3-scala2.13"
-		else
-			hadoop="hadoop3.2-scala2.13"
-		fi
-	else
-		hadoop="without-hadoop"
-	fi
-	wget --progress=dot:giga https://archive.apache.org/dist/spark/spark-$spark/spark-$spark-bin-$hadoop.tgz -O - | tar -xzC .
-	ln -s spark-$spark-bin-$hadoop spark-$spark-$scala_compat
+    if [ "$scala_compat" == "2.12" ]
+    then
+        if [[ "$spark_compat" == "3.3" || "$spark_compat" > "3.3" ]]
+        then
+            hadoop="hadoop2"
+        else
+            hadoop="hadoop2.7"
+        fi
+    elif [ "$scala_compat" == "2.13" ]
+    then
+        if [[ "$spark_compat" == "3.3" || "$spark_compat" > "3.3" ]]
+        then
+            hadoop="hadoop3-scala2.13"
+        else
+            hadoop="hadoop3.2-scala2.13"
+        fi
+    else
+        hadoop="without-hadoop"
+    fi
+    wget --progress=dot:giga https://archive.apache.org/dist/spark/spark-$spark/spark-$spark-bin-$hadoop.tgz -O - | tar -xzC .
+    ln -s spark-$spark-bin-$hadoop spark-$spark-$scala_compat
 fi
 
+echo "Testing Scala"
 spark-$spark-$scala_compat/bin/spark-shell --packages uk.co.gresearch.spark:spark-extension_$scala_compat:$version --repositories https://oss.sonatype.org/content/groups/staging/ < test-release.scala
 
-if [ ! -e "venv" ]
+echo "Testing Python with Scala package"
+spark-$spark-$scala_compat/bin/spark-submit --packages uk.co.gresearch.spark:spark-extension_$scala_compat:$version test-release.py
+
+if [ "$scala_compat" == "2.12" ]
 then
-	virtualenv -p python3 venv
-	source venv/bin/activate
-	pip install python/requirements-${spark_compat}_${scala_compat}.txt
-	deactivate
+    echo "Testing Python with whl package"
+    if [ -e "venv" ]; then rm -rf venv; fi
+    virtualenv -p python3 venv
+    source venv/bin/activate
+    pip install python/dist/pyspark_extension-${version/-*/}.$spark_compat${version/*-SNAPSHOT/.dev0}-py3-none-any.whl
+    python3 test-release.py
 fi
 
-source venv/bin/activate
-spark-$spark-$scala_compat/bin/spark-submit --packages uk.co.gresearch.spark:spark-extension_$scala_compat:$version test-release.py
-deactivate
 
 echo -e "\u001b[32;1mSUCCESS\u001b[0m"