spark-test commented on code in PR #59961:
URL: https://github.com/apache/spark/pull/59961#discussion_r4462046785
##########
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:
**Nit (P3):** Reverting this line to `new ReusedArrowColumnVector(vector)`
passes every test. `NativeColumnVectorSuite` calls `NativeColumnVector.create`
directly on vectors built with Arrow Java. The only reader-level check of the
column vectors, in 'read structs in several batches', accepts any
`ReusedArrowColumnVector`. And the test library exports no decimal128 column,
only a decimal256 one that the reader rejects. So the PR's main change, the
direct reads and the long-backed decimals with their exact integral rounding,
has no failure signal at the reader.
**Verification:**
- **Regression:** A reader-level test fails if the reader's columns are not
NativeColumnVectors, or if a native DECIMAL value of at most 18 digits read
through the reader is not long-backed.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1256,6 +1258,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.")
Review Comment:
**Nit (P3):** This assertion runs right after `close()` returns, before the
releaser task that `close()` may have queued can run, and the tests observe a
release only after the consumer's own `next()` or `close()`, or after its pool
is shut down. So three regressions pass every test. A releaser that releases
without waiting for the consumer fails here only by timing. A narrower
immediate-release condition sends the common close on the consuming thread to a
releaser, and no test closes the reader on the thread that called `next()` last
and then checks the release marker. And a releaser that keeps polling after the
consumer released the reader goes unnoticed, because every consumer pool is
shut down. These are the paths that keep the unchecked reads off freed memory
and keep closed readers from being retained.
**Recommended change:** Replace the release keyed on the thread that last
called next(). When no thread is inside next(), close() on any thread releases
the stream and the dictionaries at once, as on the base branch. If next() has
already returned false, it also frees the last batch's buffers, the native task
context and the child allocator at once. Otherwise it leaves only those to a
Cleaner registered on the batch and its column vectors, which frees them once
nothing can reach the batch; the task context stays allocated until then,
because the library may still release the last batch's arrays through their
callbacks after the stream. Remove the consumer-thread tracking, the releases
from the consumer's next() and close(), and the native-partition-reader-release
pool. Since the last batch now stays readable after close, where on the base
branch the next read failed, have script transformation destroy its script
process when its task completes, as PipedRDD does, so that its writer
thread cannot keep writing that batch to a script whose output nobody reads
and stay blocked on its pipes. State the new rule in the reader and column
vector docs, and replace the deferred-release lifecycle tests with tests of
this rule in the native data source tests.
**Why this works:** A consumer can read a batch only through a reference to
the batch or to its column vectors. Tying the buffers' release to the batch's
reachability therefore keeps every read, on any thread, inside valid memory.
Releasing the stream at close, and the buffers at once after the end of the
stream, keeps the common close prompt, whichever thread called next() and
whether or not it is still alive. Keeping the task context until the last
batch's arrays are released keeps it valid for their release callbacks.
Destroying the script at task completion stops the writer that a failed read
used to stop.
**Scope:** Rework NativePartitionReader's release rule and the lifetime of
the last batch's buffers and task context behind NativeColumnVector, remove the
consumer-thread tracking and the releaser pool, have script transformation
destroy its script at task completion, and replace the lifecycle tests with
tests of the new rule.
**Compatibility:** A close() while another thread is inside get_next still
returns without waiting and leaves the stream's release to that next() once
get_next returns.
**Risks:** The last batch's buffers of a reader closed before the end of its
stream, for example under a LIMIT or in a failed task, are freed only after a
GC, which native memory does not trigger. Until then they stay allocated
outside Spark's memory accounting, with the reader's task context and child
allocator. Neither the Cleaner's state nor the task context that it keeps may
reach the batch or its column vectors, for example through the running task
that is_interrupted consults, or they never become unreachable and the buffers
are never freed. Script transformation now destroys its script process when its
task completes, for every data source, so a script that still runs then, for
example under a LIMIT, is stopped rather than left to read the rest of its
input.
**Constraints:** A close() while another thread is inside get_next must
still return without waiting and leave the release to that next() once get_next
returns. The library must never get the release of the stream during get_next.
The native task context must stay allocated until the last batch's arrays are
released: the header keeps it valid only until the stream is released, but
every array is released before the stream on the base branch, so a library may
still use it in an array's release callback. Freeing the last batch at once
after the end of the stream is safe only once no consumer keeps native column
vectors past the reader's next batch, so this lands together with the change
that makes the columnar Python input copy what it keeps.
**Exclusions:**
- Joining PipedRDD or script-transformation writer threads before the reader
is closed is not part of this repair: the batch's buffers stay valid while
those threads can reach it.
- Making the columnar Python input copy the reused columns it keeps is a
separate, coordinated repair.
**Success:** A batch fetched on one thread and read on another stays
readable with its own values after the fetching thread closes the reader, and
its buffers are freed once nothing references the batch. When no thread is
inside next(), close() releases the stream and the dictionaries at once, on any
thread. When next() has already returned false, it also frees the buffers, the
task context and the child allocator at once. The task context stays allocated
until the last batch's arrays are released. When a script-transformation task
completes, its script process is destroyed, so its writer thread ends even when
the script's pipes are full. No thread waits for another thread to end in order
to release a reader, and the reader no longer records the thread that called
next() last. The reader and column vector docs state this release rule.
**Verification:**
- **Regression:** A batch fetched on one thread and read on another, as
script transformation reads its first batch, returns its own values without
reading freed memory when the fetching thread closes the reader during the
read. Its buffers are freed after the batch becomes unreachable.
- **Regression:** With a consumer thread that stays alive and never calls
the reader again, a close from another thread releases the stream at once,
which the library observes, and leaves no thread waiting.
- **Regression:** A script transformation over a native source under a
LIMIT, whose remaining rows exceed what the script's pipes hold, completes, and
its writer thread and script process then end.
- **Behavior:** A reader that is read to the end and then closed frees its
buffers, its task context and its child allocator at once, without waiting for
a GC, and the test fails if that release is deferred.
- **Compatibility:** A close while another thread is inside get_next still
releases the stream only after get_next returns, and the task context stays
valid until then.
- **Compatibility:** An array that the library releases after the stream,
through the Cleaner, still finds the task context allocated, which a test
library that reads the task in its array release callback observes.
- **Inspection:** No consumer-thread tracking or releaser pool remains, and
the reader and column vector docs describe the release at close and the
deferral of the last batch's buffers only while the batch can still be reached.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,316 @@
+/*
+ * 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 releases them at once only if that
thread has ended, and
+ * otherwise leaves the release to that thread, or to a thread that waits for
it to end. Spark
+ * reads a value of a nullable column only after {@link #isNullAt(int)}, and
the reader checks
+ * that the columns that the schema declares not nullable have no null.
+ *
+ * <p>The vector returns the decimals that fit their precision in a long, of a
scale from 0 to
+ * their precision, as the unscaled longs that Spark's own column vectors
hold, which Spark adds as
+ * longs, rather than as BigDecimals. The other decimals, binary values,
structs, lists and maps,
+ * and the fields of structs are read through {@link ArrowColumnVector}.
+ */
+class NativeColumnVector extends ReusedArrowColumnVector {
+
+ /** The column vector of a top-level Arrow vector of a reader. */
+ static NativeColumnVector create(ValueVector vector) {
+ if (vector instanceof BitVector) {
+ return new Booleans(vector);
+ } else if (vector instanceof TinyIntVector) {
+ return new Bytes(vector);
+ } else if (vector instanceof SmallIntVector) {
+ return new Shorts(vector);
+ } else if (vector instanceof IntVector || vector instanceof DateDayVector)
{
+ return new Ints(vector);
+ } else if (vector instanceof BigIntVector || vector instanceof
TimeStampMicroTZVector ||
+ vector instanceof TimeStampMicroVector) {
+ return new Longs(vector);
+ } else if (vector instanceof Float4Vector) {
+ return new Floats(vector);
+ } else if (vector instanceof Float8Vector) {
+ return new Doubles(vector);
+ } else if (vector instanceof DecimalVector decimalVector) {
+ return new Decimals(decimalVector);
+ } else if (vector instanceof VarCharVector) {
+ return new Strings(vector);
+ } else if (vector instanceof ViewVarCharVector viewVector) {
+ return new StringViews(viewVector);
+ } else {
+ return new NativeColumnVector(vector);
+ }
+ }
+
+ NativeColumnVector(ValueVector vector) {
+ super(vector);
+ }
+
+ /** Reads the addresses of the buffers of the batch that the reader imported
last. */
+ void refresh() {
+ }
+
+ /** A vector that reads its values and their validity at the addresses of
its buffers. */
+ private abstract static class Direct extends NativeColumnVector {
+ final ValueVector vector;
+ // The address of the validity bitmap of the batch, or 0 if the batch has
no null.
+ private long validity;
+ // The address of the data buffer of the batch.
+ long data;
+
+ Direct(ValueVector vector) {
+ super(vector);
+ this.vector = vector;
+ }
+
+ @Override
+ void refresh() {
+ validity = vector.getNullCount() == 0 ? 0 :
vector.getValidityBuffer().memoryAddress();
Review Comment:
**Nit (P3):** For a top-level column that the read schema declares not
nullable, `checkNulls` has just called `getNullCount()` on this vector and
thrown on any null, so this counts the same bitmap a second time only to reach
`validity = 0`. Arrow Java 19.0.0 does not cache the count: `getNullCount()`
rescans the bitmap on each call. That is one extra pass over the bitmap per
non-nullable column per batch. Passing the read schema's nullability to the
column, so that `refresh()` sets `validity` to 0 without counting, gives the
same nulls without the pass.
**Verification:**
- **Behavior:** A top-level column that the read schema declares not
nullable reads the same values and reports no nulls.
- **Inspection:** refresh() does not count nulls for a column that the read
schema declares not nullable.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,316 @@
+/*
+ * 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
Review Comment:
**Nit (P3):** This says strings are read at the addresses of the batch's
buffers, but `create()` gives only `VarCharVector` and `ViewVarCharVector`,
that is Utf8 and Utf8View, a direct subclass. `spark_datasource.h` also lets a
library export strings as `large_utf8`, and `checkSchema` accepts that for a
`StringType` column, since `ArrowUtils` maps `LargeUtf8` to `StringType` and
the check compares only Spark types. Such a column gets a `LargeVarCharVector`,
which falls back to `ArrowColumnVector`'s checked `LargeStringAccessor`. So the
doc overstates which reads skip Arrow Java's checks, and with them which reads
depend on the release rules. Naming Utf8 and Utf8View here, and listing large
strings with the values that read through `ArrowColumnVector`, would match
`create()`.
**Verification:**
- **Inspection:** For every Arrow type that checkSchema accepts, the class
doc lists it as read directly exactly when create() gives it a direct subclass.
--
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]
