HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4472098117
##########
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:
Thank you, I measured it on this branch, with a temporary test (not
committed) and a small C++ helper that calls `get_next` and releases each array
on a new native thread, as an engine does, on the wrapped stream and on the
stream of the library directly, and compared that with Spark's own reader on
batches of the same size. These are the ranges of two runs on Java 25 with
`native_range` (one run for 1024 rows), without the first round of each run,
which warms up:
| rows per batch | wrapped `get_next` | direct `get_next` | Spark's own read
(`get_next` + import) |
|---|---|---|---|
| 1 | 362-367 ns | 335-341 ns | 4.0-5.1 µs |
| 1024 | 3.13-3.14 µs | 3.14-3.16 µs | 7.7-8.1 µs |
| 4096 | 11.7-12.3 µs | 11.7-11.8 µs | 15.9-16.9 µs |
So the upcall and the downcall add 20 to 30 ns to a call, which is lost in
the noise from 1024 rows on, while the import of a batch on the JVM costs 4 to
5 µs. The first call on a new native thread took 40 to 90 µs more through the
wrapper once the JVM was warm, which is the attach. The JVM's thread count
rises by one while the native thread runs and falls back once it exits, so the
thread stays attached between the calls and is detached when it exits, as you
thought. Done in d40d64d: the guide now says that each call goes through the
JVM and what that adds, and both the guide and the Javadoc say that the first
call on a native thread attaches the thread until it exits. As you say, Spark
could still hand out 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:
Done in d40d64d: a state of up to 64 KiB goes with each partition, and only
a larger one is broadcast. `NativeMicroBatchPartition.state` is an `Either` of
the bytes and their broadcast, and the stream reader only keeps, and
unpersists, a broadcast. The test of a serialized partition now uses a state
above that size, with a new option `state_padding` of `native_counter`, and a
new check plans a micro-batch with the usual small state, whose bytes the
serialized partition contains, and reads it.
##########
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:
Done in d40d64d: the reads of a task share one flag, held in a map until the
task ended. The read that adds it registers the single pair of listeners, which
set the flag and remove the entry once the task failed or completed. A listener
added after the task ended runs right away, so the entry is removed then too.
The key is the `TaskContext` itself rather than its `taskAttemptId()`, so that
the contexts of `TaskContext.empty()`, whose attempt IDs are all 0, and the
tasks of a later `SparkContext` in the same JVM never share a flag. The
contexts of the reads keep their flag. In the test of the streams of a batch
scan, the task that reads both partitions has a single flag, and none is left
once it ended.
##########
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:
Done in d40d64d: both lines log with `MDC.of(LogKeys.FUNCTION_NAME,
callback)` and `MDC.of(LogKeys.DATA_SOURCE, name)`, with `release` as the
function in `engineStreamRelease`, and `sparkError` keeps the plain message.
--
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]
