From 0d3e0afc0a50ade8bb116a347a4fd210fe2d4caf Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:12:01 -0700 Subject: [PATCH 1/6] fix: depend on pyspark-client instead of the full pyspark distribution The client only ever talks to a remote Spark Connect endpoint, so it has no use for the Spark JVM jars that make up 442 MB of a 461 MB pyspark install. Apache Spark 4.0 publishes pyspark-client for exactly this case: the same Python tree, remote-only, without the jars. Installs drop to around 14 MB. sparksql-magic declares a dependency on the full pyspark distribution, which pip would install alongside pyspark-client and shadow it, so it is now installed with --no-deps wherever it is needed. --- .github/workflows/integration-tests.yaml | 3 +++ DEVELOPING.md | 7 ++++++- README.md | 23 ++++++++++++++++++++++- requirements-dev.txt | 6 ++++-- setup.py | 2 +- 5 files changed, 36 insertions(+), 5 deletions(-) diff --git a/.github/workflows/integration-tests.yaml b/.github/workflows/integration-tests.yaml index a11aa527..f901ffd5 100644 --- a/.github/workflows/integration-tests.yaml +++ b/.github/workflows/integration-tests.yaml @@ -62,6 +62,9 @@ jobs: run: | pip install -r requirements-dev.txt pip install -r requirements-test.txt + # --no-deps keeps pip from pulling the full pyspark distribution, which + # would shadow the lightweight pyspark-client package. + pip install --no-deps sparksql-magic>=0.0.3 - name: Authenticate to Google Cloud uses: google-github-actions/auth@7c6bc770dae815cd3e89ee6cdf493a5fab2cc093 # v3 diff --git a/DEVELOPING.md b/DEVELOPING.md index c9a8a5ea..5307fcc5 100644 --- a/DEVELOPING.md +++ b/DEVELOPING.md @@ -45,9 +45,14 @@ To run tests with magic functionality, install the required dependencies manuall ```sh pip install . -pip install IPython sparksql-magic +pip install IPython +pip install --no-deps sparksql-magic ``` +`sparksql-magic` declares a dependency on the full `pyspark` distribution. +Installing it with `--no-deps` keeps the lightweight `pyspark-client` package in +place; without it, pip installs `pyspark` on top and shadows the client. + Then run tests as normal. Any magic-related tests will automatically detect and use the available dependencies. ## Testing without Magic Support diff --git a/README.md b/README.md index c1117926..a1b75adb 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,22 @@ requiring additional steps. pip install google-cloud-spark-connect ``` +This depends on +[`pyspark-client`](https://pypi.org/project/pyspark-client/), the lightweight +Spark Connect client, rather than the full `pyspark` distribution — the install +is around 14 MB instead of around 460 MB, because none of the Spark JVM jars are +needed to talk to a remote session. + +If you also need a local Spark runtime, install the full distribution instead: + +```sh +pip uninstall pyspark-client +pip install 'pyspark[connect]~=4.0.0' +``` + +The two packages both provide the `pyspark` module and cannot be installed side +by side. + ## Uninstall ```sh @@ -128,9 +144,14 @@ The package supports the [sparksql-magic](https://github.com/cryeo/sparksql-magi **Installation**: To use magic commands, install the required dependencies manually: ```bash pip install google-cloud-spark-connect -pip install IPython sparksql-magic +pip install IPython +pip install --no-deps sparksql-magic ``` +`sparksql-magic` declares a dependency on the full `pyspark` distribution. +Installing it with `--no-deps` keeps the lightweight `pyspark-client` package in +place; without it, pip installs `pyspark` on top and shadows the client. + 1. Load the magic extension: ```python %load_ext sparksql_magic diff --git a/requirements-dev.txt b/requirements-dev.txt index 5cf7026e..b7bca087 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -4,8 +4,10 @@ ipython~=9.1 ipywidgets>=8.0.0 packaging>=20.0 pyink~=24.0 -pyspark[connect]~=4.0.0 +pyspark-client~=4.0.0 setuptools>=72.0 -sparksql-magic>=0.0.3 +# sparksql-magic declares a dependency on the full `pyspark` distribution, which +# would be installed alongside pyspark-client and shadow it. Install it without +# its dependencies instead: pip install --no-deps sparksql-magic>=0.0.3 tqdm>=4.67 websockets>=14.0 diff --git a/setup.py b/setup.py index 89637d54..c629120c 100644 --- a/setup.py +++ b/setup.py @@ -32,7 +32,7 @@ "google-api-core>=2.19", "google-cloud-dataproc>=5.18", "packaging>=20.0", - "pyspark[connect]~=4.0.0", + "pyspark-client~=4.0.0", "tqdm>=4.67", "websockets>=14.0", ], From 656410871be04b7229ba173355001cbbd6915fad Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:39:01 -0700 Subject: [PATCH 2/6] test: run the local Spark tests against the full pyspark distribution test_create_local_spark_session covers the Dataproc batch path, where the builder hands off to a local classic Spark session. pyspark-client ships no JVM jars and cannot start one, so the test errored with CONNECT_URL_NOT_SET once the dependency was swapped. Guard the test on pyspark.util.is_remote_only() so a pyspark-client environment skips it with a clear reason, and add a CI job that installs the full distribution so the path keeps its coverage. The test needs no GCP credentials, so it runs there rather than in the integration suite. --- .github/workflows/tests.yaml | 40 ++++++++++++++++++++++++++++++- tests/integration/test_session.py | 10 ++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index abaa39f5..ca89798e 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -49,4 +49,42 @@ jobs: pip install -r requirements-test.txt - name: Run unit tests - run: python -m pytest tests/unit/ -v --tb=short -n auto \ No newline at end of file + run: python -m pytest tests/unit/ -v --tb=short -n auto + + local-spark: + name: Run local Spark tests + runs-on: ubuntu-latest + + steps: + - name: Checkout code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Setup Python + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: "3.12" + + - name: Cache pip dependencies + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-local-spark-${{ hashFiles('requirements-dev.txt', 'requirements-test.txt') }} + restore-keys: | + ${{ runner.os }}-pip-local-spark- + ${{ runner.os }}-pip- + + - name: Install dependencies + run: | + pip install -r requirements-dev.txt + pip install -r requirements-test.txt + # The Dataproc batch code path hands off to a local classic Spark + # session, which pyspark-client cannot start. Install the full + # distribution so those tests run instead of skipping. It overlays + # pyspark-client cleanly: pyspark/core and pyspark/jars come from the + # full package, so is_remote_only() reports False. + pip install 'pyspark[connect]~=4.0.0' + + # These tests exercise the local Spark handoff only and need no GCP + # credentials, so they run here rather than in the integration suite. + - name: Run local Spark tests + run: python -m pytest tests/integration/test_session.py -v --tb=short -k test_create_local_spark_session \ No newline at end of file diff --git a/tests/integration/test_session.py b/tests/integration/test_session.py index c9217137..d6c48e0c 100644 --- a/tests/integration/test_session.py +++ b/tests/integration/test_session.py @@ -33,9 +33,18 @@ ) from pyspark.errors.exceptions import connect as connect_exceptions from pyspark.sql.types import StringType +from pyspark.util import is_remote_only _SERVICE_ACCOUNT_KEY_FILE_ = "service_account_key.json" +# The library depends on pyspark-client, which has no JVM jars and therefore +# cannot start a local Spark session. Tests that need one only run when the +# full pyspark distribution is installed instead. +requires_local_spark = pytest.mark.skipif( + is_remote_only(), + reason="requires the full pyspark distribution (a local Spark runtime)", +) + @pytest.fixture(params=[None, "3.0"]) def image_version(request): @@ -777,6 +786,7 @@ def local_spark_session(): session.stop() +@requires_local_spark def test_create_local_spark_session(batch_workload_env, local_spark_session): """Test creating a local Spark session.""" from pyspark.sql import SparkSession as PySparkSession From 9d28048d88fb0931da6b76b624de68d8e2b01951 Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:13:07 -0700 Subject: [PATCH 3/6] build: declare the local Spark test dependency in its own requirements file The full pyspark distribution is needed by one test, not by the test suite, so it does not belong in requirements-test.txt: both the unit and integration jobs install that file, and giving them a JVM Spark would stop them exercising the pyspark-client install we actually ship. Move it out of the workflow's inline pip invocation and into requirements-local-spark.txt, consumed only by the local-spark job. It now participates in the cache key and is discoverable for local development. --- .github/workflows/tests.yaml | 9 ++------- DEVELOPING.md | 11 +++++++++++ requirements-local-spark.txt | 10 ++++++++++ 3 files changed, 23 insertions(+), 7 deletions(-) create mode 100644 requirements-local-spark.txt diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index ca89798e..0a5838e4 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -68,7 +68,7 @@ jobs: uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: path: ~/.cache/pip - key: ${{ runner.os }}-pip-local-spark-${{ hashFiles('requirements-dev.txt', 'requirements-test.txt') }} + key: ${{ runner.os }}-pip-local-spark-${{ hashFiles('requirements-dev.txt', 'requirements-test.txt', 'requirements-local-spark.txt') }} restore-keys: | ${{ runner.os }}-pip-local-spark- ${{ runner.os }}-pip- @@ -77,12 +77,7 @@ jobs: run: | pip install -r requirements-dev.txt pip install -r requirements-test.txt - # The Dataproc batch code path hands off to a local classic Spark - # session, which pyspark-client cannot start. Install the full - # distribution so those tests run instead of skipping. It overlays - # pyspark-client cleanly: pyspark/core and pyspark/jars come from the - # full package, so is_remote_only() reports False. - pip install 'pyspark[connect]~=4.0.0' + pip install -r requirements-local-spark.txt # These tests exercise the local Spark handoff only and need no GCP # credentials, so they run here rather than in the integration suite. diff --git a/DEVELOPING.md b/DEVELOPING.md index 5307fcc5..83f66e64 100644 --- a/DEVELOPING.md +++ b/DEVELOPING.md @@ -11,6 +11,17 @@ pip install -r requirements-dev.txt pip install -r requirements-test.txt ``` +Tests that need a local Spark runtime are skipped unless the full `pyspark` +distribution is installed on top: + +```sh +pip install -r requirements-local-spark.txt +``` + +Install it only when you need those tests. The unit and integration suites are +meant to run against `pyspark-client` so they keep exercising the dependency +set we ship. + # Linting/formatting We use `pyink` to lint/format the code. To apply changes to your local diff --git a/requirements-local-spark.txt b/requirements-local-spark.txt new file mode 100644 index 00000000..04967a7b --- /dev/null +++ b/requirements-local-spark.txt @@ -0,0 +1,10 @@ +# Dependencies for the tests that need a local Spark runtime. +# +# The library itself depends on pyspark-client, which has no JVM jars and +# cannot start a local session. The Dataproc batch code path hands off to a +# local classic Spark session, so testing it needs the full distribution. +# +# Install this on top of requirements-dev.txt, never instead of it, and only +# for those tests: the unit and integration suites are meant to run against +# pyspark-client so they keep exercising the dependency set we ship. +pyspark[connect]~=4.0.0 From 33d3a683193d71d6ad541a8bf203f1b34807e80e Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:30:03 -0700 Subject: [PATCH 4/6] fix: warn when pyspark and pyspark-client are installed at odds Because the two distributions have different names but provide the same pyspark module, pip cannot see them as conflicting and installs both. An ordinary `pip install sparksql-magic` next to this library is enough to produce pyspark 4.2.0 alongside pyspark-client 4.0.4, and the first import then fails with "[PACKAGE_NOT_INSTALLED] zstandard >= 0.25.0 must be installed", which names neither package. Check for the mismatch before importing anything from pyspark and point at the fix. Same versions are left alone: their shared files are byte-identical, so that combination works. --- .../cloud/managed_spark_connect/__init__.py | 32 +++++++++ tests/unit/test_init.py | 68 +++++++++++++++++++ 2 files changed, 100 insertions(+) diff --git a/google/cloud/managed_spark_connect/__init__.py b/google/cloud/managed_spark_connect/__init__.py index 23567583..7603edea 100644 --- a/google/cloud/managed_spark_connect/__init__.py +++ b/google/cloud/managed_spark_connect/__init__.py @@ -14,6 +14,38 @@ import importlib.metadata import warnings + +def _check_pyspark_installation(): + """Warns when both pyspark distributions are installed at odds. + + 'pyspark-client' and 'pyspark' both provide the 'pyspark' module, so pip + installs them side by side rather than reporting a conflict. At the same + version their shared files are identical and this is harmless. At + different versions the environment ends up a mix of the two, and the + failure surfaces later as an import error naming neither package. + """ + try: + client_version = importlib.metadata.version("pyspark-client") + full_version = importlib.metadata.version("pyspark") + except importlib.metadata.PackageNotFoundError: + return + + if client_version != full_version: + warnings.warn( + f"Both 'pyspark-client' ({client_version}) and 'pyspark' " + f"({full_version}) are installed, at different versions. They " + "provide the same 'pyspark' module, so this environment holds a " + "mix of the two and imports may fail in ways that mention " + "neither. Uninstall both and reinstall only the one you need: " + "'pip uninstall pyspark pyspark-client', then " + "'pip install pyspark-client' to use remote Managed Spark " + "Sessions, or 'pip install pyspark[connect]' if you also run " + "Spark locally." + ) + + +_check_pyspark_installation() + from .session import ManagedSparkSession old_package_names = ["google-spark-connect", "dataproc-spark-connect"] diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index 794f6fa8..e60ea1b4 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -11,13 +11,81 @@ # 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. +import importlib.metadata import unittest from unittest import mock +from google.cloud.managed_spark_connect import _check_pyspark_installation from google.cloud.managed_spark_connect.session import ManagedSparkSession from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException +class TestPysparkInstallationCheck(unittest.TestCase): + + def _run_with_versions(self, versions): + """Runs the check with importlib.metadata.version stubbed out. + + `versions` maps a distribution name to its version, or to a + PackageNotFoundError to mark it as not installed. + """ + + def fake_version(name): + result = versions[name] + if isinstance(result, Exception): + raise result + return result + + with mock.patch("importlib.metadata.version", side_effect=fake_version): + with mock.patch("warnings.warn") as mock_warn: + _check_pyspark_installation() + return mock_warn + + def test_warns_when_versions_differ(self): + """Both distributions installed at different versions is a broken mix""" + mock_warn = self._run_with_versions( + {"pyspark-client": "4.0.4", "pyspark": "4.2.0"} + ) + + mock_warn.assert_called_once() + message = mock_warn.call_args[0][0] + self.assertIn("pyspark-client", message) + self.assertIn("4.0.4", message) + self.assertIn("4.2.0", message) + self.assertIn("pip uninstall pyspark pyspark-client", message) + + def test_no_warning_when_versions_match(self): + """Shared files are identical at the same version, so this is fine""" + mock_warn = self._run_with_versions( + {"pyspark-client": "4.0.4", "pyspark": "4.0.4"} + ) + + mock_warn.assert_not_called() + + def test_no_warning_without_full_pyspark(self): + """The expected install: pyspark-client alone""" + mock_warn = self._run_with_versions( + { + "pyspark-client": "4.0.4", + "pyspark": importlib.metadata.PackageNotFoundError("pyspark"), + } + ) + + mock_warn.assert_not_called() + + def test_no_warning_without_pyspark_client(self): + """The documented escape hatch: the full distribution alone""" + mock_warn = self._run_with_versions( + { + "pyspark-client": importlib.metadata.PackageNotFoundError( + "pyspark-client" + ), + "pyspark": "4.0.4", + } + ) + + mock_warn.assert_not_called() + + class TestPythonVersionCheck(unittest.TestCase): def test_python_version_mismatch_warning_for_runtime_30(self): From b5e41a18b7e06a2ddbd1552359d3c137d9fdec82 Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Thu, 17 Sep 2026 10:58:29 -0700 Subject: [PATCH 5/6] feat: choose the Spark distribution with [client] and [full] extras The base install now names no Spark distribution, so it uses whichever one the environment already provides. 'pyspark-client' and 'pyspark' both supply the 'pyspark' module but are separate distributions, so depending on either pushed a second copy onto environments that already had the other. Two extras stand in for the choice: [client] for pyspark-client, the Spark Connect client on its own, and [full] for pyspark[connect], which adds the JVM jars needed to run Spark locally. A bare install cannot start a session on its own, so the installation check now raises with the two extras named rather than letting a bare ModuleNotFoundError escape. Dropping the pinned dependency also dropped the floor it enforced, so an installed Spark older than 4.0 now warns instead of failing later inside an import. Note that [full] also makes sparksql-magic install without --no-deps, since the full distribution is what it asks for. --- DEVELOPING.md | 11 +-- README.md | 65 ++++++++++++----- .../cloud/managed_spark_connect/__init__.py | 72 +++++++++++++++---- requirements-local-spark.txt | 11 +-- setup.py | 10 ++- tests/unit/test_init.py | 42 ++++++++++- 6 files changed, 166 insertions(+), 45 deletions(-) diff --git a/DEVELOPING.md b/DEVELOPING.md index 83f66e64..1f3d89fa 100644 --- a/DEVELOPING.md +++ b/DEVELOPING.md @@ -55,23 +55,24 @@ env \ To run tests with magic functionality, install the required dependencies manually: ```sh -pip install . +pip install '.[client]' pip install IPython pip install --no-deps sparksql-magic ``` `sparksql-magic` declares a dependency on the full `pyspark` distribution. -Installing it with `--no-deps` keeps the lightweight `pyspark-client` package in -place; without it, pip installs `pyspark` on top and shadows the client. +Installing it with `--no-deps` keeps `pyspark-client` in place; without it, pip +adds `pyspark` on top and the two shadow each other. Installing `.[full]` +instead is the other way to avoid that. Then run tests as normal. Any magic-related tests will automatically detect and use the available dependencies. ## Testing without Magic Support -To run tests without the magic dependencies, simply install the base package: +To run tests without the magic dependencies, simply install the package: ```sh -pip install . +pip install '.[client]' pytest ``` diff --git a/README.md b/README.md index a1b75adb..afe92320 100644 --- a/README.md +++ b/README.md @@ -7,26 +7,44 @@ requiring additional steps. ## Install +This client needs a Spark distribution, and there are two to choose from. It +does not depend on either one directly, so that it works with whichever is +already installed. Pick with an extra: + ```sh +# Talk to remote Managed Spark Sessions. Around 14 MB. +pip install 'google-cloud-spark-connect[client]' + +# Also run Spark locally. Around 460 MB. +pip install 'google-cloud-spark-connect[full]' + +# Neither: use the Spark distribution the environment already has. pip install google-cloud-spark-connect ``` -This depends on -[`pyspark-client`](https://pypi.org/project/pyspark-client/), the lightweight -Spark Connect client, rather than the full `pyspark` distribution — the install -is around 14 MB instead of around 460 MB, because none of the Spark JVM jars are -needed to talk to a remote session. +`[client]` installs +[`pyspark-client`](https://pypi.org/project/pyspark-client/), the Spark Connect +client on its own. `[full]` installs `pyspark[connect]`, which is the same +thing plus the Spark JVM jars — those jars are the entire size difference, and +none of them are needed to talk to a remote session. + +Choose `[full]` if you also run Spark locally, or if you depend on other +packages that expect the full `pyspark` distribution. Otherwise `[client]` is +the smaller choice. -If you also need a local Spark runtime, install the full distribution instead: +The bare install is for environments that already provide Spark, such as a +Dataproc runtime image. On its own it cannot start a session, and importing the +package tells you so. Spark 4.0 or newer is required either way. + +Note that `pyspark-client` and `pyspark` both provide the `pyspark` module. +They are separate distributions, so pip will install both if asked rather than +report a conflict. Install one, not both, and switch by uninstalling the first: ```sh -pip uninstall pyspark-client -pip install 'pyspark[connect]~=4.0.0' +pip uninstall pyspark pyspark-client +pip install 'google-cloud-spark-connect[full]' ``` -The two packages both provide the `pyspark` module and cannot be installed side -by side. - ## Uninstall ```sh @@ -55,7 +73,7 @@ in your code using the builder API: 1. Install the latest version of Managed Spark Connect: ```sh - pip install -U google-cloud-spark-connect + pip install -U 'google-cloud-spark-connect[client]' ``` 2. Add the required imports into your PySpark application or notebook and start @@ -143,14 +161,23 @@ The package supports the [sparksql-magic](https://github.com/cryeo/sparksql-magi **Installation**: To use magic commands, install the required dependencies manually: ```bash -pip install google-cloud-spark-connect +pip install 'google-cloud-spark-connect[full]' +pip install IPython sparksql-magic +``` + +`sparksql-magic` declares a dependency on the full `pyspark` distribution, so +it installs cleanly next to `[full]`. If you prefer `[client]`, install it +without its dependencies, otherwise pip adds `pyspark` on top of +`pyspark-client` and the two shadow each other: + +```bash +pip install 'google-cloud-spark-connect[client]' pip install IPython pip install --no-deps sparksql-magic ``` -`sparksql-magic` declares a dependency on the full `pyspark` distribution. -Installing it with `--no-deps` keeps the lightweight `pyspark-client` package in -place; without it, pip installs `pyspark` on top and shadows the client. +It only imports `from pyspark.sql import SparkSession`, which `pyspark-client` +provides, so nothing is lost by skipping its dependencies. 1. Load the magic extension: ```python @@ -184,9 +211,9 @@ Available options: See [sparksql-magic](https://github.com/cryeo/sparksql-magic) for more examples. -**Note**: Magic commands are optional. If you only need basic ManagedSparkSession functionality without Jupyter magic support, install only the base package: +**Note**: Magic commands are optional. If you only need basic ManagedSparkSession functionality without Jupyter magic support, install the package on its own: ```bash -pip install google-cloud-spark-connect +pip install 'google-cloud-spark-connect[client]' ``` ## Migrating from dataproc-spark-connect @@ -200,7 +227,7 @@ The `dataproc-spark-connect` package has been renamed to `google-cloud-spark-con pip install dataproc-spark-connect # After -pip install google-cloud-spark-connect +pip install 'google-cloud-spark-connect[client]' ``` ### 2. Update your imports and session class diff --git a/google/cloud/managed_spark_connect/__init__.py b/google/cloud/managed_spark_connect/__init__.py index 7603edea..3ce3e322 100644 --- a/google/cloud/managed_spark_connect/__init__.py +++ b/google/cloud/managed_spark_connect/__init__.py @@ -14,23 +14,48 @@ import importlib.metadata import warnings +from packaging import version -def _check_pyspark_installation(): - """Warns when both pyspark distributions are installed at odds. +_MIN_PYSPARK_VERSION = "4.0" - 'pyspark-client' and 'pyspark' both provide the 'pyspark' module, so pip - installs them side by side rather than reporting a conflict. At the same - version their shared files are identical and this is harmless. At - different versions the environment ends up a mix of the two, and the - failure surfaces later as an import error naming neither package. - """ + +def _installed_version(distribution): + """Returns the installed version of a distribution, or None if absent.""" try: - client_version = importlib.metadata.version("pyspark-client") - full_version = importlib.metadata.version("pyspark") + return importlib.metadata.version(distribution) except importlib.metadata.PackageNotFoundError: - return + return None + + +def _check_pyspark_installation(): + """Checks the Spark distribution this package was installed alongside. + + This package depends on no Spark distribution of its own, so that it uses + whichever one is already present. 'pyspark-client' and 'pyspark' both + provide the 'pyspark' module but are separate distributions, so pip cannot + see them as alternatives and neither can be depended on without risking a + second copy landing over the first. + + That leaves three states worth reporting, since each of them otherwise + surfaces as an import error that names nothing recognizable. + """ + client_version = _installed_version("pyspark-client") + full_version = _installed_version("pyspark") - if client_version != full_version: + if client_version is None and full_version is None: + raise ImportError( + "No Spark distribution is installed. google-cloud-spark-connect " + "needs either 'pyspark-client', for remote Managed Spark Sessions " + "only, or 'pyspark', which also runs Spark locally. Install one " + "of them with 'pip install google-cloud-spark-connect[client]' or " + "'pip install google-cloud-spark-connect[full]'." + ) + + if ( + client_version is not None + and full_version is not None + and client_version != full_version + ): warnings.warn( f"Both 'pyspark-client' ({client_version}) and 'pyspark' " f"({full_version}) are installed, at different versions. They " @@ -38,10 +63,29 @@ def _check_pyspark_installation(): "mix of the two and imports may fail in ways that mention " "neither. Uninstall both and reinstall only the one you need: " "'pip uninstall pyspark pyspark-client', then " - "'pip install pyspark-client' to use remote Managed Spark " - "Sessions, or 'pip install pyspark[connect]' if you also run " + "'pip install google-cloud-spark-connect[client]' to use remote " + "Managed Spark Sessions, or " + "'pip install google-cloud-spark-connect[full]' if you also run " "Spark locally." ) + return + + installed_version = client_version or full_version + try: + too_old = version.parse(installed_version) < version.parse( + _MIN_PYSPARK_VERSION + ) + except version.InvalidVersion: + return + + if too_old: + warnings.warn( + f"Spark {installed_version} is installed, but " + "google-cloud-spark-connect uses Spark Connect APIs introduced in " + f"Spark {_MIN_PYSPARK_VERSION}. Upgrade with " + "'pip install google-cloud-spark-connect[client]' or " + "'pip install google-cloud-spark-connect[full]'." + ) _check_pyspark_installation() diff --git a/requirements-local-spark.txt b/requirements-local-spark.txt index 04967a7b..cc5213b8 100644 --- a/requirements-local-spark.txt +++ b/requirements-local-spark.txt @@ -1,10 +1,11 @@ -# Dependencies for the tests that need a local Spark runtime. +# Dependencies for the tests that need a local Spark runtime. This is what the +# [full] extra installs. # -# The library itself depends on pyspark-client, which has no JVM jars and -# cannot start a local session. The Dataproc batch code path hands off to a -# local classic Spark session, so testing it needs the full distribution. +# Development runs against pyspark-client, which has no JVM jars and cannot +# start a local session. The Dataproc batch code path hands off to a local +# classic Spark session, so testing it needs the full distribution. # # Install this on top of requirements-dev.txt, never instead of it, and only # for those tests: the unit and integration suites are meant to run against -# pyspark-client so they keep exercising the dependency set we ship. +# pyspark-client so they keep exercising the smaller of the two installs. pyspark[connect]~=4.0.0 diff --git a/setup.py b/setup.py index c629120c..ada0ea50 100644 --- a/setup.py +++ b/setup.py @@ -32,8 +32,16 @@ "google-api-core>=2.19", "google-cloud-dataproc>=5.18", "packaging>=20.0", - "pyspark-client~=4.0.0", "tqdm>=4.67", "websockets>=14.0", ], + # The base install deliberately names no Spark distribution, so it works + # with whichever one the environment already has. 'pyspark-client' and + # 'pyspark' both provide the 'pyspark' module but are separate + # distributions, so depending on either would install a second copy over + # the one already present. These extras are shorthand for picking one. + extras_require={ + "client": ["pyspark-client~=4.0.0"], + "full": ["pyspark[connect]~=4.0.0"], + }, ) diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index e60ea1b4..f9c0fe09 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -73,7 +73,7 @@ def test_no_warning_without_full_pyspark(self): mock_warn.assert_not_called() def test_no_warning_without_pyspark_client(self): - """The documented escape hatch: the full distribution alone""" + """What the [full] extra installs: the full distribution alone""" mock_warn = self._run_with_versions( { "pyspark-client": importlib.metadata.PackageNotFoundError( @@ -85,6 +85,46 @@ def test_no_warning_without_pyspark_client(self): mock_warn.assert_not_called() + def test_raises_when_no_spark_is_installed(self): + """A bare install has no Spark until an extra supplies one""" + with mock.patch( + "importlib.metadata.version", + side_effect=importlib.metadata.PackageNotFoundError, + ): + with self.assertRaises(ImportError) as context: + _check_pyspark_installation() + + message = str(context.exception) + self.assertIn("google-cloud-spark-connect[client]", message) + self.assertIn("google-cloud-spark-connect[full]", message) + + def test_warns_when_spark_is_too_old(self): + """The Spark Connect APIs this package uses arrived in Spark 4.0""" + mock_warn = self._run_with_versions( + { + "pyspark-client": importlib.metadata.PackageNotFoundError( + "pyspark-client" + ), + "pyspark": "3.5.1", + } + ) + + mock_warn.assert_called_once() + message = mock_warn.call_args[0][0] + self.assertIn("3.5.1", message) + self.assertIn("4.0", message) + + def test_no_warning_for_unparseable_version(self): + """A version we cannot read is not grounds for a warning""" + mock_warn = self._run_with_versions( + { + "pyspark-client": "not-a-version", + "pyspark": importlib.metadata.PackageNotFoundError("pyspark"), + } + ) + + mock_warn.assert_not_called() + class TestPythonVersionCheck(unittest.TestCase): From 7f01c2ab8aae8f80505cb467e53bc5a59c9b1b23 Mon Sep 17 00:00:00 2001 From: Andrew Ma <136692+ajma@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:26:00 -0700 Subject: [PATCH 6/6] fix: report a missing Spark by what is importable, not what is installed The installation check read distribution metadata, so it raised on any environment that provides Spark without installing a distribution. Runtime images commonly put SPARK_HOME/python on the path instead, which is precisely the case the bare install exists to serve, and the check rejected it. Fall back to whether pyspark can be found when neither distribution is installed, and wrap the import that would otherwise fail so a missing pyspark reports the extras rather than the module name. Other missing modules are left to surface on their own. --- .../cloud/managed_spark_connect/__init__.py | 45 ++++++++++++++--- tests/unit/test_init.py | 49 +++++++++++++++++-- 2 files changed, 82 insertions(+), 12 deletions(-) diff --git a/google/cloud/managed_spark_connect/__init__.py b/google/cloud/managed_spark_connect/__init__.py index 3ce3e322..456a96e3 100644 --- a/google/cloud/managed_spark_connect/__init__.py +++ b/google/cloud/managed_spark_connect/__init__.py @@ -12,12 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. import importlib.metadata +import importlib.util import warnings from packaging import version _MIN_PYSPARK_VERSION = "4.0" +_NO_SPARK_MESSAGE = ( + "No Spark distribution is importable. google-cloud-spark-connect needs " + "either 'pyspark-client', for remote Managed Spark Sessions only, or " + "'pyspark', which also runs Spark locally. Install one of them with " + "'pip install google-cloud-spark-connect[client]' or " + "'pip install google-cloud-spark-connect[full]'." +) + def _installed_version(distribution): """Returns the installed version of a distribution, or None if absent.""" @@ -27,6 +36,18 @@ def _installed_version(distribution): return None +def _spark_import_error(exc): + """Returns a clearer error for a missing pyspark, or None to re-raise. + + Only failures to import pyspark itself are worth rewriting. Anything else + missing is a separate problem and should surface as it is. + """ + name = exc.name or "" + if name == "pyspark" or name.startswith("pyspark."): + return ImportError(_NO_SPARK_MESSAGE) + return None + + def _check_pyspark_installation(): """Checks the Spark distribution this package was installed alongside. @@ -43,13 +64,13 @@ def _check_pyspark_installation(): full_version = _installed_version("pyspark") if client_version is None and full_version is None: - raise ImportError( - "No Spark distribution is installed. google-cloud-spark-connect " - "needs either 'pyspark-client', for remote Managed Spark Sessions " - "only, or 'pyspark', which also runs Spark locally. Install one " - "of them with 'pip install google-cloud-spark-connect[client]' or " - "'pip install google-cloud-spark-connect[full]'." - ) + # Neither distribution is installed, but Spark may still be importable: + # runtime images commonly put SPARK_HOME/python on the path instead of + # installing a distribution. Only an unimportable pyspark is a problem, + # and an unmanaged one tells us no version we can go on. + if importlib.util.find_spec("pyspark") is None: + raise ImportError(_NO_SPARK_MESSAGE) + return if ( client_version is not None @@ -90,7 +111,15 @@ def _check_pyspark_installation(): _check_pyspark_installation() -from .session import ManagedSparkSession +try: + from .session import ManagedSparkSession +except ModuleNotFoundError as e: + # The check above reads what is installed. This catches what actually + # failed to import, which covers a pyspark that is present but incomplete. + _error = _spark_import_error(e) + if _error is None: + raise + raise _error from e old_package_names = ["google-spark-connect", "dataproc-spark-connect"] current_package_name = "google-cloud-spark-connect" diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index f9c0fe09..4954091d 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -15,7 +15,10 @@ import unittest from unittest import mock -from google.cloud.managed_spark_connect import _check_pyspark_installation +from google.cloud.managed_spark_connect import ( + _check_pyspark_installation, + _spark_import_error, +) from google.cloud.managed_spark_connect.session import ManagedSparkSession from google.cloud.managed_spark_connect.exceptions import ManagedSparkConnectException @@ -85,19 +88,34 @@ def test_no_warning_without_pyspark_client(self): mock_warn.assert_not_called() - def test_raises_when_no_spark_is_installed(self): + def test_raises_when_spark_is_not_importable(self): """A bare install has no Spark until an extra supplies one""" with mock.patch( "importlib.metadata.version", side_effect=importlib.metadata.PackageNotFoundError, ): - with self.assertRaises(ImportError) as context: - _check_pyspark_installation() + with mock.patch("importlib.util.find_spec", return_value=None): + with self.assertRaises(ImportError) as context: + _check_pyspark_installation() message = str(context.exception) self.assertIn("google-cloud-spark-connect[client]", message) self.assertIn("google-cloud-spark-connect[full]", message) + def test_accepts_spark_without_distribution_metadata(self): + """Runtime images put SPARK_HOME/python on the path, not a dist""" + with mock.patch( + "importlib.metadata.version", + side_effect=importlib.metadata.PackageNotFoundError, + ): + with mock.patch( + "importlib.util.find_spec", return_value=mock.Mock() + ): + with mock.patch("warnings.warn") as mock_warn: + _check_pyspark_installation() + + mock_warn.assert_not_called() + def test_warns_when_spark_is_too_old(self): """The Spark Connect APIs this package uses arrived in Spark 4.0""" mock_warn = self._run_with_versions( @@ -114,6 +132,29 @@ def test_warns_when_spark_is_too_old(self): self.assertIn("3.5.1", message) self.assertIn("4.0", message) + def test_rewrites_a_missing_pyspark_import(self): + """A missing pyspark should name the extras, not the module""" + for missing in ["pyspark", "pyspark.sql.connect.session"]: + with self.subTest(missing=missing): + error = _spark_import_error( + ModuleNotFoundError( + f"No module named '{missing}'", name=missing + ) + ) + + self.assertIsInstance(error, ImportError) + self.assertIn("google-cloud-spark-connect[client]", str(error)) + + def test_leaves_other_missing_imports_alone(self): + """Anything else missing is a different problem entirely""" + for missing in ["tqdm", "websockets", None]: + with self.subTest(missing=missing): + error = _spark_import_error( + ModuleNotFoundError("No module named", name=missing) + ) + + self.assertIsNone(error) + def test_no_warning_for_unparseable_version(self): """A version we cannot read is not grounds for a warning""" mock_warn = self._run_with_versions(