HyukjinKwon opened a new pull request, #289: URL: https://github.com/apache/spark-connect-rust/pull/289
### What changes were proposed in this pull request? [SPARK-61084](https://issues.apache.org/jira/browse/SPARK-61084): progress handlers are now registered, and called, as PySpark Connect's are. As PySpark Connect's: - `spark.registerProgressHandler(handler)` registers a callable once and returns `None`. - `removeProgressHandler(handler)` removes it, and raises a `ValueError` if it is not registered. `clearProgressHandlers()` removes them all. - Each handler is called with the keyword arguments of `ProgressHandler`: `stages` (the `StageInfo` of each stage), `inflight_tasks`, `operation_id` and `done`. It is called for each progress of an execution that has tasks. It is called once more, with `done=True`, when the execution ends, also when it fails. - `SPARK_CONNECT_PROGRESS_BAR_ENABLED=1` prints the progress bar of each execution. `pyspark.sql.connect.shell.progress`, with `ProgressHandler`, `StageInfo` and `Progress`, is vendored from Apache Spark v4.3.0, and the drift check covers it. It imports `ExecutePlanResponse` from `pyspark.sql.connect.proto`, which now has a message class with the fields of an `ExecutionProgress` that `from_proto` reads. The core tracks the progress of each execution as PySpark's `Progress` does, from the operation ID that the client assigns before it sends the request. Every execution goes through it, including collects, local iterators, commands, catalog requests, streaming listeners and pipeline runs. The handlers of `SparkSession::register_progress_handler` now take a `ProgressUpdate` with the stages, the running tasks, the operation and whether it is done, rather than the raw `ExecutionProgress` message. One difference remains: the Rust transport calls the handlers from inside its execution. So an exception raised by a handler goes to `sys.unraisablehook`, rather than failing the execution as it does in PySpark. ### Why are the changes needed? The client: - returned an id from `registerProgressHandler`, and removed a handler by that id - called a handler with one dict of the progress, positionally, so a handler written for PySpark failed, silently - never called a handler when an execution ended - lacked `pyspark.sql.connect.shell.progress` ### Does this PR introduce _any_ user-facing change? Yes. Progress handlers, and the progress bar, work as in PySpark 4.3.0. ### How was this patch tested? - A new live test registers a handler with PySpark's keyword-only arguments, runs a query, and checks its calls: the `StageInfo`s, the running tasks, one operation ID, and one call with `done=True`, the last. It passes with official PySpark 4.3.0 and fails without this change. - A new offline test runs executions that receive a progress from a mock Connect server. It checks the calls of a handler, its registration and removal, the `ValueError`, an exception of a handler, and the progress bar. - A new lib test checks the progress of an execution in the core: progresses without tasks, the operation of the first response, and the last call, of `done`. - On a local server, the calls of a handler for fast and slow queries, a SQL command and a failed query have the same arguments as official PySpark 4.3.0's. Only the number of progress messages differs, which depends on the timing of the server. - The lib tests of both crates (593 passed), the build of all tests, also with `wasm-udf`, the offline Python suites (2040 passed, 12 skipped), and the live drop-in and Arrow suites (145 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]
