fix: support Spark Connect sessions in SparkSQLCompare - #552
Open
fdosani wants to merge 5 commits into
Open
Conversation
|
Tested the dispatch approach end-to-end locally against a Connect session — the |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
SparkSQLComparedid not work against a Spark Connect session it was handed rather than one it created (issue #535). The root cause is thatpyspark.sql.functionsandpyspark.sql.Windowonly forward to the Spark Connect implementations while the process-globalSPARK_CONNECT_MODE_ENABLEDenvironment 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 doingisinstanceagainst the Connect classes, so a classic-only PySpark install never triggers thepyspark.sql.connectimport that requires the optionalgrpciodependencywithColumns)Closes #535. Supersedes #548.
Library changes
Expression dispatch (
datacompy/comparator/utility.py)is_spark_connect_object(),get_spark_functions(), andget_spark_window()resolve the classic or Connectfunctionsmodule andWindowclass from the DataFrame or Column at handdatacompy/spark.pyand in the array, boolean, numeric, and string comparators resolvesF/Windowthis way instead of importing them at module scopeOther Connect breakages (
datacompy/spark.py,datacompy/base.py)df_to_str()checkedhasattr(df, "to_string")beforehasattr(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. ThetoPandascheck now comes first._dataframe_merge()registered fixeddf1/df2temp 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 calleddf1ordf2._intersect_compare()now guardswithColumns()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._validate_dataframe()no longer importspyspark.sql.connect.dataframe, for the samegrpcioreason as above.Tests
tests/test_spark_connect.py: end-to-endSparkSQLComparecoverage against a real Connect session. The session fixture clearsSPARK_CONNECT_MODE_ENABLED, which is what makes these tests able to fail: with the flag left set,pyspark.sql.functionsforwards 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.iniruns the existingtests/test_spark.pyandtests/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 laterSparkSession.builder.getOrCreate()in that process returns the Connect session, so classic and Connect sessions cannot coexist in one run. The regression suite is markedspark_connectand deselected from the default run viaaddopts.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 viatox-conda, so the CI matrix can be reproduced locally without a system JDK. Its header documents two traps that cost real time to find: theopenjdk>=17.0.8,<18floor (below it, conda-forge selects a GraalVM build that silently replaces CPython with GraalPy) and{envpython}substituting to the stringNonein tox 3.28.Makefile:make test,test-ansi,test-connect,test-connect-regression,test-cov,test-all, plus-no-snowflakevariants.pytest.ini/pytest-ansi.ini/pytest-connect.ini: all three now settestpaths = tests. Without it a barepytestwalks the whole repo, and a worktree checked out under the repo root contributes a secondtests/conftest.py, which aborts the run withImportPathMismatchError. The Spark driver and the Connect server are also bound to127.0.0.1so tests do not depend on the host's outward-facing network config.pyproject.toml:ruff>=0.16on theqaextra, 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 theconnectextra,cache_intermediates=Falseon runtimes that forbid caching, and Arrow always being on the wire under Connectdocs/source/install.rst: whatdatacompy[spark]installs and what a self-managed PySpark needs for ConnectREADME.md: Spark listed as classic and Spark ConnectCLAUDE.md: the testing, tox, and dispatch notes aboveValidation
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-allruns the same set.