Depend on pyspark-client instead of the full pyspark distribution

.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

.github/workflows/tests.yaml

@@ -49,4 +49,37 @@ jobs:
         pip install -r requirements-test.txt
 
     - name: Run unit tests
-      run: python -m pytest tests/unit/ -v --tb=short -n auto
+      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', 'requirements-local-spark.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
+        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.
+    - name: Run local Spark tests
+      run: python -m pytest tests/integration/test_session.py -v --tb=short -k test_create_local_spark_session

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
@@ -45,9 +56,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

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

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

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

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

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",
     ],

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

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