dongjoon-hyun opened a new pull request, #59946: URL: https://github.com/apache/spark/pull/59946
### What changes were proposed in this pull request? This PR closes the Arrow buffers of the input to Python workers when the task completes before all of the input has been written. 1. `PythonArrowInput` adds a task completion listener in `newWriter` that closes `root` and `allocator`. The listener and `close()`, which runs once the input is exhausted, share `closeRootAndAllocator()`. It closes them only once, and only if `writer`, which creates them, has been created, so that the listener does not create the lazy `root` and `allocator` (SPARK-55285). The listener does not end `writer`, since the stream that it writes to may be closed already, and nothing sends it to the Python worker anymore. As in `close()`, closing the allocator throws if any of its buffers is still allocated. 2. `ArrowWriterWrapper.createAndStartArrowWriter`, which the grouped and cogrouped runners call for each group, creates the wrapper, which adds the task completion listener that closes its `root` and `allocator`, before anything that can fail. It used to allocate the vectors and write the schema first, so a failure there, e.g. an interrupt while writing the schema, left them open. ### Why are the changes needed? `PythonArrowInput` closes `root` and its child allocator of `ArrowUtils.rootAllocator` only at the end of the input. If the task completes before that, nothing closes them, and the buffers of `root` stay allocated for the lifetime of the executor: Arrow does not release buffers on GC, and the root allocator keeps a reference to each child allocator until it is closed. A task completes before the end of its input when it - fails, e.g. when the Python UDF raises an exception or reading the input fails, - is killed, or - stops reading the output early, e.g. with a `LIMIT` above a filter on the UDF result. `ArrowWriter.reset()` keeps the capacity of the vectors, so `root` holds the buffers of a batch from the first batch on. In local mode with `spark.sql.execution.arrow.maxRecordsPerBatch=100` and 1 KiB strings, each such task left 147,456 bytes allocated and its child allocator (`stdout writer for <python>`) open. This affects the runners that use `PythonArrowInput` directly: scalar pandas and Arrow UDFs (`ArrowEvalPythonExec`), Arrow Python UDTFs, `mapInPandas` and `mapInArrow`, `applyInPandasWithState`, and `transformWithState` in PySpark, as well as the partial aggregation of `PythonIncrementalAggregateExec` and the columnar input enabled by `spark.sql.execution.arrow.pythonUDF.columnarInput.enabled`. It is a regression from SPARK-44705 (Apache Spark 4.0.0): before it, the writer thread closed them in a `finally` block. The grouped and cogrouped runners already close the `ArrowWriterWrapper` of each group in a task completion listener, but `createAndStartArrowWriter` added the listener only after allocating the vectors and writing the schema. With `spark.python.udf.pipelined.enabled`, a grouped aggregate pandas UDF under a `LIMIT` sometimes left 49,152 bytes, when the writer was interrupted there at task completion. The new listener does not race with the writer, since task completion listeners run after the writer is done. The writer runs on the task thread, except in the pipelined mode, where the pre-completion listener added by SPARK-60113 stops the writer thread and waits for it to exit before any completion listener runs. SPARK-55351 also closed them in a task completion listener but was reverted. Unlike it, this PR keeps `root` and `allocator` lazy, and the listener does not end the writer. ### Does this PR introduce _any_ user-facing change? No, except that the Arrow memory of the input to Python workers is released when a task completes before the end of its input. ### How was this patch tested? - Added a test to `PythonUDFSuite`. It runs a scalar pandas UDF over 100,000 rows of 1 KiB strings, more than the socket to the Python worker can buffer, in two jobs: one whose input fails after 1,000 rows, and one that stops reading the output early with a `LIMIT` above a filter on the UDF result. It runs both with and without `spark.python.udf.pipelined.enabled`, and checks that no allocator of the Arrow input is left open. Without this PR, it fails with two allocators left open, each holding 147,456 bytes. - Added a test to `ArrowWriterSuite` in which writing the schema fails in `ArrowWriterWrapper.createAndStartArrowWriter`. Without this PR, the allocator stays open after the task completes. - `StatefulPythonRunnerDataSentSuite`, which drives `newWriter` directly, passes `TaskContext.empty()` instead of `null`, since the writer now adds a task completion listener to it. - Also ran all suites in `org.apache.spark.sql.execution.python`, `FlatMapGroupsInPandasWithStateSuite`, `FlatMapGroupsInPandasWithStateDistributionSuite`, and `ArrowConvertersSuite`. The `PythonUDFSuite` test needs a Python with pandas and PyArrow, and is skipped otherwise. ``` PYSPARK_PYTHON=<python with pandas and pyarrow> build/sbt 'sql/testOnly org.apache.spark.sql.execution.python.PythonUDFSuite -- -z SPARK-60737' build/sbt 'sql/testOnly org.apache.spark.sql.execution.arrow.ArrowWriterSuite -- -z ArrowWriterWrapper' ``` ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Opus 5.5 -- 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]
