feat!: Move PySpark dependencies into [client] and [spark] extras

.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 }}

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'
+```

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

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"

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.
+    """

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")

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

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": [

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

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: