HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4461786162
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1186,180 @@ 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 errno values that the callbacks of a stream return.
+ private static final int EINVAL = 22;
+ private static final int EIO = 5;
+ // The message of get_last_error for a stream that Spark does not know,
which lives as long as
+ // the JVM.
+ private static final long UNKNOWN_STREAM_ERROR =
+ toCString("The stream was released, or is not a stream of Spark.");
+
+ 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 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 the current callbacks of the library's
stream, 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;
+ // The name of the data source, for the errors.
+ final String name;
+ // The message of the last error of Spark, rather than of the library, or
0 if the last call
+ // returned an error of the library, or none. The callbacks never run
concurrently.
+ long sparkError;
+
+ EngineStream(long inner, Task task, String name) {
+ this.inner = inner;
+ this.task = task;
+ this.name = name;
+ }
+
+ long callback(long offset) {
+ return Platform.getLong(null, inner + offset);
+ }
+
+ /** Records an error of Spark, which get_last_error then returns, and logs
it. */
+ int fail(String message, Throwable e) {
+ try {
+ clearError();
+ sparkError = toCString(message + ": " + rootCause(e));
+ LOGGER.error(message, e);
+ } catch (Throwable ignored) {
+ // An upcall must not throw.
+ }
+ return EIO;
+ }
+
+ void clearError() {
+ if (sparkError != 0) {
+ Platform.freeMemory(sparkError);
+ sparkError = 0;
+ }
+ }
+ }
+
+ /**
+ * Wraps the stream at the given address, which read initialized with
non-NULL callbacks, so that
+ * releasing it closes the given task context (see {@link EngineStream}). If
it fails, the stream
+ * is left as it is.
+ *
+ * @param name the name of the data source, for the errors of the stream
+ */
+ static void wrapEngineStream(long stream, Task task, String name) {
+ long inner = Platform.allocateMemory(STREAM_SIZE);
+ long id = nextEngineStreamId.getAndIncrement();
+ try {
+ Platform.copyMemory(null, stream, null, inner, STREAM_SIZE);
+ engineStreams.put(id, new EngineStream(inner, task, name));
+ } catch (Throwable e) {
+ engineStreams.remove(id);
+ Platform.freeMemory(inner);
+ throw e;
+ }
+ Platform.putLong(null, stream + STREAM_GET_SCHEMA,
Ffm.ENGINE_STREAM_GET_SCHEMA);
+ Platform.putLong(null, stream + STREAM_GET_NEXT,
Ffm.ENGINE_STREAM_GET_NEXT);
+ Platform.putLong(null, stream + STREAM_GET_LAST_ERROR,
Ffm.ENGINE_STREAM_GET_LAST_ERROR);
+ Platform.putLong(null, stream + STREAM_RELEASE, Ffm.ENGINE_STREAM_RELEASE);
+ Platform.putLong(null, stream + STREAM_PRIVATE_DATA, id);
+ }
+
+ /** The number of streams wrapped for engines that are not released yet, for
tests. */
+ static int openEngineStreams() {
+ return engineStreams.size();
+ }
+
+ private static EngineStream engineStream(long stream) {
+ return engineStreams.get(Platform.getLong(null, stream +
STREAM_PRIVATE_DATA));
+ }
+
+ // The callbacks of the streams wrapped for engines. They must not throw: an
exception thrown by
+ // an upcall crashes the JVM.
+
+ private static int engineStreamGetSchema(long stream, long out) {
+ return engineStreamCall(stream, out, STREAM_GET_SCHEMA, "get_schema");
+ }
+
+ private static int engineStreamGetNext(long stream, long out) {
+ return engineStreamCall(stream, out, STREAM_GET_NEXT, "get_next");
+ }
+
+ private static int engineStreamCall(long stream, long out, long offset,
String callback) {
+ EngineStream wrapped = null;
+ try {
+ wrapped = engineStream(stream);
+ if (wrapped == null) {
+ return EINVAL;
+ }
+ wrapped.clearError();
+ return (int) Ffm.STREAM_INT_CALLBACK.invokeExact(
+ wrapped.callback(offset), wrapped.inner, out);
+ } catch (Throwable e) {
+ return wrapped == null ? EIO : wrapped.fail(
Review Comment:
Done in 77bf7bd, as suggested: `fail` takes the name of the callback, and
builds the message in its `try`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -970,21 +1201,31 @@ class NativeDataWriter(
/** Creates the `SparkDsTaskContext` of the running task. */
private[ffi] object NativeTask {
/**
- * @param partitionId the partition that the task reads or writes, by
default the partition of
- * the task
+ * @param partitionId the `partition_id` of the context: the index of the
partition in the scan
+ * for a read, and by default the partition of the task
* @param taskAttemptId the ID of the task attempt, by default the one of
the task
+ * @param mayOutliveTask whether the library may use the context after the
task ended, as for
+ * the stream of an engine: `is_interrupted` then also
reports the end of
+ * the task, and the context does not keep the task
alive
*/
def create(
partitionId: Option[Int] = None,
- taskAttemptId: Option[Long] = None): NativeLibrary.Task = {
+ taskAttemptId: Option[Long] = None,
+ mayOutliveTask: Boolean = false): NativeLibrary.Task = {
val context = Option(TaskContext.get())
+ val interrupted: BooleanSupplier = if (mayOutliveTask) {
Review Comment:
Done in 77bf7bd: every read now gets the context of a read, also the reads
of Spark itself, so a `get_next` that waits on a thread that the task does not
wait for learns that the task ended. Data writers keep reporting only a kill,
and a read outside of a task, as in tests, is never interrupted. The header
describes `is_interrupted` that way for every read.
"release a stream only once get_next returns" now checks with
`NativeLibrary.openTaskContexts()` that the context stays open while `get_next`
waits after the task completed, and is closed once the stream is released. It
also checks that the read then gets `true` from `is_interrupted`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -886,10 +987,140 @@ class NativePartitionReader(
}
}
-/** Writes the batches of a task to the data writer of a native library. */
+object NativePartitionReader {
+ import NativeCalls._
+
+ /**
+ * Checks the arguments of `SupportsArrowStreamRead.openArrowStream` before
it looks at the
+ * partition, and returns the partition as a `P`, the class of the
partitions of the scan of the
+ * reader factory.
+ */
+ def engineStreamPartition[P <: InputPartition](
+ partition: InputPartition,
+ streamAddress: Long)(implicit tag: ClassTag[P]): P = {
+ val context = TaskContext.get()
+ if (context == null || context.isCompleted()) {
Review Comment:
Done in 77bf7bd: `openArrowStream` now rejects `context.isInterrupted() ||
context.isFailed() || context.isCompleted()`. A new case next to the one with
the completion listener calls it from a failure listener of a task that fails,
and checks the error.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1143,6 +1426,336 @@ 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 batch and micro-batch scans 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. It is the
task of the partition
+ // 1, so that the partition of the scan that a read gets is not the
partition of its task.
+ def inTask[T: ClassTag](f: => T): T =
session.artifactManager.withResources {
+ session.sparkContext.parallelize(Seq(0, 1), 2).flatMap(i => if (i == 1)
Seq(f) else Nil)
+ .collect().head
+ }
+ import TestArrowEngine.{callbacks, readIds, readIdsInPlace}
+ // 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()
+ }
+ }
+ }
+ }
+ // Checks that `f` opens `delta` more task contexts than it closes. The
count is global to the
+ // JVM, where the tasks of earlier tests may still close their contexts,
so it is read right
+ // before `f`.
+ def checkContexts[T](delta: Int)(f: => T): T = {
+ val before = NativeLibrary.openTaskContexts()
+ val result = f
+ assert(NativeLibrary.openTaskContexts() == before + delta)
+ result
+ }
+ // The size of struct ArrowArrayStream, which the layout test compares
with the header.
+ val streamSize = NativeLibrary.layout().get("ArrowArrayStream").longValue
+
+ // 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, and returns it unreleased, with what `read`
returned from it in
+ // the task.
+ def openInTask[T: ClassTag](factory: SupportsArrowStreamRead, partition:
InputPartition)(
+ read: ArrowArrayStream => T): (Long, T) = inTask {
+ val stream = Platform.allocateMemory(streamSize)
+ Platform.setMemory(stream, 0.toByte, streamSize)
+ factory.openArrowStream(partition, stream)
+ (stream, read(ArrowArrayStream.wrap(stream)))
+ }
+ val (address, copy) = checkContexts(1) {
+ // The engine reads the stream to its end in the task, but releases it
later.
+ val (address, ids) = openInTask(laterFactory,
laterPartitions.head)(readIdsInPlace)
+ assert(ids == (0L until 3L))
+ assert(NativeLibrary.openEngineStreams() == 1 && !released.exists())
+ // A copy of the stream, as an engine that moves the stream may leave
behind.
+ val copy = Platform.allocateMemory(streamSize)
+ Platform.copyMemory(null, address, null, copy, streamSize)
+ (address, copy)
+ }
+ checkContexts(-1) {
+ ArrowArrayStream.wrap(address).release()
+ assert(NativeLibrary.openEngineStreams() == 0 && released.exists())
+ }
+ Platform.freeMemory(address)
+ // Once the stream is released, the callbacks of the copy fail with
EINVAL, and its release
+ // only marks it released.
+ checkContexts(0) {
+ val array = ArrowArray.allocateNew(ArrowUtils.rootAllocator)
+ val stale = try
intercept[Exception](ArrowArrayStream.wrap(copy).getNext(array)) finally {
+ array.close()
+ }
+ assert(Utils.exceptionString(stale)
+ .contains("The stream was released, or is not a stream of Spark."))
+ ArrowArrayStream.wrap(copy).release()
+ assert(ArrowArrayStream.wrap(copy).snapshot().release == 0)
+ assert(NativeLibrary.openEngineStreams() == 0)
+ }
+ Platform.freeMemory(copy)
+
+ // Once the task ended, its context tells the library that it was
interrupted, so that the
Review Comment:
Done in 77bf7bd: the test now opens a `wait_for_interrupt` stream in a task,
and reads it there, on the thread of the task. It cancels the job group as that
test does, then checks the marker file, and that the stream and its context are
released.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -351,17 +377,30 @@ class NativeDataSourceStreamReader(
}
override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
- library.streamPartitions(handle.get, start, end)
+ if (plannedRange != (start, end)) {
+ val planned = call(name, "plan a scan of") {
+ library.streamPartitions(handle.get, start, end)
+ }
+ val partitions = NativeInputPartition.from(planned.partitions)
+ // The executors fetch the state once for the micro-batch, as a
broadcast, rather than with
+ // each of its tasks. The partitions keep the broadcast, which the
ContextCleaner removes once
+ // the driver no longer refers to them: in the plan of the micro-batch,
in an engine, or here
+ // until the next range is planned.
+ val state = Option(planned.state).filter(_.nonEmpty &&
partitions.nonEmpty).map { bytes =>
+ SparkSession.active.sparkContext.broadcast(bytes)
Review Comment:
Done in 77bf7bd: `partitions` unpersists the broadcast of the previous range
when it plans the next one, with `blocking = false`, so the driver keeps its
copy for the partitions and the reader factories that an engine keeps. The test
that reads the partitions of a micro-batch after the next one was planned still
passes.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -351,17 +377,30 @@ class NativeDataSourceStreamReader(
}
override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
- library.streamPartitions(handle.get, start, end)
+ if (plannedRange != (start, end)) {
+ val planned = call(name, "plan a scan of") {
+ library.streamPartitions(handle.get, start, end)
+ }
+ val partitions = NativeInputPartition.from(planned.partitions)
+ // The executors fetch the state once for the micro-batch, as a
broadcast, rather than with
+ // each of its tasks. The partitions keep the broadcast, which the
ContextCleaner removes once
+ // the driver no longer refers to them: in the plan of the micro-batch,
in an engine, or here
+ // until the next range is planned.
+ val state = Option(planned.state).filter(_.nonEmpty &&
partitions.nonEmpty).map { bytes =>
+ SparkSession.active.sparkContext.broadcast(bytes)
+ }
+ plannedPartitions =
partitions.map[InputPartition](NativeMicroBatchPartition(_, state))
+ plannedRange = (start, end)
}
- readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeStreamReader(handle.get))))
- NativeInputPartition.from(partitions)
+ plannedPartitions
}
override def read(partition: InputPartition): PartitionReader[ColumnarBatch]
= {
Review Comment:
Done in 77bf7bd: `ColumnarStreamReader` now declares an abstract
`readerFactory()` instead of `read`, and `TestRangeStreamReader` returns a
`TestRangeReaderFactory`. The default factory,
`ColumnarStreamPartitionReaderFactory`, the `Serializable` of the trait and the
`@transient`s of `NativeDataSourceStreamReader` are 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]
