dongjoon-hyun commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4474718389
##########
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:
This destroys the earlier broadcasts only when `end` itself has a broadcast.
If the committed range had an inline state (`<= maxInlineStateBytes`), the
broadcasts of the earlier ranges are kept on the driver until `stop()`,
although `commit(end)` means that all data up to `end` was processed.
This can happen when Spark skips the commits of intermediate batches, e.g.,
`AsyncProgressTrackingMicroBatchExecution` commits the sources only for the
batches whose progress was durably written. For example:
- batch 1, 2: large states -> broadcasts `b1`, `b2`
- batch 3: small state -> inline
- only `commit(e3)` is called -> `contains(e3)` is `false`, so `b1` and `b2`
are not destroyed.
`unpersist` in `partitions()` only removes the executor copies of the last
broadcast, so the driver blocks remain. For a long-running query whose state
size goes back and forth around 64 KiB, the driver memory keeps growing. How
about recording every planned `end` in order (including the inline ones) and
destroying all broadcasts planned up to the committed `end`?
##########
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:
This adds a fair amount of machinery to Spark (`EngineStream.project`,
`projectedRelease`, `SupportsProjectedArrowStreamRead`,
`ColumnarArrowStreamProjectionReaderFactory`, `ColumnarProjectionReader`) only
for the libraries that do not implement `prune_columns`. It also adds a native
allocation and a release upcall per `get_next`, and Spark has to keep the
tricky invariants of moving the library's struct while sharing its children.
Since we are still finalizing the v1 ABI, could we address this in the
contract instead? `read` already gets `SparkDsReadInfo.schema`, so we could
either
- require `read` to return only the columns of `info->schema`, or
- make `prune_columns` mandatory for `BATCH_READ`.
Either way, both Spark-side projection paths could be removed.
##########
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:
Is this needed? `NativeDataWriter.write` exports only the `ArrowArray` (the
data writer got the schema when it was created), and an array does not carry
the time zone of its timestamps. The only time zone that the library sees is
the one of the schema of `create_data_writer` (`NativeDataWriter.timeZoneId`).
The values of the timestamps are also the same UTC microseconds regardless of
the time zone of the Arrow type.
For the other `ColumnarWriter`s, the default is
`SQLConf.get.sessionLocalTimeZone`, which is what
`ArrowBatchingOptions(schema)` used before. So the two trait members, the
overrides in `NativeDataSourceWriter` / `NativeDataSourceStreamWriter`, and the
new parameter of `ArrowBatchingOptions.apply` seem to add API surface without a
behavior change. Could we keep `ArrowBatchingOptions(schema)` as before?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -222,9 +346,71 @@ abstract class ColumnarOnlyPartitionReaderFactory extends
PartitionReaderFactory
override def supportColumnarReads(partition: InputPartition): Boolean = true
}
-class ColumnarStreamPartitionReaderFactory(reader: ColumnarStreamReader)
+/**
+ * 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)
+ }
+}
+
+/**
+ * A reader factory that opens the Arrow C streams of the partitions for an
engine, see
+ * [[SupportsArrowStreamRead]], and can project them to some of their columns.
+ */
+trait SupportsProjectedArrowStreamRead extends SupportsArrowStreamRead {
+ /**
+ * Opens the stream of `partition`, as `openArrowStream(partition,
streamAddress)` does, but its
+ * schema and its arrays only have the columns at `projection`, in that
order.
+ */
+ def openArrowStream(partition: InputPartition, streamAddress: Long,
projection: Array[Int]): Unit
+}
+
+/**
+ * A [[ColumnarProjectionReaderFactory]] whose Arrow C streams, which an
engine opens, only have the
+ * columns that the scan reads too, as its batches do.
+ */
+class ColumnarArrowStreamProjectionReaderFactory(
+ factory: SupportsProjectedArrowStreamRead,
+ ordinals: Array[Int])
+ extends ColumnarProjectionReaderFactory(factory, ordinals) with
SupportsArrowStreamRead {
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): Unit = {
+ factory.openArrowStream(partition, streamAddress, ordinals)
}
}
+
+private class ColumnarProjectionReader(
+ reader: PartitionReader[ColumnarBatch],
+ ordinals: Array[Int])
+ extends PartitionReader[ColumnarBatch] {
+
+ // The projection of the current batch of the reader, which `get` returns
until `next`.
+ private var batch: ColumnarBatch = _
+
+ override def next(): Boolean = {
+ val hasNext = reader.next()
+ batch = if (hasNext) new ProjectedColumnarBatch(reader.get(), ordinals)
else null
Review Comment:
nit. This allocates a new `ProjectedColumnarBatch` for every batch, and the
constructor of `ColumnarBatch` allocates a `ColumnVector` array and a
`ColumnarBatchRow` too, even when the underlying reader returns the same batch
instance every time, as `NativePartitionReader` does. When `reader.get()`
returns the same instance, we could reuse the projected batch and only update
its `numRows`.
##########
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:
Since `openStream` was extracted from `NativePartitionReader`, `task` is a
local variable here, which causes a leak when `NativeLibrary.releaseStream`
fails, e.g., when `StreamCalls` fails to initialize with
`ExceptionInInitializerError` / `NoClassDefFoundError`:
1. `releaseOnFailure` skips `task.close()` on purpose because the stream was
not released.
2. The constructor of `NativePartitionReader` then calls `close()`, and
`closeResources()` releases the same stream successfully through Arrow JNI
(`stream.release()`).
3. However, the `task` field is still `null` there, so `task.close()` is
never called.
As a result, the `SparkDsTaskContext` leaks for every failed partition,
although the stream is released in the end. In `master`, `task` was a field of
the reader assigned before `read`, so `closeResources()` closed it right after
the stream. Could you make sure that the task is closed after the stream is
released on this path?
##########
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)) {
Review Comment:
This per-range dedup (`plannedRange` / `plannedPartitions`), which prevents
planning the same micro-batch more than once, e.g., for a self-union, exists
only in `NativeDataSourceStreamReader`. The other `ColumnarStreamReader`
implementations (e.g., `TestRangeStreamReader`) still get `partitions` called
once per scan node of the same micro-batch, and may return a different plan for
each. Could this be generalized, e.g., by caching by `(start, end)` in
`ColumnarMicroBatchStream.planInputPartitions`, so that it applies to all
columnar stream readers?
##########
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:
The doc says that this goes through FFM because a JVM that only reads
through engines may not have loaded the JNI library of Arrow. However, the only
caller is `releaseOnFailure`, after `openStream` called
`NativeArrow.exportSchema` (`Data.exportSchema` and `ArrowSchema.release`), so
the Arrow JNI library is always loaded on this path.
If so, `ArrowArrayStream.wrap(streamAddress).release()`, which
`NativePartitionReader.closeResources` already uses, could be used instead.
That would remove the duplicated release logic and also the new dependency of
the failure path of Spark's own reader on the initialization of `StreamCalls`
(see the comment on `releaseOnFailure`).
##########
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 =>
Review Comment:
nit. `openStream` and `NativeDataWriter` run `ArrowUtils.toArrowSchema` in
`NativeArrow.exportSchema` for every partition / data writer on the executors.
When a task reads many small partitions, e.g., after a coalesce, the same
conversion is repeated. The export itself is needed every time since the
library may move the schema, but the Arrow `Schema` could be built once, e.g.,
as a `@transient lazy val` of the reader factory / writer.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -123,53 +132,71 @@ class ColumnarScanBuilder(
override def pruneColumns(requiredSchema: StructType): Unit = {
// Only top-level columns are pruned by the data source. Spark prunes the
nested fields of the
- // columns it reads.
+ // columns it reads. The required columns are only recorded here: a batch
scan asks its planner
+ // to prune the others, and a streaming scan, which has no planner, passes
the required
+ // columns to its stream reader.
val required =
StructType(requiredSchema.fieldNames.flatMap(name => schema.find(_.name
== name)))
- if (required.length < schema.length &&
batchPlanner.pruneColumns(required)) {
- readSchema = required
+ if (required.length < schema.length) {
+ this.requiredSchema = Some(required)
}
}
// The scan takes over the planner, so it is built once.
private lazy val scan = new ColumnarScan(
- shortName, schema, readSchema, dataSource, Option(planner), pushedLimit)
+ shortName, schema, requiredSchema, dataSource, Option(planner),
pushedLimit)
override def build(): Scan = scan
}
/**
+ * A scan of a [[ColumnarDataSource]]. It reads the top-level columns that the
query needs, the
+ * columns of its read schema: a batch scan asks its planner to prune the
other columns when it
+ * plans its partitions, and drops them from the batches of the reader if the
planner does not,
+ * while the stream reader of a streaming scan reads only these columns.
+ *
+ * @param requiredSchema the top-level columns that the query needs, if fewer
than those of the
+ * full schema
* @param pushedDownPlanner the planner that the operations were pushed down
to, if any. Otherwise,
* the scan creates one when its batch is planned.
*/
class ColumnarScan(
shortName: String,
fullSchema: StructType,
- prunedSchema: StructType,
+ requiredSchema: Option[StructType],
dataSource: ColumnarDataSource,
pushedDownPlanner: Option[ColumnarBatchPlanner],
pushedLimit: Option[Int])
extends Scan with SupportsMetadata {
- // Spark may call `toBatch` more than once for the same scan; the partitions
are planned once.
- private lazy val batch = new ColumnarBatchScan(
- pushedDownPlanner.getOrElse(dataSource.batchPlanner(fullSchema)))
+ // Spark may call `toBatch` more than once for the same scan; the partitions
are planned once,
+ // after the planner prunes the columns, if it does. A failure is kept too,
so that the planner
+ // is asked at most once, as the header promises.
+ private lazy val batch: Try[ColumnarBatchScan] = Try {
+ val planner =
pushedDownPlanner.getOrElse(dataSource.batchPlanner(fullSchema))
+ // The ordinals in the full schema of the columns that the scan reads, if
the planner does
+ // not prune the others.
+ val projection = requiredSchema.filterNot(planner.pruneColumns)
Review Comment:
`planner.pruneColumns` has side effects (it updates the state of the planner
and calls `prune_columns` of the library), but it is used as the predicate of
`filterNot`. It works correctly now, but a pushdown call hidden in a filter is
easy to break in a future refactoring, e.g., evaluating it twice or replacing
it with `exists`, which would violate the "at most once" contract of the
header. How about making it explicit?
```scala
val projection = requiredSchema match {
case Some(required) if !planner.pruneColumns(required) =>
Some(required.fieldNames.map(fullSchema.fieldIndex))
case _ => None
}
```
##########
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)
Review Comment:
nit. `SparkSession.active` requires an active or default session on the
planning thread. If `planInputPartitions` is evaluated on a thread without one,
this fails with `No active or default Spark session found`, and the failure is
cached in `plannedPartitions` for the whole range. It is unlikely since most
paths propagate the session, but `SparkContext.getActive` or the `SparkContext`
captured when the reader is created would be more robust.
--
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]
