HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4461788066
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -351,17 +377,30 @@ class NativeDataSourceStreamReader(
}
override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
- library.streamPartitions(handle.get, start, end)
+ if (plannedRange != (start, end)) {
+ val planned = call(name, "plan a scan of") {
+ library.streamPartitions(handle.get, start, end)
+ }
+ val partitions = NativeInputPartition.from(planned.partitions)
+ // The executors fetch the state once for the micro-batch, as a
broadcast, rather than with
+ // each of its tasks. The partitions keep the broadcast, which the
ContextCleaner removes once
+ // the driver no longer refers to them: in the plan of the micro-batch,
in an engine, or here
+ // until the next range is planned.
+ val state = Option(planned.state).filter(_.nonEmpty &&
partitions.nonEmpty).map { bytes =>
Review Comment:
Done in 77bf7bd, as suggested: the test plans `plan(stream, 2, 2)`, which
has a state and no partition, and checks the ID of the next broadcast.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -458,17 +625,13 @@ typedef struct SparkDsLibrary {
/*
* Plans the partitions of the scan. Spark reads each partition in a
separate task.
Review Comment:
Done in 77bf7bd: L60, the doc of `partitions` and the guide now say that a
task may read several partitions, for example after a coalesce or a union.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -886,10 +987,140 @@ class NativePartitionReader(
}
}
-/** Writes the batches of a task to the data writer of a native library. */
+object NativePartitionReader {
+ import NativeCalls._
+
+ /**
+ * Checks the arguments of `SupportsArrowStreamRead.openArrowStream` before
it looks at the
+ * partition, and returns the partition as a `P`, the class of the
partitions of the scan of the
+ * reader factory.
+ */
+ def engineStreamPartition[P <: InputPartition](
+ partition: InputPartition,
+ streamAddress: Long)(implicit tag: ClassTag[P]): P = {
+ 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.")
+ }
+ partition match {
+ case expected: P => expected
+ case _ =>
+ throw new IllegalArgumentException("openArrowStream needs a partition
of the scan of " +
+ s"the reader factory, a ${tag.runtimeClass.getSimpleName}, but got "
+
+ s"${Option(partition).map(_.getClass.getName).orNull}.")
+ }
+ }
+
+ /**
+ * Initializes the stream at `streamAddress` with the data of a partition
for an engine, see
+ * [[SupportsArrowStreamRead]], once `engineStreamPartition` checked the
arguments: the context
+ * of its task is closed once the stream is released.
+ */
+ def openEngineStream(
Review Comment:
Done in 77bf7bd: one `private[ffi]` method, `openEngineStream[P]`, checks
the arguments, and then reads the partition with the state and the partition
that a function returns for it, after the checks. `openStream` and
`releaseOnFailure` are `private`.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -304,29 +312,133 @@ typedef struct SparkDsCreateInfo {
int64_t num_options;
} SparkDsCreateInfo;
+/* A batch scan, as create_reader gets it. */
+typedef struct SparkDsReaderInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+} SparkDsReaderInfo;
+
+/* A request for the partitions of a batch scan, as partitions gets it. */
+typedef struct SparkDsPartitionsInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+} SparkDsPartitionsInfo;
+
+/* 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
+ * fully qualified Hadoop path, whose characters are not escaped as in a
URI: a space stays a
+ * space, and % stays %. Its file system may be one that the library cannot
access, such as s3a
+ * or hdfs. Spark keeps what it needs there under the names that start with
_spark, such as the
+ * initial offset of the source, and the library must not use them.
+ */
+ const char* checkpoint_location;
+} SparkDsStreamReaderInfo;
+
+/* A request for the initial offset of a stream reader, as initial_offset gets
it. */
+typedef struct SparkDsInitialOffsetInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+} SparkDsInitialOffsetInfo;
+
+/* 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, for every trigger. With
Trigger.AvailableNow,
+ * Spark asks for the latest offset only once, before the first micro-batch
that it plans, and
+ * reads up to it.
+ */
+ const char* start;
+} SparkDsLatestOffsetInfo;
+
+/* A request for the partitions of a micro-batch, as stream_partitions gets
it. */
+typedef struct SparkDsStreamPartitionsInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /* The offsets, as JSON, that the micro-batch reads the data after, up to
and including end. */
+ const char* start;
+ const char* end;
+} SparkDsStreamPartitionsInfo;
+
+/* The end of the data that Spark processed, as commit_offset gets it. */
+typedef struct SparkDsCommitOffsetInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /* The offset, as JSON, up to and including which Spark processed the data.
*/
+ const char* end;
+} SparkDsCommitOffsetInfo;
+
/*
* 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 that read returns from Spark. It reads
the stream in the task
+ * that opened it, but may release it once the task ended, from another
thread: the task stays valid
+ * until then too, and is_interrupted then also returns true once the task
failed or completed, so
+ * the library may stop reading then. Such an engine may also keep the arrays
of the stream after
Review Comment:
Done in 77bf7bd: it now says that the engine may keep the arrays and the
schemas, so the release of an array or a schema must not use the stream or the
task.
--
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]
