sarutak opened a new pull request, #59914:
URL: https://github.com/apache/spark/pull/59914
### What changes were proposed in this pull request?
This PR makes `NativePartitionReader`, the reader of a partition of a native
data source (SPARK-60101), release the Arrow stream of the library only once
`get_next` returns, also when another thread closes the reader while `get_next`
runs.
- `close()` no longer releases the stream, the vectors, the task context and
the allocator while a thread is in `next()`. It marks the reader as closed, and
that `next()` releases them once it is done with the batch that `get_next`
returned: it then returns false, or throws the error of the read. A failure to
release them there is logged, since the thread that closed the reader no longer
gets it. `close()` does not wait, so it does not keep the task from completing.
- `next()` returns false once the reader is closed, so the reader never
calls a released stream, and `close()` does nothing when it is called again.
The reading thread may then close the reader again: `DataSourceRDD` closes a
reader whose `next()` returns false, as well as from the completion listener.
Before, a second `close()` failed with `NullPointerException: ArrowArrayStream
is already closed`.
- The constructor runs under the same lock, so that a `close()` on another
thread sees the task and the vectors that it created.
The test library gets an option `slow_read`: the first `get_next` waits
until the test lets it continue, and releasing the stream records whether Spark
released it during `get_next` or not. With the option `fail` set to
`slow_read`, `get_next` then fails.
### Why are the changes needed?
`spark_datasource.h` promises a library that the callbacks of the stream
never run concurrently, and that the task context passed to `read` stays valid
until the stream is released. When a task is killed, it also says that Spark
only releases the stream once `get_next` returns.
Spark did not keep these promises when another thread reads a partition.
`DataSourceRDD` closes its readers from a task completion listener, and some
operators read their input on a thread of their own that the task does not wait
for. For example, `PipedRDD` reads its parent on its stdin writer thread. On
SPARK-60101 as merged (99a1bf9a837, the same reader as its last PR commit
1f4048404be), with a library whose `get_next` takes 2 seconds,
`spark.read.format(...).load().rdd.pipe("true").count()` in the JVM returned
after 0.5 seconds, and the library got `release` of the stream while its
`get_next` was still running. Spark then also logged `IllegalStateException:
Memory was leaked by query`. After the release, Spark frees the
`SparkDsTaskContext`, which the library may still use in `get_next` through
`task->is_interrupted`. So a library that follows the header can use freed
memory, and crash the executor.
`close()` hands the release to the thread in `next()` rather than waiting
for it, since that thread can be in a `get_next` that waits for data:
`task->is_interrupted` stays false for a task that completed, so waiting could
keep the task from completing for as long as `get_next` blocks.
Besides when the library gets the release, a closed reader now reads nothing
more and can be closed again. As before, a batch that `next()` returned is only
valid until the reader is closed, so a thread that still uses it after the task
completed reads closed vectors.
### Does this PR introduce _any_ user-facing change?
No. Native data sources are not released yet (4.4.0).
### How was this patch tested?
New tests in `NativeDataSourceSuite`, which hold the library in its first
`get_next` until they let it continue, so that they do not depend on timing:
- `release a stream only once get_next returns` reads a native data source
through `rdd.pipe("cat")`, and completes the task while the stdin writer thread
of `PipedRDD` is in `get_next`. It checks that the stream is not released when
the task completes, and that the library gets the release after `get_next`
returns.
- `close a reader while another thread reads it` closes a reader while
another thread is in `next()`, and checks the same, that `next()` then returns
false, or throws when `get_next` fails, and that a closed reader reads nothing
more and can be closed again.
Without the change in `NativePartitionReader`, both fail with `The stream
was released: during get_next.` Each of the other parts of the change, except
running the constructor under the lock and logging a failure to release the
resources in `next()`, also makes a test fail when it is reverted.
`NativeDataSourceSuite` and `ColumnarDataSourceSuite` pass locally on Java
22 on macOS, except two tests that are canceled without pandas, and the new
tests passed in 5 runs.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Kiro CLI
--
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]
