gengliangwang commented on code in PR #59961:
URL: https://github.com/apache/spark/pull/59961#discussion_r4445266538
##########
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:
Agreed on both counts. Releasing based on the thread that last called next
can't work: script transformation fetches its first batch on the task thread
(`BaseScriptTransformationExec.scala:70`) and reads it on the writer, and the
reader has no way to see that. A volatile bound only narrows the window.
Waiting for in-flight reads (an epoch or lock on every access) would bring back
the per-access cost this PR removes. So I'd rather not free a batch that may
still be read than try to work out who reads it. The plan: `close()`, on any
thread unless one is inside `next()`, releases the stream, the dictionaries and
the task right away. The buffers of the last imported batch, and the child
allocator, are freed by a Cleaner only once the batch and its column vectors
become unreachable. A late read then sees the last batch's valid values, never
freed memory, and the consumer tracking and the releaser pool go away. Before
writing it I still need to confirm two things: that arrays may outliv
e the stream (the C stream interface says their lifetimes are independent),
and that a library's array release callback doesn't use the task. I'd also make
a `next()` after `close()` fail rather than return false. That way the script
writer fails and destroys the process, as it did on the base branch. I'll
rework it this way in the next revision unless you see a problem with this
direction.
<sub>Automated reply (Opus)</sub>
##########
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:
Agreed, this comes from the same thread-tracking release, and I plan to
replace it (see my reply on the thread above). With that change, `close()`
releases the stream, the dictionaries, the native task and the parked releaser
right away, on any thread. Only the last batch's buffers wait until nothing
references the batch. So a pooled pipelined-Python writer thread that stays
alive no longer pins the reader's native resources. For script transformation,
a `next()` after `close()` would fail rather than return false, so the writer
fails and destroys the script once it is done with the current batch, as on the
base branch, instead of writing to it forever.
<sub>Automated reply (Opus)</sub>
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -670,14 +685,16 @@ class NativePartitionReader(
}
/**
- * Wraps a vector of the root, into which every batch is imported. It is not
freeable between
- * batches: Spark must not close it when it moves to the next batch, as
`ColumnarToRowExec` does
- * with the column vectors that are freeable, because closing a struct
column also drops its
- * child column vectors. The reader closes the Arrow vectors at the end.
+ * Wraps a vector of the root, into which every batch is imported, in a
`NativeColumnVector`,
+ * which reads the buffers of the batch that `readNext` imported last. Like
the other
+ * `ReusedArrowColumnVector`s, it is not freeable between batches: Spark
must not close it when
+ * it moves to the next batch, as `ColumnarToRowExec` does with the column
vectors that are
+ * freeable, because closing a struct column also drops its child column
vectors. The reader
+ * closes the Arrow vectors at the end.
*/
- private def columnVector(vector: FieldVector): ColumnVector = {
+ private def columnVector(vector: FieldVector): NativeColumnVector = {
try {
- new ReusedArrowColumnVector(vector)
+ NativeColumnVector.create(vector)
Review Comment:
Agreed that this opt-in path is now worse: before, reads past a shorter
later batch or of a released batch threw, and now they read raw memory. The
underlying problem is older. Native columns were already
`ReusedArrowColumnVector`s, and with this config, rows of a batch that a later
import overwrote already read the later batch's values silently. Keeping the
last batch alive until it is unreachable (see the first thread) handles the
released-batch case. It doesn't handle a batch that is held while the next
batch is imported into the same vectors. For that I think the columnar Python
input should not keep reused vectors as pass-through columns: take the
Arrow-direct path only for an `ArrowColumnVector` that is not a
`ReusedArrowColumnVector`, and fall back to row iteration otherwise. I'll
include that with the lifecycle rework.
<sub>Automated reply (Opus)</sub>
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1256,6 +1256,78 @@ 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.")
+ 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")
+ assert(!reader.next())
Review Comment:
Done in 25616f1e1b3c: The test now also calls next() and close() on the
closing (non-consumer) thread while the release is pending and asserts that the
stream is not released. That fails if either consumerDone guard is removed. I
dropped the redundant next() on the already-released reader.
<sub>Automated reply (Opus)</sub>
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,315 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.v2.ffi;
+
+import java.nio.ByteOrder;
+import java.util.List;
+
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.vector.BaseVariableWidthViewVector;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.BitVector;
+import org.apache.arrow.vector.DateDayVector;
+import org.apache.arrow.vector.DecimalVector;
+import org.apache.arrow.vector.Float4Vector;
+import org.apache.arrow.vector.Float8Vector;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.SmallIntVector;
+import org.apache.arrow.vector.TimeStampMicroTZVector;
+import org.apache.arrow.vector.TimeStampMicroVector;
+import org.apache.arrow.vector.TinyIntVector;
+import org.apache.arrow.vector.ValueVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.ViewVarCharVector;
+
+import org.apache.spark.sql.execution.vectorized.ReusedArrowColumnVector;
+import org.apache.spark.sql.types.Decimal;
+import org.apache.spark.sql.vectorized.ArrowColumnVector;
+import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.types.UTF8String;
+
+/**
+ * The column vector of a top-level column of a {@code NativePartitionReader},
which imports every
+ * batch into the same Arrow vectors.
+ *
+ * <p>After each import, the reader calls {@link #refresh()}, and the vector
then reads booleans,
+ * integers, floats, dates, timestamps, decimals and strings at the addresses
of the buffers of the
+ * batch. {@link ArrowColumnVector} reads each value through Arrow Java, which
by default checks on
+ * each access the bounds and the reference count of the buffer, and whether
the value is null.
+ * Spark does not need these checks here: it reads only the rows of the batch,
whose buffers Arrow
+ * Java sizes by the lengths of the arrays, which the reader checks before it
imports them, and
+ * the buffers of a batch stay allocated until the thread that read it calls
next again or closes
+ * the reader: a close on another thread leaves the release to that thread.
Spark reads a value
Review Comment:
Done in 25616f1e1b3c: Rewrote the NativeColumnVector doc: it now says a
close releases at once only if the consumer thread has ended, and otherwise
leaves the release to that thread or to a thread waiting for it to end. Fixed
the NativePartitionReader class doc: a next() after close can release the
stream through the library.
<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]
