HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4417467206
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -839,10 +867,67 @@ class NativePartitionReader(
}
}
-/** Writes the batches of a task to the data writer of a native library. */
+object NativePartitionReader {
+ import NativeCalls._
+
+ /**
+ * Initializes the Arrow C stream at `streamAddress`, whose release callback
is NULL, with the
+ * data of a partition, by calling read of the library in the running task.
Returns the context
+ * of the task, which the library may use until the stream is released: it
is closed then. If
+ * this fails, the stream is released.
+ */
+ def openStream(
+ location: NativeLibraryLocation,
+ name: String,
+ kind: Int,
+ readerState: Array[Byte],
+ partition: Array[Byte],
+ schema: StructType,
+ timeZoneId: String,
+ streamAddress: Long): NativeLibrary.Task = {
+ val library = NativeLibraries.get(location)
+ val stream = ArrowArrayStream.wrap(streamAddress)
+ val task = NativeTask.create()
+ try {
+ NativeArrow.exportSchema(schema, timeZoneId) { schemaAddress =>
+ call(name, "read") {
+ library.read(task, kind, schemaAddress, readerState, partition,
streamAddress)
+ }
+ }
+ // Arrow calls the callbacks of the stream without checking them, and
calling NULL would
+ // crash the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 ||
snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream,
or a callback of " +
+ "the stream is NULL.")
+ }
+ task
+ } catch {
+ case e: Throwable =>
+ // Releases the stream, if the library initialized it, before the task
that it may use,
+ // and marks it released in case its release callback did not.
+ Utils.tryWithSafeFinally {
+ stream.release()
+ stream.markReleased()
+ }(task.close())
Review Comment:
Done in 8eae514: a failure of the cleanup is added to the error of `read`
with `addSuppressed`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1145,6 +1219,76 @@ 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)
+ }
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ // An engine opens the stream of each partition in its task, and reads it
from another thread.
+ val ids = session.sparkContext.parallelize(partitions,
partitions.length).mapPartitions {
+ partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try {
+ val handle = factory.openArrowStream(partition.next(),
stream.memoryAddress())
+ try {
+ ThreadUtils.runInNewThread("engine") {
+ // Moves the stream, which closing the reader releases.
+ val reader = Data.importArrayStream(allocator, stream)
+ try {
+ val ids = ArrayBuffer.empty[Long]
+ while (reader.loadNextBatch()) {
+ val root = reader.getVectorSchemaRoot
+ val vector = root.getVector("id").asInstanceOf[BigIntVector]
+ (0 until root.getRowCount).foreach(i => ids += vector.get(i))
+ }
+ ids.toSeq
+ } finally {
+ reader.close()
+ }
+ }.iterator
+ } finally {
+ handle.close()
Review Comment:
8eae514 removes the handle. The test moves the stream into an Arrow Java
reader right after `openArrowStream`, and closing the reader releases it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -278,37 +284,54 @@ class NativeBatchPlanner(
override def plan(): ColumnarPlannedBatch = {
try {
- val partitions = call(name, "plan a scan
of")(library.partitions(handle.get))
- val readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeReader(handle.get))))
+ val planned = call(name, "plan a scan
of")(library.partitions(handle.get))
ColumnarPlannedBatch(
- NativeInputPartition.from(partitions),
- new NativePartitionReaderFactory(executorLocation, name, readerState,
readSchema))
+ NativeInputPartition.from(planned.partitions),
+ new NativePartitionReaderFactory(executorLocation, name,
+ bytesOrEmpty(Option(planned.state)), readSchema, timeZoneId))
} finally {
- // The executors only need the serialized state.
+ // The executors only need the state that the reader returned.
handle.close()
}
}
}
-/** Reads the partitions of a planned batch scan of a native data source, on
the executors. */
+/**
+ * Reads the partitions of a planned batch scan of a native data source, on
the executors. An
+ * engine can also get the Arrow C stream of a partition from it, see
[[SupportsArrowStreamRead]].
+ */
class NativePartitionReaderFactory(
location: NativeLibraryLocation,
name: String,
readerState: Array[Byte],
- readSchema: StructType)
- extends ColumnarOnlyPartitionReaderFactory {
+ readSchema: StructType,
+ timeZoneId: String)
+ extends ColumnarOnlyPartitionReaderFactory with SupportsArrowStreamRead {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, readSchema)
+ new NativePartitionReader(location, name, NativeLibrary.READ_BATCH,
readerState,
+ partition.asInstanceOf[NativeInputPartition].bytes, readSchema,
timeZoneId)
+ }
+
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): AutoCloseable = {
+ val task = NativePartitionReader.openStream(location, name,
NativeLibrary.READ_BATCH,
+ readerState, partition.asInstanceOf[NativeInputPartition].bytes,
readSchema, timeZoneId,
+ streamAddress)
+ // In case the caller does not close it, which must release the stream
before then.
+ Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ =>
task.close()))
Review Comment:
Gone with the listener in 8eae514.
--
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]
