dongjoon-hyun commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4418419351
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -278,39 +281,52 @@ 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)
+ }
Review Comment:
**1. [Medium] Spark's reader and an engine's stream give the same partition
different `partition_id`s**
Since 8eae514, the context of an engine's stream gets the index of the
partition in the scan, `NativeInputPartition.index`, as discussed in [round
1](https://github.com/apache/spark/pull/59934#discussion_r4417274552). But
Spark's own reader, here and in `NativeDataSourceStreamReader.read`, still
creates its context with `NativeTask.create()`, which takes
`TaskContext.partitionId()`. The two agree for a plain scan, where the
`DataSourceRDD` has one partition for each input partition, but not when the
RDD of the scan is part of a larger one:
- With `df1.union(df2)` of two native scans of two partitions each, Spark's
reader gives the partitions of `df2` the `partition_id`s 2 and 3, the indexes
of the `UnionRDD`, while an engine's streams give them 0 and 1.
- With `coalesce(1)`, Spark's reader gives every partition the
`partition_id` 0, since one task reads them all.
So a library that names files or other resources after `partition_id`, as
`task_path` of the test library does, behaves differently depending on whether
an engine runs the query, and with Spark's reader two partitions of one scan
can even get the same `partition_id`. The reason given in `openEngineStream`,
that a task may read other partitions than its own, applies to Spark's reader
under `union` and `coalesce` too.
Since #59931 settles version 1 of the interface, could both paths pass the
index of the partition in the scan, for example by passing
`NativeInputPartition.index` to `NativePartitionReader`, and on to
`openStream`? If the partition of the task is intended for Spark's reader, the
doc of `partition_id` in the header, "The partition that the task reads or
writes", could say which one each path passes.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsArrowStreamRead.java:
##########
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector.read;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * A mix-in interface for a {@link PartitionReaderFactory} that also reads a
partition as an
+ * <a href="https://arrow.apache.org/docs/format/CStreamInterface.html";>Arrow
C stream</a>, for an
+ * engine that processes Arrow data in native code: the engine imports the
arrays of the stream
+ * itself, instead of reading the {@link
org.apache.spark.sql.vectorized.ColumnarBatch}es of
+ * {@link #createColumnarReader(InputPartition)}. The reader factories of the
batch and
+ * micro-batch scans of native data sources implement it, with the stream that
the library
+ * returns.
+ * <p>
+ * Each array of the stream is a struct array whose fields are the columns of
the read schema of
+ * the scan. Spark 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.
Review Comment:
**2. [Low] The API doc does not say that a stream read after its task ended
is interrupted**
This says that the caller releases the stream, also after the task ended,
and the guide says the same. An engine developer can read that as allowing the
engine to read the stream after the task too. But since 173c264, the context of
the stream reports `is_interrupted` once the task failed or completed, or once
its `TaskContext` was garbage collected. So a library that follows the header
fails `get_next` then, as `wait_for_interrupt` does in the test.
That is what I asked for in [round
2](https://github.com/apache/spark/pull/59934#discussion_r4417661673), but only
the comment of `SparkDsTaskContext` in the header says it. This Javadoc,
`docs/sql-native-data-sources.md` (L245), and the doc of the field
`is_interrupted` (`spark_datasource.h` L401, "Returns whether the task was
killed") do not. Could they say that an engine reads the stream in the task
that opened it, and that the library may stop once the task ended, so that only
the release may come later?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -189,18 +189,61 @@ class ColumnarBatchScan(planner: ColumnarBatchPlanner)
extends Batch {
case class ColumnarStreamOffset(json: String) extends Offset
-class ColumnarMicroBatchStream(reader: ColumnarStreamReader) extends
MicroBatchStream {
+/**
+ * The micro-batch stream of a [[ColumnarDataSource]]. It supports admission
control only so that
+ * the reader gets the offset that each micro-batch starts after, the end of
the last one, for
+ * every trigger: its default read limit reads all the available data. With
+ * `Trigger.AvailableNow`, Spark runs micro-batches up to the latest offset
that the first
+ * micro-batch got, rather than wrapping the stream in one that asks for the
latest offset from
+ * the initial offset.
+ */
+class ColumnarMicroBatchStream(reader: ColumnarStreamReader)
+ extends MicroBatchStream with SupportsTriggerAvailableNow {
+
+ // The offset that the reader returned last, which the progress of the query
reports.
+ @volatile private var lastLatestOffset: Offset = _
+ // Whether the trigger is Trigger.AvailableNow.
+ @volatile private var availableNow = false
+ // The end of the run with Trigger.AvailableNow: the first latest offset
asked.
+ @volatile private var availableNowEnd: Offset = _
+ // Spark asks for the initial offset for the start of the first latest
offset, and again to plan
+ // the first micro-batch, which both get the same offset.
+ private lazy val initial = ColumnarStreamOffset(reader.initialOffset())
+
+ override def initialOffset(): Offset = initial
+
+ override def latestOffset(): Offset = {
+ throw SparkException.internalError(
+ "latestOffset(Offset, ReadLimit) should be called instead of
latestOffset().")
+ }
- override def initialOffset(): Offset =
ColumnarStreamOffset(reader.initialOffset())
+ override def latestOffset(start: Offset, limit: ReadLimit): Offset = {
+ val latest = if (availableNow && availableNowEnd != null) {
+ availableNowEnd
+ } else {
+ ColumnarStreamOffset(reader.latestOffset(start.json()))
+ }
+ if (availableNow) {
+ availableNowEnd = latest
+ }
+ lastLatestOffset = latest
+ latest
+ }
- override def latestOffset(): Offset =
ColumnarStreamOffset(reader.latestOffset())
+ override def prepareForTriggerAvailableNow(): Unit = {
+ availableNow = true
+ }
+
+ override def reportLatestOffset(): Offset = lastLatestOffset
override def planInputPartitions(start: Offset, end: Offset):
Array[InputPartition] = {
reader.partitions(start.json(), end.json())
}
- override def createReaderFactory(): PartitionReaderFactory = {
- new ColumnarStreamPartitionReaderFactory(reader)
+ override def createReaderFactory(): PartitionReaderFactory = reader match {
+ case arrowStreamReader: ColumnarArrowStreamReader =>
+ new ColumnarArrowStreamPartitionReaderFactory(arrowStreamReader)
+ case _ => new ColumnarStreamPartitionReaderFactory(reader)
}
Review Comment:
**4. [Low] Could the stream reader return its reader factory, as the batch
planner does?**
For a batch scan, the data source makes the reader factory itself, as
`ColumnarPlannedBatch.readerFactory`, so `NativePartitionReaderFactory` simply
mixes in `SupportsArrowStreamRead`. For a micro-batch scan, the columnar layer
makes the factory. So 8575505 adds a trait, `ColumnarArrowStreamReader`, a
subclass, `ColumnarArrowStreamPartitionReaderFactory`, and this match on the
type of the reader. Each further capability of a reader factory would need the
same three pieces again.
Suppose `ColumnarStreamReader` had a method such as:
```scala
def readerFactory(): PartitionReaderFactory = new
ColumnarStreamPartitionReaderFactory(this)
```
Then `NativeDataSourceStreamReader` could override it with a factory that
implements `SupportsArrowStreamRead`. Both kinds of scans would work the same
way, and the trait, the subclass and the match would go away.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1162,182 @@ private static boolean isInterrupted(long task) {
}
}
+ //
---------------------------------------------------------------------------------------------
+ // Streams of engines
+ //
---------------------------------------------------------------------------------------------
+
+ // The offsets of the fields of struct ArrowArrayStream, and its size.
+ private static final long STREAM_GET_SCHEMA = 0;
+ private static final long STREAM_GET_NEXT = 8;
+ private static final long STREAM_GET_LAST_ERROR = 16;
+ private static final long STREAM_RELEASE = 24;
+ private static final long STREAM_PRIVATE_DATA = 32;
+ private static final long STREAM_SIZE = 40;
+ // The errno values that the callbacks of a stream return.
+ private static final int EINVAL = 22;
+ private static final int EIO = 5;
+ // The message of get_last_error for a stream that Spark does not know,
which lives as long as
+ // the JVM.
+ private static final long UNKNOWN_STREAM_ERROR =
+ toCString("The stream was released, or is not a stream of Spark.");
+
+ private static final AtomicLong nextEngineStreamId = new AtomicLong(1);
+ // The streams wrapped for engines that are not released yet, by their
private_data.
+ private static final ConcurrentHashMap<Long, EngineStream> engineStreams =
+ new ConcurrentHashMap<>();
+
+ /**
+ * The stream of a read that an engine got with {@code
SupportsArrowStreamRead}, which ties the
+ * task context of the read to the stream: the context is closed once the
library released the
+ * stream, whenever and on whichever thread the engine releases it, also
after the task ended.
+ * The stream of the library is moved to memory of Spark, and the stream of
the engine gets
+ * callbacks of Spark, which call the current callbacks of the library's
stream, and a
+ * private_data that identifies the wrapped stream, so that the engine can
move the stream too.
+ */
+ private static final class EngineStream {
+ final long inner;
+ final Task task;
+ // The message of the last error of Spark, rather than of the library, or
0 if the last call
+ // returned an error of the library, or none. The callbacks never run
concurrently.
+ long sparkError;
+
+ EngineStream(long inner, Task task) {
+ this.inner = inner;
+ this.task = task;
+ }
+
+ long callback(long offset) {
+ return Platform.getLong(null, inner + offset);
+ }
+
+ /** Records an error of Spark, which get_last_error then returns, and logs
it. */
+ int fail(String message, Throwable e) {
+ try {
+ clearError();
+ sparkError = toCString(message + ": " + rootCause(e));
+ LOGGER.error(message, e);
+ } catch (Throwable ignored) {
+ // An upcall must not throw.
+ }
+ return EIO;
+ }
+
+ void clearError() {
+ if (sparkError != 0) {
+ Platform.freeMemory(sparkError);
+ sparkError = 0;
+ }
+ }
+ }
+
+ /** Copies a string to native memory, as a NUL-terminated UTF-8 string,
which the caller frees. */
+ private static long toCString(String value) {
+ byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
+ long address = Platform.allocateMemory(bytes.length + 1);
+ Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address,
bytes.length);
+ Platform.putByte(null, address + bytes.length, (byte) 0);
+ return address;
+ }
+
+ /**
+ * Wraps the stream at the given address, which read initialized with
non-NULL callbacks, so that
+ * releasing it closes the given task context (see {@link EngineStream}).
+ */
+ static void wrapEngineStream(long stream, Task task) {
+ long inner = Platform.allocateMemory(STREAM_SIZE);
+ Platform.copyMemory(null, stream, null, inner, STREAM_SIZE);
+ long id = nextEngineStreamId.getAndIncrement();
+ engineStreams.put(id, new EngineStream(inner, task));
+ Platform.putLong(null, stream + STREAM_GET_SCHEMA,
Ffm.ENGINE_STREAM_GET_SCHEMA);
+ Platform.putLong(null, stream + STREAM_GET_NEXT,
Ffm.ENGINE_STREAM_GET_NEXT);
+ Platform.putLong(null, stream + STREAM_GET_LAST_ERROR,
Ffm.ENGINE_STREAM_GET_LAST_ERROR);
+ Platform.putLong(null, stream + STREAM_RELEASE, Ffm.ENGINE_STREAM_RELEASE);
+ Platform.putLong(null, stream + STREAM_PRIVATE_DATA, id);
+ }
+
+ /** The number of streams wrapped for engines that are not released yet, for
tests. */
+ static int openEngineStreams() {
+ return engineStreams.size();
+ }
+
+ /** The number of task contexts that are not closed yet, for tests. */
+ static int openTaskContexts() {
+ return tasks.size();
+ }
Review Comment:
**7. [Nit] `openTaskContexts` belongs in the section of the task contexts**
`AGENTS.md` asks: "When adding a member to an existing class or object,
follow the sectioning the file already uses and put the new member with the
code it belongs with." This only reads `tasks`, which is in the section "Task
contexts" (L1110), so it could go there, next to `isInterrupted`. For the same
reason, `toCString` could go in the section "Memory", next to `readCString`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1145,6 +1276,155 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
check(writeTasks)
}
+ nativeTest("open the partitions of a batch scan as Arrow C streams") {
+ val session = newSession(defaultPackage)
+ // The reader factory and the partitions of the batch scan of a query,
which an engine gets
+ // from its plan.
+ def plannedScan(df: DataFrame): (SupportsArrowStreamRead,
Seq[InputPartition]) = {
+ val scan = df.queryExecution.executedPlan.collectFirst { case s:
BatchScanExec => s }.get
+ (scan.readerFactory.asInstanceOf[SupportsArrowStreamRead],
scan.inputPartitions)
+ }
+ // Runs `f` in a task that only gets the files of the session.
+ def inTask[T: ClassTag](f: => T): T =
session.artifactManager.withResources {
+ session.sparkContext.parallelize(Seq(0), 1).map(_ => f).collect().head
+ }
+ import TestArrowEngine.readIds
+
+ // An engine reads both partitions of the scan in a single task.
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ val ids = inTask {
+ partitions.flatMap { partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ try {
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try factory.openArrowStream(partition, stream.memoryAddress()) catch
{
+ case e: Throwable =>
+ stream.close()
+ throw e
+ }
+ readIds(stream, allocator)
+ } finally {
+ allocator.close()
+ }
+ }
+ }
+ assert(ids.sorted == (0L until 6L))
+ // The library got the partition of the scan that it reads, not the one of
the task.
+ val tasks = taskDir.listFiles().map { file =>
+ file.getName.toInt ->
Files.readString(file.toPath).split(",").map(_.toLong).toSeq
+ }.toMap
+ assert(tasks.keySet == Set(0, 1) && tasks.values.map(_(5)).toSet.size ==
1, tasks)
+ assert(NativeLibrary.openEngineStreams() == 0)
+
+ // The engine may release the stream after its task ended, and on another
thread: Spark frees
+ // the context of the task once the library released the stream.
+ val released = new File(Utils.createTempDir(), "released")
+ val (laterFactory, laterPartitions) =
+ plannedScan(range(session, "end" -> "6", "released_path" ->
released.getPath).select("id"))
+ // Opens a stream in a task, which it returns unreleased.
+ def openInTask(factory: SupportsArrowStreamRead, partition:
InputPartition): Long = inTask {
+ val stream = Platform.allocateMemory(40)
+ Platform.setMemory(stream, 0.toByte, 40)
+ factory.openArrowStream(partition, stream)
+ stream
+ }
+ val openContexts = NativeLibrary.openTaskContexts()
+ val address = openInTask(laterFactory, laterPartitions.head)
+ assert(NativeLibrary.openEngineStreams() == 1 && !released.exists())
+ assert(NativeLibrary.openTaskContexts() == openContexts + 1)
+ assert(readIds(ArrowArrayStream.wrap(address), ArrowUtils.rootAllocator)
== (0L until 3L))
+ assert(NativeLibrary.openEngineStreams() == 0 && released.exists())
+ assert(NativeLibrary.openTaskContexts() == openContexts)
+ Platform.freeMemory(address)
+
+ // Once the task ended, its context tells the library that it was
interrupted, so that the
+ // library stops waiting.
+ val marker = new File(Utils.createTempDir(), "interrupted")
+ val (waitingFactory, waitingPartitions) = plannedScan(
+ range(session, "end" -> "6", "wait_for_interrupt" ->
marker.getPath).select("id"))
+ val waiting = openInTask(waitingFactory, waitingPartitions.head)
+ val error = intercept[Exception] {
+ readIds(ArrowArrayStream.wrap(waiting), ArrowUtils.rootAllocator)
+ }
+ assert(Utils.exceptionString(error).contains("the task was killed"))
+ assert(Files.readString(marker.toPath) == "interrupted")
+ assert(NativeLibrary.openEngineStreams() == 0)
Review Comment:
**8. [Nit] The test does not check that the context is closed after an
interrupted or failed stream**
The test checks `openTaskContexts()` only around the stream that is read to
its end (L1334-L1340). Here, and after each failure at L1424, it only checks
`openEngineStreams()`. So the test would still pass if a change left the
context open after a failed `get_next`, or after a failure in
`openEngineStream`. `assert(NativeLibrary.openTaskContexts() == openContexts)`
in both places would cover these cases.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1162,182 @@ private static boolean isInterrupted(long task) {
}
}
+ //
---------------------------------------------------------------------------------------------
+ // Streams of engines
+ //
---------------------------------------------------------------------------------------------
+
+ // The offsets of the fields of struct ArrowArrayStream, and its size.
+ private static final long STREAM_GET_SCHEMA = 0;
+ private static final long STREAM_GET_NEXT = 8;
+ private static final long STREAM_GET_LAST_ERROR = 16;
+ private static final long STREAM_RELEASE = 24;
+ private static final long STREAM_PRIVATE_DATA = 32;
+ private static final long STREAM_SIZE = 40;
+ // The errno values that the callbacks of a stream return.
+ private static final int EINVAL = 22;
+ private static final int EIO = 5;
+ // The message of get_last_error for a stream that Spark does not know,
which lives as long as
+ // the JVM.
+ private static final long UNKNOWN_STREAM_ERROR =
+ toCString("The stream was released, or is not a stream of Spark.");
+
+ private static final AtomicLong nextEngineStreamId = new AtomicLong(1);
+ // The streams wrapped for engines that are not released yet, by their
private_data.
+ private static final ConcurrentHashMap<Long, EngineStream> engineStreams =
+ new ConcurrentHashMap<>();
+
+ /**
+ * The stream of a read that an engine got with {@code
SupportsArrowStreamRead}, which ties the
+ * task context of the read to the stream: the context is closed once the
library released the
+ * stream, whenever and on whichever thread the engine releases it, also
after the task ended.
+ * The stream of the library is moved to memory of Spark, and the stream of
the engine gets
+ * callbacks of Spark, which call the current callbacks of the library's
stream, and a
+ * private_data that identifies the wrapped stream, so that the engine can
move the stream too.
+ */
+ private static final class EngineStream {
+ final long inner;
+ final Task task;
+ // The message of the last error of Spark, rather than of the library, or
0 if the last call
+ // returned an error of the library, or none. The callbacks never run
concurrently.
+ long sparkError;
+
+ EngineStream(long inner, Task task) {
+ this.inner = inner;
+ this.task = task;
+ }
+
+ long callback(long offset) {
+ return Platform.getLong(null, inner + offset);
+ }
+
+ /** Records an error of Spark, which get_last_error then returns, and logs
it. */
+ int fail(String message, Throwable e) {
+ try {
+ clearError();
+ sparkError = toCString(message + ": " + rootCause(e));
+ LOGGER.error(message, e);
+ } catch (Throwable ignored) {
+ // An upcall must not throw.
+ }
+ return EIO;
+ }
+
+ void clearError() {
+ if (sparkError != 0) {
+ Platform.freeMemory(sparkError);
+ sparkError = 0;
+ }
+ }
+ }
+
+ /** Copies a string to native memory, as a NUL-terminated UTF-8 string,
which the caller frees. */
+ private static long toCString(String value) {
+ byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
+ long address = Platform.allocateMemory(bytes.length + 1);
+ Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address,
bytes.length);
+ Platform.putByte(null, address + bytes.length, (byte) 0);
+ return address;
Review Comment:
**5. [Low] `toCString` repeats the copy of `Memory.string` and `Ffm.host`**
This is the third copy of the same code: `Memory.string` (L1073) and the
copy of the Spark version in `Ffm.host` (L562-L566) also copy a string as UTF-8
followed by a NUL. The copies also compute the length slightly differently,
`bytes.length + 1` here and `bytes.length + 1L` there. `toCString` could move
next to `readCString`, with the other memory helpers, and `Ffm.host` could use
it. `Memory.string` could call it too, since it only adds the allocation to the
list to free after the call.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1162,182 @@ private static boolean isInterrupted(long task) {
}
}
+ //
---------------------------------------------------------------------------------------------
+ // Streams of engines
+ //
---------------------------------------------------------------------------------------------
+
+ // The offsets of the fields of struct ArrowArrayStream, and its size.
+ private static final long STREAM_GET_SCHEMA = 0;
+ private static final long STREAM_GET_NEXT = 8;
+ private static final long STREAM_GET_LAST_ERROR = 16;
+ private static final long STREAM_RELEASE = 24;
+ private static final long STREAM_PRIVATE_DATA = 32;
+ private static final long STREAM_SIZE = 40;
+ // The errno values that the callbacks of a stream return.
+ private static final int EINVAL = 22;
+ private static final int EIO = 5;
+ // The message of get_last_error for a stream that Spark does not know,
which lives as long as
+ // the JVM.
+ private static final long UNKNOWN_STREAM_ERROR =
+ toCString("The stream was released, or is not a stream of Spark.");
+
+ private static final AtomicLong nextEngineStreamId = new AtomicLong(1);
+ // The streams wrapped for engines that are not released yet, by their
private_data.
+ private static final ConcurrentHashMap<Long, EngineStream> engineStreams =
+ new ConcurrentHashMap<>();
+
+ /**
+ * The stream of a read that an engine got with {@code
SupportsArrowStreamRead}, which ties the
+ * task context of the read to the stream: the context is closed once the
library released the
+ * stream, whenever and on whichever thread the engine releases it, also
after the task ended.
+ * The stream of the library is moved to memory of Spark, and the stream of
the engine gets
+ * callbacks of Spark, which call the current callbacks of the library's
stream, and a
+ * private_data that identifies the wrapped stream, so that the engine can
move the stream too.
+ */
+ private static final class EngineStream {
+ final long inner;
+ final Task task;
+ // The message of the last error of Spark, rather than of the library, or
0 if the last call
+ // returned an error of the library, or none. The callbacks never run
concurrently.
+ long sparkError;
+
+ EngineStream(long inner, Task task) {
+ this.inner = inner;
+ this.task = task;
+ }
+
+ long callback(long offset) {
+ return Platform.getLong(null, inner + offset);
+ }
+
+ /** Records an error of Spark, which get_last_error then returns, and logs
it. */
+ int fail(String message, Throwable e) {
+ try {
+ clearError();
+ sparkError = toCString(message + ": " + rootCause(e));
+ LOGGER.error(message, e);
+ } catch (Throwable ignored) {
+ // An upcall must not throw.
+ }
+ return EIO;
+ }
+
+ void clearError() {
+ if (sparkError != 0) {
+ Platform.freeMemory(sparkError);
+ sparkError = 0;
+ }
+ }
+ }
+
+ /** Copies a string to native memory, as a NUL-terminated UTF-8 string,
which the caller frees. */
+ private static long toCString(String value) {
+ byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
+ long address = Platform.allocateMemory(bytes.length + 1);
+ Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address,
bytes.length);
+ Platform.putByte(null, address + bytes.length, (byte) 0);
+ return address;
+ }
+
+ /**
+ * Wraps the stream at the given address, which read initialized with
non-NULL callbacks, so that
+ * releasing it closes the given task context (see {@link EngineStream}).
+ */
+ static void wrapEngineStream(long stream, Task task) {
+ long inner = Platform.allocateMemory(STREAM_SIZE);
+ Platform.copyMemory(null, stream, null, inner, STREAM_SIZE);
+ long id = nextEngineStreamId.getAndIncrement();
+ engineStreams.put(id, new EngineStream(inner, task));
Review Comment:
**3. [Low] `inner` leaks if the wrapping fails after it was allocated**
173c264 dropped the `try` that freed `inner` when the wrapping failed,
together with "If it fails, the stream is left as it is" in the Javadoc. Now
`new EngineStream(...)` or `engineStreams.put` can throw, for example an
`OutOfMemoryError`. `ConcurrentHashMap.put` can also throw it while it resizes
the table, after it inserted the entry. Then `openEngineStream` releases the
library's stream through the caller's struct, which still has the library's
callbacks, and closes the task. But `inner` is never freed. In the second case,
`engineStreams` also keeps an entry whose task is closed, so
`openEngineStreams()` stays one too high.
This is unlikely. Still, restoring the `try` around the copy and the `put`,
with `engineStreams.remove(id)` and `Platform.freeMemory(inner)` on failure,
would keep the behavior that the Javadoc used to state.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -399,6 +515,25 @@ static MethodHandle downcall(Object address, C result,
C... parameters) throws T
return handle;
}
+ /**
+ * Returns a method handle that calls the C function at the address of its
first argument, a
+ * {@code long}, with the given types, as {@link #downcall} does for a
fixed address.
+ */
+ static MethodHandle unboundDowncall(C result, C... parameters) throws
Throwable {
+ MethodHandle handle = (MethodHandle) UNBOUND_DOWNCALL.invoke(
+ LINKER, descriptor(result, parameters), NO_OPTIONS);
+ handle = MethodHandles.filterArguments(handle, 0, OF_ADDRESS);
+ for (int i = 0; i < parameters.length; i++) {
+ if (parameters[i] == C.PTR) {
+ handle = MethodHandles.filterArguments(handle, i + 1, OF_ADDRESS);
+ }
+ }
+ if (result == C.PTR) {
+ handle = MethodHandles.filterReturnValue(handle, ADDRESS_OF);
+ }
Review Comment:
**6. [Nit] `unboundDowncall` repeats the filters of `downcall`**
This repeats the loop of `downcall` with the indexes shifted by one, so the
conversion of pointers is now written three times, counting `upcall`. One copy
would be enough: either a helper that takes the index of the first C parameter,
or `downcall` written as `MethodHandles.insertArguments(unboundDowncall(result,
parameters), 0, address)`, with the address as a `long`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1145,6 +1276,155 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
check(writeTasks)
}
+ nativeTest("open the partitions of a batch scan as Arrow C streams") {
+ val session = newSession(defaultPackage)
+ // The reader factory and the partitions of the batch scan of a query,
which an engine gets
+ // from its plan.
+ def plannedScan(df: DataFrame): (SupportsArrowStreamRead,
Seq[InputPartition]) = {
+ val scan = df.queryExecution.executedPlan.collectFirst { case s:
BatchScanExec => s }.get
+ (scan.readerFactory.asInstanceOf[SupportsArrowStreamRead],
scan.inputPartitions)
+ }
+ // Runs `f` in a task that only gets the files of the session.
+ def inTask[T: ClassTag](f: => T): T =
session.artifactManager.withResources {
+ session.sparkContext.parallelize(Seq(0), 1).map(_ => f).collect().head
+ }
+ import TestArrowEngine.readIds
+
+ // An engine reads both partitions of the scan in a single task.
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ val ids = inTask {
+ partitions.flatMap { partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ try {
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try factory.openArrowStream(partition, stream.memoryAddress()) catch
{
+ case e: Throwable =>
+ stream.close()
+ throw e
+ }
+ readIds(stream, allocator)
+ } finally {
+ allocator.close()
+ }
+ }
+ }
+ assert(ids.sorted == (0L until 6L))
+ // The library got the partition of the scan that it reads, not the one of
the task.
+ val tasks = taskDir.listFiles().map { file =>
+ file.getName.toInt ->
Files.readString(file.toPath).split(",").map(_.toLong).toSeq
+ }.toMap
+ assert(tasks.keySet == Set(0, 1) && tasks.values.map(_(5)).toSet.size ==
1, tasks)
+ assert(NativeLibrary.openEngineStreams() == 0)
+
+ // The engine may release the stream after its task ended, and on another
thread: Spark frees
+ // the context of the task once the library released the stream.
+ val released = new File(Utils.createTempDir(), "released")
+ val (laterFactory, laterPartitions) =
+ plannedScan(range(session, "end" -> "6", "released_path" ->
released.getPath).select("id"))
+ // Opens a stream in a task, which it returns unreleased.
+ def openInTask(factory: SupportsArrowStreamRead, partition:
InputPartition): Long = inTask {
+ val stream = Platform.allocateMemory(40)
+ Platform.setMemory(stream, 0.toByte, 40)
+ factory.openArrowStream(partition, stream)
+ stream
+ }
+ val openContexts = NativeLibrary.openTaskContexts()
+ val address = openInTask(laterFactory, laterPartitions.head)
+ assert(NativeLibrary.openEngineStreams() == 1 && !released.exists())
+ assert(NativeLibrary.openTaskContexts() == openContexts + 1)
+ assert(readIds(ArrowArrayStream.wrap(address), ArrowUtils.rootAllocator)
== (0L until 3L))
+ assert(NativeLibrary.openEngineStreams() == 0 && released.exists())
+ assert(NativeLibrary.openTaskContexts() == openContexts)
+ Platform.freeMemory(address)
+
+ // Once the task ended, its context tells the library that it was
interrupted, so that the
+ // library stops waiting.
+ val marker = new File(Utils.createTempDir(), "interrupted")
+ val (waitingFactory, waitingPartitions) = plannedScan(
+ range(session, "end" -> "6", "wait_for_interrupt" ->
marker.getPath).select("id"))
+ val waiting = openInTask(waitingFactory, waitingPartitions.head)
+ val error = intercept[Exception] {
+ readIds(ArrowArrayStream.wrap(waiting), ArrowUtils.rootAllocator)
+ }
+ assert(Utils.exceptionString(error).contains("the task was killed"))
+ assert(Files.readString(marker.toPath) == "interrupted")
+ assert(NativeLibrary.openEngineStreams() == 0)
+ Platform.freeMemory(waiting)
+
+ // So does the reader factory of a micro-batch scan, with the state of its
stream reader.
+ val query =
session.readStream.format("native_counter").option("max_offset", "4").load()
+ .writeStream.format("noop").start()
+ val microBatchScan = try {
+ query.processAllAvailable()
+
query.asInstanceOf[StreamingQueryWrapper].streamingQuery.lastExecution.executedPlan
+ .collectFirst { case s: MicroBatchScanExec => s }.get
+ } finally {
+ query.stop()
+ }
+ val streamFactory =
microBatchScan.readerFactory.asInstanceOf[SupportsArrowStreamRead]
+ val streamPartitions = microBatchScan.inputPartitions
+ val streamIds = inTask {
+ streamPartitions.flatMap { partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ try {
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try streamFactory.openArrowStream(partition, stream.memoryAddress())
catch {
+ case e: Throwable =>
+ stream.close()
+ throw e
+ }
+ readIds(stream, allocator)
+ } finally {
+ allocator.close()
+ }
+ }
+ }
Review Comment:
**10. [Nit] The batch and micro-batch parts of the test repeat the same
block**
These lines repeat L1298-L1313 with another factory and other partitions. A
helper such as `def readAll(factory: SupportsArrowStreamRead, partitions:
Seq[InputPartition]): Seq[Long]` would keep a single copy of the allocation,
the cleanup on failure and the read.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -926,18 +1063,28 @@ private[ffi] object NativeTask {
* @param partitionId the partition that the task reads or writes, by
default the partition of
* the task
* @param taskAttemptId the ID of the task attempt, by default the one of
the task
+ * @param mayOutliveTask whether the library may use the context after the
task ended, as for
+ * the stream of an engine: `is_interrupted` then also
reports the end of
+ * the task, and the context does not keep the task
alive
*/
def create(
partitionId: Option[Int] = None,
- taskAttemptId: Option[Long] = None): NativeLibrary.Task = {
+ taskAttemptId: Option[Long] = None,
+ mayOutliveTask: Boolean = false): NativeLibrary.Task = {
val context = Option(TaskContext.get())
+ val interrupted: () => Boolean = if (mayOutliveTask) {
+ val task = new WeakReference(context.orNull)
+ () => Option(task.get()).forall(t => t.isInterrupted() || t.isFailed()
|| t.isCompleted())
+ } else {
+ () => context.exists(_.isInterrupted())
+ }
new NativeLibrary.Task(
context.map(_.stageId()).getOrElse(-1),
context.map(_.stageAttemptNumber()).getOrElse(-1),
partitionId.orElse(context.map(_.partitionId())).getOrElse(-1),
context.map(_.attemptNumber()).getOrElse(0),
taskAttemptId.orElse(context.map(_.taskAttemptId())).getOrElse(-1L),
- () => context.exists(_.isInterrupted()))
+ () => interrupted())
Review Comment:
**11. [Nit] `interrupted` could be a `BooleanSupplier`**
`interrupted` is a Scala function, which is wrapped again in `() =>
interrupted()` to make the `BooleanSupplier` of `NativeLibrary.Task`. It could
be declared as `val interrupted: BooleanSupplier = if (mayOutliveTask) { ... }
else { ... }` and passed as it is, since the two lambdas convert to
`BooleanSupplier` just as the wrapper does now. That would drop the extra
closure from every call to `is_interrupted`.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -903,6 +1162,182 @@ private static boolean isInterrupted(long task) {
}
}
+ //
---------------------------------------------------------------------------------------------
+ // Streams of engines
+ //
---------------------------------------------------------------------------------------------
+
+ // The offsets of the fields of struct ArrowArrayStream, and its size.
+ private static final long STREAM_GET_SCHEMA = 0;
+ private static final long STREAM_GET_NEXT = 8;
+ private static final long STREAM_GET_LAST_ERROR = 16;
+ private static final long STREAM_RELEASE = 24;
+ private static final long STREAM_PRIVATE_DATA = 32;
+ private static final long STREAM_SIZE = 40;
+ // The errno values that the callbacks of a stream return.
+ private static final int EINVAL = 22;
+ private static final int EIO = 5;
+ // The message of get_last_error for a stream that Spark does not know,
which lives as long as
+ // the JVM.
+ private static final long UNKNOWN_STREAM_ERROR =
+ toCString("The stream was released, or is not a stream of Spark.");
+
+ private static final AtomicLong nextEngineStreamId = new AtomicLong(1);
+ // The streams wrapped for engines that are not released yet, by their
private_data.
+ private static final ConcurrentHashMap<Long, EngineStream> engineStreams =
+ new ConcurrentHashMap<>();
+
+ /**
+ * The stream of a read that an engine got with {@code
SupportsArrowStreamRead}, which ties the
+ * task context of the read to the stream: the context is closed once the
library released the
+ * stream, whenever and on whichever thread the engine releases it, also
after the task ended.
+ * The stream of the library is moved to memory of Spark, and the stream of
the engine gets
+ * callbacks of Spark, which call the current callbacks of the library's
stream, and a
+ * private_data that identifies the wrapped stream, so that the engine can
move the stream too.
+ */
+ private static final class EngineStream {
+ final long inner;
+ final Task task;
+ // The message of the last error of Spark, rather than of the library, or
0 if the last call
+ // returned an error of the library, or none. The callbacks never run
concurrently.
+ long sparkError;
+
+ EngineStream(long inner, Task task) {
+ this.inner = inner;
+ this.task = task;
+ }
+
+ long callback(long offset) {
+ return Platform.getLong(null, inner + offset);
+ }
+
+ /** Records an error of Spark, which get_last_error then returns, and logs
it. */
+ int fail(String message, Throwable e) {
+ try {
+ clearError();
+ sparkError = toCString(message + ": " + rootCause(e));
+ LOGGER.error(message, e);
+ } catch (Throwable ignored) {
+ // An upcall must not throw.
+ }
+ return EIO;
+ }
+
+ void clearError() {
+ if (sparkError != 0) {
+ Platform.freeMemory(sparkError);
+ sparkError = 0;
+ }
+ }
+ }
+
+ /** Copies a string to native memory, as a NUL-terminated UTF-8 string,
which the caller frees. */
+ private static long toCString(String value) {
+ byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
+ long address = Platform.allocateMemory(bytes.length + 1);
+ Platform.copyMemory(bytes, Platform.BYTE_ARRAY_OFFSET, null, address,
bytes.length);
+ Platform.putByte(null, address + bytes.length, (byte) 0);
+ return address;
+ }
+
+ /**
+ * Wraps the stream at the given address, which read initialized with
non-NULL callbacks, so that
+ * releasing it closes the given task context (see {@link EngineStream}).
+ */
+ static void wrapEngineStream(long stream, Task task) {
+ long inner = Platform.allocateMemory(STREAM_SIZE);
+ Platform.copyMemory(null, stream, null, inner, STREAM_SIZE);
+ long id = nextEngineStreamId.getAndIncrement();
+ engineStreams.put(id, new EngineStream(inner, task));
+ Platform.putLong(null, stream + STREAM_GET_SCHEMA,
Ffm.ENGINE_STREAM_GET_SCHEMA);
+ Platform.putLong(null, stream + STREAM_GET_NEXT,
Ffm.ENGINE_STREAM_GET_NEXT);
+ Platform.putLong(null, stream + STREAM_GET_LAST_ERROR,
Ffm.ENGINE_STREAM_GET_LAST_ERROR);
+ Platform.putLong(null, stream + STREAM_RELEASE, Ffm.ENGINE_STREAM_RELEASE);
+ Platform.putLong(null, stream + STREAM_PRIVATE_DATA, id);
+ }
+
+ /** The number of streams wrapped for engines that are not released yet, for
tests. */
+ static int openEngineStreams() {
+ return engineStreams.size();
+ }
+
+ /** The number of task contexts that are not closed yet, for tests. */
+ static int openTaskContexts() {
+ return tasks.size();
+ }
+
+ private static EngineStream engineStream(long stream) {
+ return engineStreams.get(Platform.getLong(null, stream +
STREAM_PRIVATE_DATA));
+ }
+
+ // The callbacks of the streams wrapped for engines. They must not throw: an
exception thrown by
+ // an upcall crashes the JVM.
+
+ private static int engineStreamGetSchema(long stream, long out) {
+ return engineStreamCall(stream, out, STREAM_GET_SCHEMA, "get_schema");
+ }
+
+ private static int engineStreamGetNext(long stream, long out) {
+ return engineStreamCall(stream, out, STREAM_GET_NEXT, "get_next");
+ }
+
+ private static int engineStreamCall(long stream, long out, long offset,
String callback) {
+ EngineStream wrapped = null;
+ try {
+ wrapped = engineStream(stream);
+ if (wrapped == null) {
+ return EINVAL;
+ }
+ wrapped.clearError();
+ return (int) Ffm.STREAM_INT_CALLBACK.invokeExact(
+ wrapped.callback(offset), wrapped.inner, out);
+ } catch (Throwable e) {
+ return wrapped == null ? EIO : wrapped.fail(
+ "The " + callback + " of the stream of a native data source failed",
e);
Review Comment:
**9. [Nit] Spark's errors on an engine's stream do not name the data source**
When a callback fails on Spark's side, the engine gets "The get_next of the
stream of a native data source failed: ..." from `get_last_error`, and the
executor log has the same message. Neither names the data source. An engine
that reads several native data sources in one query then cannot tell which one
failed, while Spark's own reader reports the name in
`NATIVE_DATA_SOURCE_ERROR`. It would help to pass the name to `EngineStream`
and include it in this message, and in the one of `release`.
--
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]
