HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4417795613
##########
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:
Done in 173c264: the logging in the four callbacks has its own `try`.
##########
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:
Done in 173c264: `EngineStream` keeps the message of an error of Spark,
which `get_last_error` returns until the next call, and a stream that Spark
does not know gets a static message.
##########
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:
Done in 173c264: the Javadoc now says that the stream is released if the
method fails after it checked its arguments, and that a stream whose release
callback is not NULL is left as it is.
##########
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:
Done in 173c264: the doc of `read` now says that this holds when Spark reads
the stream itself, and that an engine that gets the stream from Spark reads and
releases it itself, pointing to `SparkDsTaskContext`.
--
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]
