sarutak commented on code in PR #59961:
URL: https://github.com/apache/spark/pull/59961#discussion_r4445437083
##########
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 direction makes sense to me. One concern about the Cleaner: if every
close leaves the last imported batch to it, every reader, including one that
was read to the end, gives its native memory back only after a GC. Off-heap
memory does not trigger a GC, and a batch that lived as long as the task has
often been promoted, so it may wait for an old-generation collection. Could
`close()` release the buffers right away when no batch can still be read, e.g.
once `next()` has returned false or thrown, and keep the Cleaner only for a
close while a batch that `next()` returned may still be in use?
##########
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:
The Scala linter fails here: "Public method must have explicit type".
--
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]
