From 21cb40ec62cde1f632223e434f2d45977d8f4552 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Fri, 2 Oct 2026 08:48:51 -0700 Subject: [PATCH] feat!: Move PySpark dependencies into [client] and [spark] extras pyspark is an ~800 MB install that most Spark Connect users do not need, while pyspark-client is a lightweight (~2 MB) pure-Python Spark Connect client. However, pyspark and pyspark-client write to the same pyspark module directory, so installing both in a single environment risks immediate version skew or broken imports if either package is later uninstalled. Because Python package extras are strictly additive, move the Spark dependency out of install_requires into two mutually exclusive extras: - [client]: depends on pyspark-client~=4.0.0 - [spark]: depends on pyspark[connect]~=4.0.0 - no Spark extra: relies on whichever pyspark provider is already present in the environment If pyspark is missing or incomplete (such as a leftover pyspark/ directory after uninstalling one of the two packages) when google.cloud.managed_spark_connect is imported, raise PySparkNotInstalledError recommending google-cloud-spark-connect[client,interactive] when inside an IPython shell, or google-cloud-spark-connect[client] otherwise. It subclasses both ManagedSparkConnectException and ImportError, so existing `except ImportError` guards keep working. As install_requires no longer pins pyspark, also reject a pyspark older than 4.0 on import with UnsupportedPySparkVersionError, a PySparkNotInstalledError subclass, recommending an in-place upgrade via [spark]. %%sparksql gives the same install hint when pyspark is missing. Update README.md to recommend [client,interactive] as the primary installation command and [spark,interactive] when upgrading an existing dataproc-spark-connect environment in-place. Switch requirements-dev.txt to pyspark-client~=4.0.0, and document how to switch a dev environment to full pyspark. CI runs unit tests against each of the [client] and [spark] extras; the local Spark tests need a JVM, so they only run against [spark]. BREAKING CHANGE: pyspark[connect] is no longer installed by default. Install google-cloud-spark-connect[client] (or [client,interactive] for notebooks), google-cloud-spark-connect[spark], or provide pyspark / pyspark-client in the environment. Release-As: 2.0.0 --- .github/workflows/tests.yaml | 20 ++- DEVELOPING.md | 24 ++- README.md | 55 ++++++- .../cloud/managed_spark_connect/__init__.py | 67 +++++++- .../cloud/managed_spark_connect/exceptions.py | 20 +++ google/cloud/managed_spark_magics/sparksql.py | 5 +- requirements-dev.txt | 2 +- setup.py | 12 +- tests/integration/test_session.py | 5 + .../managed_spark_magics/test_sparksql.py | 8 + tests/unit/test_init.py | 153 ++++++++++++++++++ tests/unit/test_session.py | 16 ++ 12 files changed, 370 insertions(+), 17 deletions(-) diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index d298beb..94b9722 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -23,8 +23,14 @@ permissions: jobs: test: - name: Run unit tests + name: Run unit tests ([${{ matrix.spark-extra }}]) runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + # Mutually exclusive: [client] installs pyspark-client, [spark] full + # pyspark[connect]. + spark-extra: ["client", "spark"] steps: - name: Checkout code @@ -39,20 +45,26 @@ jobs: uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: path: ~/.cache/pip - key: ${{ runner.os }}-pip-${{ hashFiles('requirements-dev.txt', 'requirements-test.txt') }} + key: ${{ runner.os }}-pip-${{ matrix.spark-extra }}-${{ hashFiles('setup.py', 'requirements-test.txt') }} restore-keys: | + ${{ runner.os }}-pip-${{ matrix.spark-extra }}- ${{ runner.os }}-pip- + # Installs through the package's own extras, rather than + # requirements-dev.txt (which pins pyspark-client), so that both extras are + # exercised as users would install them. - name: Install dependencies run: | - pip install -r requirements-dev.txt + pip install -e '.[${{ matrix.spark-extra }},interactive]' pip install -r requirements-test.txt - name: Run unit tests run: python -m pytest tests/unit/ -v --tb=short -n auto - # Starts Spark in-process, so it runs serially, outside of xdist. + # Starts Spark in-process, so it runs serially, outside of xdist. Needs + # the JVM that only full pyspark ships. - name: Run local Spark tests + if: matrix.spark-extra == 'spark' run: python -m pytest tests/local/ -v --tb=short env: JAVA_HOME: ${{ env.JAVA_HOME_17_X64 }} \ No newline at end of file diff --git a/DEVELOPING.md b/DEVELOPING.md index 06a1cae..23b5a02 100644 --- a/DEVELOPING.md +++ b/DEVELOPING.md @@ -42,8 +42,9 @@ env \ ## Testing the Interactive Extras The dependencies of the interactive extras (`explore_dataframe`, `%dpip`, -`%%sparksql`) live in the `interactive` extra (`pip install '.[interactive]'`), -and `requirements-dev.txt` already includes them. The tests in +`%%sparksql`) live in the `interactive` extra +(`pip install '.[client,interactive]'`), and `requirements-dev.txt` already +includes them along with `pyspark-client`. The tests in `tests/unit/test_init.py` assume `google-colabsqlviz`, `ipython`, `sqlparse` and `traitlets` are importable and will fail, not skip, if they are not. Missing dependencies are simulated within the tests themselves, by setting @@ -71,3 +72,22 @@ env \ MANAGED_SPARK_CONNECT_SERVICE_ACCOUNT='service@account.test' \ pytest -n auto --tb=auto -v ``` + +## Testing Against Full PySpark + +`requirements-dev.txt` installs the lightweight `pyspark-client`, under which +`tests/local` and the integration tests that need a local JVM skip. To run them, +switch the environment to full `pyspark`. Both packages write to the same +`pyspark` directory, so always uninstall both before installing either: + +```sh +pip uninstall -y pyspark-client pyspark +pip install 'pyspark[connect]~=4.0.0' +``` + +To switch back to `pyspark-client`: + +```sh +pip uninstall -y pyspark-client pyspark +pip install 'pyspark-client~=4.0.0' +``` diff --git a/README.md b/README.md index 135812c..196f036 100644 --- a/README.md +++ b/README.md @@ -7,16 +7,44 @@ requiring additional steps. ## Install +For notebooks and interactive development, install with the `client` and +`interactive` extras to get the lightweight `pyspark-client` package along with +the interactive extras (`explore_dataframe()`, `%dpip`, `%%sparksql`): + ```sh -pip install google-cloud-spark-connect +pip install 'google-cloud-spark-connect[client,interactive]' ``` +Other installation options: + +* `pip install 'google-cloud-spark-connect[client]'`: Installs `pyspark-client` + without the interactive extras (for standalone scripts or applications). +* `pip install 'google-cloud-spark-connect[spark]'` (or + `'[spark,interactive]'`): Installs the full `pyspark[connect]` distribution + instead of `pyspark-client`, for environments that also need a local Spark JVM + or already have full `pyspark` installed. +* `pip install google-cloud-spark-connect` (or `'[interactive]'`): Does not + install a Spark package, relying on whichever `pyspark` or `pyspark-client` is + already present in your environment (such as pre-installed Dataproc or Colab + runtime images). Use this in Dataproc batch and Serverless runtimes, which + must keep their bundled `pyspark`. + +> **Note:** Avoid installing both `pyspark` and `pyspark-client` in the same +> Python environment. Both packages write to the `pyspark` module namespace, so +> mixing them can cause version skew or break the remaining package if one is +> later uninstalled. In particular, never combine the `client` and `spark` +> extras (`'[client,spark]'`). + ## Uninstall ```sh -pip uninstall google-cloud-spark-connect +pip uninstall google-cloud-spark-connect pyspark-client ``` +`pip uninstall` never removes dependencies, so this also names the +`pyspark-client` that the `client` extra installed. If you installed the `spark` +extra instead, uninstall `pyspark` too if nothing else needs it. + ## Setup This client requires permissions to @@ -39,7 +67,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,interactive]' ``` 2. Add the required imports into your PySpark application or notebook and start @@ -123,10 +151,11 @@ To create or connect to a named session: ### Interactive Extras -For notebooks and other interactive use, install the `interactive` extra: +For notebooks and other interactive use, include the `interactive` extra (along +with `client` or `spark`, unless PySpark is already in your environment): ```sh -pip install 'google-cloud-spark-connect[interactive]' +pip install 'google-cloud-spark-connect[client,interactive]' ``` When you import the package inside an IPython kernel, it then automatically sets @@ -207,12 +236,24 @@ The `dataproc-spark-connect` package has been renamed to `google-cloud-spark-con ### 1. Update the package you install +`dataproc-spark-connect` previously pulled in full `pyspark` by default. When +upgrading an existing environment in-place, use the `spark` extra so `pip` +reuses your existing `pyspark` installation without installing `pyspark-client` +alongside it (along with `interactive` for the interactive extras): + ```sh # Before pip install dataproc-spark-connect -# After -pip install google-cloud-spark-connect +# After (upgrading an existing environment with pyspark already installed) +pip install 'google-cloud-spark-connect[spark,interactive]' +``` + +In a fresh environment (or after running `pip uninstall dataproc-spark-connect pyspark`), +prefer the lightweight `client` extra instead: + +```sh +pip install 'google-cloud-spark-connect[client,interactive]' ``` ### 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 06082a2..6738db2 100644 --- a/google/cloud/managed_spark_connect/__init__.py +++ b/google/cloud/managed_spark_connect/__init__.py @@ -15,7 +15,72 @@ import importlib.util import warnings -from .session import ManagedSparkSession + +def _in_ipython() -> bool: + # Not environment.is_interactive(), which is also true in a plain + # `python -i` REPL, where the [interactive] extras do nothing. + try: + from IPython import get_ipython + + return get_ipython() is not None + except Exception: + return False + + +def _install_extras(spark_extra: str) -> str: + """The extras to recommend: also [interactive] inside IPython.""" + return f"{spark_extra},interactive" if _in_ipython() else spark_extra + + +def _check_pyspark_version(pyspark_version: str) -> None: + from packaging.version import InvalidVersion, Version + + try: + too_old = Version(pyspark_version).major < 4 + except InvalidVersion: + # E.g. a custom build. Let the imports decide whether it works. + return + if too_old: + from .exceptions import UnsupportedPySparkVersionError + + # pyspark-client starts at 4.0, so this is full pyspark: upgrade it in + # place with [spark] rather than adding pyspark-client next to it. + raise UnsupportedPySparkVersionError( + f"PySpark {pyspark_version} is installed, but " + "google-cloud-spark-connect requires PySpark 4.0 or later. " + "Upgrade it with: pip install -U " + f"'google-cloud-spark-connect[{_install_extras('spark')}]'" + ) + + +# Needs nothing beyond the standard library, so it is safe to import before +# pyspark has been checked. +from .exceptions import PySparkNotInstalledError + +try: + # pyspark.version first, so that an old pyspark gets a clear error rather + # than whichever ImportError session.py happens to hit first. + from pyspark.version import __version__ as _pyspark_version + + _check_pyspark_version(_pyspark_version) + + from .session import ManagedSparkSession +except PySparkNotInstalledError: + # From _check_pyspark_version(): already an ImportError with a clear + # message, not one to rewrite as "missing or incomplete". + raise +except ImportError as e: + # Only translate errors about pyspark itself (missing, a leftover + # namespace directory, a partial install), not about its dependencies, + # which pyspark already reports clearly. + if (e.name or "").partition(".")[0] != "pyspark": + raise + + raise PySparkNotInstalledError( + f"PySpark is missing or incomplete in this environment ({e}). " + "Install it with: pip install " + f"'google-cloud-spark-connect[{_install_extras('client')}]'" + ) from e old_package_names = ["google-spark-connect", "dataproc-spark-connect"] current_package_name = "google-cloud-spark-connect" diff --git a/google/cloud/managed_spark_connect/exceptions.py b/google/cloud/managed_spark_connect/exceptions.py index afdea9e..16c5595 100644 --- a/google/cloud/managed_spark_connect/exceptions.py +++ b/google/cloud/managed_spark_connect/exceptions.py @@ -25,3 +25,23 @@ def __init__(self, message): def _render_traceback_(self): return [self.message] + + +class PySparkNotInstalledError(ManagedSparkConnectException, ImportError): + """Raised on import when pyspark is missing or incomplete. + + Subclasses ImportError so that the usual optional-dependency guard, + ``try: import ... except ImportError:``, still works. + """ + + def __init__(self, message): + super().__init__(message) + self.name = "pyspark" + + +class UnsupportedPySparkVersionError(PySparkNotInstalledError): + """Raised on import when the installed pyspark is older than 4.0. + + A PySparkNotInstalledError, and so an ImportError, too: a pyspark this + old is as unusable as a missing one. + """ diff --git a/google/cloud/managed_spark_magics/sparksql.py b/google/cloud/managed_spark_magics/sparksql.py index fd4c1f0..122148e 100644 --- a/google/cloud/managed_spark_magics/sparksql.py +++ b/google/cloud/managed_spark_magics/sparksql.py @@ -566,7 +566,10 @@ def resolve_session(user_ns: Mapping[str, Any], name: Optional[str]) -> Any: """ classes = _session_classes() if not classes: - raise UsageError("%%sparksql needs pyspark, which is not installed.") + raise UsageError( + "%%sparksql needs pyspark, which is not installed. Install it with:" + " pip install 'google-cloud-spark-connect[client,interactive]'" + ) if name is not None: if name not in user_ns: raise UsageError(f"-s {name}: no such variable") diff --git a/requirements-dev.txt b/requirements-dev.txt index f3c15fc..772be7c 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -5,7 +5,7 @@ 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 sqlparse>=0.5.0 tqdm>=4.67 diff --git a/setup.py b/setup.py index ed02808..f971bb7 100644 --- a/setup.py +++ b/setup.py @@ -32,11 +32,21 @@ "google-api-core>=2.19", "google-cloud-dataproc>=5.18", "packaging>=20.0", - "pyspark[connect]~=4.0.0", "tqdm>=4.67", "websockets>=14.0", ], extras_require={ + # client and spark are mutually exclusive: install at most one. Both + # packages write to the same pyspark module directory, so installing + # both risks version skew, and uninstalling either breaks the other. + # Extras are additive, so pip cannot enforce this. With neither, the + # package uses whichever pyspark the environment already provides. + "client": [ + "pyspark-client~=4.0.0", + ], + "spark": [ + "pyspark[connect]~=4.0.0", + ], # Everything here is optional at runtime: managed_spark_connect._ipython # only initializes the extras whose dependencies are importable. "interactive": [ diff --git a/tests/integration/test_session.py b/tests/integration/test_session.py index b2cf110..77bf133 100644 --- a/tests/integration/test_session.py +++ b/tests/integration/test_session.py @@ -33,6 +33,7 @@ ) 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" @@ -741,6 +742,10 @@ def local_spark_session(): session.stop() +@pytest.mark.skipif( + is_remote_only(), + reason="Local SparkSession requires full pyspark (not pyspark-client).", +) 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 diff --git a/tests/unit/managed_spark_magics/test_sparksql.py b/tests/unit/managed_spark_magics/test_sparksql.py index 6183c67..bd150f0 100644 --- a/tests/unit/managed_spark_magics/test_sparksql.py +++ b/tests/unit/managed_spark_magics/test_sparksql.py @@ -309,6 +309,14 @@ def test_ambiguous_and_missing(self): with self.assertRaisesRegex(UsageError, "No SparkSession found"): sparksql.resolve_session({}, None) + def test_no_pyspark_recommends_install(self): + with mock.patch.object(sparksql, "_session_classes", return_value=()): + with self.assertRaisesRegex( + UsageError, + r"pip install 'google-cloud-spark-connect\[client,interactive\]'", + ): + sparksql.resolve_session({}, None) + class FakeContext: diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index cf164ad..ea92d8e 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -1047,5 +1047,158 @@ def test_import_without_interactive_deps(self): self.assertEqual(result.stdout.split(), ["True", "False", "False"]) +_NO_PYSPARK_PROBE = """ +import importlib.machinery, importlib.util, json, os, sys, types + +ipython_mode, pyspark_mode, tmp = sys.argv[1:] + +if ipython_mode == "in_ipython": + from IPython.core.interactiveshell import InteractiveShell + InteractiveShell.instance() +elif ipython_mode == "no_ipython": + sys.modules["IPython"] = None +elif ipython_mode == "broken_ipython": + fake_ipython = types.ModuleType("IPython") + def _boom(): + raise RuntimeError("broken IPython") + fake_ipython.get_ipython = _boom + sys.modules["IPython"] = fake_ipython + +if pyspark_mode == "absent": + sys.modules["pyspark"] = None +elif pyspark_mode == "namespace": + # What a leftover pyspark/ directory without __init__.py imports as. + spec = importlib.machinery.ModuleSpec("pyspark", None, is_package=True) + spec.submodule_search_locations = [os.path.join(tmp, "pyspark")] + sys.modules["pyspark"] = importlib.util.module_from_spec(spec) +else: + # A fake pyspark package, created by the test, shadows the real one. + sys.path.insert(0, tmp) + +try: + import google.cloud.managed_spark_connect +except Exception as e: + print(json.dumps({ + "type": f"{type(e).__module__}.{type(e).__qualname__}", + "message": str(e), + "is_import_error": isinstance(e, ImportError), + "name": getattr(e, "name", None), + })) +""" + +_PYSPARK_MISSING_PREFIX = ( + "PySpark is missing or incomplete in this environment (" +) + + +class TestWithoutPySpark(unittest.TestCase): + """Importing without a usable pyspark must fail fast with a clear error. + + Runs out of process so that the real pyspark, already imported by the test + runner, can be hidden. + """ + + def _run(self, ipython_mode, pyspark_mode, pyspark_version=None): + import json + import os + import shutil + import subprocess + import sys + import tempfile + + tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, tmp, ignore_errors=True) + os.mkdir(os.path.join(tmp, "pyspark")) + if pyspark_version is not None: + with open(os.path.join(tmp, "pyspark", "__init__.py"), "w"): + pass + with open(os.path.join(tmp, "pyspark", "version.py"), "w") as f: + f.write(f"__version__ = {pyspark_version!r}\n") + + env = dict(os.environ) + env["IPYTHONDIR"] = os.path.join(tmp, "ipython") + result = subprocess.run( + [sys.executable, "-c", _NO_PYSPARK_PROBE] + + [ipython_mode, pyspark_mode, tmp], + capture_output=True, + text=True, + env=env, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertNotIn("Traceback", result.stderr) + self.assertTrue(result.stdout.strip(), "import unexpectedly succeeded") + return json.loads(result.stdout) + + def _assert_missing_pyspark(self, error, extras): + self.assertEqual( + error["type"], + "google.cloud.managed_spark_connect.exceptions.PySparkNotInstalledError", + ) + self.assertTrue( + error["message"].startswith(_PYSPARK_MISSING_PREFIX), + error["message"], + ) + self.assertTrue( + error["message"].endswith( + "). Install it with: pip install " + f"'google-cloud-spark-connect[{extras}]'" + ), + error["message"], + ) + + def test_missing_pyspark_outside_ipython_recommends_client(self): + error = self._run("plain", "absent") + self._assert_missing_pyspark(error, "client") + + def test_missing_pyspark_in_ipython_recommends_client_interactive(self): + error = self._run("in_ipython", "absent") + self._assert_missing_pyspark(error, "client,interactive") + + def test_missing_pyspark_with_missing_or_broken_ipython_recommends_client( + self, + ): + for mode in ("no_ipython", "broken_ipython"): + with self.subTest(mode=mode): + error = self._run(mode, "absent") + self._assert_missing_pyspark(error, "client") + + def test_missing_pyspark_is_an_import_error_named_pyspark(self): + error = self._run("plain", "absent") + self.assertTrue(error["is_import_error"]) + self.assertEqual(error["name"], "pyspark") + + def test_leftover_pyspark_namespace_dir_recommends_client(self): + error = self._run("plain", "namespace") + self._assert_missing_pyspark(error, "client") + self.assertIn("pyspark.version", error["message"]) + + def test_partial_pyspark_recommends_client(self): + error = self._run("plain", "partial", pyspark_version="4.0.0") + self._assert_missing_pyspark(error, "client") + self.assertIn("pyspark.sql", error["message"]) + + def test_old_pyspark_recommends_upgrading_via_spark_extra(self): + for ipython_mode, extras in ( + ("plain", "spark"), + ("in_ipython", "spark,interactive"), + ): + with self.subTest(ipython_mode=ipython_mode): + error = self._run(ipython_mode, "old", pyspark_version="3.5.1") + self.assertEqual( + error["type"], + "google.cloud.managed_spark_connect.exceptions.UnsupportedPySparkVersionError", + ) + # Caught by the same `except ImportError:` guards as a + # missing pyspark. + self.assertTrue(error["is_import_error"]) + self.assertEqual(error["name"], "pyspark") + self.assertEqual( + error["message"], + "PySpark 3.5.1 is installed, but google-cloud-spark-connect " + "requires PySpark 4.0 or later. Upgrade it with: " + f"pip install -U 'google-cloud-spark-connect[{extras}]'", + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index dbd4811..46814a6 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -2667,6 +2667,22 @@ def test_wait_for_session_available_timeout( str(context.exception), ) + @mock.patch.dict(os.environ, {"DATAPROC_WORKLOAD_TYPE": "batch"}) + @mock.patch("pyspark.sql.SparkSession") + def test_get_or_create_dataproc_batch_returns_local_session( + self, mock_pyspark_sql_session + ): + """Test that getOrCreate delegates to pyspark.sql.SparkSession in batch mode.""" + mock_local_session = mock.Mock() + mock_pyspark_sql_session.builder.getOrCreate.return_value = ( + mock_local_session + ) + + session = ManagedSparkSession.builder.getOrCreate() + + self.assertIs(session, mock_local_session) + mock_pyspark_sql_session.builder.getOrCreate.assert_called_once_with() + class SessionIdValidationTests(unittest.TestCase): """Test cases for session ID validation and custom session ID functionality."""