HyukjinKwon opened a new pull request, #325: URL: https://github.com/apache/spark-connect-rust/pull/325
### What changes were proposed in this pull request? [SPARK-61129](https://issues.apache.org/jira/browse/SPARK-61129): The client's `SparkSession` and `DataFrame` now pickle with cloudpickle, as PySpark Connect's do: - **`SparkSession`** pickles as a function of its ID, which cloudpickle pickles by value. The function gives back the active session of that ID through the new `SparkSession._getActiveSessionIfMatches`: - Where no session is active, it raises `PySparkRuntimeError` `NO_ACTIVE_SESSION`, as on a Python worker of the server. - Where another session is active, it raises `PySparkAssertionError`. The function reads the class of `pyspark.sql.connect.session` when it is called, so the server's worker, which runs PySpark itself, raises the same error. - **`DataFrame`** pickles of its session and the bytes of its relation, PySpark's plan. The session unpickles first, so on the server's worker a pickled DataFrame raises `NO_ACTIVE_SESSION`, as PySpark's does. - **The core** rebuilds the DataFrame of a relation that its session sent with the plan ID of that relation (`DataFrame::from_sent_relation`). So the copy of a DataFrame is sent exactly as the DataFrame was, and the columns of the original resolve against it. PySpark's copies keep their plan IDs the same way. Plain `pickle` still cannot pickle either, as with PySpark: the function is a local one. ### Why are the changes needed? The client could pickle neither: `TypeError: cannot pickle 'pyspark.sql.session.SparkSession' object`, and the same for its DataFrame. So a UDF or a UDTF that captured the session or a DataFrame failed on the client, while PySpark sends it and its worker raises `NO_ACTIVE_SESSION`. Apache Spark's `test_udtf_access_spark_session_connect` failed on this client. ### Does this PR introduce _any_ user-facing change? Yes: - A session and a DataFrame can be pickled with cloudpickle, as with PySpark 4.3.0. - A UDF or a UDTF that captures either raises `PythonException` `NO_ACTIVE_SESSION` from the server, rather than a pickling error on the client. ### How was this patch tested? - A new offline test checks: - the round trips of a session, which is the same object, and of a DataFrame, whose relation is the same - the `NO_ACTIVE_SESSION` of either in a thread without an active session - the `PySparkAssertionError` of another session ID - that plain `pickle` refuses both - A new live test runs a UDTF that captures a DataFrame and a UDF that captures the session, which both raise `PythonException` `NO_ACTIVE_SESSION`. It passes with official PySpark 4.3.0 and fails without this change. - A new Rust test (`raw_relation.rs`) checks that the DataFrame of a relation that its session sent is sent as the relation was, and that the columns of the original refer to it. - Apache Spark's `test_udtf_access_spark_session_connect`, of `ArrowUDTFParityTests` and `LegacyArrowUDTFParityTests` in `test_parity_udtf.py`, now passes on this client against a local server. - The lib tests of both crates (590 passed), the build of all tests, also with `wasm-udf`, the offline Python suites (2189 passed, 8 skipped), and the live drop-in and Arrow suites (178 passed) and live cases (10 passed) against a local Spark 4.3.0 server. Clippy reports no diagnostics in the changed lines. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Opus 5.5 This pull request and its description were written by Isaac. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
