dongjoon-hyun commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4420249991
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -453,6 +493,15 @@ object NativeInputPartition {
}
}
+/**
+ * A partition of a micro-batch, with the state that the stream reader
returned with the
+ * partitions of the micro-batch. The state changes with each micro-batch, so
it goes with the
+ * partitions rather than with the reader factory, which an engine may keep
while Spark plans the
+ * next micro-batch.
+ */
+case class NativeMicroBatchPartition(partition: NativeInputPartition,
readerState: Array[Byte])
Review Comment:
**1. [Low] The reader state of a micro-batch is now serialized with every
task**
d81e619 moves the state that `stream_partitions` returned from the stream
reader into each `NativeMicroBatchPartition`, which answers [round
4](https://github.com/apache/spark/pull/59934#discussion_r4419050240): a reader
factory that an engine keeps now reads each partition with the state of its
micro-batch. It also changes how often the state is sent. `DataSourceRDD` keeps
the reader factory as a field, so the factory travels in the task binary, which
Spark broadcasts once per stage, while each task serializes its own
`DataSourceRDDPartition` with its input partitions. Before, the state went with
the reader inside the factory, so each executor fetched it once per
micro-batch. Now every task of the micro-batch carries its own copy, so a state
of S bytes over N partitions costs N times S per micro-batch. For example, a
state of 1 MiB over 1,000 partitions sends about 1 GiB, and each task goes over
the 1000 KiB at which `TaskSetManager` warns about large tasks. The batch scan
still
keeps its state in `NativePartitionReaderFactory`.
I agree that the partitions are the right place to tie the state to its
micro-batch, since the order of `planInputPartitions` and `createReaderFactory`
is not defined. Two ways to keep the cost down:
- Broadcast the state of each micro-batch once, in `partitions` on the
driver, and keep the `Broadcast` in the partitions.
- Or keep it as is, and say in the doc of `stream_partitions` in the header
that the state is sent with every task, so that a library keeps it small.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1143,6 +1353,261 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
.mode("append")
.save()
check(writeTasks)
+
+ // A read gets the index of its partition in the scan, also when a task
reads several
+ // partitions, or the partitions of several scans: the task attempt that
read each partition.
+ def readers(dir: File): Map[Int, Long] = dir.listFiles().map { file =>
+ file.getName.toInt -> Files.readString(file.toPath).split(",")(5).toLong
+ }.toMap
+ val (first, second) = (Utils.createTempDir(), Utils.createTempDir())
+ checkRows(
+ range(session, "end" -> "6", "task_path" -> first.getPath).select("id")
+ .union(range(session, "end" -> "6", "task_path" ->
second.getPath).select("id")),
+ (0L until 6L).flatMap(i => Seq(Row(i), Row(i))))
+ assert(readers(first).keySet == Set(0, 1) && readers(second).keySet ==
Set(0, 1))
+ val coalesced = Utils.createTempDir()
+ checkRows(range(session, "end" -> "6", "task_path" ->
coalesced.getPath).select("id")
+ .coalesce(1), (0L until 6L).map(Row(_)))
+ assert(readers(coalesced).keySet == Set(0, 1) &&
readers(coalesced).values.toSet.size == 1)
+ // So does a read of a micro-batch.
+ val (firstStream, secondStream) = (Utils.createTempDir(),
Utils.createTempDir())
+ def counter(dir: File): DataFrame =
session.readStream.format("native_counter")
+ .option("max_offset", "4").option("task_path", dir.getPath).load()
+ val query =
counter(firstStream).union(counter(secondStream)).writeStream.format("noop").start()
+ try query.processAllAvailable() finally query.stop()
+ assert(readers(firstStream).keySet == Set(0) &&
readers(secondStream).keySet == Set(0))
+ }
+
+ nativeTest("open the partitions of batch and micro-batch scans 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)
+ }
+ // Runs `f` in a task that only gets the files of the session.
+ def inTask[T: ClassTag](f: => T): T =
session.artifactManager.withResources {
+ session.sparkContext.parallelize(Seq(0), 1).map(_ => f).collect().head
+ }
+ import TestArrowEngine.{callbacks, readIds}
+ // Reads the given partitions in a single task, as an engine does.
+ def readAll(factory: SupportsArrowStreamRead, partitions:
Seq[InputPartition]): Seq[Long] = {
+ inTask {
+ partitions.flatMap { partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ try {
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try factory.openArrowStream(partition, stream.memoryAddress())
catch {
+ case e: Throwable =>
+ stream.close()
+ throw e
+ }
+ readIds(stream, allocator)
+ } finally {
+ allocator.close()
+ }
+ }
+ }
+ }
+ // Checks that `f` opens `delta` more task contexts than it closes. The
count is global to the
+ // JVM, where the tasks of earlier tests may still close their contexts,
so it is read right
+ // before `f`.
+ def checkContexts[T](delta: Int)(f: => T): T = {
+ val before = NativeLibrary.openTaskContexts()
+ val result = f
+ assert(NativeLibrary.openTaskContexts() == before + delta)
+ result
+ }
+
+ // An engine reads both partitions of the scan in a single task.
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ assert(readAll(factory, partitions).sorted == (0L until 6L))
+ // The library got the partition of the scan that it reads, not the one of
the task.
+ val tasks = taskDir.listFiles().map { file =>
+ file.getName.toInt ->
Files.readString(file.toPath).split(",").map(_.toLong).toSeq
+ }.toMap
+ assert(tasks.keySet == Set(0, 1) && tasks.values.map(_(5)).toSet.size ==
1, tasks)
+ assert(NativeLibrary.openEngineStreams() == 0)
+
+ // The engine may release the stream after its task ended, and on another
thread: Spark frees
+ // the context of the task once the library released the stream.
+ val released = new File(Utils.createTempDir(), "released")
+ val (laterFactory, laterPartitions) =
+ plannedScan(range(session, "end" -> "6", "released_path" ->
released.getPath).select("id"))
+ // Opens a stream in a task, which it returns unreleased.
+ def openInTask(factory: SupportsArrowStreamRead, partition:
InputPartition): Long = inTask {
+ val stream = Platform.allocateMemory(40)
+ Platform.setMemory(stream, 0.toByte, 40)
+ factory.openArrowStream(partition, stream)
+ stream
+ }
+ val (address, copy) = checkContexts(1) {
+ val address = openInTask(laterFactory, laterPartitions.head)
+ assert(NativeLibrary.openEngineStreams() == 1 && !released.exists())
+ // A copy of the stream, as an engine that moves the stream may leave
behind.
+ val copy = Platform.allocateMemory(40)
+ Platform.copyMemory(null, address, null, copy, 40)
+ (address, copy)
+ }
+ checkContexts(-1) {
+ assert(readIds(ArrowArrayStream.wrap(address), ArrowUtils.rootAllocator)
== (0L until 3L))
Review Comment:
**2. [Low] The test reads a stream after its task ended, which the API lets
the source refuse**
This part answers [round
2](https://github.com/apache/spark/pull/59934#discussion_r4417661698), and
checks that the context outlives the task until the stream is released. But
`openInTask` returns once its job finished, so this line reads the whole stream
on the driver after the task completed. The context then reports
`is_interrupted` as true, and the Javadoc of `openArrowStream` says that "once
the task ended, the source may stop returning data and fail". The read only
succeeds because `get_next` of `native_range` checks `is_interrupted` only with
the option `wait_for_interrupt`. If the test library followed the guide and
checked `is_interrupted` in every `get_next`, this part would fail for a reason
that has nothing to do with the release that it checks. It also shows an engine
reading a stream after its task ended, which the API asks engines not to rely
on.
Since the release is what this part checks, could `openInTask` read the
stream to its end in the task, or not read it at all, and return it unreleased?
The driver would then only release it, for example with
`ArrowArrayStream.wrap(address).release()`, and check the context count and
`released` as now.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,17 +364,31 @@ class NativeDataSourceStreamReader(
}
override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
+ val planned = call(name, "plan a scan of") {
library.streamPartitions(handle.get, start, end)
}
- readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeStreamReader(handle.get))))
- NativeInputPartition.from(partitions)
+ val readerState = bytesOrEmpty(Option(planned.state))
+ NativeInputPartition.from(planned.partitions).map { partition =>
+ NativeMicroBatchPartition(partition.asInstanceOf[NativeInputPartition],
readerState)
+ }
}
override def read(partition: InputPartition): PartitionReader[ColumnarBatch]
= {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, schema)
+ val microBatchPartition = partition.asInstanceOf[NativeMicroBatchPartition]
+ new NativePartitionReader(location, name, NativeLibrary.READ_MICRO_BATCH,
+ microBatchPartition.readerState, microBatchPartition.partition, schema,
timeZoneId)
+ }
+
+ // An engine can also get the Arrow C stream of a partition, see
[[SupportsArrowStreamRead]].
+ override def readerFactory(): PartitionReaderFactory = {
+ new ColumnarStreamPartitionReaderFactory(this) with
SupportsArrowStreamRead {
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): Unit = {
+ val microBatchPartition =
partition.asInstanceOf[NativeMicroBatchPartition]
Review Comment:
**3. [Nit] The casts of the partitions now run before the checks of
`openEngineStream`**
d81e619 moved the casts into the factories, here and in
`NativePartitionReaderFactory.openArrowStream`, so they run before
`openEngineStream` checks the task and the stream. A partition of another kind,
for example a `NativeInputPartition` of a batch scan passed to the factory of a
micro-batch scan, or the other way around, now fails with a
`ClassCastException` such as `NativeInputPartition cannot be cast to
NativeMicroBatchPartition`. That also happens outside of a task, where the
Javadoc says that the call fails because it does not run in the task. In
b98fd19, the cast came after the checks. Nothing leaks, since the stream is not
touched yet.
Could the factories run the checks of `openEngineStream` first, for example
through a helper that both call before they look at the partition, and reject a
partition of another kind with an `IllegalArgumentException` that names the
expected class?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1143,6 +1353,261 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
.mode("append")
.save()
check(writeTasks)
+
+ // A read gets the index of its partition in the scan, also when a task
reads several
+ // partitions, or the partitions of several scans: the task attempt that
read each partition.
+ def readers(dir: File): Map[Int, Long] = dir.listFiles().map { file =>
+ file.getName.toInt -> Files.readString(file.toPath).split(",")(5).toLong
+ }.toMap
+ val (first, second) = (Utils.createTempDir(), Utils.createTempDir())
+ checkRows(
+ range(session, "end" -> "6", "task_path" -> first.getPath).select("id")
+ .union(range(session, "end" -> "6", "task_path" ->
second.getPath).select("id")),
+ (0L until 6L).flatMap(i => Seq(Row(i), Row(i))))
+ assert(readers(first).keySet == Set(0, 1) && readers(second).keySet ==
Set(0, 1))
+ val coalesced = Utils.createTempDir()
+ checkRows(range(session, "end" -> "6", "task_path" ->
coalesced.getPath).select("id")
+ .coalesce(1), (0L until 6L).map(Row(_)))
+ assert(readers(coalesced).keySet == Set(0, 1) &&
readers(coalesced).values.toSet.size == 1)
+ // So does a read of a micro-batch.
+ val (firstStream, secondStream) = (Utils.createTempDir(),
Utils.createTempDir())
+ def counter(dir: File): DataFrame =
session.readStream.format("native_counter")
+ .option("max_offset", "4").option("task_path", dir.getPath).load()
+ val query =
counter(firstStream).union(counter(secondStream)).writeStream.format("noop").start()
+ try query.processAllAvailable() finally query.stop()
+ assert(readers(firstStream).keySet == Set(0) &&
readers(secondStream).keySet == Set(0))
+ }
+
+ nativeTest("open the partitions of batch and micro-batch scans 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)
+ }
+ // Runs `f` in a task that only gets the files of the session.
+ def inTask[T: ClassTag](f: => T): T =
session.artifactManager.withResources {
+ session.sparkContext.parallelize(Seq(0), 1).map(_ => f).collect().head
+ }
+ import TestArrowEngine.{callbacks, readIds}
+ // Reads the given partitions in a single task, as an engine does.
+ def readAll(factory: SupportsArrowStreamRead, partitions:
Seq[InputPartition]): Seq[Long] = {
+ inTask {
+ partitions.flatMap { partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ try {
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try factory.openArrowStream(partition, stream.memoryAddress())
catch {
+ case e: Throwable =>
+ stream.close()
+ throw e
+ }
+ readIds(stream, allocator)
+ } finally {
+ allocator.close()
+ }
+ }
+ }
+ }
+ // Checks that `f` opens `delta` more task contexts than it closes. The
count is global to the
+ // JVM, where the tasks of earlier tests may still close their contexts,
so it is read right
+ // before `f`.
+ def checkContexts[T](delta: Int)(f: => T): T = {
+ val before = NativeLibrary.openTaskContexts()
+ val result = f
+ assert(NativeLibrary.openTaskContexts() == before + delta)
+ result
+ }
+
+ // An engine reads both partitions of the scan in a single task.
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ assert(readAll(factory, partitions).sorted == (0L until 6L))
+ // The library got the partition of the scan that it reads, not the one of
the task.
+ val tasks = taskDir.listFiles().map { file =>
+ file.getName.toInt ->
Files.readString(file.toPath).split(",").map(_.toLong).toSeq
+ }.toMap
+ assert(tasks.keySet == Set(0, 1) && tasks.values.map(_(5)).toSet.size ==
1, tasks)
+ assert(NativeLibrary.openEngineStreams() == 0)
+
+ // The engine may release the stream after its task ended, and on another
thread: Spark frees
+ // the context of the task once the library released the stream.
+ val released = new File(Utils.createTempDir(), "released")
+ val (laterFactory, laterPartitions) =
+ plannedScan(range(session, "end" -> "6", "released_path" ->
released.getPath).select("id"))
+ // Opens a stream in a task, which it returns unreleased.
+ def openInTask(factory: SupportsArrowStreamRead, partition:
InputPartition): Long = inTask {
+ val stream = Platform.allocateMemory(40)
+ Platform.setMemory(stream, 0.toByte, 40)
+ factory.openArrowStream(partition, stream)
+ stream
+ }
+ val (address, copy) = checkContexts(1) {
+ val address = openInTask(laterFactory, laterPartitions.head)
+ assert(NativeLibrary.openEngineStreams() == 1 && !released.exists())
+ // A copy of the stream, as an engine that moves the stream may leave
behind.
+ val copy = Platform.allocateMemory(40)
+ Platform.copyMemory(null, address, null, copy, 40)
+ (address, copy)
+ }
+ checkContexts(-1) {
+ assert(readIds(ArrowArrayStream.wrap(address), ArrowUtils.rootAllocator)
== (0L until 3L))
+ assert(NativeLibrary.openEngineStreams() == 0 && released.exists())
+ }
+ Platform.freeMemory(address)
+ // Once the stream is released, the callbacks of the copy fail with
EINVAL, and its release
+ // only marks it released.
+ checkContexts(0) {
+ val array = ArrowArray.allocateNew(ArrowUtils.rootAllocator)
+ val stale = try
intercept[Exception](ArrowArrayStream.wrap(copy).getNext(array)) finally {
+ array.close()
+ }
+ assert(Utils.exceptionString(stale)
+ .contains("The stream was released, or is not a stream of Spark."))
+ ArrowArrayStream.wrap(copy).release()
+ assert(ArrowArrayStream.wrap(copy).snapshot().release == 0)
+ assert(NativeLibrary.openEngineStreams() == 0)
+ }
+ Platform.freeMemory(copy)
+
+ // Once the task ended, its context tells the library that it was
interrupted, so that the
+ // library stops waiting.
+ val marker = new File(Utils.createTempDir(), "interrupted")
+ val (waitingFactory, waitingPartitions) = plannedScan(
+ range(session, "end" -> "6", "wait_for_interrupt" ->
marker.getPath).select("id"))
+ checkContexts(0) {
+ val waiting = openInTask(waitingFactory, waitingPartitions.head)
+ val error = intercept[Exception] {
+ readIds(ArrowArrayStream.wrap(waiting), ArrowUtils.rootAllocator)
+ }
+ assert(Utils.exceptionString(error).contains("the task was killed"))
+ assert(Files.readString(marker.toPath) == "interrupted")
+ assert(NativeLibrary.openEngineStreams() == 0)
+ Platform.freeMemory(waiting)
+ }
+
+ // So does the reader factory of a micro-batch scan, with the state of its
stream reader.
+ val query =
session.readStream.format("native_counter").option("max_offset", "4").load()
+ .writeStream.format("noop").start()
+ val microBatchScan = try {
+ query.processAllAvailable()
+
query.asInstanceOf[StreamingQueryWrapper].streamingQuery.lastExecution.executedPlan
+ .collectFirst { case s: MicroBatchScanExec => s }.get
+ } finally {
+ query.stop()
+ }
+ val streamFactory =
microBatchScan.readerFactory.asInstanceOf[SupportsArrowStreamRead]
+ assert(readAll(streamFactory, microBatchScan.inputPartitions).sorted ==
(0L until 4L))
Review Comment:
**5. [Nit] The test does not check the `partition_id` or the contexts of the
streams of a micro-batch**
The batch part checks, with `task_path`, that the library got the index of
the partition in the scan, and, with `checkContexts`, that every context is
closed. The micro-batch part only checks the ids that it reads. So a change
that passed another `partition_id` from the factory of a micro-batch scan, or
that left its contexts open, would still pass. Since `native_counter` now
accepts `task_path` too, could this part pass it, and wrap the read in
`checkContexts(0)`?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1177,187 @@ private static boolean isInterrupted(long task) {
}
}
+ /** The number of task contexts that are not closed yet, for tests. */
+ static int openTaskContexts() {
+ return tasks.size();
+ }
+
+ //
---------------------------------------------------------------------------------------------
+ // Streams of engines
+ //
---------------------------------------------------------------------------------------------
+
+ // The offsets of the fields of struct ArrowArrayStream, and its size.
+ private static final long STREAM_GET_SCHEMA = 0;
Review Comment:
**10. [Nit] The offsets of `ArrowArrayStream` are outside the section of the
layouts**
The offsets of every other struct, such as `TASK_*` and `READ_INFO_*`, are
in the block of "The layouts of the structs of the header" near the top of the
class, and `layout()` now lists these offsets with them. `AGENTS.md` asks to
"follow the sectioning the file already uses and put the new member with the
code it belongs with". Could `STREAM_*` move to that block, while the errno
values and `UNKNOWN_STREAM_ERROR`, which only the streams of engines use, stay
here?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,17 +364,31 @@ class NativeDataSourceStreamReader(
}
override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
+ val planned = call(name, "plan a scan of") {
library.streamPartitions(handle.get, start, end)
}
- readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeStreamReader(handle.get))))
- NativeInputPartition.from(partitions)
+ val readerState = bytesOrEmpty(Option(planned.state))
+ NativeInputPartition.from(planned.partitions).map { partition =>
+ NativeMicroBatchPartition(partition.asInstanceOf[NativeInputPartition],
readerState)
Review Comment:
**7. [Nit] `partitions` casts back what `NativeInputPartition.from` built**
`from` builds `NativeInputPartition`s, but returns them as an
`Array[InputPartition]`, so this line casts each one back to the class that it
was created with. If `from` returned an `Array[NativeInputPartition]`, the cast
would go away, and the batch planner could widen the array once, for example
with `.toArray[InputPartition]`, since arrays are invariant.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1143,6 +1353,261 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
.mode("append")
.save()
check(writeTasks)
+
+ // A read gets the index of its partition in the scan, also when a task
reads several
+ // partitions, or the partitions of several scans: the task attempt that
read each partition.
+ def readers(dir: File): Map[Int, Long] = dir.listFiles().map { file =>
+ file.getName.toInt -> Files.readString(file.toPath).split(",")(5).toLong
+ }.toMap
+ val (first, second) = (Utils.createTempDir(), Utils.createTempDir())
+ checkRows(
+ range(session, "end" -> "6", "task_path" -> first.getPath).select("id")
+ .union(range(session, "end" -> "6", "task_path" ->
second.getPath).select("id")),
+ (0L until 6L).flatMap(i => Seq(Row(i), Row(i))))
+ assert(readers(first).keySet == Set(0, 1) && readers(second).keySet ==
Set(0, 1))
+ val coalesced = Utils.createTempDir()
+ checkRows(range(session, "end" -> "6", "task_path" ->
coalesced.getPath).select("id")
+ .coalesce(1), (0L until 6L).map(Row(_)))
+ assert(readers(coalesced).keySet == Set(0, 1) &&
readers(coalesced).values.toSet.size == 1)
+ // So does a read of a micro-batch.
+ val (firstStream, secondStream) = (Utils.createTempDir(),
Utils.createTempDir())
+ def counter(dir: File): DataFrame =
session.readStream.format("native_counter")
+ .option("max_offset", "4").option("task_path", dir.getPath).load()
+ val query =
counter(firstStream).union(counter(secondStream)).writeStream.format("noop").start()
+ try query.processAllAvailable() finally query.stop()
+ assert(readers(firstStream).keySet == Set(0) &&
readers(secondStream).keySet == Set(0))
+ }
+
+ nativeTest("open the partitions of batch and micro-batch scans 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)
+ }
+ // Runs `f` in a task that only gets the files of the session.
+ def inTask[T: ClassTag](f: => T): T =
session.artifactManager.withResources {
+ session.sparkContext.parallelize(Seq(0), 1).map(_ => f).collect().head
+ }
+ import TestArrowEngine.{callbacks, readIds}
+ // Reads the given partitions in a single task, as an engine does.
+ def readAll(factory: SupportsArrowStreamRead, partitions:
Seq[InputPartition]): Seq[Long] = {
+ inTask {
+ partitions.flatMap { partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ try {
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try factory.openArrowStream(partition, stream.memoryAddress())
catch {
+ case e: Throwable =>
+ stream.close()
+ throw e
+ }
+ readIds(stream, allocator)
+ } finally {
+ allocator.close()
+ }
+ }
+ }
+ }
+ // Checks that `f` opens `delta` more task contexts than it closes. The
count is global to the
+ // JVM, where the tasks of earlier tests may still close their contexts,
so it is read right
+ // before `f`.
+ def checkContexts[T](delta: Int)(f: => T): T = {
+ val before = NativeLibrary.openTaskContexts()
+ val result = f
+ assert(NativeLibrary.openTaskContexts() == before + delta)
+ result
+ }
+
+ // An engine reads both partitions of the scan in a single task.
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ assert(readAll(factory, partitions).sorted == (0L until 6L))
+ // The library got the partition of the scan that it reads, not the one of
the task.
+ val tasks = taskDir.listFiles().map { file =>
+ file.getName.toInt ->
Files.readString(file.toPath).split(",").map(_.toLong).toSeq
+ }.toMap
+ assert(tasks.keySet == Set(0, 1) && tasks.values.map(_(5)).toSet.size ==
1, tasks)
+ assert(NativeLibrary.openEngineStreams() == 0)
+
+ // The engine may release the stream after its task ended, and on another
thread: Spark frees
+ // the context of the task once the library released the stream.
+ val released = new File(Utils.createTempDir(), "released")
+ val (laterFactory, laterPartitions) =
+ plannedScan(range(session, "end" -> "6", "released_path" ->
released.getPath).select("id"))
+ // Opens a stream in a task, which it returns unreleased.
+ def openInTask(factory: SupportsArrowStreamRead, partition:
InputPartition): Long = inTask {
+ val stream = Platform.allocateMemory(40)
Review Comment:
**9. [Nit] The test hard-codes the size of `ArrowArrayStream`**
`40` appears four times here and below. d81e619 added `ArrowArrayStream` to
`NativeLibrary.layout()`, which the layout test compares with the header, so
could the test read the size from there once, as it already does for
`SparkDsTaskContext`? For example, `val streamSize =
NativeLibrary.layout().get("ArrowArrayStream").longValue`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,17 +364,31 @@ class NativeDataSourceStreamReader(
}
override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
+ val planned = call(name, "plan a scan of") {
library.streamPartitions(handle.get, start, end)
}
- readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeStreamReader(handle.get))))
- NativeInputPartition.from(partitions)
+ val readerState = bytesOrEmpty(Option(planned.state))
+ NativeInputPartition.from(planned.partitions).map { partition =>
+ NativeMicroBatchPartition(partition.asInstanceOf[NativeInputPartition],
readerState)
+ }
}
override def read(partition: InputPartition): PartitionReader[ColumnarBatch]
= {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, schema)
+ val microBatchPartition = partition.asInstanceOf[NativeMicroBatchPartition]
+ new NativePartitionReader(location, name, NativeLibrary.READ_MICRO_BATCH,
+ microBatchPartition.readerState, microBatchPartition.partition, schema,
timeZoneId)
+ }
+
+ // An engine can also get the Arrow C stream of a partition, see
[[SupportsArrowStreamRead]].
+ override def readerFactory(): PartitionReaderFactory = {
+ new ColumnarStreamPartitionReaderFactory(this) with
SupportsArrowStreamRead {
Review Comment:
**6. [Nit] The reader factory of a micro-batch scan is anonymous, and
repeats `read`**
`openArrowStream` here repeats the unpacking of `NativeMicroBatchPartition`
in `read` above, with the same arguments. The factory also shows up as
`NativeDataSourceStreamReader$$anon$1` in logs and stack traces. A named class
next to `NativePartitionReaderFactory`, for example
`NativeMicroBatchPartitionReaderFactory(location, name, schema, timeZoneId)`,
could implement both `createColumnarReader` and `openArrowStream` from the
partition, and `read` could delegate to it. Or `NativePartitionReaderFactory`
could take the kind of the scan, and read the state from the partition when it
is a `NativeMicroBatchPartition`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1143,6 +1353,261 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
.mode("append")
.save()
check(writeTasks)
+
+ // A read gets the index of its partition in the scan, also when a task
reads several
+ // partitions, or the partitions of several scans: the task attempt that
read each partition.
+ def readers(dir: File): Map[Int, Long] = dir.listFiles().map { file =>
+ file.getName.toInt -> Files.readString(file.toPath).split(",")(5).toLong
+ }.toMap
+ val (first, second) = (Utils.createTempDir(), Utils.createTempDir())
+ checkRows(
+ range(session, "end" -> "6", "task_path" -> first.getPath).select("id")
+ .union(range(session, "end" -> "6", "task_path" ->
second.getPath).select("id")),
+ (0L until 6L).flatMap(i => Seq(Row(i), Row(i))))
+ assert(readers(first).keySet == Set(0, 1) && readers(second).keySet ==
Set(0, 1))
+ val coalesced = Utils.createTempDir()
+ checkRows(range(session, "end" -> "6", "task_path" ->
coalesced.getPath).select("id")
+ .coalesce(1), (0L until 6L).map(Row(_)))
+ assert(readers(coalesced).keySet == Set(0, 1) &&
readers(coalesced).values.toSet.size == 1)
+ // So does a read of a micro-batch.
+ val (firstStream, secondStream) = (Utils.createTempDir(),
Utils.createTempDir())
+ def counter(dir: File): DataFrame =
session.readStream.format("native_counter")
+ .option("max_offset", "4").option("task_path", dir.getPath).load()
+ val query =
counter(firstStream).union(counter(secondStream)).writeStream.format("noop").start()
+ try query.processAllAvailable() finally query.stop()
+ assert(readers(firstStream).keySet == Set(0) &&
readers(secondStream).keySet == Set(0))
+ }
+
+ nativeTest("open the partitions of batch and micro-batch scans 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)
+ }
+ // Runs `f` in a task that only gets the files of the session.
+ def inTask[T: ClassTag](f: => T): T =
session.artifactManager.withResources {
+ session.sparkContext.parallelize(Seq(0), 1).map(_ => f).collect().head
+ }
+ import TestArrowEngine.{callbacks, readIds}
+ // Reads the given partitions in a single task, as an engine does.
+ def readAll(factory: SupportsArrowStreamRead, partitions:
Seq[InputPartition]): Seq[Long] = {
+ inTask {
+ partitions.flatMap { partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ try {
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try factory.openArrowStream(partition, stream.memoryAddress())
catch {
+ case e: Throwable =>
+ stream.close()
+ throw e
+ }
+ readIds(stream, allocator)
+ } finally {
+ allocator.close()
+ }
+ }
+ }
+ }
+ // Checks that `f` opens `delta` more task contexts than it closes. The
count is global to the
+ // JVM, where the tasks of earlier tests may still close their contexts,
so it is read right
+ // before `f`.
+ def checkContexts[T](delta: Int)(f: => T): T = {
+ val before = NativeLibrary.openTaskContexts()
+ val result = f
+ assert(NativeLibrary.openTaskContexts() == before + delta)
+ result
+ }
+
+ // An engine reads both partitions of the scan in a single task.
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ assert(readAll(factory, partitions).sorted == (0L until 6L))
+ // The library got the partition of the scan that it reads, not the one of
the task.
+ val tasks = taskDir.listFiles().map { file =>
+ file.getName.toInt ->
Files.readString(file.toPath).split(",").map(_.toLong).toSeq
+ }.toMap
+ assert(tasks.keySet == Set(0, 1) && tasks.values.map(_(5)).toSet.size ==
1, tasks)
+ assert(NativeLibrary.openEngineStreams() == 0)
+
+ // The engine may release the stream after its task ended, and on another
thread: Spark frees
+ // the context of the task once the library released the stream.
+ val released = new File(Utils.createTempDir(), "released")
+ val (laterFactory, laterPartitions) =
+ plannedScan(range(session, "end" -> "6", "released_path" ->
released.getPath).select("id"))
+ // Opens a stream in a task, which it returns unreleased.
+ def openInTask(factory: SupportsArrowStreamRead, partition:
InputPartition): Long = inTask {
+ val stream = Platform.allocateMemory(40)
+ Platform.setMemory(stream, 0.toByte, 40)
+ factory.openArrowStream(partition, stream)
+ stream
+ }
+ val (address, copy) = checkContexts(1) {
+ val address = openInTask(laterFactory, laterPartitions.head)
+ assert(NativeLibrary.openEngineStreams() == 1 && !released.exists())
+ // A copy of the stream, as an engine that moves the stream may leave
behind.
+ val copy = Platform.allocateMemory(40)
+ Platform.copyMemory(null, address, null, copy, 40)
+ (address, copy)
+ }
+ checkContexts(-1) {
+ assert(readIds(ArrowArrayStream.wrap(address), ArrowUtils.rootAllocator)
== (0L until 3L))
+ assert(NativeLibrary.openEngineStreams() == 0 && released.exists())
+ }
+ Platform.freeMemory(address)
+ // Once the stream is released, the callbacks of the copy fail with
EINVAL, and its release
+ // only marks it released.
+ checkContexts(0) {
+ val array = ArrowArray.allocateNew(ArrowUtils.rootAllocator)
+ val stale = try
intercept[Exception](ArrowArrayStream.wrap(copy).getNext(array)) finally {
+ array.close()
+ }
+ assert(Utils.exceptionString(stale)
+ .contains("The stream was released, or is not a stream of Spark."))
+ ArrowArrayStream.wrap(copy).release()
+ assert(ArrowArrayStream.wrap(copy).snapshot().release == 0)
+ assert(NativeLibrary.openEngineStreams() == 0)
+ }
+ Platform.freeMemory(copy)
+
+ // Once the task ended, its context tells the library that it was
interrupted, so that the
+ // library stops waiting.
+ val marker = new File(Utils.createTempDir(), "interrupted")
+ val (waitingFactory, waitingPartitions) = plannedScan(
+ range(session, "end" -> "6", "wait_for_interrupt" ->
marker.getPath).select("id"))
+ checkContexts(0) {
+ val waiting = openInTask(waitingFactory, waitingPartitions.head)
+ val error = intercept[Exception] {
+ readIds(ArrowArrayStream.wrap(waiting), ArrowUtils.rootAllocator)
+ }
+ assert(Utils.exceptionString(error).contains("the task was killed"))
+ assert(Files.readString(marker.toPath) == "interrupted")
+ assert(NativeLibrary.openEngineStreams() == 0)
+ Platform.freeMemory(waiting)
+ }
+
+ // So does the reader factory of a micro-batch scan, with the state of its
stream reader.
+ val query =
session.readStream.format("native_counter").option("max_offset", "4").load()
+ .writeStream.format("noop").start()
+ val microBatchScan = try {
+ query.processAllAvailable()
+
query.asInstanceOf[StreamingQueryWrapper].streamingQuery.lastExecution.executedPlan
+ .collectFirst { case s: MicroBatchScanExec => s }.get
+ } finally {
+ query.stop()
+ }
+ val streamFactory =
microBatchScan.readerFactory.asInstanceOf[SupportsArrowStreamRead]
+ assert(readAll(streamFactory, microBatchScan.inputPartitions).sorted ==
(0L until 4L))
+ assert(NativeLibrary.openEngineStreams() == 0)
+ // The partitions of a micro-batch keep the state that the stream reader
returned with them:
+ // a reader factory that an engine keeps reads them with it, also once
Spark planned the next
+ // micro-batch, and so does the reader of Spark.
+ val relation =
session.readStream.format("native_counter").option("max_offset", "4").load()
+ .queryExecution.analyzed.collectFirst { case relation:
StreamingRelationV2 => relation }.get
+ val stream =
relation.table.asInstanceOf[SupportsRead].newScanBuilder(relation.extraOptions)
+ .build().toMicroBatchStream(Utils.createTempDir().getPath)
+ try {
+ def offset(i: Int): Offset =
stream.deserializeOffset(s"""{"offset":$i}""")
+ val first = stream.planInputPartitions(offset(0), offset(2)).toSeq
+ val keptFactory = stream.createReaderFactory()
+ val second = stream.planInputPartitions(offset(2), offset(4)).toSeq
+ assert(readAll(keptFactory.asInstanceOf[SupportsArrowStreamRead], first)
== Seq(0L, 1L))
+ assert(readAll(keptFactory.asInstanceOf[SupportsArrowStreamRead],
second) == Seq(2L, 3L))
+ val read = inTask {
+ first.flatMap { partition =>
+ val reader = keptFactory.createColumnarReader(partition)
+ try {
+ val ids = ArrayBuffer.empty[Long]
+ while (reader.next()) {
+ val batch = reader.get()
+ (0 until batch.numRows()).foreach(i => ids +=
batch.column(0).getLong(i))
+ }
+ ids.toSeq
+ } finally {
+ reader.close()
+ }
+ }
+ }
+ assert(read == Seq(0L, 1L))
+ } finally {
+ stream.stop()
+ }
+
+ // It runs in the running task, with a stream whose release callback is
NULL.
+ val outside = intercept[IllegalStateException] {
+ factory.openArrowStream(partitions.head, 0)
+ }
+ assert(outside.getMessage.contains("must be called in the running task"))
+ val completed = intercept[Exception] {
+ session.artifactManager.withResources {
+ session.sparkContext.parallelize(Seq(0), 1).foreach { _ =>
+ TaskContext.get().addTaskCompletionListener[Unit] { _ =>
+ factory.openArrowStream(partitions.head, 0)
+ }
+ }
+ }
+ }
+ assert(Utils.exceptionString(completed).contains("must be called in the
running task"))
+ val noStream =
intercept[Exception](inTask(factory.openArrowStream(partitions.head, 0)))
+ assert(Utils.exceptionString(noStream).contains("whose release callback is
NULL"))
+ // A stream whose release callback is not NULL, such as an open one, is
left as it is.
+ val (rejected, unchanged, ids) = checkContexts(0) {
+ inTask {
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ try {
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ factory.openArrowStream(partitions.head, stream.memoryAddress())
+ val before = callbacks(stream)
+ val rejected = try {
+ factory.openArrowStream(partitions.head, stream.memoryAddress())
+ None
+ } catch {
+ case e: IllegalArgumentException => Some(e.getMessage)
+ }
+ (rejected, callbacks(stream) == before, readIds(stream, allocator))
+ } finally {
+ allocator.close()
+ }
+ }
+ }
+ assert(rejected.exists(_.contains("whose release callback is NULL")),
rejected)
+ assert(unchanged && ids == (0L until 3L))
+
+ // If it fails, the stream is released, including the stream of the
library that Spark rejects.
+ Seq(
+ "open" -> "injected failure in open",
+ "uninitialized" ->
+ "The library did not initialize the stream, or a callback of the
stream is NULL.",
+ "null_get_next" ->
+ "The library did not initialize the stream, or a callback of the
stream is NULL."
+ ).foreach { case (fail, msg) =>
+ val failReleased = new File(Utils.createTempDir(), "released")
+ val (failing, failingPartitions) =
+ plannedScan(range(session, "fail" -> fail, "released_path" ->
failReleased.getPath))
+ val e = checkContexts(0) {
+ intercept[Exception] {
+ inTask {
+ val stream = ArrowArrayStream.allocateNew(ArrowUtils.rootAllocator)
+ try {
+ failing.openArrowStream(failingPartitions.head,
stream.memoryAddress())
+ } finally {
+ val notReleased = stream.snapshot().release != 0
+ stream.close()
+ if (notReleased) throw new IllegalStateException("The stream is
not released.")
Review Comment:
**4. [Nit] The `finally` of the failure cases hides the original failure**
If a change left the stream unreleased, the `IllegalStateException` thrown
here would replace the exception of `openArrowStream`. `findError` would then
rethrow it, since it has no `NATIVE_DATA_SOURCE_ERROR` among its causes, and
the failure would show neither the original error nor which of the three cases
failed. Could the check keep the original exception, for example by catching it
before the `finally` and adding it to the new exception with `addSuppressed`,
and name the case, as in `s"The stream is not released after the failure
$fail."`? Adding it as the cause instead would let `findError` find the
expected error and pass.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -879,7 +1153,7 @@ static final class Task implements AutoCloseable {
}
@Override
- public void close() {
+ public synchronized void close() {
Review Comment:
**8. [Nit] `synchronized` does not guard against a concurrent close**
The description says that closing a `Task` is now synchronized, since the
release of an engine's stream may run on any thread. But a single close on
another thread does not race with anything: `address` is final, and only the
thread that closes reads `closed`. What makes the task of an engine's stream
close once is `engineStreams.remove` in `engineStreamRelease`, which lets only
one release of the stream, or of a copy of it, reach `task.close()`.
`releaseOnFailure` only closes a task that was never wrapped. The keyword is
harmless, but a reader will look for a second closer that does not exist. Could
it be dropped, or could the comment say that the removal from `engineStreams`
is what makes the close happen once?
--
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]
