HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4418954837
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -399,6 +515,25 @@ static MethodHandle downcall(Object address, C result,
C... parameters) throws T
return handle;
}
+ /**
+ * 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 #downcall} does for a
fixed address.
+ */
+ static MethodHandle unboundDowncall(C result, C... parameters) throws
Throwable {
+ MethodHandle handle = (MethodHandle) UNBOUND_DOWNCALL.invoke(
+ LINKER, descriptor(result, parameters), NO_OPTIONS);
+ handle = MethodHandles.filterArguments(handle, 0, OF_ADDRESS);
+ for (int i = 0; i < parameters.length; i++) {
+ if (parameters[i] == C.PTR) {
+ handle = MethodHandles.filterArguments(handle, i + 1, OF_ADDRESS);
+ }
+ }
+ if (result == C.PTR) {
+ handle = MethodHandles.filterReturnValue(handle, ADDRESS_OF);
+ }
Review Comment:
Done in b98fd19: `convertPointers` applies the filters once, for `downcall`,
`unboundDowncall` (from the parameter 1) and `upcall` (with the filters the
other way).
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1162,182 @@ private static boolean isInterrupted(long task) {
}
}
+ //
---------------------------------------------------------------------------------------------
+ // Streams of engines
+ //
---------------------------------------------------------------------------------------------
+
+ // The offsets of the fields of struct ArrowArrayStream, and its size.
+ private static final long STREAM_GET_SCHEMA = 0;
+ private static final long STREAM_GET_NEXT = 8;
+ private static final long STREAM_GET_LAST_ERROR = 16;
+ private static final long STREAM_RELEASE = 24;
+ private static final long STREAM_PRIVATE_DATA = 32;
+ private static final long STREAM_SIZE = 40;
+ // The errno values that the callbacks of a stream return.
+ private static final int EINVAL = 22;
+ private static final int EIO = 5;
+ // 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 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) {
+ this.inner = inner;
+ this.task = task;
+ }
+
+ 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;
+ }
+ }
+ }
+
+ /** Copies a string to native memory, as a NUL-terminated UTF-8 string,
which the caller frees. */
+ private static long toCString(String value) {
+ byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
+ long address = Platform.allocateMemory(bytes.length + 1);
+ Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address,
bytes.length);
+ Platform.putByte(null, address + bytes.length, (byte) 0);
+ return address;
+ }
+
+ /**
+ * 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}).
+ */
+ static void wrapEngineStream(long stream, Task task) {
+ long inner = Platform.allocateMemory(STREAM_SIZE);
+ Platform.copyMemory(null, stream, null, inner, STREAM_SIZE);
+ long id = nextEngineStreamId.getAndIncrement();
+ engineStreams.put(id, new EngineStream(inner, task));
+ 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();
+ }
+
+ /** The number of task contexts that are not closed yet, for tests. */
+ static int openTaskContexts() {
+ return tasks.size();
+ }
Review Comment:
Done in b98fd19: `openTaskContexts` is in the Task contexts section, after
`isInterrupted`, and `toCString` in the Memory section.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1145,6 +1276,155 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
check(writeTasks)
}
+ nativeTest("open the partitions of a batch scan as Arrow C streams") {
+ val session = newSession(defaultPackage)
+ // The reader factory and the partitions of the batch scan of a query,
which an engine gets
+ // from its plan.
+ def plannedScan(df: DataFrame): (SupportsArrowStreamRead,
Seq[InputPartition]) = {
+ val scan = df.queryExecution.executedPlan.collectFirst { case s:
BatchScanExec => s }.get
+ (scan.readerFactory.asInstanceOf[SupportsArrowStreamRead],
scan.inputPartitions)
+ }
+ // Runs `f` in a task that only gets the files of the session.
+ def inTask[T: ClassTag](f: => T): T =
session.artifactManager.withResources {
+ session.sparkContext.parallelize(Seq(0), 1).map(_ => f).collect().head
+ }
+ import TestArrowEngine.readIds
+
+ // An engine reads both partitions of the scan in a single task.
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ val ids = inTask {
+ partitions.flatMap { partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ try {
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try factory.openArrowStream(partition, stream.memoryAddress()) catch
{
+ case e: Throwable =>
+ stream.close()
+ throw e
+ }
+ readIds(stream, allocator)
+ } finally {
+ allocator.close()
+ }
+ }
+ }
+ assert(ids.sorted == (0L until 6L))
+ // The library got the partition of the scan that it reads, not the one of
the task.
+ val tasks = taskDir.listFiles().map { file =>
+ file.getName.toInt ->
Files.readString(file.toPath).split(",").map(_.toLong).toSeq
+ }.toMap
+ assert(tasks.keySet == Set(0, 1) && tasks.values.map(_(5)).toSet.size ==
1, tasks)
+ assert(NativeLibrary.openEngineStreams() == 0)
+
+ // The engine may release the stream after its task ended, and on another
thread: Spark frees
+ // the context of the task once the library released the stream.
+ val released = new File(Utils.createTempDir(), "released")
+ val (laterFactory, laterPartitions) =
+ plannedScan(range(session, "end" -> "6", "released_path" ->
released.getPath).select("id"))
+ // Opens a stream in a task, which it returns unreleased.
+ def openInTask(factory: SupportsArrowStreamRead, partition:
InputPartition): Long = inTask {
+ val stream = Platform.allocateMemory(40)
+ Platform.setMemory(stream, 0.toByte, 40)
+ factory.openArrowStream(partition, stream)
+ stream
+ }
+ val openContexts = NativeLibrary.openTaskContexts()
+ val address = openInTask(laterFactory, laterPartitions.head)
+ assert(NativeLibrary.openEngineStreams() == 1 && !released.exists())
+ assert(NativeLibrary.openTaskContexts() == openContexts + 1)
+ assert(readIds(ArrowArrayStream.wrap(address), ArrowUtils.rootAllocator)
== (0L until 3L))
+ assert(NativeLibrary.openEngineStreams() == 0 && released.exists())
+ assert(NativeLibrary.openTaskContexts() == openContexts)
+ Platform.freeMemory(address)
+
+ // Once the task ended, its context tells the library that it was
interrupted, so that the
+ // library stops waiting.
+ val marker = new File(Utils.createTempDir(), "interrupted")
+ val (waitingFactory, waitingPartitions) = plannedScan(
+ range(session, "end" -> "6", "wait_for_interrupt" ->
marker.getPath).select("id"))
+ val waiting = openInTask(waitingFactory, waitingPartitions.head)
+ val error = intercept[Exception] {
+ readIds(ArrowArrayStream.wrap(waiting), ArrowUtils.rootAllocator)
+ }
+ assert(Utils.exceptionString(error).contains("the task was killed"))
+ assert(Files.readString(marker.toPath) == "interrupted")
+ assert(NativeLibrary.openEngineStreams() == 0)
Review Comment:
Done in b98fd19: the test checks `openTaskContexts()` after the interrupted
stream and after each failure too.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1162,182 @@ private static boolean isInterrupted(long task) {
}
}
+ //
---------------------------------------------------------------------------------------------
+ // Streams of engines
+ //
---------------------------------------------------------------------------------------------
+
+ // The offsets of the fields of struct ArrowArrayStream, and its size.
+ private static final long STREAM_GET_SCHEMA = 0;
+ private static final long STREAM_GET_NEXT = 8;
+ private static final long STREAM_GET_LAST_ERROR = 16;
+ private static final long STREAM_RELEASE = 24;
+ private static final long STREAM_PRIVATE_DATA = 32;
+ private static final long STREAM_SIZE = 40;
+ // The errno values that the callbacks of a stream return.
+ private static final int EINVAL = 22;
+ private static final int EIO = 5;
+ // 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 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) {
+ this.inner = inner;
+ this.task = task;
+ }
+
+ 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;
+ }
+ }
+ }
+
+ /** Copies a string to native memory, as a NUL-terminated UTF-8 string,
which the caller frees. */
+ private static long toCString(String value) {
+ byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
+ long address = Platform.allocateMemory(bytes.length + 1);
+ Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address,
bytes.length);
+ Platform.putByte(null, address + bytes.length, (byte) 0);
+ return address;
+ }
+
+ /**
+ * 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}).
+ */
+ static void wrapEngineStream(long stream, Task task) {
+ long inner = Platform.allocateMemory(STREAM_SIZE);
+ Platform.copyMemory(null, stream, null, inner, STREAM_SIZE);
+ long id = nextEngineStreamId.getAndIncrement();
+ engineStreams.put(id, new EngineStream(inner, task));
+ 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();
+ }
+
+ /** The number of task contexts that are not closed yet, for tests. */
+ static int openTaskContexts() {
+ return tasks.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(
+ "The " + callback + " of the stream of a native data source failed",
e);
Review Comment:
Done in b98fd19: `EngineStream` gets the name of the data source, and the
errors of its callbacks and of its release name it.
--
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]
