gengliangwang commented on code in PR #59961:
URL: https://github.com/apache/spark/pull/59961#discussion_r4484969139
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -858,9 +897,21 @@ class NativePartitionReader(
override def close(): Unit = synchronized {
if (!closed) {
closed = true
- if (!reading) {
+ // Another thread may be in next, or still read the batch that next
returned to it.
+ if (!reading && consumerDone) {
Review Comment:
The rework is in d56b02ddb24. `close()` now frees everything at once if
`next()` was called again since it returned the batch: at the end of the
stream, or after a failed `next()`. Otherwise it releases the stream, the
dictionaries and the task at once, and a Cleaner frees the last batch's buffers
and the allocator once the reader is unreachable. Each `NativeColumnVector`
refers to its reader, so this only happens once the batch and its column
vectors are unreachable too. A thread that is inside `next()` still gets the
release once `get_next` returns, as before. A `next()` called after `close()`
now throws an `IllegalStateException` rather than return false. The consumer
tracking and the `releaser` pool are gone.
On the two points I had to check:
1. The C stream interface says that the lifetimes of what `get_schema` and
`get_next` return are not tied to that of the stream.
2. The header already requires that the release of an array or a schema not
use the stream or the task (`SparkDsTaskContext`). The Rust SDK exports each
batch with `FFI_ArrowArray::new`, whose `ArrayData` is `'static` and so cannot
borrow the `'task` context. The C++ SDK uses `arrow::ExportRecordBatch`, whose
release only drops the references to the buffers. I wrote this into the
`SupportsArrowStreamRead` contract. The header now also says that Spark itself
may release the last array after the stream, from a JVM Cleaner thread.
The class doc notes the GC trade-off. For the tests, `free the buffers of a
reader closed at the end of its stream without a GC` is the end-of-stream test,
with a reader whose `next()` failed as well. `close a reader while another
thread reads the batch that it returned` now checks four things: the stream and
the task are released at once, the other thread still reads valid values,
`next()` fails on both threads, and a GC frees the buffers.
`NativeDataSourceSuite` also waits in `afterAll` until no reader allocator is
left, so that later suites in the same JVM don't see that memory freed while
they run.
<sub>Automated reply (Opus)</sub>
--
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]
