HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4417793209
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -278,37 +284,69 @@ 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): Unit = {
+ val context = TaskContext.get()
+ if (context == null || context.isCompleted()) {
+ throw new IllegalStateException(
+ "openArrowStream must be called in the running task that reads the
partition.")
+ }
+ if (streamAddress == 0 ||
ArrowArrayStream.wrap(streamAddress).snapshot().release != 0) {
+ throw new IllegalArgumentException(
+ "openArrowStream needs the address of a stream whose release callback
is NULL.")
+ }
+ val nativePartition = partition.asInstanceOf[NativeInputPartition]
+ // The partition of the scan: the task of an engine may read others than
its own.
+ val task = NativePartitionReader.openStream(location, name,
NativeLibrary.READ_BATCH,
+ readerState, nativePartition.bytes, readSchema, timeZoneId,
streamAddress,
+ Some(nativePartition.index))
Review Comment:
Done in 173c264: the context of an engine's stream reports `is_interrupted`
as `isInterrupted() || isFailed() || isCompleted()` of its task. Spark's own
reader keeps reporting only a kill. The test opens a `wait_for_interrupt`
stream in a task and reads it after the task completed: `get_next` sees the
interruption right away.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1053,163 @@ private static boolean isInterrupted(long task) {
}
}
+ //
---------------------------------------------------------------------------------------------
+ // 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;
+ // The errno values that the callbacks of a stream return.
+ private static final int EINVAL = 22;
+ private static final int EIO = 5;
+
+ private static final AtomicLong nextEngineStreamId = new AtomicLong(1);
+ // The streams wrapped for engines that are not released yet, by their
private_data.
+ private static final ConcurrentHashMap<Long, EngineStream> engineStreams =
+ new ConcurrentHashMap<>();
+ // The callbacks of the streams of the libraries, by address, for each
callback of a stream: a
+ // library typically has one of each, for all its streams.
+ private static final ConcurrentHashMap<Long, MethodHandle>
getSchemaCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> getNextCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle>
getLastErrorCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> releaseCallbacks =
+ new ConcurrentHashMap<>();
+
+ /**
+ * The stream of a read that an engine got with {@code
SupportsArrowStreamRead}, which ties the
+ * task context of the read to the stream: the context is closed once the
library released the
+ * stream, whenever and on whichever thread the engine releases it, also
after the task ended.
+ * The stream of the library is moved to memory of Spark, and the stream of
the engine gets
+ * callbacks of Spark, which call those of the library, and a private_data
that identifies the
+ * wrapped stream, so that the engine can move the stream too.
+ */
+ private static final class EngineStream {
+ final long inner;
+ final Task task;
+ final MethodHandle getSchema;
+ final MethodHandle getNext;
+ final MethodHandle getLastError;
+ final MethodHandle release;
+
+ EngineStream(long inner, Task task) throws Throwable {
+ this.inner = inner;
+ this.task = task;
+ getSchema = callback(getSchemaCallbacks, inner + STREAM_GET_SCHEMA,
C.I32, C.PTR, C.PTR);
+ getNext = callback(getNextCallbacks, inner + STREAM_GET_NEXT, C.I32,
C.PTR, C.PTR);
+ getLastError = callback(getLastErrorCallbacks, inner +
STREAM_GET_LAST_ERROR, C.PTR, C.PTR);
+ release = callback(releaseCallbacks, inner + STREAM_RELEASE, C.VOID,
C.PTR);
+ }
+
+ private static MethodHandle callback(
+ ConcurrentHashMap<Long, MethodHandle> callbacks,
+ long field,
+ C result,
+ C... parameters) throws Throwable {
+ long address = Platform.getLong(null, field);
+ MethodHandle handle = callbacks.get(address);
+ if (handle == null) {
+ handle = Ffm.downcall(Ffm.OF_ADDRESS.invoke(address), result,
parameters);
+ callbacks.putIfAbsent(address, handle);
+ }
Review Comment:
Done in 173c264: `Ffm` has three downcall handles that take the address of
the function first (`Linker.downcallHandle(FunctionDescriptor,
Linker.Option...)`), and each upcall reads the current pointer from `inner`.
The four maps, `callback` and the handle fields are gone.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1145,6 +1222,105 @@ 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"))
+ val address = inTask {
+ val stream = Platform.allocateMemory(40)
+ Platform.setMemory(stream, 0.toByte, 40)
+ laterFactory.openArrowStream(laterPartitions.head, stream)
+ stream
+ }
+ assert(NativeLibrary.openEngineStreams() == 1 && !released.exists())
+ assert(readIds(ArrowArrayStream.wrap(address), ArrowUtils.rootAllocator)
== (0L until 3L))
+ assert(NativeLibrary.openEngineStreams() == 0 && released.exists())
Review Comment:
Done in 173c264: `NativeLibrary.openTaskContexts()` counts the contexts in
`tasks`. The test checks that the count is one higher after the task ended with
the stream unreleased, and back after `readIds` released it.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -304,9 +308,36 @@ typedef struct SparkDsCreateInfo {
int64_t num_options;
} SparkDsCreateInfo;
+/* A micro-batch streaming scan, as create_stream_reader gets it. */
+typedef struct SparkDsStreamReaderInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /*
+ * The location of a directory in the checkpoint of the streaming query that
is reserved for this
+ * source, such as file:/tmp/checkpoint/sources/0, where the library may
keep what it needs to
+ * recover from a failure. It stays the same when the query restarts from
its checkpoint. It is a
+ * URI, possibly of a file system that the library cannot access, such as
s3a or hdfs.
+ */
+ const char* checkpoint_location;
+} SparkDsStreamReaderInfo;
+
+/* A request for the latest offset of a stream reader, as latest_offset gets
it. */
+typedef struct SparkDsLatestOffsetInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /*
+ * The offset, as JSON, that the next micro-batch starts after: the end of
the last micro-batch,
+ * or the initial offset before the first one. With Trigger.AvailableNow,
Spark asks for the
+ * latest offset only once, when the query starts, with the initial offset.
+ */
+ const char* start;
+} SparkDsLatestOffsetInfo;
+
/*
* The task that calls read or create_data_writer, on an executor. It stays
valid until the stream
- * that read returns is released, or until the data writer is closed.
+ * that read returns is released, or until the data writer is closed. An
engine that runs a query
+ * in native code may get the stream of read from Spark, and release it once
the task ended, from
+ * another thread: the task stays valid until then too.
Review Comment:
Done in 173c264: the comment of `SparkDsTaskContext` says that an engine may
keep the arrays of the stream after it released the stream, so that the release
of an array must not use the stream or the task.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1053,163 @@ private static boolean isInterrupted(long task) {
}
}
+ //
---------------------------------------------------------------------------------------------
+ // 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;
+ // The errno values that the callbacks of a stream return.
+ private static final int EINVAL = 22;
+ private static final int EIO = 5;
+
+ private static final AtomicLong nextEngineStreamId = new AtomicLong(1);
+ // The streams wrapped for engines that are not released yet, by their
private_data.
+ private static final ConcurrentHashMap<Long, EngineStream> engineStreams =
+ new ConcurrentHashMap<>();
+ // The callbacks of the streams of the libraries, by address, for each
callback of a stream: a
+ // library typically has one of each, for all its streams.
+ private static final ConcurrentHashMap<Long, MethodHandle>
getSchemaCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> getNextCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle>
getLastErrorCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> releaseCallbacks =
+ new ConcurrentHashMap<>();
+
+ /**
+ * The stream of a read that an engine got with {@code
SupportsArrowStreamRead}, which ties the
+ * task context of the read to the stream: the context is closed once the
library released the
+ * stream, whenever and on whichever thread the engine releases it, also
after the task ended.
+ * The stream of the library is moved to memory of Spark, and the stream of
the engine gets
+ * callbacks of Spark, which call those of the library, and a private_data
that identifies the
+ * wrapped stream, so that the engine can move the stream too.
+ */
+ private static final class EngineStream {
+ final long inner;
+ final Task task;
Review Comment:
Done in 173c264: the supplier holds the `TaskContext` through a
`WeakReference`, and reports `true` once it is gone.
--
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]
