HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4418953409
##########
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:
Done in b98fd19: Spark's reader of batch and micro-batch scans now passes
`NativeInputPartition.index` too, through `NativePartitionReader` and
`openStream`. The header and the guide say that `partition_id` is the index of
the partition in the scan for `read`, and the partition of the task for
`create_data_writer`. "pass the task to reads and data writers" checks the ids
that the library gets after a union, a coalesce and a union of two streams.
##########
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:
Done in b98fd19: the Javadoc, the guide and the doc of `SparkDsTaskContext`
and `is_interrupted` now say that the engine reads the stream in the task that
opened it, that `is_interrupted` returns true once the task ended so the
library may stop then, and that only the release may come later.
##########
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:
Done in b98fd19: `wrapEngineStream` copies the stream and adds the entry in
a `try` that removes the entry and frees `inner` on failure. The Javadoc says
again that the stream is left as it is then.
##########
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:
Done in b98fd19: `ColumnarStreamReader.readerFactory()` returns a
`ColumnarStreamPartitionReaderFactory` by default, and
`NativeDataSourceStreamReader` overrides it with one that implements
`SupportsArrowStreamRead`. `ColumnarArrowStreamReader`,
`ColumnarArrowStreamPartitionReaderFactory` and the match are gone.
##########
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:
Done in b98fd19: a single `toCString`, in the Memory section next to
`readCString`, which `Memory.string` and `Ffm.host` now use.
--
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]
