HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4465714734
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1158,234 @@ 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.
+ /**
+ * The callbacks of the streams that Spark wraps for engines, see {@link
EngineStream}, and the
+ * method handles that call the callbacks of the wrapped streams. 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;
+
+ static {
+ try {
+ MethodHandles.Lookup lookup = MethodHandles.lookup();
+ 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 = Ffm.unboundDowncall(C.I32, C.PTR, C.PTR);
+ POINTER_CALLBACK = Ffm.unboundDowncall(C.PTR, C.PTR);
+ VOID_CALLBACK = Ffm.unboundDowncall(C.VOID, C.PTR);
+ } 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);
+ }
+ }
+ }
+
+ 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 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) {
+ 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,
EngineStreamCallbacks.GET_SCHEMA);
Review Comment:
Fixed in d93a1ab: `wrapEngineStream` reads the four addresses of
`EngineStreamCallbacks` into locals first, before it allocates or records
anything, so a failure to create the callbacks leaves the stream as it is:
`releaseOnFailure` releases the stream of the library and closes the task, and
no `EngineStream` stays in `engineStreams`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -351,17 +378,29 @@ 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. 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 d93a1ab: `stop()` also unpersists the broadcast of the last
micro-batch (`blocking = false`, as when the next range is planned) before it
closes the handle, in `Utils.tryWithSafeFinally` so that the handle is closed
even if the unpersist fails.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1158,234 @@ 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.
+ /**
+ * The callbacks of the streams that Spark wraps for engines, see {@link
EngineStream}, and the
+ * method handles that call the callbacks of the wrapped streams. 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;
+
+ static {
+ try {
+ MethodHandles.Lookup lookup = MethodHandles.lookup();
+ 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 = Ffm.unboundDowncall(C.I32, C.PTR, C.PTR);
+ POINTER_CALLBACK = Ffm.unboundDowncall(C.PTR, C.PTR);
+ VOID_CALLBACK = Ffm.unboundDowncall(C.VOID, C.PTR);
+ } 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);
+ }
+ }
+ }
+
+ 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;
Review Comment:
Done in d93a1ab: `sparkError` is `volatile`, and its comment says that the
callbacks may run on different threads one after the other.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -886,10 +986,132 @@ 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 = {
+ // A task that failed or was killed is over too, also while its listeners
run.
+ val context = TaskContext.get()
+ if (context == null || context.isInterrupted() || context.isFailed() ||
Review Comment:
Done in d93a1ab: after the check of `null`, `openEngineStream` calls
`context.killTaskIfInterrupted()`, so a killed task gets the
`TaskKilledException` of its kill, with its reason, while `isFailed()` and
`isCompleted()` stay in the check. The test sets the context of a killed task
on the thread and checks the reason.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsArrowStreamRead.java:
##########
@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector.read;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * A mix-in interface for a {@link PartitionReaderFactory} that also reads a
partition as an
+ * <a href="https://arrow.apache.org/docs/format/CStreamInterface.html";>Arrow
C stream</a>, for an
+ * engine that processes Arrow data in native code: the engine imports the
arrays of the stream
+ * itself, instead of reading the {@link
org.apache.spark.sql.vectorized.ColumnarBatch}es of
+ * {@link #createColumnarReader(InputPartition)}. The reader factories of the
batch and
+ * micro-batch scans of native data sources implement it, with the stream that
the library
+ * returns.
+ * <p>
+ * Each array of the stream is a struct array whose fields are the columns of
the read schema of
+ * the scan. Spark only checks that the source initialized the callbacks of
the stream, not its
+ * arrays, while it checks the arrays of the streams that it reads itself, so
the engine checks
+ * what its import of the arrays requires, such as the schema of the stream
against the read
+ * schema, and the offsets and the lengths of the arrays. The callbacks of the
stream can be called
+ * from any thread, but never concurrently. They call into the JVM, so the
engine releases the
+ * stream before the JVM exits.
+ *
+ * @since 4.4.0
+ */
+@Evolving
+public interface SupportsArrowStreamRead extends PartitionReaderFactory {
+
+ /**
+ * Initializes an Arrow C stream with the data of a partition. It is called
on the thread of the
+ * running task that reads the partition, and fails otherwise. The caller
reads the stream in
Review Comment:
Done in d93a1ab: the Javadoc now says "It is called in the running task that
reads the partition, on a thread that has its `TaskContext`, and fails
otherwise", and the error of the check says it the same way.
--
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]
