HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4418955888
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1145,6 +1276,155 @@ 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)
+ }
+ // 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.readIds
+
+ // 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)
+ val ids = 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()
+ }
+ }
+ }
+ assert(ids.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 openContexts = NativeLibrary.openTaskContexts()
+ val address = openInTask(laterFactory, laterPartitions.head)
+ assert(NativeLibrary.openEngineStreams() == 1 && !released.exists())
+ assert(NativeLibrary.openTaskContexts() == openContexts + 1)
+ assert(readIds(ArrowArrayStream.wrap(address), ArrowUtils.rootAllocator)
== (0L until 3L))
+ assert(NativeLibrary.openEngineStreams() == 0 && released.exists())
+ assert(NativeLibrary.openTaskContexts() == openContexts)
+ Platform.freeMemory(address)
+
+ // 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"))
+ 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]
+ val streamPartitions = microBatchScan.inputPartitions
+ val streamIds = inTask {
+ streamPartitions.flatMap { partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ try {
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try streamFactory.openArrowStream(partition, stream.memoryAddress())
catch {
+ case e: Throwable =>
+ stream.close()
+ throw e
+ }
+ readIds(stream, allocator)
+ } finally {
+ allocator.close()
+ }
+ }
+ }
Review Comment:
Done in b98fd19: a `readAll(factory, partitions)` helper reads the
partitions of the batch and micro-batch scans.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -926,18 +1063,28 @@ private[ffi] object NativeTask {
* @param partitionId the partition that the task reads or writes, by
default the partition of
* the task
* @param taskAttemptId the ID of the task attempt, by default the one of
the task
+ * @param mayOutliveTask whether the library may use the context after the
task ended, as for
+ * the stream of an engine: `is_interrupted` then also
reports the end of
+ * the task, and the context does not keep the task
alive
*/
def create(
partitionId: Option[Int] = None,
- taskAttemptId: Option[Long] = None): NativeLibrary.Task = {
+ taskAttemptId: Option[Long] = None,
+ mayOutliveTask: Boolean = false): NativeLibrary.Task = {
val context = Option(TaskContext.get())
+ val interrupted: () => Boolean = if (mayOutliveTask) {
+ val task = new WeakReference(context.orNull)
+ () => Option(task.get()).forall(t => t.isInterrupted() || t.isFailed()
|| t.isCompleted())
+ } else {
+ () => context.exists(_.isInterrupted())
+ }
new NativeLibrary.Task(
context.map(_.stageId()).getOrElse(-1),
context.map(_.stageAttemptNumber()).getOrElse(-1),
partitionId.orElse(context.map(_.partitionId())).getOrElse(-1),
context.map(_.attemptNumber()).getOrElse(0),
taskAttemptId.orElse(context.map(_.taskAttemptId())).getOrElse(-1L),
- () => context.exists(_.isInterrupted()))
+ () => interrupted())
Review Comment:
Done in b98fd19: `interrupted` is a `BooleanSupplier`, passed to
`NativeLibrary.Task` as it is.
--
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]
