HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4417465365
##########
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 {
Review Comment:
Done in 8eae514: the jobs of the test run in
`session.artifactManager.withResources`.
##########
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()
+ }
+ } finally {
+ stream.close()
+ allocator.close()
+ }
+ }.collect()
+ assert(ids.sorted.toSeq == (0L until 6L))
+ // The library got the context of each task.
+ assert(taskDir.list().toSet == Set("0", "1"))
+
+ // If it fails, the stream is released.
+ Seq(
+ "open" -> "injected failure in open",
+ "uninitialized" ->
+ "The library did not initialize the stream, or a callback of the
stream is NULL."
+ ).foreach { case (fail, msg) =>
+ val (failing, failingPartitions) = plannedScan(range(session, "fail" ->
fail))
+ val stream = ArrowArrayStream.allocateNew(ArrowUtils.rootAllocator)
+ try {
+ val e = intercept[Exception] {
+ failing.openArrowStream(failingPartitions.head,
stream.memoryAddress())
+ }
+ checkError(
+ exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"),
+ condition = "NATIVE_DATA_SOURCE_ERROR",
+ parameters = Map("action" -> "read", "name" -> "native_range", "msg"
-> msg))
+ assert(stream.snapshot().release == 0)
Review Comment:
Done in 8eae514: a `null_get_next` failure initializes the stream with a
NULL `get_next`, and the `released_path` option of the test library checks that
Spark calls its `release`. The listener and its race are gone.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsArrowStreamRead.java:
##########
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector.read;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * A mix-in interface for a {@link PartitionReaderFactory} that also reads a
partition as an
+ * <a href="https://arrow.apache.org/docs/format/CStreamInterface.html";>Arrow
C stream</a>, for an
+ * engine that processes Arrow data in native code: the engine imports the
arrays of the stream
+ * itself, instead of reading the {@link
org.apache.spark.sql.vectorized.ColumnarBatch}es of
+ * {@link #createColumnarReader(InputPartition)}. The reader factories of the
batch scans of
+ * native data sources implement it, with the stream that the library returns.
+ * <p>
+ * Each array of the stream is a struct array whose fields are the columns of
the read schema of
+ * the scan. Spark does not check the stream, while it checks the streams that
it reads itself, so
+ * the engine checks what its import of the arrays requires, such as the
schema of the stream
+ * against the read schema, and the offsets and the lengths of the arrays. The
callbacks of the
+ * stream can be called from any thread, but never concurrently.
+ *
+ * @since 4.4.0
+ */
+@Evolving
+public interface SupportsArrowStreamRead extends PartitionReaderFactory {
+
+ /**
+ * Initializes an Arrow C stream with the data of a partition. It is called
in the task that
+ * reads the partition, whose context the stream may use until it is
released.
+ *
+ * @param partition a partition of the scan, as the scan planned it
+ * @param streamAddress the address of a {@code struct ArrowArrayStream}
that the caller
+ * allocated, and whose release callback is NULL, as it
is in a
+ * zero-initialized struct. If this method fails, the
stream is released.
+ * @return a handle that the caller closes once the stream is released,
which frees what the
+ * stream uses until then. Spark also closes it when the task
completes, so the stream
+ * must be released before then.
Review Comment:
Done in 8eae514: it fails with an `IllegalStateException` without a
`TaskContext` or once the task completed.
##########
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)
Review Comment:
Done in 8eae514: an address of 0, or a stream whose `release` is not NULL,
fails with an `IllegalArgumentException` before `read`.
--
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]
