dongjoon-hyun commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4471627269
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -304,29 +313,135 @@ 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. The
stream of a read may be
+ * released once the task ended, from another thread: an engine that runs a
query in native code
+ * may get the stream from Spark, read it in the task that opened it and
release it later, and
+ * Spark itself may read a stream on a thread that the task does not wait for,
for example the one
+ * of a PipedRDD. The task stays valid until then too, and for a read,
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 and the schemas of the stream after it
released the stream, so
+ * the release of an array or a schema must not use the stream or the task.
*/
typedef struct SparkDsTaskContext {
/* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
int64_t size;
/* The ID of the stage of the task, and the attempt number of the stage. */
int32_t stage_id;
int32_t stage_attempt_number;
- /* The partition that the task reads or writes, and the attempt number of
the task. */
+ /*
+ * The partition: for read, the index of the partition in the list that
partitions or
+ * stream_partitions returned, which is not the partition of the task when
Spark reads several
+ * scans or partitions in one task, for example after a union or a coalesce.
It is unique within
+ * a scan, but a task that reads several scans may get the same index for
each. For
+ * create_data_writer, the partition of the task that writes. Then the
attempt number of the
+ * task.
+ */
int32_t partition_id;
Review Comment:
**2. [Low] Could the index of the partition in the scan be a field of
`SparkDsReadInfo`?**
Since b98fd19, which answered [my item 1 of round
3](https://github.com/apache/spark/pull/59934#discussion_r4418419351),
`partition_id` is the index of the partition in the scan for `read`, and the
partition of the task for `create_data_writer`. This comment documents both,
and that a task that reads several scans may get the same index for each. That
leaves a read with no way to get the partition of its task, and gives the same
field of the same struct two meanings, while `SparkDsReadInfo` already
describes the partition being read (`kind`, `partition`).
Since #59931 settles version 1 of the interface, could `partition_id` stay
the partition of the task for both, and `SparkDsReadInfo` get the index, for
example as `partition_index` after `partition_length`? Then `(stage_id,
stage_attempt_number, partition_id, attempt_number)` identifies the task in a
read as it does in a data writer, and the index stays unique within a scan, as
now. If you prefer the current shape, since `task_attempt_id` already gives a
unique key, that is fine with me too.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -178,30 +204,89 @@ class ColumnarScan(
* The batch of a scan, which plans it once, when Spark first asks for its
partitions or its reader
* factory. A failure is planned once too: the physical plan can ask again,
for example from a lazy
* value, and gets the same error.
+ *
+ * @param projection the ordinals of the columns that the scan reads, if the
partitions return
+ * more columns
*/
-class ColumnarBatchScan(planner: ColumnarBatchPlanner) extends Batch {
+class ColumnarBatchScan(planner: ColumnarBatchPlanner, projection:
Option[Array[Int]] = None)
+ extends Batch {
private lazy val planned: Try[ColumnarPlannedBatch] = Try(planner.plan())
override def planInputPartitions(): Array[InputPartition] =
planned.get.partitions
- override def createReaderFactory(): PartitionReaderFactory =
planned.get.readerFactory
+ override def createReaderFactory(): PartitionReaderFactory = {
Review Comment:
**1. [Medium] A batch scan whose library does not prune the columns loses
`SupportsArrowStreamRead`**
Since 870e3fa, which came from [my review of
#59943](https://github.com/apache/spark/pull/59943#discussion_r4463353297), a
batch scan that reads fewer columns than the full schema wraps the factory of
the planner in a `ColumnarProjectionReaderFactory` when the planner does not
prune them. That wrapper does not implement `SupportsArrowStreamRead`. So for a
library that leaves `prune_columns` NULL, or returns `pruned = false`, a query
such as `select("id")` gives `BatchScanExec.readerFactory` a factory without
the interface.
The Javadoc (`SupportsArrowStreamRead.java` L27-L29), the guide
(`sql-native-data-sources.md` L254-L255) and the PR description say that the
reader factories of native batch scans implement it. An engine that casts, as
the test does at `NativeDataSourceSuite.scala` L1549, then gets a
`ClassCastException`, and one that checks falls back to Spark's import without
telling why. The test library always prunes (`test_native_datasource.cc`
L1188), so no test reaches this case.
A wrapper that only delegated `openArrowStream` would return a stream with
all the columns of the partitions, not those of the read schema that the
Javadoc promises (L31-L32). So could the docs say that the factory implements
it only when the library pruned the columns, or could the wrapper implement it
and project the stream too, as it projects the batches? Either way, a test with
a library that does not prune would cover it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -886,10 +998,135 @@ class NativePartitionReader(
}
}
-/** Writes the batches of a task to the data writer of a native library. */
+object NativePartitionReader {
+ import NativeCalls._
+
+ /**
+ * Initializes the stream at `streamAddress` with the data of a partition
for an engine, see
+ * [[SupportsArrowStreamRead]]: the context of its task is closed once the
stream is released.
+ * It first checks the arguments of `openArrowStream`, before it looks at
the partition, which
+ * must be a `P`, the class of the partitions of the scan of the reader
factory. It then reads
+ * the partition with the state and the partition of the library that
`stateAndPartition`
+ * returns for it.
+ */
+ 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 = {
+ val context = TaskContext.get()
+ if (context != null) {
+ // A task that was killed gets the error of its kill, as Spark reports
it.
+ context.killTaskIfInterrupted()
+ }
+ // A task that failed is over too, also while its listeners run.
+ if (context == null || context.isFailed() || context.isCompleted()) {
+ throw new IllegalStateException("openArrowStream must be called in the
running task that " +
+ "reads the partition, on a thread that has its TaskContext.")
+ }
+ 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 (readerState, nativePartition) = partition match {
+ case expected: P => stateAndPartition(expected)
+ case _ =>
+ throw new IllegalArgumentException("openArrowStream needs a partition
of the scan of " +
+ s"the reader factory, a ${classTag[P].runtimeClass.getSimpleName},
but got " +
+ s"${Option(partition).map(_.getClass.getName).orNull}.")
+ }
+ val task = openStream(location, name, kind, readerState, nativePartition,
schema, timeZoneId,
+ streamAddress)
+ try {
+ NativeLibrary.wrapEngineStream(streamAddress, task, name)
Review Comment:
**3. [Nit] A failure to create the callbacks is not a
`NATIVE_DATA_SOURCE_ERROR`**
`call` rethrows the link errors as `NATIVE_DATA_SOURCE_ERROR` too, "such as
when Arrow cannot load its own JNI library" (L1312-L1327). This
`wrapEngineStream` runs outside of a `call`, though. So if
`EngineStreamCallbacks` cannot be created, the engine gets the
`ExceptionInInitializerError` of its first use, and a `NoClassDefFoundError`
from each later one, rather than an error that names the data source.
The stream is still released by `releaseOnFailure`, so only the error
differs from the one of a failing `read`, which the test checks. Could this
line run in `call(name, "read")`, as `library.read` does in `openStream`?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1180,6 +1964,7 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
val marker = new File(Utils.createTempDir(), "released")
val reading = marker.getPath + ".reading"
val df = range(session, "end" -> "10", "partitions" -> "1", "slow_read" ->
marker.getPath)
+ val contexts = NativeLibrary.openTaskContexts()
Review Comment:
**4. [Nit] The killed task of the previous test may close its context after
`contexts` is read**
"tell a read that its task was killed" ends once the library wrote its
marker (L1956-L1958), but its task only closes the context after `read`
returned and Spark closed the reader, which may happen after this line.
`contexts` would then count that context, and `== contexts + 1` at L1986 would
fail, since the count went down by one meanwhile. Reading the count right
before each step, which answered [my item 5 of round
4](https://github.com/apache/spark/pull/59934#discussion_r4419050277), narrows
the window but does not close it, and `checkContexts` (L1577-L1585) would have
the same exposure if a test that kills a task ran before it.
Could the test of the kill also wait for its context, as the kill part of
the stream test does at L1702? For example, it could read
`NativeLibrary.openTaskContexts()` before the query and check it in its last
`eventually`.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1143,260 @@ 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 callbacks of the streams that Spark wraps for engines, see {@link
EngineStream}, the
+ * method handles that call the callbacks of the wrapped streams, and what
the callbacks return
+ * for a stream that Spark does not know. They are created when the first
stream is wrapped, so
+ * that a JVM that wraps no stream does not create them, and a failure to
create them only fails
+ * the streams of engines.
+ */
+ private static final class EngineStreamCallbacks {
+ // The addresses of the C functions that call Java, which live as long as
the JVM.
+ static final long GET_SCHEMA;
+ static final long GET_NEXT;
+ static final long GET_LAST_ERROR;
+ static final long RELEASE;
+ // The callbacks of struct ArrowArrayStream, called with the address of
the function first:
+ // get_schema and get_next, get_last_error, and release.
+ static final MethodHandle INT_CALLBACK;
+ static final MethodHandle POINTER_CALLBACK;
+ static final MethodHandle VOID_CALLBACK;
+ // The message of get_last_error for a stream that Spark does not know,
which lives as long as
+ // the JVM.
+ static final long UNKNOWN_STREAM_ERROR;
+
+ static {
+ try {
+ MethodHandles.Lookup lookup = MethodHandles.lookup();
+ // Linker.downcallHandle(FunctionDescriptor, Linker.Option...), whose
method handles take
+ // the address of the function to call as their first argument.
+ MethodHandle unboundDowncall = lookup.findVirtual(
+ Class.forName("java.lang.foreign.Linker"), "downcallHandle",
+ MethodType.methodType(MethodHandle.class,
+ Class.forName("java.lang.foreign.FunctionDescriptor"),
Ffm.NO_OPTIONS.getClass()))
+ .asFixedArity();
+ GET_SCHEMA = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "engineStreamGetSchema",
+ MethodType.methodType(int.class, long.class, long.class)),
+ C.I32, C.PTR, C.PTR);
+ GET_NEXT = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "engineStreamGetNext",
+ MethodType.methodType(int.class, long.class, long.class)),
+ C.I32, C.PTR, C.PTR);
+ GET_LAST_ERROR = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "engineStreamGetLastError",
+ MethodType.methodType(long.class, long.class)),
+ C.PTR, C.PTR);
+ RELEASE = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "engineStreamRelease",
+ MethodType.methodType(void.class, long.class)),
+ C.VOID, C.PTR);
+ INT_CALLBACK = callback(unboundDowncall, C.I32, C.PTR, C.PTR);
+ POINTER_CALLBACK = callback(unboundDowncall, C.PTR, C.PTR);
+ VOID_CALLBACK = callback(unboundDowncall, C.VOID, C.PTR);
+ UNKNOWN_STREAM_ERROR = toCString("The stream was released, or is not a
stream of Spark.");
+ } catch (Throwable e) {
+ // The message names the cause, because the JVM only keeps this
message for the later
+ // uses of the class.
+ throw new IllegalStateException(
+ "Cannot create the callbacks of the streams of engines: " +
rootCause(e), e);
+ }
+ }
+
+ /**
+ * Returns a method handle that calls the C function at the address of its
first argument, a
+ * {@code long}, with the given types, as {@link Ffm#downcall} does for a
fixed address.
+ */
+ private static MethodHandle callback(MethodHandle unboundDowncall, C
result, C... parameters)
+ throws Throwable {
+ MethodHandle handle = (MethodHandle) unboundDowncall.invoke(
+ Ffm.LINKER, Ffm.descriptor(result, parameters), Ffm.NO_OPTIONS);
+ handle = MethodHandles.filterArguments(handle, 0, Ffm.OF_ADDRESS);
+ return Ffm.convertPointers(handle, 1, Ffm.OF_ADDRESS, Ffm.ADDRESS_OF,
result, parameters);
+ }
+ }
+
+ 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, but may
+ // run on different threads one after the other.
+ volatile 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 in the given callback, which get_last_error
then returns, and logs
+ * it. The message is built here, since building it can throw too, as a
StackOverflowError on
+ * the small stack of a thread of an engine can.
+ */
+ int fail(String callback, Throwable e) {
+ try {
+ clearError();
+ String message =
+ "The " + callback + " of the stream of the native data source " +
name + " failed";
+ 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) {
Review Comment:
**5. [Nit] Each call of an engine's stream goes through an upcall and a
downcall**
In [round
1](https://github.com/apache/spark/pull/59934#discussion_r4417274537) I
suggested replacing only `release`, and noted that wrapping all four callbacks
costs an upcall per `get_next`. I still think that the wrapping is the more
robust choice, since it needs no key from the library. Its cost is a bit larger
than I wrote, though: each `get_schema`, `get_next` and `get_last_error` of the
engine is an upcall into `engineStreamCall` and then a downcall back into the
library. The first call on a thread of the engine also attaches that thread to
the JVM, which, as far as I know, keeps it attached until the thread exits.
The Javadoc and the guide say that the callbacks call into the JVM, but not
what that costs an engine whose point is to stay out of it. Did you measure it
against the import of a batch? If it is small, as I expect for batches of
thousands of rows, could the guide say that each call goes through the JVM, and
attaches a native thread to it? Since the wrapping is internal to Spark, it
could still keep the callbacks of the library later, without a change of the
interface.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,25 +382,44 @@ class NativeDataSourceStreamReader(
offset
}
- override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
- library.streamPartitions(handle.get, start, end)
+ override def partitions(start: String, end: String): Array[InputPartition] =
synchronized {
+ if (plannedRange != (start, end)) {
+ plannedPartitions = Try {
+ 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. Spark only plans the next range once the
micro-batch of the previous
+ // one finished, so the copies of its state on the executors are
removed then. The driver
+ // keeps its own copy, so that the partitions and the reader factories
that an engine keeps
+ // can still fetch it, until the ContextCleaner removes it once
nothing refers to them.
+ plannedState.foreach(_.unpersist(blocking = false))
Review Comment:
**6. [Nit] A small state still gets a broadcast for each micro-batch**
The broadcast answered [my item 1 of round
5](https://github.com/apache/spark/pull/59934#discussion_r4420249991), which
was about a large state sent with every task. For the state of a typical stream
reader, such as a few offsets, each micro-batch now writes a broadcast to the
block manager of the driver, registers it with the ContextCleaner, makes each
executor fetch it once, and sends an unpersist to every executor when the next
range is planned. With a short trigger interval, that may cost more than the
bytes that it saves in each task.
Could a state below some size, for example 64 KiB, stay in the partitions,
with `NativeMicroBatchPartition` holding either the bytes or the broadcast? The
test of a serialized partition would then use a state above that size.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -222,9 +338,52 @@ abstract class ColumnarOnlyPartitionReaderFactory extends
PartitionReaderFactory
override def supportColumnarReads(partition: InputPartition): Boolean = true
}
-class ColumnarStreamPartitionReaderFactory(reader: ColumnarStreamReader)
+object ColumnarProjection {
+ /**
+ * Returns `factory`, or a factory whose readers only return the columns at
`projection` of the
+ * batches of the readers of `factory`, if any.
+ */
+ def apply(
+ factory: PartitionReaderFactory,
+ projection: Option[Array[Int]]): PartitionReaderFactory = projection
match {
+ case Some(ordinals) => new ColumnarProjectionReaderFactory(factory,
ordinals)
+ case None => factory
+ }
+}
+
+/**
+ * Reads only the columns at `ordinals` of the batches of the readers of
`factory`, for a scan
+ * whose partitions return more columns than the scan reads.
+ */
+class ColumnarProjectionReaderFactory(factory: PartitionReaderFactory,
ordinals: Array[Int])
extends ColumnarOnlyPartitionReaderFactory {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- reader.read(partition)
+ new ColumnarProjectionReader(factory.createColumnarReader(partition),
ordinals)
}
}
+
+private class ColumnarProjectionReader(
+ reader: PartitionReader[ColumnarBatch],
+ ordinals: Array[Int])
+ extends PartitionReader[ColumnarBatch] {
+
+ override def next(): Boolean = reader.next()
+
+ override def get(): ColumnarBatch = new ProjectedColumnarBatch(reader.get(),
ordinals)
Review Comment:
**9. [Nit] `ColumnarProjectionReader.get()` returns a new batch on each
call**
This comes from 870e3fa of #59931. `PartitionReader.get()` "should return
same value until `next` is called", but this allocates a new
`ProjectedColumnarBatch`, and an array of its columns, on each call, so two
calls return different batches. Spark calls `get()` once for each `next()`, so
this is only an allocation today, but could `next()` build the batch once, and
`get()` return it?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1143,260 @@ 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 callbacks of the streams that Spark wraps for engines, see {@link
EngineStream}, the
+ * method handles that call the callbacks of the wrapped streams, and what
the callbacks return
+ * for a stream that Spark does not know. They are created when the first
stream is wrapped, so
+ * that a JVM that wraps no stream does not create them, and a failure to
create them only fails
+ * the streams of engines.
+ */
+ private static final class EngineStreamCallbacks {
+ // The addresses of the C functions that call Java, which live as long as
the JVM.
+ static final long GET_SCHEMA;
+ static final long GET_NEXT;
+ static final long GET_LAST_ERROR;
+ static final long RELEASE;
+ // The callbacks of struct ArrowArrayStream, called with the address of
the function first:
+ // get_schema and get_next, get_last_error, and release.
+ static final MethodHandle INT_CALLBACK;
+ static final MethodHandle POINTER_CALLBACK;
+ static final MethodHandle VOID_CALLBACK;
+ // The message of get_last_error for a stream that Spark does not know,
which lives as long as
+ // the JVM.
+ static final long UNKNOWN_STREAM_ERROR;
+
+ static {
+ try {
+ MethodHandles.Lookup lookup = MethodHandles.lookup();
+ // Linker.downcallHandle(FunctionDescriptor, Linker.Option...), whose
method handles take
+ // the address of the function to call as their first argument.
+ MethodHandle unboundDowncall = lookup.findVirtual(
+ Class.forName("java.lang.foreign.Linker"), "downcallHandle",
+ MethodType.methodType(MethodHandle.class,
+ Class.forName("java.lang.foreign.FunctionDescriptor"),
Ffm.NO_OPTIONS.getClass()))
+ .asFixedArity();
+ GET_SCHEMA = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "engineStreamGetSchema",
+ MethodType.methodType(int.class, long.class, long.class)),
+ C.I32, C.PTR, C.PTR);
+ GET_NEXT = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "engineStreamGetNext",
+ MethodType.methodType(int.class, long.class, long.class)),
+ C.I32, C.PTR, C.PTR);
+ GET_LAST_ERROR = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "engineStreamGetLastError",
+ MethodType.methodType(long.class, long.class)),
+ C.PTR, C.PTR);
+ RELEASE = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "engineStreamRelease",
+ MethodType.methodType(void.class, long.class)),
+ C.VOID, C.PTR);
+ INT_CALLBACK = callback(unboundDowncall, C.I32, C.PTR, C.PTR);
+ POINTER_CALLBACK = callback(unboundDowncall, C.PTR, C.PTR);
+ VOID_CALLBACK = callback(unboundDowncall, C.VOID, C.PTR);
+ UNKNOWN_STREAM_ERROR = toCString("The stream was released, or is not a
stream of Spark.");
+ } catch (Throwable e) {
+ // The message names the cause, because the JVM only keeps this
message for the later
+ // uses of the class.
+ throw new IllegalStateException(
+ "Cannot create the callbacks of the streams of engines: " +
rootCause(e), e);
+ }
+ }
+
+ /**
+ * Returns a method handle that calls the C function at the address of its
first argument, a
+ * {@code long}, with the given types, as {@link Ffm#downcall} does for a
fixed address.
+ */
+ private static MethodHandle callback(MethodHandle unboundDowncall, C
result, C... parameters)
+ throws Throwable {
+ MethodHandle handle = (MethodHandle) unboundDowncall.invoke(
+ Ffm.LINKER, Ffm.descriptor(result, parameters), Ffm.NO_OPTIONS);
+ handle = MethodHandles.filterArguments(handle, 0, Ffm.OF_ADDRESS);
+ return Ffm.convertPointers(handle, 1, Ffm.OF_ADDRESS, Ffm.ADDRESS_OF,
result, parameters);
+ }
+ }
+
+ 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, but may
+ // run on different threads one after the other.
+ volatile 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 in the given callback, which get_last_error
then returns, and logs
+ * it. The message is built here, since building it can throw too, as a
StackOverflowError on
+ * the small stack of a thread of an engine can.
+ */
+ int fail(String callback, Throwable e) {
+ try {
+ clearError();
+ String message =
+ "The " + callback + " of the stream of the native data source " +
name + " failed";
+ sparkError = toCString(message + ": " + rootCause(e));
+ LOGGER.error(message, e);
Review Comment:
**8. [Nit] The errors of an engine's stream are logged without MDC**
`log` logs the messages of the library with `MDC.of(LogKeys.PATH, ...)` and
`MDC.of(LogKeys.MESSAGE, ...)` (L597-L604), but this line, and the one in
`engineStreamRelease` (L1389-L1390), put the name of the data source into the
message itself. Could they pass it as `MDC.of(LogKeys.DATA_SOURCE, name)`, and
the callback as `MDC.of(LogKeys.FUNCTION_NAME, callback)`? For example:
```java
LOGGER.error("The {} of the stream of the native data source {} failed", e,
MDC.of(LogKeys.FUNCTION_NAME, callback), MDC.of(LogKeys.DATA_SOURCE,
name));
```
`sparkError` can keep the plain message.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -970,21 +1207,43 @@ 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 read whether the context is the one of a read, which the library
may use after the
+ * task ended: an engine may read the stream of a partition
then, and a thread that
+ * the task does not wait for may be in get_next once it
completes, as with
+ * `PipedRDD`. `is_interrupted` then also reports that the task
failed or completed,
+ * and the context does not keep the task alive. A data writer
only runs within its
+ * task, whose kill `is_interrupted` reports. Outside of a task,
it reports nothing.
*/
def create(
partitionId: Option[Int] = None,
- taskAttemptId: Option[Long] = None): NativeLibrary.Task = {
+ taskAttemptId: Option[Long] = None,
+ read: Boolean = false): NativeLibrary.Task = {
val context = Option(TaskContext.get())
+ val interrupted: BooleanSupplier = context match {
+ case Some(taskContext) if read =>
+ // Set once the task failed or completed, by listeners, so that a
check only reads
+ // volatile fields, unlike `isFailed()` and `isCompleted()`, which
lock the task.
+ val ended = new AtomicBoolean(false)
+ taskContext.addTaskFailureListener((_, _) => ended.set(true))
Review Comment:
**7. [Nit] Each read registers two listeners that stay until the end of its
task**
Thank you for the flag, which answers [item 6 of round
7](https://github.com/apache/spark/pull/59934#discussion_r4463375838). Since
every read now gets the context of a read, also Spark's own (L1074), a task
that reads many partitions keeps two listeners and a flag for each until it
ends, and runs them all then. For example, `coalesce(1)` of a scan with 100,000
partitions keeps 200,000 listeners. That is the cost of [my item 11 of round
1](https://github.com/apache/spark/pull/59934#discussion_r4417274621), now for
Spark's own reads too.
Since the flag only depends on the task, could the reads of a task share
one? For example, a map by `taskAttemptId()` could hold the flag of each
running task, whose single pair of listeners sets it and removes the entry once
the task ended. The contexts of the reads keep their flag, so they still see it
after that.
--
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]
