Skip to content

fix: support Spark Connect sessions in SparkSQLCompare - #552

Open
fdosani wants to merge 5 commits into
mainfrom
spark-connect-sessions
Open

fix: support Spark Connect sessions in SparkSQLCompare#552
fdosani wants to merge 5 commits into
mainfrom
spark-connect-sessions

Conversation

@fdosani

@fdosani fdosani commented Aug 6, 2026

Copy link
Copy Markdown
Member

Summary

SparkSQLCompare did not work against a Spark Connect session it was handed rather than one it created (issue #535). The root cause is that pyspark.sql.functions and pyspark.sql.Window only forward to the Spark Connect implementations while the process-global SPARK_CONNECT_MODE_ENABLED environment variable is set. A notebook runtime, a serverless runtime, or another framework that hands over an already-built Connect session does not necessarily set it, so datacompy built classic expressions against Connect DataFrames.

This branch selects the expression API from the object being operated on instead of from the global flag, and then backs that up with a Connect test session, a regression suite, and CI/tox coverage so the path stays green.

Credit

The dispatch approach here is the one proposed by @patrickswedish in #548, and that PR is what identified the process-global flag as the real cause. This branch keeps that design and extends it:

  • is_spark_connect_object() walks the MRO comparing module names rather than doing isinstance against the Connect classes, so a classic-only PySpark install never triggers the pyspark.sql.connect import that requires the optional grpcio dependency
  • fixes to the non-comparator Connect breakages found while building an end-to-end Connect test run (report rendering, temp views, empty withColumns)
  • a Spark Connect test session, a regression suite, and CI + tox coverage

Closes #535. Supersedes #548.

Library changes

Expression dispatch (datacompy/comparator/utility.py)

  • is_spark_connect_object(), get_spark_functions(), and get_spark_window() resolve the classic or Connect functions module and Window class from the DataFrame or Column at hand
  • every call site in datacompy/spark.py and in the array, boolean, numeric, and string comparators resolves F / Window this way instead of importing them at module scope

Other Connect breakages (datacompy/spark.py, datacompy/base.py)

  • df_to_str() checked hasattr(df, "to_string") before hasattr(df, "toPandas"). A Connect DataFrame synthesizes a Column for any unknown attribute, so the first check is always True for it and every Spark report took the pandas branch. The toPandas check now comes first.
  • _dataframe_merge() registered fixed df1 / df2 temp views. Connect resolves views lazily, so a second comparison on the same session replaced the views an earlier plan still referred to; the names are now unique per merge, which also stops datacompy clobbering a user view called df1 or df2.
  • _intersect_compare() now guards withColumns() on a non-empty mapping. Connect asserts on an empty one where classic Spark treats it as a no-op; this is reached when the two frames share nothing but the join columns.
  • the type check in _validate_dataframe() no longer imports pyspark.sql.connect.dataframe, for the same grpcio reason as above.

Tests

  • tests/test_spark_connect.py: end-to-end SparkSQLCompare coverage against a real Connect session. The session fixture clears SPARK_CONNECT_MODE_ENABLED, which is what makes these tests able to fail: with the flag left set, pyspark.sql.functions forwards to Connect on its own and the suite would pass against the unfixed library.
  • tests/comparator/test_utility_spark.py: unit coverage for the three new helpers.
  • pytest-connect.ini runs the existing tests/test_spark.py and tests/comparator/ against a Connect session, so the fix is checked against the full existing suite rather than only the new one.

The Connect suites must each run in their own pytest process: starting a local Connect server sets SPARK_LOCAL_REMOTE, after which every later SparkSession.builder.getOrCreate() in that process returns the Connect session, so classic and Connect sessions cannot coexist in one run. The regression suite is marked spark_connect and deselected from the default run via addopts.

CI and tooling

  • .github/workflows/test-package.yml: the Spark matrix was a cross product where the Spark jobs are the entire cost of a run (~20 min each) and the no-Spark jobs finish in under a minute. It now covers the axes independently: Python 3.10-3.13 breadth from the four cheap basic-install jobs, pandas 2 vs 3 and Spark 3.5 vs 4 each present once, and ANSI mode plus the two Spark Connect steps on the 3.12 baseline only. The known gap (no Spark job on 3.11 or 3.13) is recorded in the file's header along with what closing it costs.
  • tox.ini (new): mirrors that workflow job for job via tox-conda, so the CI matrix can be reproduced locally without a system JDK. Its header documents two traps that cost real time to find: the openjdk>=17.0.8,<18 floor (below it, conda-forge selects a GraalVM build that silently replaces CPython with GraalPy) and {envpython} substituting to the string None in tox 3.28.
  • Makefile: make test, test-ansi, test-connect, test-connect-regression, test-cov, test-all, plus -no-snowflake variants.
  • pytest.ini / pytest-ansi.ini / pytest-connect.ini: all three now set testpaths = tests. Without it a bare pytest walks the whole repo, and a worktree checked out under the repo root contributes a second tests/conftest.py, which aborts the run with ImportPathMismatchError. The Spark driver and the Connect server are also bound to 127.0.0.1 so tests do not depend on the host's outward-facing network config.
  • pyproject.toml: ruff>=0.16 on the qa extra, since the config uses selectors that older ruff fails to parse; dropped the Python 3.14 classifier, which claimed support for a version nothing in the test matrix covers.

Docs

  • docs/source/spark_usage.rst: a Spark Connect section covering the handed-over-session case and the reason dispatch is per-DataFrame, plus notes on the connect extra, cache_intermediates=False on runtimes that forbid caching, and Arrow always being on the wire under Connect
  • docs/source/install.rst: what datacompy[spark] installs and what a self-managed PySpark needs for Connect
  • README.md: Spark listed as classic and Spark Connect
  • CLAUDE.md: the testing, tox, and dispatch notes above

Validation

CI covers this: the 3.12 Spark 4 job runs the default suite, ANSI mode, the existing Spark suite against a Connect session, and the Connect regression suite. Locally, make test-all runs the same set.

@fdosani
fdosani requested review from a team as code owners August 6, 2026 15:45
@patrickswedish

Copy link
Copy Markdown

Tested the dispatch approach end-to-end locally against a Connect session — the is_spark_connect_object() MRO walk is clean, the grpcio import guard works, and the temp-view uniqueness fix is solid. Happy to help with any follow-up review items if the maintainers have questions. 🚀

@fdosani
fdosani requested review from ak-gupta and rhaffar August 7, 2026 02:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

SparkSQLCompare fails with Spark Connect sessions (pyspark.sql.connect.session.SparkSession)

2 participants