dongjoon-hyun commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4417661687
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1053,163 @@ 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;
+
+ 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 callbacks of the streams of the libraries, by address, for each
callback of a stream: a
+ // library typically has one of each, for all its streams.
+ private static final ConcurrentHashMap<Long, MethodHandle>
getSchemaCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> getNextCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle>
getLastErrorCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> releaseCallbacks =
+ 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 those of the library, 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;
+ final MethodHandle getSchema;
+ final MethodHandle getNext;
+ final MethodHandle getLastError;
+ final MethodHandle release;
+
+ EngineStream(long inner, Task task) throws Throwable {
+ this.inner = inner;
+ this.task = task;
+ getSchema = callback(getSchemaCallbacks, inner + STREAM_GET_SCHEMA,
C.I32, C.PTR, C.PTR);
+ getNext = callback(getNextCallbacks, inner + STREAM_GET_NEXT, C.I32,
C.PTR, C.PTR);
+ getLastError = callback(getLastErrorCallbacks, inner +
STREAM_GET_LAST_ERROR, C.PTR, C.PTR);
+ release = callback(releaseCallbacks, inner + STREAM_RELEASE, C.VOID,
C.PTR);
+ }
+
+ private static MethodHandle callback(
+ ConcurrentHashMap<Long, MethodHandle> callbacks,
+ long field,
+ C result,
+ C... parameters) throws Throwable {
+ long address = Platform.getLong(null, field);
+ MethodHandle handle = callbacks.get(address);
+ if (handle == null) {
+ handle = Ffm.downcall(Ffm.OF_ADDRESS.invoke(address), result,
parameters);
+ callbacks.putIfAbsent(address, handle);
+ }
Review Comment:
**2. [Low] The library's callbacks are bound once, and cached by address
without bound**
`EngineStream` resolves the four function pointers of the library's stream
when it wraps it, and keeps their downcall handles in four static maps keyed by
address, which are never cleared. This has two consequences:
- A library that creates its callbacks per stream, such as a binding that
makes a C closure for each stream with libffi, adds four entries to these maps
for each stream, and each `openArrowStream` builds four downcall handles.
- Each callback gets the struct as its argument, and nothing in the C stream
interface stops a producer from updating its own callbacks in `inner` after
`read` returns, for example to switch `get_next` once the stream is exhausted.
The wrapper keeps calling the functions it found at wrap time.
The FFM API can also bind a signature without an address:
`Linker.downcallHandle(FunctionDescriptor, Linker.Option...)` returns a handle
whose first argument is the function to call. With three such handles in `Ffm`,
since `get_schema` and `get_next` have the same signature, each upcall could
read the current pointer from `inner` and call it. That removes the four maps,
`callback`, and the handle fields of `EngineStream`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -278,37 +284,69 @@ class NativeBatchPlanner(
override def plan(): ColumnarPlannedBatch = {
try {
- val partitions = call(name, "plan a scan
of")(library.partitions(handle.get))
- val readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeReader(handle.get))))
+ val planned = call(name, "plan a scan
of")(library.partitions(handle.get))
ColumnarPlannedBatch(
- NativeInputPartition.from(partitions),
- new NativePartitionReaderFactory(executorLocation, name, readerState,
readSchema))
+ NativeInputPartition.from(planned.partitions),
+ new NativePartitionReaderFactory(executorLocation, name,
+ bytesOrEmpty(Option(planned.state)), readSchema, timeZoneId))
} finally {
- // The executors only need the serialized state.
+ // The executors only need the state that the reader returned.
handle.close()
}
}
}
-/** Reads the partitions of a planned batch scan of a native data source, on
the executors. */
+/**
+ * Reads the partitions of a planned batch scan of a native data source, on
the executors. An
+ * engine can also get the Arrow C stream of a partition from it, see
[[SupportsArrowStreamRead]].
+ */
class NativePartitionReaderFactory(
location: NativeLibraryLocation,
name: String,
readerState: Array[Byte],
- readSchema: StructType)
- extends ColumnarOnlyPartitionReaderFactory {
+ readSchema: StructType,
+ timeZoneId: String)
+ extends ColumnarOnlyPartitionReaderFactory with SupportsArrowStreamRead {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, readSchema)
+ new NativePartitionReader(location, name, NativeLibrary.READ_BATCH,
readerState,
+ partition.asInstanceOf[NativeInputPartition].bytes, readSchema,
timeZoneId)
+ }
+
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): Unit = {
+ val context = TaskContext.get()
+ if (context == null || context.isCompleted()) {
+ throw new IllegalStateException(
+ "openArrowStream must be called in the running task that reads the
partition.")
+ }
+ 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 nativePartition = partition.asInstanceOf[NativeInputPartition]
+ // The partition of the scan: the task of an engine may read others than
its own.
+ val task = NativePartitionReader.openStream(location, name,
NativeLibrary.READ_BATCH,
+ readerState, nativePartition.bytes, readSchema, timeZoneId,
streamAddress,
+ Some(nativePartition.index))
Review Comment:
**1. [Medium] `is_interrupted` no longer tells the library that the task is
over**
The context of an engine's stream reports `is_interrupted` only through
`TaskContext.isInterrupted()`, the supplier of `NativeTask.create`, which is
true only when the task is killed. So once the engine's task fails, or
completes, the context keeps reporting `false` until the engine releases the
stream. In 9bc17e2, the completion listener removed the context from
`NativeLibrary.tasks`, so `isInterrupted` reported `true` after the end of the
task. 8eae514 rightly drops the listener, but that signal goes with it.
This matters more now that the stream may be read on another thread than the
task's. For example, the engine's scan thread is blocked in the library's
`get_next`, waiting for remote data and polling `task->is_interrupted` as the
header asks. Another operator of the task throws, so the task fails without
being killed. `get_next` keeps waiting until the library's own timeout, and
since the callbacks never run concurrently, the engine cannot release the
stream before `get_next` returns. An engine whose cleanup on the task thread
waits for its scan thread then holds the task, and the task reaper does not
help, since the task was not killed. With Spark's own reader this cannot
happen, because `get_next` runs on the task thread.
Could the supplier of these contexts also report the end of the task, for
example `context.isInterrupted() || context.isFailed() ||
context.isCompleted()`? Both new calls are synchronized in `TaskContextImpl`,
so the library's threads can make them. Spark's own reader releases its stream
from a completion listener, after `completed` is set, so applying this only to
`openArrowStream` would leave that reader as it is.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -304,9 +308,36 @@ typedef struct SparkDsCreateInfo {
int64_t num_options;
} SparkDsCreateInfo;
+/* A micro-batch streaming scan, as create_stream_reader gets it. */
+typedef struct SparkDsStreamReaderInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /*
+ * The location of a directory in the checkpoint of the streaming query that
is reserved for this
+ * source, such as file:/tmp/checkpoint/sources/0, where the library may
keep what it needs to
+ * recover from a failure. It stays the same when the query restarts from
its checkpoint. It is a
+ * URI, possibly of a file system that the library cannot access, such as
s3a or hdfs.
+ */
+ const char* checkpoint_location;
+} SparkDsStreamReaderInfo;
+
+/* A request for the latest offset of a stream reader, as latest_offset gets
it. */
+typedef struct SparkDsLatestOffsetInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /*
+ * The offset, as JSON, that the next micro-batch starts after: the end of
the last micro-batch,
+ * or the initial offset before the first one. With Trigger.AvailableNow,
Spark asks for the
+ * latest offset only once, when the query starts, with the initial offset.
+ */
+ const char* start;
+} SparkDsLatestOffsetInfo;
+
/*
* The task that calls read or create_data_writer, on an executor. It stays
valid until the stream
- * that read returns is released, or until the data writer is closed.
+ * that read returns is released, or until the data writer is closed. An
engine that runs a query
+ * in native code may get the stream of read from Spark, and release it once
the task ended, from
+ * another thread: the task stays valid until then too.
Review Comment:
**4. [Low] The header could say that an engine may release the arrays after
the stream**
This tells the library that an engine may read and release the stream, but
not that an engine may also keep the arrays of `get_next` after it released the
stream. The C stream interface allows that, and engines do it: for example, a
record batch that arrow-rs imports from a stream owns its FFI arrays,
independently of the reader. Spark's own reader never does it, since its
`close()` closes the vectors, which releases every array, before it releases
the stream and closes the context.
So a library tested only with Spark can free the memory of its arrays in the
release of its stream, or use `task` in the release of its arrays, and work
until an engine reads it. Then the engine reads freed buffers, or the release
of an array uses the context that Spark freed when the stream was released. A
sentence here, or in the doc of `read`, saying that the arrays of the stream
may be released after the stream, and must not use the stream or `task`, would
warn library authors.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1053,163 @@ 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;
+
+ 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 callbacks of the streams of the libraries, by address, for each
callback of a stream: a
+ // library typically has one of each, for all its streams.
+ private static final ConcurrentHashMap<Long, MethodHandle>
getSchemaCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> getNextCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle>
getLastErrorCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> releaseCallbacks =
+ 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 those of the library, 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;
+ final MethodHandle getSchema;
+ final MethodHandle getNext;
+ final MethodHandle getLastError;
+ final MethodHandle release;
+
+ EngineStream(long inner, Task task) throws Throwable {
+ this.inner = inner;
+ this.task = task;
+ getSchema = callback(getSchemaCallbacks, inner + STREAM_GET_SCHEMA,
C.I32, C.PTR, C.PTR);
+ getNext = callback(getNextCallbacks, inner + STREAM_GET_NEXT, C.I32,
C.PTR, C.PTR);
+ getLastError = callback(getLastErrorCallbacks, inner +
STREAM_GET_LAST_ERROR, C.PTR, C.PTR);
+ release = callback(releaseCallbacks, inner + STREAM_RELEASE, C.VOID,
C.PTR);
+ }
+
+ private static MethodHandle callback(
+ ConcurrentHashMap<Long, MethodHandle> callbacks,
+ long field,
+ C result,
+ C... parameters) throws Throwable {
+ long address = Platform.getLong(null, field);
+ MethodHandle handle = callbacks.get(address);
+ if (handle == null) {
+ handle = Ffm.downcall(Ffm.OF_ADDRESS.invoke(address), result,
parameters);
+ callbacks.putIfAbsent(address, handle);
+ }
+ return handle;
+ }
+ }
+
+ /**
+ * 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.
+ */
+ static void wrapEngineStream(long stream, Task task) throws Throwable {
+ long inner = Platform.allocateMemory(STREAM_SIZE);
+ EngineStream wrapped;
+ try {
+ Platform.copyMemory(null, stream, null, inner, STREAM_SIZE);
+ wrapped = new EngineStream(inner, task);
+ } catch (Throwable e) {
+ Platform.freeMemory(inner);
+ throw e;
+ }
+ long id = nextEngineStreamId.getAndIncrement();
+ engineStreams.put(id, wrapped);
+ 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();
+ }
+
+ 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) {
+ try {
+ EngineStream wrapped = engineStream(stream);
+ return wrapped == null ? EINVAL : (int)
wrapped.getSchema.invokeExact(wrapped.inner, out);
+ } catch (Throwable e) {
+ LOGGER.error("The get_schema of the stream of a native data source
failed.", e);
+ return EIO;
+ }
+ }
+
+ private static int engineStreamGetNext(long stream, long out) {
+ try {
+ EngineStream wrapped = engineStream(stream);
+ return wrapped == null ? EINVAL : (int)
wrapped.getNext.invokeExact(wrapped.inner, out);
+ } catch (Throwable e) {
+ LOGGER.error("The get_next of the stream of a native data source
failed.", e);
Review Comment:
**6. [Nit] A failure to log in these upcalls would crash the JVM**
The comment above says that these callbacks must not throw, but their
`catch` blocks call `LOGGER.error`. An exception from the logger, such as an
`OutOfMemoryError` while formatting the stack trace, or an appender that
throws, leaves the upcall and terminates the executor instead of returning
`EIO`. The other upcalls avoid this: `log` swallows everything, `isInterrupted`
returns `false` without logging, and `collect` only records the error. Wrapping
the logging in its own `try`, in the four callbacks, would keep the guarantee.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsArrowStreamRead.java:
##########
@@ -0,0 +1,54 @@
+/*
+ * 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 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 does not check the stream, while it checks 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.
+ *
+ * @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
owns the stream, which
+ * it may move, and releases it, also after the task ended: Spark frees what
the stream uses,
+ * such as the context of the task, once the stream is released. A stream
that is never released
+ * leaks.
+ *
+ * @param partition a partition of the scan, as the scan planned it
+ * @param streamAddress the address of a {@code struct ArrowArrayStream}
that the caller
+ * allocated, and whose release callback is NULL, as it
is in a
+ * zero-initialized struct. If this method fails, the
stream is released.
Review Comment:
**8. [Nit] "If this method fails, the stream is released" does not hold for
a rejected stream**
When the release callback of the stream is not NULL, `openArrowStream`
throws `IllegalArgumentException` and leaves the stream as it is. That is
right, since it is not a stream of Spark, but an engine that relies on this
sentence skips its own release on that error path, and leaks the stream that it
passed by mistake. Something like "If this method fails after checking its
arguments, the stream is released; a stream whose release callback is not NULL
is left as it is" would match the code.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -525,34 +594,23 @@ typedef struct SparkDsLibrary {
* Reads a partition of a batch or streaming scan.
*
* task: the task that reads the partition, valid until the stream is
released
- * reader_state: the state returned by serialize_reader or
serialize_stream_reader, empty if
- * there is none
- * partition: a partition returned by partitions or stream_partitions
+ * info: the partition, the kind of the scan, the schema to read and the
state of the reader
* out: the stream to initialize. Each array of the stream is a struct array
whose fields are the
- * columns to read: the columns of the schema to read, or the pruned
columns. The
- * callbacks of the stream can be called from any thread, but never
concurrently. get_next
- * may block, for example on network I/O, but Spark cannot interrupt
native code: when the
- * task is killed, Spark only stops reading once get_next returns, and
then releases the
- * stream. So a get_next that blocks forever keeps the task running,
unless the task reaper
- * (spark.task.reaper.enabled and spark.task.reaper.killTimeout) kills
the executor. Use
- * timeouts for blocking operations, and check task->is_interrupted.
- */
- int (*read)(const SparkDsTaskContext* task, const uint8_t* reader_state,
- int64_t reader_state_length, const uint8_t* partition, int64_t
partition_length,
+ * columns of the schema to read, info->schema. The callbacks of the
stream can be called
+ * from any thread, but never concurrently. get_next may block, for
example on network
+ * I/O, but Spark cannot interrupt native code: when the task is
killed, Spark only stops
+ * reading once get_next returns, and then releases the stream. So a
get_next that blocks
+ * forever keeps the task running, unless the task reaper
(spark.task.reaper.enabled and
Review Comment:
**9. [Nit] The doc of `read` still says that Spark reads and releases the
stream**
This says that when the task is killed, Spark stops reading once `get_next`
returns and then releases the stream, and that a `get_next` that blocks forever
keeps the task running. Both hold for Spark's own reader. Under
`SupportsArrowStreamRead`, the engine reads and releases the stream, possibly
on another thread and after the task ended, so neither is guaranteed.
Qualifying this with "When Spark reads the stream itself", and pointing to the
comment of `SparkDsTaskContext`, would keep the two consistent.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1145,6 +1222,105 @@ 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"))
+ val address = inTask {
+ val stream = Platform.allocateMemory(40)
+ Platform.setMemory(stream, 0.toByte, 40)
+ laterFactory.openArrowStream(laterPartitions.head, stream)
+ stream
+ }
+ assert(NativeLibrary.openEngineStreams() == 1 && !released.exists())
+ assert(readIds(ArrowArrayStream.wrap(address), ArrowUtils.rootAllocator)
== (0L until 3L))
+ assert(NativeLibrary.openEngineStreams() == 0 && released.exists())
Review Comment:
**3. [Low] The test does not check that the context outlives the task**
This part tests what 8eae514 fixes, but nothing in it fails if Spark frees
the context at the end of the task again. The test library uses `task` in
`get_next` only with `wait_for_interrupt`, and `openEngineStreams()` counts the
`EngineStream`s, not the contexts. If a completion listener that closes the
`Task` came back, `readIds` would still return the ids 0 to 2, the count would
still go from 1 to 0, and `released` would still be created.
A hook like `openEngineStreams()` for the contexts, the size of
`NativeLibrary.tasks`, checked before and after `readIds`, would cover it. A
test library whose `get_next` and `release` call `task->is_interrupted(task)`
would exercise the context too, although reading a freed context does not
reliably fail.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1053,163 @@ 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;
+
+ 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 callbacks of the streams of the libraries, by address, for each
callback of a stream: a
+ // library typically has one of each, for all its streams.
+ private static final ConcurrentHashMap<Long, MethodHandle>
getSchemaCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> getNextCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle>
getLastErrorCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> releaseCallbacks =
+ 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 those of the library, 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;
+ final MethodHandle getSchema;
+ final MethodHandle getNext;
+ final MethodHandle getLastError;
+ final MethodHandle release;
+
+ EngineStream(long inner, Task task) throws Throwable {
+ this.inner = inner;
+ this.task = task;
+ getSchema = callback(getSchemaCallbacks, inner + STREAM_GET_SCHEMA,
C.I32, C.PTR, C.PTR);
+ getNext = callback(getNextCallbacks, inner + STREAM_GET_NEXT, C.I32,
C.PTR, C.PTR);
+ getLastError = callback(getLastErrorCallbacks, inner +
STREAM_GET_LAST_ERROR, C.PTR, C.PTR);
+ release = callback(releaseCallbacks, inner + STREAM_RELEASE, C.VOID,
C.PTR);
+ }
+
+ private static MethodHandle callback(
+ ConcurrentHashMap<Long, MethodHandle> callbacks,
+ long field,
+ C result,
+ C... parameters) throws Throwable {
+ long address = Platform.getLong(null, field);
+ MethodHandle handle = callbacks.get(address);
+ if (handle == null) {
+ handle = Ffm.downcall(Ffm.OF_ADDRESS.invoke(address), result,
parameters);
+ callbacks.putIfAbsent(address, handle);
+ }
+ return handle;
+ }
+ }
+
+ /**
+ * 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.
+ */
+ static void wrapEngineStream(long stream, Task task) throws Throwable {
+ long inner = Platform.allocateMemory(STREAM_SIZE);
+ EngineStream wrapped;
+ try {
+ Platform.copyMemory(null, stream, null, inner, STREAM_SIZE);
+ wrapped = new EngineStream(inner, task);
+ } catch (Throwable e) {
+ Platform.freeMemory(inner);
+ throw e;
+ }
+ long id = nextEngineStreamId.getAndIncrement();
+ engineStreams.put(id, wrapped);
+ 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();
+ }
+
+ 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) {
+ try {
+ EngineStream wrapped = engineStream(stream);
+ return wrapped == null ? EINVAL : (int)
wrapped.getSchema.invokeExact(wrapped.inner, out);
+ } catch (Throwable e) {
+ LOGGER.error("The get_schema of the stream of a native data source
failed.", e);
+ return EIO;
+ }
+ }
+
+ private static int engineStreamGetNext(long stream, long out) {
+ try {
+ EngineStream wrapped = engineStream(stream);
+ return wrapped == null ? EINVAL : (int)
wrapped.getNext.invokeExact(wrapped.inner, out);
+ } catch (Throwable e) {
+ LOGGER.error("The get_next of the stream of a native data source
failed.", e);
+ return EIO;
+ }
+ }
+
+ private static long engineStreamGetLastError(long stream) {
+ try {
+ EngineStream wrapped = engineStream(stream);
+ return wrapped == null ? 0 : (long)
wrapped.getLastError.invokeExact(wrapped.inner);
Review Comment:
**7. [Nit] `get_last_error` does not describe the errors that Spark returns**
When the wrapper fails by itself, with `EINVAL` for a stream that it does
not know, or with `EIO` after an exception, this still returns the library's
`get_last_error`, which describes an earlier error of the library, if any, or
is NULL. Arrow C++ and Arrow Java call `get_last_error` after a non-zero code,
so the engine reports an unrelated message, or none. A message kept in
`EngineStream` for the errors of Spark, and returned here instead, would make
them accurate. For a stream that is not known, a static message would do.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1053,163 @@ 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;
+
+ 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 callbacks of the streams of the libraries, by address, for each
callback of a stream: a
+ // library typically has one of each, for all its streams.
+ private static final ConcurrentHashMap<Long, MethodHandle>
getSchemaCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> getNextCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle>
getLastErrorCallbacks =
+ new ConcurrentHashMap<>();
+ private static final ConcurrentHashMap<Long, MethodHandle> releaseCallbacks =
+ 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 those of the library, 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;
Review Comment:
**5. [Low] An unreleased stream keeps the whole `TaskContext` of its task**
`task` keeps its entry in `NativeLibrary.tasks`, whose supplier from
`NativeTask.create` captures the `TaskContext`. So each stream that the engine
has not released yet keeps the `TaskContextImpl` of its task after the task
ended: its `TaskMemoryManager`, whose page table alone has 8192 entries, its
`TaskMetrics` with the accumulators of the task, and its local properties, such
as the job description. Before this PR, the context of a read was dropped with
its task. An engine that keeps streams after its tasks, as the API now allows,
or leaks some on an error path, then keeps much more than the 40 bytes of the
context.
A supplier that holds the `TaskContext` through a `WeakReference`, and
reports `true` once it is gone, as in item 1, would drop the reference.
--
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]
