dongjoon-hyun commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4419050240
##########
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:
**1. [Low] The reader factory of a micro-batch scan uses the reader state of
the last planned micro-batch**
The factory of a batch scan, `NativePartitionReaderFactory`, gets the state
of its scan as a constructor argument, so it is self-contained. This factory is
not: it reads `readerState` of the stream reader (L335), which `partitions`
overwrites on the driver for every micro-batch. So the state that reaches the
executor is the one of the micro-batch that was planned last when the factory
was serialized, not necessarily the one of the partitions that the engine
passes with it.
Spark's own reader has the same structure since #59931, and it is safe
there, because Spark submits the job of a micro-batch before it plans the next
one. But this PR hands the factory to engines as an object of a public API, and
the PR description says that it reads "with the state of the stream reader of
the micro-batch". An engine that keeps the factory and the partitions of
micro-batch N, and serializes them after micro-batch N+1 was planned, would
pass the partitions of N to `read` with the reader state of N+1, and Spark
would not notice.
Could the factory keep the state of its micro-batch? As far as I can see,
`MicroBatchScanExec` creates its factory after it planned its partitions, so
`readerFactory()` could copy `readerState` into the factory, and both `read`
and `openArrowStream` would use that copy. Passing the state with the
partitions would work as well. Otherwise, the Javadoc of
`SupportsArrowStreamRead` could say that the factory of a micro-batch scan is
only valid until the next micro-batch is planned.
##########
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:
**2. [Nit] The docs could say that the callbacks of an engine's stream call
into the JVM**
"can be called from any thread" holds, but since the four callbacks of the
stream are upcalls of Spark (`EngineStream`), each call enters the JVM. A
native thread of the engine, such as a thread of the Velox executor or a tokio
worker, is attached to the JVM by its first call, and each batch costs an
upcall and a downcall. A call once the JVM is exiting, such as a release from a
C++ static destructor or an `atexit` handler that cleans up leaked streams,
cannot enter the JVM and may abort the process.
An engine developer cannot tell this from the guide, nor from the Javadoc of
`SupportsArrowStreamRead` (L34-L35), which says the same. A sentence in both,
such as "The callbacks of the stream call into the JVM, so the engine releases
the stream before the JVM exits", would make it explicit.
##########
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:
**4. [Nit] The test does not cover a stream whose release callback is not
NULL**
Both calls pass the address 0, so they stop at `streamAddress == 0`
(`NativeDataSource.scala` L896). The other half of that check,
`snapshot().release != 0`, is not covered, and neither is a task that already
completed (L892). The Javadoc says that "A stream whose release callback is not
NULL is left as it is" (`SupportsArrowStreamRead.java` L54-L55), but the test
would still pass if that branch released the stream or overwrote its callbacks.
For example, passing a stream that was already opened, before reading it, could
check the `IllegalArgumentException`, that its callbacks and `private_data` did
not change, and that it can still be read to its end.
##########
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()
Review Comment:
**5. [Nit] The checks of the task contexts compare a count that is global to
the JVM**
Thank you for adding the checks of [item
8](https://github.com/apache/spark/pull/59934#discussion_r4418419426). Now that
the count is read once here and compared at L1421, L1424, L1439 and L1493,
though, the test depends on other tests more than before. `openTaskContexts()`
is the size of a static map, which also counts the contexts of the tasks of
earlier tests, and a task that Spark killed when its job failed, or a task of a
streaming query that a test stopped, closes its context after the job returned,
as the comment of `checkDataWritersClosed` (L152-L153) says for data writers.
If such a close happens while this test runs, the next comparison fails. The
test that runs just before only has jobs that succeed, so this is unlikely
today, but it depends on the order of the tests.
Reading the count again right before each operation, or waiting until it
stops changing before reading it here, would make these checks independent of
the earlier tests.
##########
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:
**3. [Nit] `layout()` could list the offsets of `ArrowArrayStream`**
The Javadoc of the class (L50-L52) says that Spark allocates the memory that
it passes to the library at the offsets of the structs of the header, which
`layout()` lists for the tests, and "the bindings match the layout of
spark_datasource.h" compares that list with `offsetof` in C++. Spark now reads
and writes these fields of `struct ArrowArrayStream` too, which the header
defines (`spark_datasource.h` L251), but they are not in the list, so the test
would not catch a wrong offset here. Adding them to `layout()` would let the
existing test check them:
```java
layout.put("ArrowArrayStream", STREAM_SIZE);
layout.put("ArrowArrayStream.get_schema", STREAM_GET_SCHEMA);
layout.put("ArrowArrayStream.get_next", STREAM_GET_NEXT);
layout.put("ArrowArrayStream.get_last_error", STREAM_GET_LAST_ERROR);
layout.put("ArrowArrayStream.release", STREAM_RELEASE);
layout.put("ArrowArrayStream.private_data", STREAM_PRIVATE_DATA);
```
##########
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") {
Review Comment:
**6. [Nit] The name of the test only mentions batch scans**
Since 8575505, this test also reads the partitions of a micro-batch scan
(L1442-L1454), and "How was this patch tested?" in the PR description quotes
this name. A name such as "open the partitions of batch and micro-batch scans
as Arrow C streams" would say what it covers.
--
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]
