HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4475108066
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,37 +401,117 @@ 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)
+ val bytes = bytesOrEmpty(Option(planned.state))
+ // A small state goes with each task, while the executors fetch a
larger one once for the
+ // micro-batch, as a broadcast. Spark only plans the next range once
the micro-batch of the
+ // previous one finished, so the copies of the previous broadcast on
the executors are
+ // removed then. The broadcast itself is destroyed once its
micro-batch is committed,
+ // whose partitions are not read anymore then, also by an engine.
+ uncommittedBroadcasts.lastOption.foreach(_._2.unpersist(blocking =
false))
+ val state = if (bytes.length >
NativeDataSourceStreamReader.maxInlineStateBytes &&
+ partitions.nonEmpty) {
+ val broadcast = SparkSession.active.sparkContext.broadcast(bytes)
+ uncommittedBroadcasts.remove(end).foreach(_.destroy(blocking =
false))
+ uncommittedBroadcasts.put(end, broadcast)
+ Right(broadcast)
+ } else {
+ Left(bytes)
+ }
+ 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.get
}
- override def read(partition: InputPartition): PartitionReader[ColumnarBatch]
= {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, schema)
+ override def readerFactory(): PartitionReaderFactory = {
+ new NativeMicroBatchPartitionReaderFactory(location, name, schema,
timeZoneId)
}
- override def commit(end: String): Unit = {
+ override def commit(end: String): Unit = synchronized {
+ // The broadcasts of the micro-batches up to the committed one, which are
committed too.
+ if (uncommittedBroadcasts.contains(end)) {
Review Comment:
Fixed in 9711532: the stream reader records every range that it plans by its
end, in the order in which they were planned, with the broadcast of its state
if it has one, and `commit(end)` destroys the broadcasts of all the ranges up
to and including `end`, also when the committed range kept its state inline. As
you say, `AsyncProgressTrackingMicroBatchExecution` only commits the sources of
the batches whose offsets reached durable storage. The test plans three ranges
with broadcasts, commits the first one, and then commits a later range without
partitions, which destroys the broadcasts of the other two.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -886,10 +1047,150 @@ 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.
+ *
+ * @param projection the columns of `schema` that the engine gets, in their
order, if not all of
+ * them: the schema and the arrays of the stream then only
have these children
+ */
+ private[ffi] def openEngineStream[P <: InputPartition: ClassTag](
+ location: NativeLibraryLocation,
+ name: String,
+ kind: Int,
+ partition: InputPartition,
+ schema: StructType,
+ timeZoneId: String,
+ streamAddress: Long,
+ projection: Option[Array[Int]])(
+ 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 {
+ // As a failure of read, a failure to create the callbacks of the
streams of engines, such as
+ // a link error, is an error of the data source.
+ call(name, "read") {
+ NativeLibrary.wrapEngineStream(streamAddress, task, name,
projection.orNull)
+ }
+ } catch {
+ case e: Throwable =>
+ releaseOnFailure(e, streamAddress, task)
+ throw e
+ }
+ }
+
+ /**
+ * Initializes the Arrow C stream at `streamAddress`, whose release callback
is NULL, with the
+ * data of a partition, by calling read of the library in the running task.
Returns the context
+ * of the task, which the library may use until the stream is released: it
is closed then. If
+ * this fails, the stream is released.
+ *
+ * @param partition the partition, whose index in the scan the library gets
as `partition_id`,
+ * rather than the partition of the task: a task may read
several partitions of
+ * a scan, or of several scans, as after a coalesce or a
union
+ */
+ private def openStream(
+ location: NativeLibraryLocation,
+ name: String,
+ kind: Int,
+ readerState: Array[Byte],
+ partition: NativeInputPartition,
+ schema: StructType,
+ timeZoneId: String,
+ streamAddress: Long): NativeLibrary.Task = {
+ val library = NativeLibraries.get(location)
+ val stream = ArrowArrayStream.wrap(streamAddress)
+ val task = NativeTask.create(Some(partition.index), read = true)
+ try {
+ NativeArrow.exportSchema(schema, timeZoneId) { schemaAddress =>
+ call(name, "read") {
+ library.read(task, kind, schemaAddress, readerState,
partition.bytes, streamAddress)
+ }
+ }
+ // Arrow calls the callbacks of the stream without checking them, and
calling NULL would
+ // crash the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 ||
snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream,
or a callback of " +
+ "the stream is NULL.")
+ }
+ task
+ } catch {
+ case e: Throwable =>
+ releaseOnFailure(e, streamAddress, task)
+ throw e
+ }
+ }
+
+ /**
+ * After `failure`, releases the stream at `streamAddress`, if the library
initialized it, before
+ * the task that it may use, and marks it released in case its release
callback did not. The
+ * stream is released through FFM, as the streams of engines are, without
the JNI library of
+ * Arrow. If it cannot be released, the task is not closed either, since the
library may still
+ * use it: both then leak, rather than letting the release of the stream use
a freed task. A
+ * failure to clean up is added to `failure`.
+ */
+ private def releaseOnFailure(
+ failure: Throwable,
+ streamAddress: Long,
+ task: NativeLibrary.Task): Unit = {
+ val released = try {
+ NativeLibrary.releaseStream(streamAddress)
+ true
+ } catch {
+ case e: Throwable =>
+ failure.addSuppressed(e)
+ false
+ }
+ if (released) {
Review Comment:
Fixed in 9711532: `openStream` hands the task to its caller before `read`,
as the reader of master set its field, and no longer releases the stream
itself. So when it fails, the reader closes as in master: `closeResources`
releases the stream, and then closes the task. The path of engines keeps the
task in `openEngineStream` for `releaseOnFailure`, which closes it once it
released the stream.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1169,432 @@ 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}, of the
+ * schemas and the arrays that it projects for them, 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;
+ static final long SCHEMA_RELEASE;
+ static final long ARRAY_RELEASE;
+ // 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();
+ 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);
+ SCHEMA_RELEASE = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "projectedSchemaRelease",
+ MethodType.methodType(void.class, long.class)),
+ C.VOID, C.PTR);
+ ARRAY_RELEASE = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "projectedArrayRelease",
+ MethodType.methodType(void.class, long.class)),
+ 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);
+ }
+ }
+ }
+
+ /**
+ * The method handles that call the callbacks of the structs of the Arrow C
data and stream
+ * interfaces, with the address of the function first: get_schema and
get_next, get_last_error,
+ * and release. They are apart from {@link EngineStreamCallbacks}, so that
Spark can still release
+ * a stream if the upcalls of the callbacks cannot be created.
+ */
+ private static final class StreamCalls {
+ static final MethodHandle INT_CALLBACK;
+ static final MethodHandle POINTER_CALLBACK;
+ static final MethodHandle VOID_CALLBACK;
+
+ static {
+ try {
+ // Linker.downcallHandle(FunctionDescriptor, Linker.Option...), whose
method handles take
+ // the address of the function to call as their first argument.
+ MethodHandle unboundDowncall = MethodHandles.lookup().findVirtual(
+ Class.forName("java.lang.foreign.Linker"), "downcallHandle",
+ MethodType.methodType(MethodHandle.class,
+ Class.forName("java.lang.foreign.FunctionDescriptor"),
Ffm.NO_OPTIONS.getClass()))
+ .asFixedArity();
+ 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);
+ } catch (Throwable e) {
+ throw new IllegalStateException(
+ "Cannot call the callbacks of Arrow C streams: " + 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);
+ }
+ }
+
+ /**
+ * Releases the stream at {@code stream} with its release callback, if it is
not released yet,
+ * and marks it released, as the C stream interface requires of release, in
case the callback
+ * did not. It calls the callback through FFM, as the rest of the streams of
engines do, rather
+ * than through the JNI library of Arrow, which a JVM that only reads
through engines may not
+ * have loaded.
+ */
+ static void releaseStream(long stream) throws Throwable {
Review Comment:
Right, thank you: `exportSchema` loads the JNI library of Arrow before
`read`. In 9711532, `releaseOnFailure` releases the stream with
`ArrowArrayStream.wrap(...).release()` again, as `closeResources` does, and
`NativeLibrary.releaseStream` and the separate holder of the downcalls are
removed. If the release fails, the task still stays open, as in your round 9
comment.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala:
##########
@@ -132,6 +158,12 @@ trait ColumnarWriter extends Serializable {
def commit(messages: Array[WriterCommitMessage]): Unit
def abort(messages: Array[WriterCommitMessage]): Unit
+
+ /**
+ * The time zone of the timestamps of the Arrow batches that the data
writers get, read on the
+ * driver: by default, the session time zone when the data writers are
planned.
+ */
+ def timeZoneId: String = SQLConf.get.sessionLocalTimeZone
Review Comment:
Done in #59931 (4a76b34), which is merged here in d5b1e3d:
`ColumnarWriter.timeZoneId` and `ColumnarStreamWriter.timeZoneId` are removed,
and `ArrowBatchingOptions(schema)` is as in master. I had added them for item
11 of your third round on #59931, to give the arrays of `write` and the schema
of `create_writer` a single source, but as you say, an exported array carries
no time zone, and the values are the same.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1169,432 @@ 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}, of the
+ * schemas and the arrays that it projects for them, 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;
+ static final long SCHEMA_RELEASE;
+ static final long ARRAY_RELEASE;
+ // 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();
+ 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);
+ SCHEMA_RELEASE = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "projectedSchemaRelease",
+ MethodType.methodType(void.class, long.class)),
+ C.VOID, C.PTR);
+ ARRAY_RELEASE = Ffm.upcall(
+ lookup.findStatic(NativeLibrary.class, "projectedArrayRelease",
+ MethodType.methodType(void.class, long.class)),
+ 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);
+ }
+ }
+ }
+
+ /**
+ * The method handles that call the callbacks of the structs of the Arrow C
data and stream
+ * interfaces, with the address of the function first: get_schema and
get_next, get_last_error,
+ * and release. They are apart from {@link EngineStreamCallbacks}, so that
Spark can still release
+ * a stream if the upcalls of the callbacks cannot be created.
+ */
+ private static final class StreamCalls {
+ static final MethodHandle INT_CALLBACK;
+ static final MethodHandle POINTER_CALLBACK;
+ static final MethodHandle VOID_CALLBACK;
+
+ static {
+ try {
+ // Linker.downcallHandle(FunctionDescriptor, Linker.Option...), whose
method handles take
+ // the address of the function to call as their first argument.
+ MethodHandle unboundDowncall = MethodHandles.lookup().findVirtual(
+ Class.forName("java.lang.foreign.Linker"), "downcallHandle",
+ MethodType.methodType(MethodHandle.class,
+ Class.forName("java.lang.foreign.FunctionDescriptor"),
Ffm.NO_OPTIONS.getClass()))
+ .asFixedArity();
+ 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);
+ } catch (Throwable e) {
+ throw new IllegalStateException(
+ "Cannot call the callbacks of Arrow C streams: " + 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);
+ }
+ }
+
+ /**
+ * Releases the stream at {@code stream} with its release callback, if it is
not released yet,
+ * and marks it released, as the C stream interface requires of release, in
case the callback
+ * did not. It calls the callback through FFM, as the rest of the streams of
engines do, rather
+ * than through the JNI library of Arrow, which a JVM that only reads
through engines may not
+ * have loaded.
+ */
+ static void releaseStream(long stream) throws Throwable {
+ long release = Platform.getLong(null, stream + STREAM_RELEASE);
+ if (release != 0) {
+ StreamCalls.VOID_CALLBACK.invokeExact(release, stream);
+ Platform.putLong(null, stream + STREAM_RELEASE, 0);
+ }
+ }
+
+ 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 children of the schema and of the arrays of the stream of the
library that the stream of
+ // the engine has, in their order, or null for all of them, see project.
+ final int[] projection;
+ // 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, int[] projection) {
+ this.inner = inner;
+ this.task = task;
+ this.name = name;
+ this.projection = projection;
+ }
+
+ 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();
+ sparkError = toCString("The " + callback + " of the stream of the
native data source " +
+ name + " failed: " + rootCause(e));
+ 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));
+ } 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
+ * @param projection the children of the schema and of the arrays of the
stream that the engine
+ * gets, in their order, or null for all of them
+ */
+ static void wrapEngineStream(long stream, Task task, String name, int[]
projection) {
+ // Creates the callbacks first, so that the stream is left as it is if
that fails.
+ long getSchema;
+ long getNext;
+ long getLastError;
+ long release;
+ try {
+ getSchema = EngineStreamCallbacks.GET_SCHEMA;
+ getNext = EngineStreamCallbacks.GET_NEXT;
+ getLastError = EngineStreamCallbacks.GET_LAST_ERROR;
+ release = EngineStreamCallbacks.RELEASE;
+ // Creates the method handles of the callbacks too.
+ Objects.requireNonNull(StreamCalls.INT_CALLBACK);
+ } catch (ExceptionInInitializerError | NoClassDefFoundError e) {
+ // The first use of a class that failed to initialize gets the error,
and the later ones a
+ // NoClassDefFoundError, whose cause has its message: both name the root
cause here, as the
+ // errors of load do.
+ throw new IllegalStateException(
+ "Cannot create the callbacks of the streams of engines: " +
rootCause(e), e);
+ }
+ 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, projection));
+ } catch (Throwable e) {
+ engineStreams.remove(id);
+ Platform.freeMemory(inner);
+ throw e;
+ }
+ Platform.putLong(null, stream + STREAM_GET_SCHEMA, getSchema);
+ Platform.putLong(null, stream + STREAM_GET_NEXT, getNext);
+ Platform.putLong(null, stream + STREAM_GET_LAST_ERROR, getLastError);
+ Platform.putLong(null, stream + STREAM_RELEASE, 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();
+ int result = (int) StreamCalls.INT_CALLBACK.invokeExact(
+ wrapped.callback(offset), wrapped.inner, out);
+ if (result == 0 && wrapped.projection != null) {
+ project(out, offset == STREAM_GET_SCHEMA, wrapped.projection);
+ }
+ return result;
+ } catch (Throwable e) {
+ return wrapped == null ? EIO : wrapped.fail(callback, e);
+ }
+ }
+
+ /**
+ * Projects the struct at {@code out}, the schema or an array that the
stream of a library
+ * returned, in place, to its children at {@code projection}, for a scan
that reads only some of
+ * the columns from a library that does not prune the others. The struct of
the library is moved
+ * to memory of Spark, and {@code out} keeps its fields, but has the
projected children and a
+ * release of Spark, which releases the moved struct of the library, and so
all its children.
+ * The children are not copied: an engine that moves a projected child marks
it released where
+ * the library keeps it, as the Arrow C data interface requires. A released
array, the end of
+ * the stream, stays as it is. If the projection fails, the struct of the
library is released,
+ * so that the engine gets a released struct with the error.
+ */
+ private static void project(long out, boolean schema, int[] projection)
throws Throwable {
Review Comment:
Done with your first option, in #59931 (4a76b34), which is merged here:
`read` returns exactly the columns of `SparkDsReadInfo.schema`, which are the
columns that the scan reads, also when the library does not implement
`prune_columns`. `prune_columns` only lets the library plan to read them, so
its `pruned` out-parameter is removed while v1 is still open. Both projections
of Spark are gone: `ColumnarProjectionReader`, `ProjectedColumnarBatch` and
`ColumnarProjectionReaderFactory` in #59931, and the projection of the streams
of engines with `SupportsProjectedArrowStreamRead`,
`ColumnarArrowStreamProjectionReaderFactory` and `EngineStream.project` here,
in 9711532. Spark's own reader already rejects a stream with other columns
("The data has the schema ..., but the expected schema is ..."), which a new
test checks with a library that ignores its schema to read.
--
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]
