Add Spark Connect compatibility guards by zhengruifeng · Pull Request #501 · databrickslabs/dbldatagen · GitHub
Skip to content

Add Spark Connect compatibility guards - #501

Draft
zhengruifeng wants to merge 2 commits into
databrickslabs:masterfrom
zhengruifeng:spark-connect-compatibility
Draft

zhengruifeng wants to merge 2 commits into
databrickslabs:masterfrom
zhengruifeng:spark-connect-compatibility

Conversation

@zhengruifeng

Copy link
Copy Markdown

Changes

  • Add Spark Connect detection helpers using PySpark's remote-session predicate where available, with environment and Spark 3 connect-class fallbacks.
  • Avoid forcing local Spark session creation when a caller does not pass a session; prefer the active/default session path so Spark Connect environments can supply a remote session.
  • Avoid SparkContext.defaultParallelism access for Spark Connect sessions and fall back to the existing default parallelism constant.
  • Skip Arrow Spark conf mutation for Spark Connect sessions and tolerate restricted conf writes in shared/remote environments.
  • Add unit coverage for Spark Connect-style sessions without requiring a live Connect server.

Linked issues

Related to #255.

Requirements

  • manually tested
  • updated documentation
  • updated demos
  • updated tests

Tests run:

  • JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 conda run -n spark-connect-tck-test python -m pytest tests/test_shared_env.py tests/test_serverless.py tests/test_spark_singleton.py
  • JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 conda run -n spark-connect-tck-test python -m pytest tests/test_basic_test.py
  • conda run -n spark-connect-tck-test python -m ruff check dbldatagen/spark_compat.py dbldatagen/spark_singleton.py dbldatagen/data_generator.py dbldatagen/datasets_object.py

@zhengruifeng
zhengruifeng requested review from a team as code owners August 19, 2026 01:01
@zhengruifeng
zhengruifeng requested review from renardeinside and removed request for a team August 19, 2026 01:01
@CLAassistant

Copy link
Copy Markdown

@zhengruifeng
zhengruifeng marked this pull request as draft August 19, 2026 07:34

@ghanse ghanse left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for your contribution, I left a few small comments.

Comment thread dbldatagen/spark_singleton.py Outdated
Comment on lines +56 to +58
active_session = cls.getActiveSession()
if active_session is not None:
return active_session

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we getInstance here to avoid duplication?


def _get_pyspark_is_remote() -> Callable[[], bool] | None:
"""Returns PySpark's public remote-session predicate when available."""
for import_path in ("pyspark.sql.utils", "pyspark.sql"):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe add a small comment about why we use these modules to test if PySpark is remote?

Comment thread dbldatagen/spark_compat.py Outdated
Comment on lines +45 to +47

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we simplify this to only test for isinstance(sparkSession, "pyspark.sql.connect.session.SparkSession")? Or are there some edge cases?

This branch has not been deployed

No deployments
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.

3 participants