HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4419222580
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,17 +368,26 @@ 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)
+ readerState = bytesOrEmpty(Option(planned.state))
+ NativeInputPartition.from(planned.partitions)
}
override def read(partition: InputPartition): PartitionReader[ColumnarBatch]
= {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, schema)
+ new NativePartitionReader(location, name, NativeLibrary.READ_MICRO_BATCH,
readerState,
+ partition.asInstanceOf[NativeInputPartition], 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 = {
+ NativePartitionReader.openEngineStream(location, name,
NativeLibrary.READ_MICRO_BATCH,
+ readerState, partition, schema, timeZoneId, streamAddress)
+ }
+ }
Review Comment:
Done in d81e619: the order of `planInputPartitions` and
`createReaderFactory` is not guaranteed (`MicroBatchStream` does not define it,
and an engine may get `MicroBatchScanExec.readerFactory` before
`inputPartitions`), so the state now goes with the partitions:
`stream_partitions` returns `NativeMicroBatchPartition`s that hold the state of
their micro-batch, and both `read` and `openArrowStream` use it. The test
library now puts the range of the micro-batch in its state, and `read` fails on
a partition of another micro-batch. The test plans two micro-batches and reads
the partitions of both with the factory of the first, through `openArrowStream`
and Spark's reader; it fails with "read got the partition 0,2 with the state of
the micro-batch 2,4" without the change.
##########
docs/sql-native-data-sources.md:
##########
@@ -226,6 +243,21 @@ their memory, and at least 384 MiB, by default. Increase
it for libraries that u
bound the memory that a library uses, for example with the size of its
batches: on YARN or
Kubernetes, an executor that uses more memory than it requested is killed.
+### Native Execution Engines
+
+An engine that runs queries in native code within Spark, such as Apache Gluten
or Apache
+DataFusion Comet, can read a native data source without Spark importing its
data on the JVM. The
+reader factory of a batch or micro-batch scan of a native data source
implements
+`org.apache.spark.sql.connector.read.SupportsArrowStreamRead`: in the task
that reads a partition,
+`openArrowStream` initializes an Arrow C stream that the engine allocated with
the stream that
+`read` returns. The engine reads the stream in the task that opened it, and
owns and releases the
+stream, which it may do after the task ended: Spark frees the
`SparkDsTaskContext` of the read
+once the library released its stream. Once the task ended, `is_interrupted`
returns true, so the
+library may stop reading then. The callbacks
+of the stream can be called from any thread, but never concurrently.
Review Comment:
Done in d81e619: the Javadoc of `SupportsArrowStreamRead` and the guide now
say that the callbacks of the stream call into the JVM, so the engine releases
the stream before the JVM exits.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1170,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;
+ private static final long STREAM_GET_NEXT = 8;
+ private static final long STREAM_GET_LAST_ERROR = 16;
+ private static final long STREAM_RELEASE = 24;
+ private static final long STREAM_PRIVATE_DATA = 32;
+ private static final long STREAM_SIZE = 40;
Review Comment:
Done in d81e619: `layout()` now lists the size and the five fields of
`ArrowArrayStream`, which the layout test compares with the header.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1143,6 +1332,166 @@ 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 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
+ // 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()
+ }
+ }
+ }
+ }
+ val openContexts = NativeLibrary.openTaskContexts()
+
+ // 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 = 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)
+ assert(NativeLibrary.openTaskContexts() == openContexts)
+ 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)
+
+ // It runs in a 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 noStream =
intercept[Exception](inTask(factory.openArrowStream(partitions.head, 0)))
+ assert(Utils.exceptionString(noStream).contains("whose release callback is
NULL"))
Review Comment:
Done in d81e619: the test now passes an open stream again, and checks the
`IllegalArgumentException`, that its callbacks and `private_data` did not
change, and that it is still read to its end. It also calls `openArrowStream`
in the completion listener of a task, and, after a stream is released, checks
that a copy of it gets the error of an unknown stream from `get_next`, and that
releasing the copy closes nothing and marks it released.
--
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]
