viirya opened a new pull request, #59647: URL: https://github.com/apache/spark/pull/59647
### What changes were proposed in this pull request? This is a backport of #59332 to `branch-4.3`. In pipelined Python UDF mode (`spark.python.udf.pipelined.enabled`), the listener that stops the writer thread now waits until the writer has actually exited, and it runs before every task completion listener: `TaskContextImpl.markTaskCompleted` runs the new pre-completion listeners on its own thread before it marks the task completed. See #59332 for the full description. The cherry-pick of the merged commit (ca1f5c5752e) applied cleanly except for the imports of `BasePythonRunnerSuite`: `branch-4.3` does not have the SPARK-59773 worker metrics tests that use some of master's imports, so only the imports that the backported tests use are added. `CloseShieldOutputStream` (SPARK-60305) and the per-buffer `DataOutputStream` of SPARK-60343, which this change builds on, are already in `branch-4.3`. ### Why are the changes needed? The pipelined mode was added in 4.3.0 by SPARK-56642, so `branch-4.3` has the same use-after-free: once `cancel(true)` succeeds, the old listener's `Future.get()` returns without waiting for the writer, so the listeners that run after it can free memory the writer still uses, e.g. the pages of the `HybridRowQueue`. ### Does this PR introduce _any_ user-facing change? No. This fixes a bug in the opt-in pipelined mode (default `false`). For tasks with no pre-completion listener, listener behavior is unchanged. ### How was this patch tested? On this branch, with JDK 17: - `build/sbt 'core/testOnly *BasePythonRunnerSuite *TaskContextSuite *CloseShieldOutputStreamSuite *BarrierTaskContextSuite *ExecutorSuite *StreamingShuffleSuite'` - `build/sbt core/scalastyle core/Test/scalastyle core/mimaReportBinaryIssues` (scalastyle passes; MiMa reports no problem from this change) - `python/run-tests --testnames pyspark.sql.tests.pandas.test_pipelined_udf` after `build/sbt -Phive package` ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (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]
