spark-test commented on code in PR #59961:
URL: https://github.com/apache/spark/pull/59961#discussion_r4439604925
##########
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:
**Non-blocking (P2):** When the recorded consumer thread is alive, a
`close()` from another thread only marks the reader closed, and parks a
releaser here until that thread ends. Before then, only that thread's next
`next()` or `close()` releases the reader. A live thread that has stopped
reading calls neither.
With `spark.python.udf.pipelined.enabled=true`, a reused thread of the
shared `python-udf-pipelined-writer` pool pulls the native scan in
`writeNextInputToStream`, so it is the consumer. When the task ends early (a
`show()` or `LIMIT` above the UDF, a failing UDF, or a killed attempt), the
pre-completion listener stops the writer and returns the thread to the pool
alive. `DataSourceRDD` then closes the reader on the task thread, and the
release waits until the cached pool reaps that thread after 60 s idle. That
does not happen while pipelined Python tasks keep reusing the thread. Until
then, the reader of each such task keeps its last batch, the library's stream
and whatever the library holds for it, the dictionaries, the native task and
the child allocator, and one `native-partition-reader-release` thread stays
parked. The base branch released all of them at close.
Script transformation reaches the same state by default once its writer
thread has fetched a batch itself. After a `LIMIT`, the batch stays readable,
so the writer keeps writing to a script whose output nobody reads. Once the
pipes fill, the writer, the script process, the reader and a releaser thread
stay blocked for good. Before this change, the writer's next read failed, and
the writer destroyed the script.
See **Shared repair plan 1** in the review body.
##########
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:
**Blocking (P1):** This releases at once whenever the closing thread is the
one that called `next()` last, which assumes that no other thread is still
reading the batch that `next()` returned. Script transformation breaks that
assumption. `BaseScriptTransformationExec.doExecute` calls `iter.hasNext` on
the task thread, which fetches the first batch of the native scan.
`processIterator` then hands the same iterator to a writer thread, which reads
the remaining rows of that batch. Nothing stops or joins the writer at task
completion.
So with `SELECT TRANSFORM(id) USING 'cat' AS (x) FROM native_t LIMIT 1`, or
with a script such as `head -1` that exits early, the task can complete while
the writer is still inside the first batch. `DataSourceRDD`'s completion
listener then calls `close()` on the task thread. `consumerDone` holds, and
`closeResources()` hands the imported buffers back to the library, while the
writer keeps calling `isNullAt`, `getLong` and `getUTF8String` at the addresses
that `refresh()` cached. The writer is then reading freed native memory. It can
send garbage to the script. Garbage `Utf8` offsets make `UnsafeRowWriter` copy
arbitrary memory, and reading a page that the library has unmapped crashes the
executor JVM, which fails every task running on it.
Before this change, the same late reads went through Arrow Java's bounds and
reference-count checks on the cleared vectors and threw, so the writer failed
and destroyed the script process. `HiveScriptTransformationExec` shares this
`doExecute`, and the same holds for any consumer that reads a batch on a thread
other than the one that called `next()` for it.
See **Shared repair plan 1** in the review body.
##########
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:
**Non-blocking (P2):** Every native column now reads through
`NativeColumnVector`, which is still an `ArrowColumnVector`. So with
`spark.sql.execution.arrow.pythonUDF.columnarInput.enabled=true`,
`ColumnarArrowEvalPythonEvaluatorFactory` takes its Arrow columnar path for
native batches. That path keeps each batch's column vectors as pass-through
columns until Python returns that batch's result. Meanwhile, the reader imports
later batches into the same vectors and releases them at the end of the scan.
The config's documentation says that this path needs upstream batches that
stay valid, which a columnar reader does not promise. On the base branch, the
rows of an overwritten batch that lie within the later batch already read the
later batch's values. But a read past a shorter later batch, or a read of the
released last batch, threw `IndexOutOfBoundsException` through Arrow Java's
checks. Now these reads go to raw addresses past the buffers or in freed
memory, so the combined output silently carries garbage values, or the executor
JVM crashes. Only that opt-in config with a native source is affected.
See **Shared repair plan 1** in the review body.
##########
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:
**Nit (P3):** This says that a close on another thread leaves the release to
the thread that read the batch. But `close()` also releases at once when that
thread has already ended, and otherwise a thread of the
`native-partition-reader-release` pool releases once it ends.
`NativePartitionReader`'s class doc is also out of date: it still says that
a `next()` called once the reader is closed returns false without calling the
library. After a deferred close, though, a `next()` on the consumer thread runs
`releaseClosed()`, which releases the stream through the library's release
callback, as the doc's own new paragraph says. These docs state the rules that
keep the unchecked reads off freed memory, so a maintainer who relies on them
misjudges which thread frees the buffers and when the library gets called.
See **Shared repair plan 1** in the review body.
##########
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:
**Nit (P3):** By this point, the pool thread's `next()` has already released
the reader, so this `next()`, like the `close()` in the `finally` block, only
exercises a released reader. No test calls `next()` or `close()` on a
non-consumer thread while the release is still pending. So replacing `if
(consumerDone) releaseClosed()` in `startReading` with an unconditional
`releaseClosed()`, or dropping `&& consumerDone` from the repeated-close branch
of `close()`, keeps every test green. Yet either change lets another thread
free a batch that the consumer thread is still reading at raw addresses.
See **Shared repair plan 1** in the review body.
--
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]
