spark-test commented on code in PR #59961:
URL: https://github.com/apache/spark/pull/59961#discussion_r4488995108
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -2416,14 +2535,110 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
}
}
try {
- assert(closeWhileReading(failing = false) == Right(false))
+ checkClosed(closeWhileReading(failing = false).swap.getOrElse(fail("next
did not fail")))
val error = closeWhileReading(failing = true).swap.getOrElse(fail("next
did not fail"))
assert(error.getMessage.contains("injected failure after a slow
get_next"))
} finally {
pool.shutdown()
}
}
+ nativeTest("close a reader while another thread reads the batch that it
returned") {
+ val session = newSession(defaultPackage)
+ val pool = ThreadUtils.newDaemonSingleThreadExecutor("native-reader")
+ def onPool[T](f: => T): T = pool.submit(new Callable[T] {
+ override def call(): T = f
+ }).get(1, TimeUnit.MINUTES)
+ // The contexts of the tasks of earlier tests are closed first, so that
the counts below are
+ // exactly those of this test.
+ eventually(timeout(1.minute))(assert(NativeLibrary.openTaskContexts() ==
0))
+ // Closes a reader on this thread once the other thread read a batch, and
returns the allocator
+ // of the reader, which is no longer referenced then. With `nextAgain`,
the other thread then
+ // calls next, which tells that it no longer reads the batch.
+ def closeWhileRead(nextAgain: Boolean): BufferAllocator = {
+ val marker = new File(Utils.createTempDir(), "released")
+ // get_next does not wait: the file .continue exists.
+ new File(marker.getPath + ".continue").createNewFile()
+ val (reader, allocator) = readerAndAllocator(
+ range(session, "end" -> "10", "partitions" -> "1", "slow_read" ->
marker.getPath)
+ .select("id"))
+ try {
+ assert(onPool(reader.next()))
+ assert(NativeLibrary.openTaskContexts() == 1)
+ reader.close()
+ // The stream and the task are released at once, but not the buffers
of the batch, which
+ // the other thread may still read.
+ assert(Files.readString(marker.toPath) == "after get_next")
+ assert(NativeLibrary.openTaskContexts() == 0)
+ assert(readerAllocators().contains(allocator) &&
allocator.getAllocatedMemory > 0)
+ val rows = onPool((0 until
reader.get().numRows()).map(reader.get().column(0).getLong))
Review Comment:
**Nit (P3):** This read goes through `reader.get()`, and the test keeps
`reader` in a local until `closeWhileRead` returns, so the reader stays
strongly reachable while the other thread reads. What keeps a closed reader's
buffers allocated for a consumer that holds only the batch is the `reader`
field of `NativeColumnVector`, which no code reads. Deleting `column.reader =
reader;` in `NativeColumnVector.create`, as a cleanup of an unused field might,
still passes every test: `NativeColumnVectorSuite` passes a null reader, and
`checkReadersFreed` only waits for frees. In production,
`PartitionIterator.close()` drops `DataSourceRDD`'s reference to the reader, so
script transformation's writer holds only the batch, and without the field a
garbage collection would run the Cleaner and free the buffers that the writer
still reads at their addresses.
**Verification:**
- **Regression:** With the column vectors' reference to their reader
removed, a garbage collection frees the allocator of a closed reader whose
batch is still held, and the test fails before it reads the batch.
- **Behavior:** With the reference in place, a closed reader's batch held
without the reader keeps its allocator across garbage collections and reads its
values, and the allocator is freed once the batch is dropped.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -155,11 +157,15 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
val inputColumnIndices = resolveColumnIndices(allInputs.toSeq)
- // Peek at first batch to check if Arrow-backed.
+ // Peek at first batch to check if Arrow-backed. Path 1 holds the
pass-through columns past
+ // the next batch, which a ReusedArrowColumnVector does not allow: its
producer overwrites or
+ // releases it, so such batches take path 2, which copies them.
val peekIter = new PeekableIterator(inputIter)
val isArrow = peekIter.peek().exists { batch =>
batch.numCols() > 0 &&
- batch.column(0).isInstanceOf[ArrowColumnVector]
+ batch.column(0).isInstanceOf[ArrowColumnVector] &&
+ !(0 until batch.numCols()).exists(
Review Comment:
**Nit (P3):** This checks every column, but the only test that reaches it,
'Arrow UDFs on the columns of a native data source', reads batches whose
columns are all `NativeColumnVector`s, so narrowing it back to
`batch.column(0)` passes every test. The case that tells the two apart is a
batch that mixes plain and reused Arrow columns. Path 1 builds its output from
the input's pass-through columns plus the result columns that
`ArrowOutputProcessorImpl` wraps as `ReusedArrowColumnVector`s, so an Arrow UDF
stacked directly on another, as `ExtractPythonUDFs` stacks nested UDFs of
different eval types, gets the source's plain vectors first and the inner UDF's
reused result vector last. Under a first-column check the outer node would hold
that result vector past the next batch and pair its results with the inner
results of a later batch, with no failure. The Arrow-backed sources of
`ArrowColumnarPythonUDFSuite` build plain `ArrowColumnVector`s, so no test
there reaches this check, and th
e native test is skipped in pull request CI, which runs Java 21.
This extends an [existing
discussion](https://github.com/apache/spark/pull/59961#discussion_r4488238424).
**Verification:**
- **Regression:** With the guard narrowed back to the first column, the
outer UDF keeps the inner UDF's reused result vector past the next batch, so
over many small batches its output carries the inner results of later batches,
and the test fails.
- **Behavior:** Stacked columnar Arrow UDFs of different eval types over an
Arrow-backed source return correct pass-through and result values over many
small batches, with both nodes taking columnar input.
##########
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) {
+ released = true
closeResources()
+ } else if (!reading) {
+ // That thread may not call next or close again, so the reader is
released once it ends.
+ val thread = consumer
+ NativePartitionReader.releaser.execute(() => {
+ thread.join()
Review Comment:
The rework settles the pipelined Python case: close() now releases the
stream, the task context and the dictionaries at once, and no thread stays
parked. The script transformation case still holds at this head, though.
Because the closed reader keeps its last batch readable, the writer goes on
writing the rest of that batch to the script, and it fails only when it asks
for the next batch. When nobody reads the script's output any more, as under a
LIMIT, and the rest of the batch is more than the two pipes hold, the writer
blocks in the write before that. Nothing destroys the script when the task
completes, so the writer thread, the script process and the batch's native
buffers stay alive for good. On the base branch, the writer's next read of the
batch after the close failed Arrow Java's checks, and the writer destroyed the
script. Destroying the script when its task completes, as PipedRDD does, would
end the writer and let the batch be freed.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4439604925","thread_id":"inline:4439604925","verdict_sha256":"fd564c19e251e7f8c8c414ff2dd2288dbad9d6be08b14411f2b2dfdb9b160e0a"}
-->
--
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]
