dongjoon-hyun commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4424626654
##########
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:
**2. [Low] Spark's own reader does not tell a read that its task ended**
This flag gives the signal of [round
2](https://github.com/apache/spark/pull/59934#discussion_r4417661673) to the
streams of engines only. A stream that Spark reads itself can outlive its task
too, as the doc of `NativePartitionReader` describes (L567-L574): with
`PipedRDD`, a thread of its own is in `get_next` when the task completes, and
`close` leaves the release to `next`.
For that read, `context.exists(_.isInterrupted())` stays false, since a
normal completion is not a kill. So a `get_next` that waits for data and polls
`is_interrupted`, as the header asks, keeps waiting until its own timeout, and
the thread, the stream, the allocator and the context stay until then. "release
a stream only once get_next returns" asserts this `false`
(`NativeDataSourceSuite.scala` L1822).
The behavior predates this PR, so it need not block it. This PR does make
the difference part of the interface, though, since the header now describes
`is_interrupted` differently for the stream of an engine. Could every read use
the first branch, so that reads have one meaning of it? Data writers run within
their task, so they could keep the second branch. That test would then need
another way to tell that the context was not freed, for example
`NativeLibrary.openTaskContexts()` while `get_next` waits.
##########
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:
**1. [Low] The message of this `catch` is built outside the guard of `fail`**
`fail` guards its work with a `try`, but its argument is evaluated before
the call, here in the `catch` block. Besides a NULL callback, what the downcall
can throw is mostly `StackOverflowError` or `OutOfMemoryError`, and this
concatenation can then throw the same error again. On its first run, it also
has to link its `StringConcatFactory` call site. An exception that leaves an
upcall aborts the JVM, as the comment at L1292 says. The native thread of an
engine may have a small stack, so a stack overflow is the more likely of the
two.
Could `fail` take the name of the callback, and build the message inside its
`try`?
```java
return wrapped == null ? EIO : wrapped.fail(callback, e);
```
##########
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 =>
Review Comment:
**7. [Nit] The check of `partitions.nonEmpty` is not tested**
`plan(stateless, 0, 2)` covers a micro-batch without a state
(`NativeDataSourceSuite.scala` L1648), but not one with a state and without
partitions. If this check went away, such a micro-batch would add a broadcast
that nothing reads, and no test would fail, since there is no partition to look
at. The test library returns exactly that for `plan(stream, 2, 2)`, and the ID
of the next broadcast would show it:
```scala
val before = session.sparkContext.broadcast(0).id
assert(plan(stream, 2, 2).isEmpty)
assert(session.sparkContext.broadcast(0).id == before + 1)
```
##########
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:
**6. [Low] `read` is no longer used outside of a test**
This stream reader is the only `ColumnarStreamReader` in the main code.
Since it overrides `readerFactory()`, Spark never calls this `read`. That
leaves the default `readerFactory()` and `ColumnarStreamPartitionReaderFactory`
only for `TestRangeStreamReader` in `ColumnarDataSourceSuite`, and the same
goes for the `Serializable` of the trait. The `@transient`s of this class
(L335, L343, L344, L346) have no effect now that nothing serializes it.
Could the trait declare an abstract `readerFactory()` instead of `read`?
`TestRangeStreamReader` could return `new TestRangeReaderFactory(schema)`,
which already has the code of its `read`, and the rest could go.
##########
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:
**3. [Low] `openArrowStream` also opens a stream in a task that failed**
The Javadoc says that it is called on the thread of the running task, "and
fails otherwise", but this only rejects a completed task.
`TaskContextImpl.markTaskCompleted` only sets `completed` after the failure
listeners and the pre-completion listeners ran, so a call from one of them
passes. `read` of the library then runs for a task that is over, the context
reports `is_interrupted` as true at once, and an engine that is cleaning up may
lose the stream, which then leaks.
Could this reject what `NativeTask.create` reports for such a context,
`context.isInterrupted() || context.isFailed() || context.isCompleted()`? A
test with `addTaskFailureListener`, next to the one with the completion
listener, would cover it.
##########
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:
**4. [Low] The kill of the task of an engine's stream is not tested**
Of the three conditions of this context (`NativeDataSource.scala` L1218),
the test only covers `isCompleted`, since this part reads after the task ended.
"tell a read that its task was killed" reads with Spark's own reader, which
takes the other branch of `NativeTask.create`. So from reading the suite, a
change that dropped `t.isInterrupted()` or `t.isFailed()` would still pass,
although a canceled query is the main case for `is_interrupted`.
Could a test open a `wait_for_interrupt` stream in a task, read it there,
cancel the job group as that test does, and check the marker file?
##########
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()) {
+ 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.")
+ }
+ partition match {
+ case expected: P => expected
+ case _ =>
+ throw new IllegalArgumentException("openArrowStream needs a partition
of the scan of " +
+ s"the reader factory, a ${tag.runtimeClass.getSimpleName}, but got "
+
+ s"${Option(partition).map(_.getClass.getName).orNull}.")
+ }
+ }
+
+ /**
+ * Initializes the stream at `streamAddress` with the data of a partition
for an engine, see
+ * [[SupportsArrowStreamRead]], once `engineStreamPartition` checked the
arguments: the context
+ * of its task is closed once the stream is released.
+ */
+ def openEngineStream(
Review Comment:
**9. [Nit] The checks and the opening of an engine's stream are two steps**
Both factories call `engineStreamPartition` and then this method, and only
its doc says that the checks come first. A new caller of this public method
alone gets no check, and `read` then overwrites a stream that is still open,
which leaks it. Could one `private[ffi]` method do both, with a function that
reads the state and the partition after the checks? For example:
```scala
private[ffi] def openEngineStream[P <: InputPartition: ClassTag](
location: NativeLibraryLocation,
name: String,
kind: Int,
partition: InputPartition,
schema: StructType,
timeZoneId: String,
streamAddress: Long)(stateAndPartition: P => (Array[Byte],
NativeInputPartition)): Unit
```
`openStream` and `releaseOnFailure` are `private[ffi]`, but only this object
and its companion class use them, so they could be `private`.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -853,7 +1133,10 @@ private static byte[] readCString(long address) {
/**
* The {@code SparkDsTaskContext} of a task that reads a partition or writes
data. The reader or
Review Comment:
**11. [Nit] The comment of `tasks` only mentions a kill**
L1131, which is outside of the diff, says "Whether the task of each
SparkDsTaskContext was killed". Since this PR, the context of an engine's
stream also reports that its task failed or completed (`NativeDataSource.scala`
L1218). Could it say something like "Whether the library should stop reading or
writing"?
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -304,29 +312,133 @@ typedef struct SparkDsCreateInfo {
int64_t num_options;
} SparkDsCreateInfo;
+/* A batch scan, as create_reader gets it. */
+typedef struct SparkDsReaderInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+} SparkDsReaderInfo;
+
+/* A request for the partitions of a batch scan, as partitions gets it. */
+typedef struct SparkDsPartitionsInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+} SparkDsPartitionsInfo;
+
+/* 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
+ * fully qualified Hadoop path, whose characters are not escaped as in a
URI: a space stays a
+ * space, and % stays %. Its file system may be one that the library cannot
access, such as s3a
+ * or hdfs. Spark keeps what it needs there under the names that start with
_spark, such as the
+ * initial offset of the source, and the library must not use them.
+ */
+ const char* checkpoint_location;
+} SparkDsStreamReaderInfo;
+
+/* A request for the initial offset of a stream reader, as initial_offset gets
it. */
+typedef struct SparkDsInitialOffsetInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+} SparkDsInitialOffsetInfo;
+
+/* 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, for every trigger. With
Trigger.AvailableNow,
+ * Spark asks for the latest offset only once, before the first micro-batch
that it plans, and
+ * reads up to it.
+ */
+ const char* start;
+} SparkDsLatestOffsetInfo;
+
+/* A request for the partitions of a micro-batch, as stream_partitions gets
it. */
+typedef struct SparkDsStreamPartitionsInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /* The offsets, as JSON, that the micro-batch reads the data after, up to
and including end. */
+ const char* start;
+ const char* end;
+} SparkDsStreamPartitionsInfo;
+
+/* The end of the data that Spark processed, as commit_offset gets it. */
+typedef struct SparkDsCommitOffsetInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /* The offset, as JSON, up to and including which Spark processed the data.
*/
+ const char* end;
+} SparkDsCommitOffsetInfo;
+
/*
* 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 that read returns from Spark. It reads
the stream in the task
+ * that opened it, but may release it once the task ended, from another
thread: the task stays valid
+ * until then too, and is_interrupted then also returns true once the task
failed or completed, so
+ * the library may stop reading then. Such an engine may also keep the arrays
of the stream after
Review Comment:
**10. [Nit] The schema of a stream can outlive the stream too**
The Arrow C stream interface says this of the results of both `get_schema`
and `get_next` ([Result
lifetimes](https://arrow.apache.org/docs/format/CStreamInterface.html#result-lifetimes)),
but this sentence only names the arrays. A library may then take it that the
release of the schema can still use the stream or the task. Could it say "the
release of an array or a schema"?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -347,10 +467,33 @@ private static final class Ffm {
lookup.findStatic(NativeLibrary.class, "sink",
MethodType.methodType(void.class, long.class, long.class,
long.class)),
C.VOID, C.PTR, C.PTR, C.I64);
+ STATE_SINK = upcall(
+ lookup.findStatic(NativeLibrary.class, "stateSink",
+ MethodType.methodType(void.class, long.class, long.class,
long.class)),
+ C.VOID, C.PTR, C.PTR, C.I64);
IS_INTERRUPTED = upcall(
lookup.findStatic(NativeLibrary.class, "isInterrupted",
MethodType.methodType(boolean.class, long.class)),
C.BOOL, C.PTR);
+ ENGINE_STREAM_GET_SCHEMA = upcall(
Review Comment:
**12. [Nit] Every JVM creates the callbacks of the streams of engines**
These four upcall stubs and the three downcall handles after them
(L494-L496) are created in the static initializer of `Ffm`. So each executor
that reads or writes a native data source pays for them, also without an
engine, and a failure here fails all of `Ffm`. I did not measure the cost,
which may well be small. Only `wrapEngineStream` and the callbacks of the
wrapped streams use them, so a small holder class in the section of the streams
of engines would create them on first use.
##########
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:
**5. [Low] The broadcasts of past micro-batches stay until a GC of the
driver**
Thank you for moving the state to a broadcast, which answers [round
5](https://github.com/apache/spark/pull/59934#discussion_r4420249991). Each
micro-batch with a state now adds one, and the ContextCleaner only removes it
once the `Broadcast` object was collected on the driver. One that reached the
old generation waits for a GC of that generation, which
`spark.cleaner.periodicGC.interval` only forces every 30 minutes by default.
With a trigger of one second, that is up to about 1800 broadcasts for each
source, cached as `MEMORY_AND_DISK` on the driver and on the executors that
read them.
Could `partitions` unpersist the broadcast of the previous range when it
plans the next one? `unpersist(blocking = false)` only removes the copies on
the executors, and the driver keeps its own (`TorrentBroadcast.doUnpersist`),
so a factory and partitions that an engine keeps can still fetch the state.
Spark only plans the next range once the previous micro-batch finished. It
takes one more field here, so if the states are expected to stay small, leaving
it as it is would be fine too.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -458,17 +625,13 @@ typedef struct SparkDsLibrary {
/*
* Plans the partitions of the scan. Spark reads each partition in a
separate task.
Review Comment:
**8. [Nit] The header and the docs still say that each partition gets its
own task**
The new text of `partition_id` (L395-L399) says that a task may read several
partitions, but this line says "Spark reads each partition in a separate task",
and L60 and `docs/sql-native-data-sources.md` L105 say "in one task for each
partition". A library that believes them may keep a thread-local or a buffer
for each task, and assume that it serves one partition, which breaks after a
coalesce or a union. Could these three say that a task may read several
partitions?
--
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]
