Add support for running Spark Connect client inside the Dataproc s8s batch

[ZhiweiLinnn]

For batch, we want initiated sessions.
Add integration test for batch environ

[gemini-code-assist]

Summary of Changes

Hello @ZhiweiLinnn, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This PR enables the Spark Connect client to correctly integrate with Dataproc serverless batch workloads. It introduces logic to detect batch environments and, when detected, ensures the client connects to an existing Spark session rather than provisioning a new interactive one. This is achieved by adding environment detection, a session availability polling mechanism, and modifying the session creation flow. Comprehensive unit and integration tests are included to validate these changes.

Highlights

• Intent: This pull request adds support for running the Spark Connect client within a Dataproc s8s (serverless) batch environment. The primary goal is to ensure that when the client is used in a batch context, it connects to an already initiated Spark session rather than attempting to create a new interactive Dataproc session. An integration test is also added to validate this new batch environment behavior.
• Changes: The changes introduce a mechanism to detect Dataproc batch workloads and modify the session creation logic accordingly. A new utility function is_dataproc_batch() is added to environment.py to check for the DATAPROC_WORKLOAD_TYPE environment variable. In session.py, a new private method _wait_for_session_available is introduced to poll for the Spark Connect endpoint readiness of a session. Crucially, the DataprocSparkSession.Builder.getOrCreate() method is updated to conditionally return a standard pyspark.sql.SparkSession directly if is_dataproc_batch() is true, bypassing the Dataproc interactive session creation flow. Integration tests in test_session.py now include fixtures and a test case to verify this batch-specific behavior, ensuring that getOrCreate() yields a local PySpark session when in a simulated batch environment. Unit tests for the new _wait_for_session_available method are also added to test_session.py to cover success and timeout scenarios.
• Reviewer Activity: No reviewer activity has been recorded for this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

Thank you for this contribution. The changes to support Dataproc batch workloads are a great addition. I've identified a couple of areas for improvement, one of which is critical regarding the return type of getOrCreate in batch mode. Please see my detailed comments below.

[gemini-code-assist]

Catching a generic Exception is too broad and can mask unexpected issues or even catch system-exiting exceptions like KeyboardInterrupt. It's better to catch more specific exceptions that you expect from the get_session API call and want to retry on.

Given that this loop is waiting for a session to become available, it would be safer to only catch exceptions that indicate a transient state, such as NotFound (if the session is not yet created) or other retryable API errors. The google-api-core library provides specific exception types for this. Using the already imported exceptions would make this more robust.

Suggested change

	except Exception as e:
	logger.warning(
	f"Error while polling for Spark Connect endpoint: {e}"
	)
	time.sleep(5)
	except (NotFound, Aborted, FailedPrecondition) as e:
	logger.warning(
	f"API error while polling for Spark Connect endpoint, will retry: {e}"
	)
	time.sleep(5)