diff --git a/.github/workflows/integration-tests.yaml b/.github/workflows/integration-tests.yaml index a11aa52..f901ffd 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/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index abaa39f..0a5838e 100644 --- a/.github/workflows/tests.yaml +++ b/.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 \ 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', '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 \ No newline at end of file diff --git a/DEVELOPING.md b/DEVELOPING.md index c9a8a5e..1f3d89f 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 @@ -44,18 +55,24 @@ env \ To run tests with magic functionality, install the required dependencies manually: ```sh -pip install . -pip install IPython sparksql-magic +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 `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 c111792..afe9232 100644 --- a/README.md +++ b/README.md @@ -7,10 +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 ``` +`[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. + +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 pyspark-client +pip install 'google-cloud-spark-connect[full]' +``` + ## Uninstall ```sh @@ -39,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 @@ -127,10 +161,24 @@ 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 +``` + +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 %load_ext sparksql_magic @@ -163,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 @@ -179,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 2356758..456a96e 100644 --- a/google/cloud/managed_spark_connect/__init__.py +++ b/google/cloud/managed_spark_connect/__init__.py @@ -12,9 +12,114 @@ # See the License for the specific language governing permissions and # limitations under the License. import importlib.metadata +import importlib.util import warnings -from .session import ManagedSparkSession +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.""" + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + 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. + + 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 is None and full_version is None: + # 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 + 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 " + "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 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() + +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/requirements-dev.txt b/requirements-dev.txt index 5cf7026..b7bca08 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/requirements-local-spark.txt b/requirements-local-spark.txt new file mode 100644 index 0000000..cc5213b --- /dev/null +++ b/requirements-local-spark.txt @@ -0,0 +1,11 @@ +# Dependencies for the tests that need a local Spark runtime. This is what the +# [full] extra installs. +# +# 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 smaller of the two installs. +pyspark[connect]~=4.0.0 diff --git a/setup.py b/setup.py index 89637d5..ada0ea5 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[connect]~=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/integration/test_session.py b/tests/integration/test_session.py index c921713..d6c48e0 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 diff --git a/tests/unit/test_init.py b/tests/unit/test_init.py index 794f6fa..4954091 100644 --- a/tests/unit/test_init.py +++ b/tests/unit/test_init.py @@ -11,13 +11,162 @@ # 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, + _spark_import_error, +) 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): + """What the [full] extra installs: 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() + + 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 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( + { + "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_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( + { + "pyspark-client": "not-a-version", + "pyspark": importlib.metadata.PackageNotFoundError("pyspark"), + } + ) + + mock_warn.assert_not_called() + + class TestPythonVersionCheck(unittest.TestCase): def test_python_version_mismatch_warning_for_runtime_30(self):