dongjoon-hyun commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4475570426
##########
docs/sql-native-data-sources.md:
##########
@@ -161,28 +179,33 @@ to, but it cannot check that the buffers of the arrays
hold those values, becaus
interface does not carry the sizes of buffers: a library that exports a buffer
that is too short
makes Spark read past its end.
-`read` and `create_data_writer` get a `SparkDsTaskContext`, with the IDs of
the task, and its
-partition and attempt, and an `is_interrupted` function that returns whether
the task was killed,
-for example because its query was canceled. Spark cannot interrupt native
code: when a task is
-killed, Spark stops reading only once the `get_next` function of the stream
returns. So a read or
-a write that runs for long, or blocks, for example on network I/O, should
check `is_interrupted`
-regularly and use timeouts. The task reaper, `spark.task.reaper.enabled` and
-`spark.task.reaper.killTimeout`, kills an executor whose killed task does not
end.
+`read` and `create_data_writer` get a `SparkDsTaskContext`, with the IDs of
the task, its
+attempt, and the partition: for `read`, the index of the partition in the list
that the planning
+returned, which is unique within a scan but not across the scans that a task
reads, and for
+`create_data_writer`, the partition of the task. It also has an
`is_interrupted` function that
Review Comment:
**3. [Low] The guide still says that `is_interrupted` returns whether the
task was killed**
For a read, `is_interrupted` now also returns true once the task failed or
completed, and every read gets the context of a read, also the reads of Spark
itself (NativeDataSource.scala:1098). So a library that no engine reads also
sees true in a task that succeeded, for example while the stdin writer thread
of a `PipedRDD` is still in get_next after the task completed.
`spark_datasource.h` says so (L384-L387), but the guide only says it in the
section on engines (L261). Could this paragraph say it too, for example
"returns whether the task was killed, or for `read`, also whether the task
failed or completed"?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,37 +384,104 @@ class NativeDataSourceStreamReader(
offset
}
- override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
- library.streamPartitions(handle.get, start, end)
+ override def partitions(start: String, end: String): Array[InputPartition] =
synchronized {
+ val planned = call(name, "plan a scan
of")(library.streamPartitions(handle.get, start, end))
+ val partitions = NativeInputPartition.from(planned.partitions)
+ val bytes = bytesOrEmpty(Option(planned.state))
+ // A small state goes with each task, while the executors fetch a larger
one once for the
+ // micro-batch, as a broadcast. Spark only plans the next range once the
micro-batch of the
+ // previous one finished, so the copies of the previous broadcast on the
executors are removed
+ // then. The broadcast itself is destroyed once its micro-batch is
committed, whose partitions
+ // are not read anymore then, also by an engine.
+ uncommitted.values.lastOption.flatten.foreach(_.unpersist(blocking =
false))
+ val state = if (bytes.length >
NativeDataSourceStreamReader.maxInlineStateBytes &&
Review Comment:
**2. [Low] Correction to my item 6 of round 8: the size limit of an inline
state does not count the partitions**
This follows [my item 6 of round
8](https://github.com/apache/spark/pull/59934#discussion_r4471627311), which
suggested a limit on the size of the state alone. But the cost that [item 1 of
round 5](https://github.com/apache/spark/pull/59934#discussion_r4420249991)
described also grows with the number of partitions: each task serializes its
own `NativeMicroBatchPartition`, so a state just below 64 KiB over 1,000
partitions sends about 64 MB with the tasks of each micro-batch, while its
broadcast would go once to each executor. Could the limit apply to the total,
for example `bytes.length.toLong * partitions.length > maxInlineStateBytes`?
The comments of `maxInlineStateBytes` (L440-L443) and of
`NativeMicroBatchPartition` (L569-L571) would then say so.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSource.scala:
##########
@@ -101,18 +111,37 @@ case class ColumnarPlannedBatch(
/**
* Reads the data of a [[ColumnarDataSource]] for a micro-batch streaming
scan, with offsets in
- * JSON. Spark plans the partitions of each micro-batch on the driver, then
serializes the reader
- * to the executors, where `read` is called for each partition.
+ * JSON. Spark plans the partitions of each micro-batch on the driver with
this reader, which stays
+ * there, and serializes the reader factory that `readerFactory` returns to
the executors, which
+ * read the partitions with it.
*/
-trait ColumnarStreamReader extends Serializable {
+trait ColumnarStreamReader {
+ /**
+ * Returns the offset to start from when the streaming query has no
checkpoint. Spark keeps it in
+ * the checkpoint of the source, and only asks for it when the checkpoint
does not have it yet.
+ */
def initialOffset(): String
- def latestOffset(): String
+ /**
+ * Returns the latest offset available.
+ *
+ * @param start the offset that the next micro-batch starts after: the end
of the last
+ * micro-batch, or the initial offset before the first one
+ */
+ def latestOffset(start: String): String
- /** Plans the partitions that read the data after `start` up to and
including `end`. */
+ /**
+ * Plans the partitions that read the data after `start` up to and including
`end`. Spark plans
+ * each range once, also for a micro-batch that reads the source several
times.
+ */
def partitions(start: String, end: String): Array[InputPartition]
- def read(partition: InputPartition): PartitionReader[ColumnarBatch]
+ /**
+ * Returns the reader factory of a micro-batch, which Spark serializes to
the executors. It may
Review Comment:
**4. [Low] The doc of `readerFactory()` does not say that its factory reads
the partitions of every micro-batch**
"the reader factory of a micro-batch" reads as one factory for each
micro-batch, but `readerFactory()` gets no range, and an engine may keep the
factory and read the partitions of later micro-batches with it, as
NativeDataSourceSuite.scala:1774-1776 says. The native stream reader handles
this by keeping the state in the partitions, which answered [item 1 of round
4](https://github.com/apache/spark/pull/59934#discussion_r4419050240), but the
contract is not written here, so another `ColumnarStreamReader` could put the
state of the last planned range in its factory and bring that bug back. Could
this say that the factory reads the partitions of any micro-batch, so what
changes with each micro-batch goes with its partitions?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,37 +384,104 @@ class NativeDataSourceStreamReader(
offset
}
- override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
- library.streamPartitions(handle.get, start, end)
+ override def partitions(start: String, end: String): Array[InputPartition] =
synchronized {
+ val planned = call(name, "plan a scan
of")(library.streamPartitions(handle.get, start, end))
+ val partitions = NativeInputPartition.from(planned.partitions)
+ val bytes = bytesOrEmpty(Option(planned.state))
+ // A small state goes with each task, while the executors fetch a larger
one once for the
+ // micro-batch, as a broadcast. Spark only plans the next range once the
micro-batch of the
+ // previous one finished, so the copies of the previous broadcast on the
executors are removed
+ // then. The broadcast itself is destroyed once its micro-batch is
committed, whose partitions
+ // are not read anymore then, also by an engine.
+ uncommitted.values.lastOption.flatten.foreach(_.unpersist(blocking =
false))
Review Comment:
**5. [Nit] The `unpersist` in `partitions()` only does something when
commits are skipped**
`MicroBatchExecution.constructNextBatch` commits the previous micro-batch
(`cleanUpLastExecutedMicroBatch`, MicroBatchExecution.scala:1157) before
`runBatch` plans the next range, so `commit(end)` has already destroyed and
removed its entry, and `uncommitted.values.lastOption` is empty here. This line
only finds a broadcast when Spark skipped the commit of the previous
micro-batch, as with async progress tracking, while the comment above reads as
if it always removes the previous copies. Could the comment say so, or the line
go, since the next commit destroys the broadcast anyway?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsArrowStreamRead.java:
##########
@@ -0,0 +1,62 @@
+/*
+ * 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 only checks that the source initialized the callbacks of
the stream, not its
+ * arrays, while it checks the arrays of 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. Each call goes through the JVM,
which then calls the
+ * callback of the stream of the source, so the engine releases the stream
before the JVM exits.
+ * The first call on a native thread attaches the thread to the JVM, until the
thread exits.
+ *
+ * @since 4.4.0
+ */
+@Evolving
+public interface SupportsArrowStreamRead extends PartitionReaderFactory {
+
+ /**
+ * Initializes an Arrow C stream with the data of a partition. It is called
in the running task
+ * that reads the partition, on a thread that has its {@code TaskContext},
and fails otherwise.
+ * The caller reads the stream in that task: once the task ended, the source
may stop returning
+ * data and fail. The caller owns
Review Comment:
**8. [Nit] A line of the Javadoc breaks early**
This line ends after "The caller owns", while the lines around it wrap near
100 characters. Could L48-L52 be rewrapped?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -267,65 +276,90 @@ class NativeBatchPlanner(
call(name, "plan a scan of")(optional(library.pushLimit(handle.get,
limit))).getOrElse(false)
}
- override def pruneColumns(requiredSchema: StructType): Boolean = {
- val pruned = call(name, "plan a scan of") {
+ override def pruneColumns(requiredSchema: StructType): Unit = {
+ // read gets the columns to read as its schema, and returns exactly them,
so a library that
+ // does not implement prune_columns, which only lets it plan to read them,
reads them too.
+ call(name, "plan a scan of") {
optional(library.pruneColumns(handle.get, requiredSchema.fieldNames))
- }.getOrElse(false)
- if (pruned) {
- readSchema = requiredSchema
}
- pruned
+ readSchema = requiredSchema
}
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).toArray[InputPartition],
+ 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], readSchema, timeZoneId)
+ }
+
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): Unit = {
+ NativePartitionReader.openEngineStream[NativeInputPartition](location,
name,
+ NativeLibrary.READ_BATCH, partition, readSchema, timeZoneId,
streamAddress) { partition =>
+ (readerState, partition)
+ }
}
}
-/** Reads a native data source for a micro-batch streaming scan. */
+/**
+ * Reads a native data source for a micro-batch streaming scan.
+ *
+ * @param timeZoneId the time zone of the timestamps of the schema that the
stream reader got
+ */
class NativeDataSourceStreamReader(
location: NativeLibraryLocation,
name: String,
schema: StructType,
- @transient private val handle: NativeHandle)
+ timeZoneId: String,
+ handle: NativeHandle)
extends ColumnarStreamReader {
import NativeCalls._
- // Set on the driver whenever the partitions of a micro-batch are planned,
and serialized to
- // the executors with the partitions of that micro-batch.
- private var readerState: Array[Byte] = Array.emptyByteArray
+ // The ends of the planned ranges that are not committed yet, in the order
in which they were
+ // planned, with the broadcast of the state of their micro-batch, if it was
broadcast. Spark plans
+ // each range once (see ColumnarMicroBatchStream), so the scans of a
micro-batch get the
+ // partitions of a single stream_partitions call, with the state that it
returned.
+ private val uncommitted = mutable.LinkedHashMap.empty[String,
Option[Broadcast[Array[Byte]]]]
+ // The context that broadcasts the states, taken on the driver when the
reader is created, so
+ // that a range can also be planned on a thread without an active session.
+ private val sparkContext =
+ SparkContext.getActive.getOrElse(SparkSession.active.sparkContext)
Review Comment:
**7. [Nit] The fallback to `SparkSession.active.sparkContext` cannot give a
running context**
Thank you for taking the context when the reader is created, which answers
[item 10 of round
10](https://github.com/apache/spark/pull/59934#discussion_r4474718480).
`SparkContext.getActive` is only empty when no `SparkContext` runs in the JVM,
and then the context of any session is stopped too, or `SparkSession.active`
throws, so `broadcast` (L399) fails either way. Could this keep only one of
them, such as `SparkSession.active.sparkContext`? `toMicroBatchStream` runs on
the thread of the query, whose session is active, which
`ColumnarMicroBatchStream` also relies on for its initial offset log
(ColumnarTable.scala:251).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -447,13 +553,31 @@ class NativeDataSourceStreamWriter(
case class NativeInputPartition(index: Int, bytes: Array[Byte]) extends
InputPartition
object NativeInputPartition {
- def from(partitions: Array[Array[Byte]]): Array[InputPartition] = {
+ def from(partitions: Array[Array[Byte]]): Array[NativeInputPartition] = {
Option(partitions).getOrElse(Array.empty).zipWithIndex.map { case (bytes,
i) =>
NativeInputPartition(i, Option(bytes).getOrElse(Array.emptyByteArray))
}
}
}
+/**
+ * A partition of a micro-batch, with the state that the stream reader
returned with the
+ * partitions of the micro-batch, if any. The state changes with each
micro-batch, so it goes with
+ * the partitions rather than with the reader factory, which an engine may
keep while Spark plans
+ * the next micro-batch.
+ *
+ * @param state the bytes of the state, empty if the stream reader returned
none, or for a state
+ * larger than 64 KiB, their broadcast, so that each executor
fetches it once for the
+ * micro-batch rather than with each task
+ */
+case class NativeMicroBatchPartition(
+ partition: NativeInputPartition,
+ state: Either[Array[Byte], Broadcast[Array[Byte]]])
+ extends InputPartition {
+ /** The state to read the partition with, empty if the stream reader
returned none. */
+ def readerState: Array[Byte] = state.fold(identity, _.value)
Review Comment:
**1. [Low] The partitions of a committed or stopped micro-batch fail to read
only when their state was broadcast**
`commit(end)` (L412-L418) and `stop()` (L427-L432) destroy the broadcasts of
the micro-batches that they end. A partition of such a micro-batch still reads
when its state stayed inline, but with a state above 64 KiB, `_.value` here, or
the serialization of the broadcast with a task, fails with the internal error
"Attempted to use Broadcast(N) after it was destroyed". So the same read
succeeds or fails depending on the size of the state. The error also comes
before `call` (L463 for the reader of Spark, L469 for an engine), so it is not
a `NATIVE_DATA_SOURCE_ERROR` and does not name the data source, and the Javadoc
of `SupportsArrowStreamRead` does not say how long the partitions of a
micro-batch can be read.
NativeDataSourceSuite.scala:1757-1769 relies on the small case: it reads the
partitions of the last micro-batch after `query.stop()`. `MicroBatchExecution`
commits a micro-batch only when it constructs the next one
(MicroBatchExecution.scala:1157), so that micro-batch is not committed, and
`stop()` destroys its broadcast. The test passes because the state of
`native_counter` stays inline. ([Item 2 of round
7](https://github.com/apache/spark/pull/59934#discussion_r4463375800) asked
`stop()` for an `unpersist`, after which the executors could still fetch the
broadcast from the driver.)
Could the Javadoc say that the partitions of a micro-batch can be read until
the micro-batch is committed or its query stops, and the test read them before
`query.stop()`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -969,22 +1234,70 @@ class NativeDataWriter(
/** Creates the `SparkDsTaskContext` of the running task. */
private[ffi] object NativeTask {
+ // The flag of each running task that reads, set once the task failed or
completed by a single
+ // pair of listeners for all its reads, which then also remove it: a task
that reads many
+ // partitions, as after a coalesce, registers them once. The contexts of the
reads keep their
+ // flag, so they still see it once it is removed. A task is the key rather
than its attempt ID,
+ // which is 0 for each context of `TaskContext.empty()`. The keys are weak,
compared by identity,
+ // so that a context that never ends, such as one that a test sets on its
thread, does not stay.
+ private val endedTasks = new MapMaker().weakKeys().makeMap[TaskContext,
AtomicBoolean]()
+
+ /** Returns the flag that is set once the given task failed or completed. */
+ private def endedFlag(taskContext: TaskContext): AtomicBoolean = {
+ Option(endedTasks.get(taskContext)).getOrElse {
+ val ended = new AtomicBoolean(false)
+ Option(endedTasks.putIfAbsent(taskContext, ended)).getOrElse {
+ // A listener runs right away if the task already failed or completed.
+ def end(): Unit = {
+ ended.set(true)
+ endedTasks.remove(taskContext, ended)
+ }
+ taskContext.addTaskFailureListener((_, _) => end())
Review Comment:
**6. [Nit] A read that starts after its task completed registers another
failure listener**
Thank you for the shared flag, which answers [item 7 of round
8](https://github.com/apache/spark/pull/59934#discussion_r4471627323). Once the
task completed, `end()` has removed its entry, so a read that starts after
that, such as the next partition of a coalesced scan read by the stdin writer
thread of a `PipedRDD`, registers a new pair. The completion listener runs
right away and removes the entry again, but `addTaskFailureListener` only runs
a listener right away if the task failed (TaskContextImpl.scala:179-185), so
the failure listener of each such read stays in the context. That is bounded by
the reads after the end of the task, but the comment above says that a task
registers them once. Could `endedFlag` return a set flag without listeners when
the task already failed or completed?
--
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]
