dongjoon-hyun commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4417274524
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -278,37 +284,54 @@ class NativeBatchPlanner(
override def plan(): ColumnarPlannedBatch = {
try {
- val partitions = call(name, "plan a scan
of")(library.partitions(handle.get))
- val readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeReader(handle.get))))
+ val planned = call(name, "plan a scan
of")(library.partitions(handle.get))
ColumnarPlannedBatch(
- NativeInputPartition.from(partitions),
- new NativePartitionReaderFactory(executorLocation, name, readerState,
readSchema))
+ NativeInputPartition.from(planned.partitions),
+ new NativePartitionReaderFactory(executorLocation, name,
+ bytesOrEmpty(Option(planned.state)), readSchema, timeZoneId))
} finally {
- // The executors only need the serialized state.
+ // The executors only need the state that the reader returned.
handle.close()
}
}
}
-/** Reads the partitions of a planned batch scan of a native data source, on
the executors. */
+/**
+ * Reads the partitions of a planned batch scan of a native data source, on
the executors. An
+ * engine can also get the Arrow C stream of a partition from it, see
[[SupportsArrowStreamRead]].
+ */
class NativePartitionReaderFactory(
location: NativeLibraryLocation,
name: String,
readerState: Array[Byte],
- readSchema: StructType)
- extends ColumnarOnlyPartitionReaderFactory {
+ readSchema: StructType,
+ timeZoneId: String)
+ extends ColumnarOnlyPartitionReaderFactory with SupportsArrowStreamRead {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, readSchema)
+ new NativePartitionReader(location, name, NativeLibrary.READ_BATCH,
readerState,
+ partition.asInstanceOf[NativeInputPartition].bytes, readSchema,
timeZoneId)
+ }
+
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): AutoCloseable = {
+ val task = NativePartitionReader.openStream(location, name,
NativeLibrary.READ_BATCH,
+ readerState, partition.asInstanceOf[NativeInputPartition].bytes,
readSchema, timeZoneId,
+ streamAddress)
+ // In case the caller does not close it, which must release the stream
before then.
+ Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ =>
task.close()))
Review Comment:
**1. [High] The completion listener can free the task context while the
stream is still open**
`task.close()` frees the `SparkDsTaskContext` whether or not the engine has
released the stream. `spark_datasource.h` tells the library that `task` is
"valid until the stream is released", and the test library keeps it in its
`Stream` and calls `task->is_interrupted(task)` from `get_next`. This listener
can run first in two cases:
- Task completion listeners run in reverse order of registration
(`TaskContextImpl` keeps them on a stack). An engine that registers its own
cleanup, which releases the stream, before it calls `openArrowStream` gets this
listener first: the context is freed, and then the library's `release` runs
against it.
- The stream may be read on another thread, as the new test does. When the
task is killed, the task thread unwinds and runs the listeners while `get_next`
may still be blocked on the engine's thread, polling `is_interrupted` as the
header asks libraries to do. That reads freed memory, which may already be
reused.
In both cases the result is a crash of the executor instead of a leak. Could
this listener only make the context report an interruption, by removing it from
`NativeLibrary.tasks` (which `isInterrupted` already treats as a read that is
over), and log a warning, without freeing the memory? A leaked 40-byte struct
is cheaper than a use-after-free. Item 2 suggests a way to avoid the trade-off.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -278,37 +284,54 @@ class NativeBatchPlanner(
override def plan(): ColumnarPlannedBatch = {
try {
- val partitions = call(name, "plan a scan
of")(library.partitions(handle.get))
- val readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeReader(handle.get))))
+ val planned = call(name, "plan a scan
of")(library.partitions(handle.get))
ColumnarPlannedBatch(
- NativeInputPartition.from(partitions),
- new NativePartitionReaderFactory(executorLocation, name, readerState,
readSchema))
+ NativeInputPartition.from(planned.partitions),
+ new NativePartitionReaderFactory(executorLocation, name,
+ bytesOrEmpty(Option(planned.state)), readSchema, timeZoneId))
} finally {
- // The executors only need the serialized state.
+ // The executors only need the state that the reader returned.
handle.close()
}
}
}
-/** Reads the partitions of a planned batch scan of a native data source, on
the executors. */
+/**
+ * Reads the partitions of a planned batch scan of a native data source, on
the executors. An
+ * engine can also get the Arrow C stream of a partition from it, see
[[SupportsArrowStreamRead]].
+ */
class NativePartitionReaderFactory(
location: NativeLibraryLocation,
name: String,
readerState: Array[Byte],
- readSchema: StructType)
- extends ColumnarOnlyPartitionReaderFactory {
+ readSchema: StructType,
+ timeZoneId: String)
+ extends ColumnarOnlyPartitionReaderFactory with SupportsArrowStreamRead {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, readSchema)
+ new NativePartitionReader(location, name, NativeLibrary.READ_BATCH,
readerState,
+ partition.asInstanceOf[NativeInputPartition].bytes, readSchema,
timeZoneId)
+ }
+
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): AutoCloseable = {
Review Comment:
**2. [Medium] Could Spark free the context when the stream is released,
instead of returning a handle?**
The handle and the listener give the context a lifetime of its own, next to
the one of the stream, so every engine has to get their order right, and
getting it wrong crashes the executor (item 1). The header already ties the
context to the stream, and Spark could make that hold for any consumer:
`openArrowStream` would replace `release` with an upcall, like `IS_INTERRUPTED`
and the sinks, and remember the library's `release` and the `Task` by the
stream's `private_data`. The upcall would restore the library's `release`, call
it, and then close the `Task`. The other callbacks stay those of the library,
and a moved stream keeps working, since `private_data` moves with it. A library
that leaves `private_data` NULL would need another key; wrapping `private_data`
with trampolines for all four callbacks avoids that, at the cost of an upcall
per `get_next`.
Then `openArrowStream` would not need to return a handle or register a
listener, the context would be freed exactly when the library is done with it,
and a stream that an engine never releases would leak 40 bytes instead of
crashing. `NativePartitionReader` could use the same path.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1145,6 +1219,76 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
check(writeTasks)
}
+ nativeTest("open the partitions of a batch scan as Arrow C streams") {
+ val session = newSession(defaultPackage)
+ // The reader factory and the partitions of the batch scan of a query,
which an engine gets
+ // from its plan.
+ def plannedScan(df: DataFrame): (SupportsArrowStreamRead,
Seq[InputPartition]) = {
+ val scan = df.queryExecution.executedPlan.collectFirst { case s:
BatchScanExec => s }.get
+ (scan.readerFactory.asInstanceOf[SupportsArrowStreamRead],
scan.inputPartitions)
+ }
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ // An engine opens the stream of each partition in its task, and reads it
from another thread.
+ val ids = session.sparkContext.parallelize(partitions,
partitions.length).mapPartitions {
Review Comment:
**5. [Medium] This job runs outside of a SQL execution, which `checkRows`
says to avoid in this suite**
The comment on `checkRows` explains that a job outside of a SQL execution
"gets the files of all the sessions", and that the streaming queries of the
tests create sessions whose files are removed when they are garbage collected.
This `parallelize(...).collect()` is such a job: without a `JobArtifactState`,
`JobArtifactSet.getActiveOrDefault` ships `sc.allAddedFiles` with the tasks. If
a session of an earlier streaming test is collected around this point, the
executor may fail to fetch one of its files, and the test fails for a reason
that has nothing to do with the stream.
Wrapping the job in `session.artifactManager.withResources { ... }` gives it
only the files of `session`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -839,10 +867,67 @@ class NativePartitionReader(
}
}
-/** Writes the batches of a task to the data writer of a native library. */
+object NativePartitionReader {
+ import NativeCalls._
+
+ /**
+ * Initializes the Arrow C stream at `streamAddress`, whose release callback
is NULL, with the
+ * data of a partition, by calling read of the library in the running task.
Returns the context
+ * of the task, which the library may use until the stream is released: it
is closed then. If
+ * this fails, the stream is released.
+ */
+ def openStream(
+ location: NativeLibraryLocation,
+ name: String,
+ kind: Int,
+ readerState: Array[Byte],
+ partition: Array[Byte],
+ schema: StructType,
+ timeZoneId: String,
+ streamAddress: Long): NativeLibrary.Task = {
+ val library = NativeLibraries.get(location)
+ val stream = ArrowArrayStream.wrap(streamAddress)
+ val task = NativeTask.create()
+ try {
+ NativeArrow.exportSchema(schema, timeZoneId) { schemaAddress =>
+ call(name, "read") {
+ library.read(task, kind, schemaAddress, readerState, partition,
streamAddress)
+ }
+ }
+ // Arrow calls the callbacks of the stream without checking them, and
calling NULL would
+ // crash the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 ||
snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream,
or a callback of " +
+ "the stream is NULL.")
+ }
+ task
+ } catch {
+ case e: Throwable =>
+ // Releases the stream, if the library initialized it, before the task
that it may use,
+ // and marks it released in case its release callback did not.
+ Utils.tryWithSafeFinally {
+ stream.release()
+ stream.markReleased()
+ }(task.close())
Review Comment:
**9. [Nit] The error of `read` is dropped if `stream.release()` throws**
If `stream.release()` throws, `Utils.tryWithSafeFinally` rethrows that
exception after `task.close()`, so `throw e` never runs, and the original error
is lost without even being added as suppressed. `release()` only throws when
the JNI library of Arrow is not usable, which would usually have failed
`exportSchema` first, so this is unlikely; still, adding a failure of the
cleanup to `e` with `addSuppressed` keeps the error that matters.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsArrowStreamRead.java:
##########
@@ -0,0 +1,54 @@
+/*
+ * 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.connector.read;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * A mix-in interface for a {@link PartitionReaderFactory} that also reads a
partition as an
+ * <a href="https://arrow.apache.org/docs/format/CStreamInterface.html";>Arrow
C stream</a>, for an
+ * engine that processes Arrow data in native code: the engine imports the
arrays of the stream
+ * itself, instead of reading the {@link
org.apache.spark.sql.vectorized.ColumnarBatch}es of
+ * {@link #createColumnarReader(InputPartition)}. The reader factories of the
batch scans of
+ * native data sources implement it, with the stream that the library returns.
+ * <p>
+ * Each array of the stream is a struct array whose fields are the columns of
the read schema of
+ * the scan. Spark does not check the stream, while it checks the streams that
it reads itself, so
+ * the engine checks what its import of the arrays requires, such as the
schema of the stream
+ * against the read schema, and the offsets and the lengths of the arrays. The
callbacks of the
+ * stream can be called from any thread, but never concurrently.
+ *
+ * @since 4.4.0
+ */
+@Evolving
+public interface SupportsArrowStreamRead extends PartitionReaderFactory {
+
+ /**
+ * Initializes an Arrow C stream with the data of a partition. It is called
in the task that
+ * reads the partition, whose context the stream may use until it is
released.
+ *
+ * @param partition a partition of the scan, as the scan planned it
+ * @param streamAddress the address of a {@code struct ArrowArrayStream}
that the caller
+ * allocated, and whose release callback is NULL, as it
is in a
+ * zero-initialized struct. If this method fails, the
stream is released.
+ * @return a handle that the caller closes once the stream is released,
which frees what the
+ * stream uses until then. Spark also closes it when the task
completes, so the stream
+ * must be released before then.
Review Comment:
**7. [Low] Nothing checks that `openArrowStream` runs in a task that has not
completed**
"Spark also closes it when the task completes" only holds when
`openArrowStream` runs on the task thread before the task completes. The
Javadoc asks for that ("It is called in the task that reads the partition"),
but nothing checks it:
- Without a `TaskContext`, for example on a thread of the engine, the
context gets -1 for every ID, `is_interrupted` never reports a kill, and no
listener is registered, so a handle that the engine does not close is never
freed.
- After the task completed, for example from a completion listener of the
engine, `addTaskCompletionListener` runs the new listener right away, so
`openArrowStream` returns a handle that is already closed, for a stream that is
still open.
Should `openArrowStream` fail when `TaskContext.get()` is null or
`isCompleted()`? The failure cases of the new test call it on the driver, so
they would need to run in a task.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -278,37 +284,54 @@ class NativeBatchPlanner(
override def plan(): ColumnarPlannedBatch = {
try {
- val partitions = call(name, "plan a scan
of")(library.partitions(handle.get))
- val readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeReader(handle.get))))
+ val planned = call(name, "plan a scan
of")(library.partitions(handle.get))
ColumnarPlannedBatch(
- NativeInputPartition.from(partitions),
- new NativePartitionReaderFactory(executorLocation, name, readerState,
readSchema))
+ NativeInputPartition.from(planned.partitions),
+ new NativePartitionReaderFactory(executorLocation, name,
+ bytesOrEmpty(Option(planned.state)), readSchema, timeZoneId))
} finally {
- // The executors only need the serialized state.
+ // The executors only need the state that the reader returned.
handle.close()
}
}
}
-/** Reads the partitions of a planned batch scan of a native data source, on
the executors. */
+/**
+ * Reads the partitions of a planned batch scan of a native data source, on
the executors. An
+ * engine can also get the Arrow C stream of a partition from it, see
[[SupportsArrowStreamRead]].
+ */
class NativePartitionReaderFactory(
location: NativeLibraryLocation,
name: String,
readerState: Array[Byte],
- readSchema: StructType)
- extends ColumnarOnlyPartitionReaderFactory {
+ readSchema: StructType,
+ timeZoneId: String)
+ extends ColumnarOnlyPartitionReaderFactory with SupportsArrowStreamRead {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, readSchema)
+ new NativePartitionReader(location, name, NativeLibrary.READ_BATCH,
readerState,
+ partition.asInstanceOf[NativeInputPartition].bytes, readSchema,
timeZoneId)
+ }
+
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): AutoCloseable = {
+ val task = NativePartitionReader.openStream(location, name,
NativeLibrary.READ_BATCH,
+ readerState, partition.asInstanceOf[NativeInputPartition].bytes,
readSchema, timeZoneId,
+ streamAddress)
+ // In case the caller does not close it, which must release the stream
before then.
+ Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ =>
task.close()))
Review Comment:
**11. [Nit] One completion listener per stream stays registered until the
end of the task**
A `TaskContext` cannot remove a listener, so each `openArrowStream` adds one
that keeps its `Task` until the task ends, even when the engine closed the
handle long before. An engine that opens many streams in one task accumulates
them, and the end of the task then runs one no-op close for each. One listener
per task, which closes the contexts that are still open from a set that
`Task.close()` removes them from, would keep this constant. Item 2 would remove
the listener altogether.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -304,6 +308,31 @@ typedef struct SparkDsCreateInfo {
int64_t num_options;
} SparkDsCreateInfo;
+/* A micro-batch streaming scan, as create_stream_reader gets it. */
+typedef struct SparkDsStreamReaderInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /*
+ * The location of a directory in the checkpoint of the streaming query that
is reserved for this
+ * source, such as file:/tmp/checkpoint/sources/0, where the library may
keep what it needs to
+ * recover from a failure. It stays the same when the query restarts from
its checkpoint. It is a
+ * URI, possibly of a file system that the library cannot access, such as
s3a or hdfs.
+ */
+ const char* checkpoint_location;
+} SparkDsStreamReaderInfo;
+
+/* A request for the latest offset of a stream reader, as latest_offset gets
it. */
+typedef struct SparkDsLatestOffsetInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /*
+ * The offset, as JSON, that the next micro-batch starts after: the end of
the last micro-batch,
+ * or the initial offset before the first one. With Trigger.AvailableNow,
Spark asks for the
+ * latest offset only once, when the query starts, with the initial offset.
+ */
+ const char* start;
+} SparkDsLatestOffsetInfo;
+
/*
* The task that calls read or create_data_writer, on an executor. It stays
valid until the stream
* that read returns is released, or until the data writer is closed.
Review Comment:
**3. [Medium] The header still says that the context outlives the stream**
This comment says the context "stays valid until the stream that read
returns is released", and the doc of `read` (lines 594 to 603) says that when
the task is killed, Spark only stops reading once `get_next` returns, and then
releases the stream. Both hold for `NativePartitionReader`, but not for a
stream read through `SupportsArrowStreamRead`: the engine owns the stream, may
read it on another thread, and Spark frees the context when the task completes,
whatever the state of the stream (item 1).
A library author reads only this header, so it should say what changes under
an engine, for example that `task` must not be used after the end of the task
even if the stream is not released yet. If item 2 is taken, the current wording
becomes true again.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -278,37 +284,54 @@ class NativeBatchPlanner(
override def plan(): ColumnarPlannedBatch = {
try {
- val partitions = call(name, "plan a scan
of")(library.partitions(handle.get))
- val readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeReader(handle.get))))
+ val planned = call(name, "plan a scan
of")(library.partitions(handle.get))
ColumnarPlannedBatch(
- NativeInputPartition.from(partitions),
- new NativePartitionReaderFactory(executorLocation, name, readerState,
readSchema))
+ NativeInputPartition.from(planned.partitions),
+ new NativePartitionReaderFactory(executorLocation, name,
+ bytesOrEmpty(Option(planned.state)), readSchema, timeZoneId))
} finally {
- // The executors only need the serialized state.
+ // The executors only need the state that the reader returned.
handle.close()
}
}
}
-/** Reads the partitions of a planned batch scan of a native data source, on
the executors. */
+/**
+ * Reads the partitions of a planned batch scan of a native data source, on
the executors. An
+ * engine can also get the Arrow C stream of a partition from it, see
[[SupportsArrowStreamRead]].
+ */
class NativePartitionReaderFactory(
location: NativeLibraryLocation,
name: String,
readerState: Array[Byte],
- readSchema: StructType)
- extends ColumnarOnlyPartitionReaderFactory {
+ readSchema: StructType,
+ timeZoneId: String)
+ extends ColumnarOnlyPartitionReaderFactory with SupportsArrowStreamRead {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, readSchema)
+ new NativePartitionReader(location, name, NativeLibrary.READ_BATCH,
readerState,
+ partition.asInstanceOf[NativeInputPartition].bytes, readSchema,
timeZoneId)
+ }
+
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): AutoCloseable = {
+ val task = NativePartitionReader.openStream(location, name,
NativeLibrary.READ_BATCH,
+ readerState, partition.asInstanceOf[NativeInputPartition].bytes,
readSchema, timeZoneId,
+ streamAddress)
Review Comment:
**4. [Medium] `partition_id` is the partition of the engine's task, not the
partition being read**
`NativeTask.create()` takes the partition from `TaskContext.partitionId()`.
In `DataSourceRDD` that is the partition of the scan, but an engine runs the
scan in its own RDD, where a task may read another partition of the scan, or
several of them, for example after coalescing small partitions. The library
then gets the same `partition_id` for different partitions, while the header
documents it as "The partition that the task reads or writes".
`NativeDataWriter` already passes its partition explicitly
(`NativeTask.create(Some(partitionId), Some(taskId))`), and here
`NativeInputPartition.index` is at hand. The new test does not catch this,
because the partitions 0 and 1 of the engine's RDD happen to read the
partitions 0 and 1 of the scan.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1145,6 +1219,76 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
check(writeTasks)
}
+ nativeTest("open the partitions of a batch scan as Arrow C streams") {
+ val session = newSession(defaultPackage)
+ // The reader factory and the partitions of the batch scan of a query,
which an engine gets
+ // from its plan.
+ def plannedScan(df: DataFrame): (SupportsArrowStreamRead,
Seq[InputPartition]) = {
+ val scan = df.queryExecution.executedPlan.collectFirst { case s:
BatchScanExec => s }.get
+ (scan.readerFactory.asInstanceOf[SupportsArrowStreamRead],
scan.inputPartitions)
+ }
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ // An engine opens the stream of each partition in its task, and reads it
from another thread.
+ val ids = session.sparkContext.parallelize(partitions,
partitions.length).mapPartitions {
+ partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try {
+ val handle = factory.openArrowStream(partition.next(),
stream.memoryAddress())
+ try {
+ ThreadUtils.runInNewThread("engine") {
+ // Moves the stream, which closing the reader releases.
+ val reader = Data.importArrayStream(allocator, stream)
+ try {
+ val ids = ArrayBuffer.empty[Long]
+ while (reader.loadNextBatch()) {
+ val root = reader.getVectorSchemaRoot
+ val vector = root.getVector("id").asInstanceOf[BigIntVector]
+ (0 until root.getRowCount).foreach(i => ids += vector.get(i))
+ }
+ ids.toSeq
+ } finally {
+ reader.close()
+ }
+ }.iterator
+ } finally {
+ handle.close()
+ }
+ } finally {
+ stream.close()
+ allocator.close()
+ }
+ }.collect()
+ assert(ids.sorted.toSeq == (0L until 6L))
+ // The library got the context of each task.
+ assert(taskDir.list().toSet == Set("0", "1"))
+
+ // If it fails, the stream is released.
+ Seq(
+ "open" -> "injected failure in open",
+ "uninitialized" ->
+ "The library did not initialize the stream, or a callback of the
stream is NULL."
+ ).foreach { case (fail, msg) =>
+ val (failing, failingPartitions) = plannedScan(range(session, "fail" ->
fail))
+ val stream = ArrowArrayStream.allocateNew(ArrowUtils.rootAllocator)
+ try {
+ val e = intercept[Exception] {
+ failing.openArrowStream(failingPartitions.head,
stream.memoryAddress())
+ }
+ checkError(
+ exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"),
+ condition = "NATIVE_DATA_SOURCE_ERROR",
+ parameters = Map("action" -> "read", "name" -> "native_range", "msg"
-> msg))
+ assert(stream.snapshot().release == 0)
Review Comment:
**6. [Medium] `release == 0` holds whether or not Spark released the stream**
Neither injected failure initializes the stream: with `"open"`,
`readPartition` throws before it touches the stream, and with `"uninitialized"`
it returns without touching it. `ArrowArrayStream.allocateNew` already sets
`release` to NULL, and `openStream` calls `markReleased()` in any case, so this
assertion passes even if the `stream.release()` call in `openStream` is removed.
A case where the library does initialize the stream but Spark rejects it,
for example a non-NULL `release` with a NULL `get_next`, would check that Spark
calls the library's `release`; the test library could record that in a file, as
it does with `closed_path`. The listener that closes the context at the end of
the task, and a close of the handle that races with it, which is why
`Task.close()` is now `synchronized`, are not exercised either.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1145,6 +1219,76 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
check(writeTasks)
}
+ nativeTest("open the partitions of a batch scan as Arrow C streams") {
+ val session = newSession(defaultPackage)
+ // The reader factory and the partitions of the batch scan of a query,
which an engine gets
+ // from its plan.
+ def plannedScan(df: DataFrame): (SupportsArrowStreamRead,
Seq[InputPartition]) = {
+ val scan = df.queryExecution.executedPlan.collectFirst { case s:
BatchScanExec => s }.get
+ (scan.readerFactory.asInstanceOf[SupportsArrowStreamRead],
scan.inputPartitions)
+ }
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ // An engine opens the stream of each partition in its task, and reads it
from another thread.
+ val ids = session.sparkContext.parallelize(partitions,
partitions.length).mapPartitions {
+ partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try {
+ val handle = factory.openArrowStream(partition.next(),
stream.memoryAddress())
+ try {
+ ThreadUtils.runInNewThread("engine") {
+ // Moves the stream, which closing the reader releases.
+ val reader = Data.importArrayStream(allocator, stream)
+ try {
+ val ids = ArrayBuffer.empty[Long]
+ while (reader.loadNextBatch()) {
+ val root = reader.getVectorSchemaRoot
+ val vector = root.getVector("id").asInstanceOf[BigIntVector]
+ (0 until root.getRowCount).foreach(i => ids += vector.get(i))
+ }
+ ids.toSeq
+ } finally {
+ reader.close()
+ }
+ }.iterator
+ } finally {
+ handle.close()
Review Comment:
**10. [Nit] As an example of the protocol, the test could release the stream
before it closes the handle**
If anything fails between `openArrowStream` and `Data.importArrayStream`,
which moves the stream, this `finally` closes the handle while the stream of
the library is still open, and `stream.close()` below only frees the struct.
Since the test also shows engine authors how to use the API, calling
`stream.release()` before `handle.close()` (a no-op once the stream was moved)
would follow the order that the Javadoc requires.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -839,10 +867,67 @@ class NativePartitionReader(
}
}
-/** Writes the batches of a task to the data writer of a native library. */
+object NativePartitionReader {
+ import NativeCalls._
+
+ /**
+ * Initializes the Arrow C stream at `streamAddress`, whose release callback
is NULL, with the
+ * data of a partition, by calling read of the library in the running task.
Returns the context
+ * of the task, which the library may use until the stream is released: it
is closed then. If
+ * this fails, the stream is released.
+ */
+ def openStream(
+ location: NativeLibraryLocation,
+ name: String,
+ kind: Int,
+ readerState: Array[Byte],
+ partition: Array[Byte],
+ schema: StructType,
+ timeZoneId: String,
+ streamAddress: Long): NativeLibrary.Task = {
+ val library = NativeLibraries.get(location)
+ val stream = ArrowArrayStream.wrap(streamAddress)
Review Comment:
**8. [Low] The preconditions of `streamAddress` are not checked**
The Javadoc requires a stream whose `release` is NULL, but `openStream`
passes the address to the library as it is. A 0 makes the library write to
NULL, which crashes the JVM instead of failing the task. A struct that still
holds a stream of the engine is overwritten by the library, which leaks that
stream, and if `read` fails before it initializes the struct, the `catch` below
releases the engine's stream. Checking `streamAddress != 0` and
`stream.snapshot().release == 0` before calling `read` would turn both into
errors.
--
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]
