gengliangwang commented on code in PR #59961:
URL: https://github.com/apache/spark/pull/59961#discussion_r4445873980
##########
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:
Good point, and I agree: the Cleaner should only cover the case it exists
for. In the next revision, if `next()` has returned false (the end of the
stream), `close()` frees everything right away: the stream, the dictionaries,
the task, the last batch's buffers and the child allocator. The Cleaner only
takes the buffers when `close()` runs while a batch that `next()` returned may
still be read, for example with a LIMIT or when a script exits early. That
covers the normal fully-read close from `DataSourceRDD`. Once a reader has
handed back the end of the stream, `PartitionReader` gives `get()` no valid
batch, and every consumer I know of stops reading the batch by then. Script
transformation is one example: after its first batch, the writer thread is the
one that calls `hasNext`. The exception is the one you name.
`NativeColumnVector` extends `ReusedArrowColumnVector` and therefore
`ArrowColumnVector`, so `ColumnarArrowEvalPythonEvaluatorFactory` picks path 1
for it (with `spark.sql
.execution.arrow.pythonUDF.columnarInput.enabled`), and path 1 holds
references to the pass-through columns until Python returns the results. Path 1
is already wrong for this reader before the end of the stream, though, because
the vectors are reused: the pass-through columns of batch N would show batch
N+1. So the fallback I plan sends batches whose columns are
`ReusedArrowColumnVector`s down path 2. Path 2 copies the pass-through rows
into the `HybridRowQueue` before the writer asks for the next batch, so nothing
reads the reader's buffers after `next()` returns false. With that in place,
freeing at the end of the stream is safe. I'll add a test that a reader read to
the end and then closed frees its buffers without waiting for a GC.
<sub>Automated reply (Opus)</sub>
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1256,6 +1256,81 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
}
}
+ nativeTest("close a reader while another thread reads the batch that it
returned") {
+ val session = newSession(defaultPackage)
+ val pool = ThreadUtils.newDaemonSingleThreadExecutor("native-reader")
+ val marker = new File(Utils.createTempDir(), "released")
+ // get_next does not wait: the file .continue exists.
+ new File(marker.getPath + ".continue").createNewFile()
+ val df = range(session, "end" -> "10", "partitions" -> "1", "slow_read" ->
marker.getPath)
+ val batch = scan(df.select("id")).toBatch
+ val reader =
batch.createReaderFactory().createColumnarReader(batch.planInputPartitions().head)
+ try {
+ def onPool[T](f: => T): T = pool.submit(new Callable[T] {
+ override def call(): T = f
+ }).get(1, TimeUnit.MINUTES)
+ assert(onPool(reader.next()))
+ // The other thread may still read the batch, whose buffers the reader
keeps until that
+ // thread calls next again.
+ reader.close()
+ assert(!marker.isFile, "The stream was released.")
+ // Nor does a next or another close on this thread, which did not read
the batch.
+ assert(!reader.next())
+ reader.close()
+ assert(!marker.isFile, "The stream was released.")
+ val rows = onPool((0 until
reader.get().numRows()).map(reader.get().column(0).getLong))
+ assert(rows === (0L until 10L))
+ assert(!onPool(reader.next()))
+ assert(Files.readString(marker.toPath) == "after get_next")
+ } finally {
+ // Once the other thread ends, a close that the test did not reach
releases the reader.
+ pool.shutdown()
+ pool.awaitTermination(1, TimeUnit.MINUTES)
+ Utils.tryLogNonFatalError(reader.close())
+ }
+ }
+
+ nativeTest("release a closed reader when the thread that read its batch
closes it or ends") {
+ val session = newSession(defaultPackage)
+ def onPool[T](pool: ExecutorService)(f: => T): T = pool.submit(new
Callable[T] {
+ override def call(): T = f
+ }).get(1, TimeUnit.MINUTES)
+ // Returns a reader closed while a thread of the pool may still read its
batch, and the file
+ // that the library creates when the stream is released.
+ def closedReader(pool: ExecutorService) = {
Review Comment:
Done in 9b9dfb979474: Gave `closedReader` the explicit return type
`(PartitionReader[ColumnarBatch], File)` and added the two imports it needs.
<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]
