HyukjinKwon opened a new pull request, #312: URL: https://github.com/apache/spark-connect-rust/pull/312
### What changes were proposed in this pull request? [SPARK-61117](https://issues.apache.org/jira/browse/SPARK-61117): The hooks of `SparkSession.Hook` are called, as PySpark does. In PySpark 4.3.0, `SparkSession.Builder._registerHook(factory)` registers a factory of `SparkSession.Hook`. A session that the builder creates calls each factory with itself. Before it sends each `ExecutePlanRequest`, the client calls the `on_execute_plan` of each hook in order, with the request of the hook before. It sends the request that the hooks return, with the operation ID of the original. This PR does the same: - `_registerHook` keeps the factories. `create()`, and `getOrCreate()` when it creates a session, call them with the new session. The factories run after the lock of session creation is released, since Python code may release the GIL that the lock is held with. - Before each execution, the Rust transport calls the hooks with the request as a `pb2.ExecutePlanRequest` of the generated `pyspark.sql.connect.proto`, and sends the request that they return. Clones of the session keep its hooks, as PySpark's clones do. - The exception of a hook is raised by the execution, as it is. - The core gains `SparkSession::set_execute_plan_hook` and the `ExecutePlanHook` type. - A subclass of `Hook` can take arguments in its `__init__`, as a Python class does. ### Why are the changes needed? `_registerHook` registered nothing: the client sent its requests from Rust and never called the hooks. A hook that adds tags to a request or changes its plan had no effect. ### Does this PR introduce _any_ user-facing change? Yes. Hooks run as with PySpark 4.3.0. ### How was this patch tested? - New offline tests, against a mock Connect server, check: - that the factories are called with a new session, and not with a session that `getOrCreate()` gets - the tags that two hooks add, in order, the operation ID of the original, and the command of the request sent - that the exception of a hook is raised by the execution, and nothing is sent - A new live test checks that the server runs the request that a hook returns, whose plan the hook replaced. It passes with official PySpark 4.3.0. - New Rust tests (`execute_plan_hook.rs`): the error of a hook, a cleared or replaced hook, and, against the server, a request of a hook sent with the operation ID of the original. - All of the new tests fail without this change. - The lib tests of both crates (590 passed), the build of all tests, also with `wasm-udf`, the offline Python suites (2162 passed, 8 skipped), and the live drop-in and Arrow suites (172 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]
