HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4423390179
##########
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:
Done in 22cf88d: `partitions` now broadcasts a non-empty state once for each
planned micro-batch, on the driver, and the partitions keep the `Broadcast`
(`NativeMicroBatchPartition.state`). Each executor then fetches the state once
for the micro-batch, as it fetched the reader with its state in the task binary
before d81e619, and each task only carries the ID of the broadcast. A
micro-batch without a state, or without partitions, gets no broadcast. The
ContextCleaner removes the broadcast once the driver no longer refers to the
partitions, neither in the plan of the micro-batch nor in an engine that keeps
them, so a factory and partitions that an engine keeps still read with the
state of their micro-batch. The test checks that a serialized partition does
not contain the state, and that a micro-batch without a state is not broadcast.
##########
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:
Done in 22cf88d: `openInTask` now reads the stream to its end in the task,
through the callbacks of the stream and without moving it
(`TestArrowEngine.readIdsInPlace`), and returns it unreleased. The driver then
only releases it, with `ArrowArrayStream.wrap(address).release()`, and checks
the count of the contexts and `released` as before. The part with
`wait_for_interrupt` still reads on the driver after the task ended, since it
checks that `is_interrupted` is true then.
##########
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:
Done in 22cf88d: both factories first call
`NativePartitionReader.engineStreamPartition`, which checks the task and the
stream as `openEngineStream` did, and only then the class of the partition,
which it rejects with an `IllegalArgumentException` that names the expected
class, for example "openArrowStream needs a partition of the scan of the reader
factory, a NativeMicroBatchPartition, but got
org.apache.spark.sql.execution.datasources.v2.ffi.NativeInputPartition." The
test passes a partition of each kind to the factory of the other kind, outside
of a task, where it fails because it does not run in the task, and in a task.
##########
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:
Done in 22cf88d: the failure cases now catch the failure, and if the stream
is not released, throw `s"The stream is not released after the failure $fail."`
with the failure added as suppressed rather than as the cause, so that
`findError` does not find the expected error then. The stream is closed in the
`finally`, also when `openArrowStream` does not fail.
##########
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:
Done in 22cf88d: the micro-batch part passes `task_path` to
`native_counter`, removes the files of the tasks of the query once it stopped,
and checks that the library got the partition 0 of the scan, the only one of
the micro-batch. `inTask` now runs in the task of the partition 1 of two, so
the index of the partition in the scan is not the partition of the task here
either. The reads of both micro-batch parts are wrapped in `checkContexts(0)`.
--
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]
