HyukjinKwon commented on code in PR #59942:
URL: https://github.com/apache/spark/pull/59942#discussion_r4418943149
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -189,11 +189,64 @@ 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]]. The reader gets the
offset that each
+ * micro-batch starts after, the end of the last one, for every trigger, and
the limit on the data
+ * of a micro-batch, its default read limit, all the available data unless it
limits it. With
+ * `Trigger.AvailableNow`, Spark runs micro-batches up to the latest offset
without a limit that
+ * the first micro-batch got, rather than wrapping the stream in one that asks
for the latest
+ * offset from the initial offset: in one micro-batch, or within the limit in
several.
+ */
+class ColumnarMicroBatchStream(reader: ColumnarStreamReader)
+ extends MicroBatchStream with SupportsTriggerAvailableNow {
+
+ // The latest offset without a limit that the reader returned last, which
the progress of the
+ // query reports.
+ @volatile private var lastLatestOffset: Offset = _
+ // The limit on the data of each micro-batch, asked once.
+ private lazy val readLimit = reader.defaultReadLimit()
+ // Whether the trigger is Trigger.AvailableNow.
+ @volatile private var availableNow = false
+ // The end of the run with Trigger.AvailableNow: the first latest offset
without a limit 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 getDefaultReadLimit: ReadLimit = readLimit
+
+ override def latestOffset(start: Offset, limit: ReadLimit): Offset = {
+ def ask(limit: ReadLimit, end: Option[Offset]): Offset =
+ ColumnarStreamOffset(reader.latestOffset(start.json(), limit,
end.map(_.json())))
+ val allAvailable = ReadLimit.allAvailable()
+ if (availableNow && availableNowEnd == null) {
+ availableNowEnd = ask(allAvailable, None)
+ }
+ if (limit == allAvailable) {
+ val latest = if (availableNow) availableNowEnd else ask(allAvailable,
None)
+ lastLatestOffset = latest
+ latest
+ } else if (availableNow) {
+ lastLatestOffset = availableNowEnd
+ ask(limit, Some(availableNowEnd))
+ } else {
+ lastLatestOffset = ask(allAvailable, None)
+ ask(limit, None)
Review Comment:
Done in bcb8f8d: Kept the two calls, and documented them: the header
(Streaming reads and `latest_offset`) and the guide now say that with a limit
Spark may call `latest_offset` more than once for a micro-batch, and that a
call with `SPARK_DS_READ_ALL_AVAILABLE` may only be for the progress, so a
library must not count the calls as micro-batches or reserve data in them.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -56,47 +56,55 @@
*
* - Batch reads. On the driver, Spark creates a reader with create_reader,
pushes operations down
* by calling, each at most once and in this order, push_limit and
prune_columns, plans the
- * partitions with partitions, serializes the state of the reader with
serialize_reader, and
- * closes it. Then it calls read on the executors, in one task for each
partition. If Spark never
- * plans the partitions, for example to only explain the plan of the query,
it closes the reader
- * once the reader is no longer referenced.
- * - Streaming reads. On the driver, Spark creates a stream reader with
create_stream_reader and
- * keeps it until the streaming query stops. For each micro-batch, it asks
for the latest_offset,
- * plans the partitions with stream_partitions and serializes the state with
- * serialize_stream_reader, then calls read on the executors. Offsets are
defined by the library,
- * as compact JSON on a single line: Spark stores each offset as a line of
the offset log of the
- * checkpoint, so it rejects an offset that is not valid JSON or contains a
line break. When a
- * streaming query restarts, from its checkpoint, Spark creates a new stream
reader and may call
- * stream_partitions again for the last range that it planned before, so the
partitions of a
- * range must read the same data each time.
+ * partitions with partitions, which also returns the state of the reader,
and closes it. Then it
+ * calls read on the executors, in one task for each partition, with the
schema to read and the
+ * state of the reader. If Spark never plans the partitions, for example to
only explain the plan
+ * of the query, it closes the reader once the reader is no longer
referenced.
+ * - Streaming reads. On the driver, Spark creates a stream reader with
create_stream_reader, with
+ * the location of the checkpoint of the source, and keeps it until the
streaming query stops. For
+ * each micro-batch, or only before the first one with Trigger.AvailableNow,
it asks for the
+ * latest_offset, with the offset that the micro-batch starts after, plans
the partitions with
+ * stream_partitions, which also returns the state of the stream reader,
then calls read on the
+ * executors. A data source with the capability SPARK_DS_ADMISSION_CONTROL
limits the data of each
+ * micro-batch: Spark asks for the limit once with default_read_limit, and
passes it to each
+ * latest_offset, also with Trigger.AvailableNow, whose run then reads the
data that was available
+ * when it started in several micro-batches. Offsets are defined by the
library, as compact JSON
Review Comment:
Done in bcb8f8d: The Streaming reads section now says that Spark asks for
the `latest_offset` for each micro-batch (once for a run of
`Trigger.AvailableNow`), passes the limit to the one that ends each
micro-batch, and may call it more than once with a limit. The comment of
`latest_offset` now says that it returns the latest offset within the limit and
up to the end of the run, and the guide says the same.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -304,6 +315,107 @@ 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.
+ */
+ const char* checkpoint_location;
+} SparkDsStreamReaderInfo;
+
+/*
+ * A request for the initial offset of a stream reader, as initial_offset gets
it. Spark may call
+ * initial_offset more than once for a query, for example when it restarts
before its first
+ * micro-batch is committed, so it must return the same offset each time: a
library whose initial
+ * offset depends on when it is asked, such as the latest data, keeps it under
checkpoint_location.
+ */
+typedef struct SparkDsInitialOffsetInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+} SparkDsInitialOffsetInfo;
+
+/* The kinds of the limits on the data of a micro-batch, see SparkDsReadLimit.
*/
+#define SPARK_DS_READ_ALL_AVAILABLE 0
+#define SPARK_DS_READ_MAX_ROWS 1
+
+/*
+ * A limit on the data of each micro-batch of a streaming read, which
default_read_limit returns
+ * for a data source with the capability SPARK_DS_ADMISSION_CONTROL. The
library sets the fields
+ * that fit in the size that Spark set.
+ */
+typedef struct SparkDsReadLimit {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /* SPARK_DS_READ_ALL_AVAILABLE, as Spark initializes it, or
SPARK_DS_READ_MAX_ROWS. */
+ int32_t kind;
+ /* With SPARK_DS_READ_MAX_ROWS, the maximum number of rows of a micro-batch,
at least 1. */
+ int64_t max_rows;
Review Comment:
Done in bcb8f8d: The doc of `limit` in `SparkDsLatestOffsetInfo` and the
guide now say that when there is data after `start`, the latest offset must be
after `start`, even if the next unit of data is larger than the limit.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -813,6 +816,133 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
}
}
+ nativeTest("pass the checkpoint location and the start of a micro-batch to
the stream reader") {
+ val session = newSession(defaultPackage)
+ val checkpoint = Utils.createTempDir()
+ val dir = Utils.createTempDir()
+ // Runs the query until it processed the rows [0, maxOffset), 5 more rows
in each micro-batch,
+ // and returns the checkpoint location and the start offsets that the
stream reader got.
+ def run(maxOffset: Int): (String, Seq[Long]) = {
+ val checkpointPath = new File(dir, s"checkpoint_$maxOffset")
+ val startsPath = new File(dir, s"starts_$maxOffset")
+ val query = session.readStream.format("native_counter")
+ .option("max_offset", maxOffset.toString).option("step", "5")
+ .option("checkpoint_path", checkpointPath.getPath)
+ .option("starts_path", startsPath.getPath)
+ .load()
+ .writeStream
+ .format("noop")
+ .option("checkpointLocation", checkpoint.getPath)
+ .start()
+ try query.processAllAvailable() finally query.stop()
+ val starts = Files.readAllLines(startsPath.toPath).asScala.toSeq
+ .map(_.stripPrefix("{\"offset\":").stripSuffix("}").toLong)
+ (Files.readString(checkpointPath.toPath), starts)
+ }
+ // The first micro-batch starts after the initial offset, and the next
ones after the end of
+ // the previous one, until no more rows are available.
+ val (location, starts) = run(12)
+ assert(location.endsWith(s"${checkpoint.getCanonicalPath}/sources/0"),
location)
+ assert(starts.take(4) == Seq(0L, 5L, 10L, 12L) && starts.drop(4).forall(_
== 12L), starts)
+ // A query that restarts from its checkpoint starts after the end of its
last micro-batch, with
+ // the same checkpoint location.
+ val (restartedLocation, restartedStarts) = run(15)
+ assert(restartedLocation == location)
+ assert(restartedStarts.take(2) == Seq(12L, 15L), restartedStarts)
+ }
+
+ nativeTest("read up to the latest offset with Trigger.AvailableNow, also
after a restart") {
+ val session = newSession(defaultPackage)
+ // Spark does not wrap the stream of a native data source, which supports
the trigger itself,
+ // even with the internal flag that wraps the others.
+ session.conf.set("spark.sql.streaming.triggerAvailableNowWrapper.enabled",
"true")
+ val checkpoint = Utils.createTempDir()
+ val dir = Utils.createTempDir()
+ val out = Utils.createTempDir()
+ // Runs the query with Trigger.AvailableNow, and returns the start offsets
that the stream
+ // reader got.
+ def run(name: String): Seq[String] = {
+ val startsPath = new File(dir, name)
+ val query = session.readStream.format("native_counter")
+ .option("max_offset", "12").option("step", "5").option("starts_path",
startsPath.getPath)
+ .load()
+ .writeStream
+ .format("native_sink")
+ .option("path", out.getPath)
+ .option("checkpointLocation", checkpoint.getPath)
+ .trigger(Trigger.AvailableNow())
+ .start()
+ try query.awaitTermination() finally query.stop()
+ Files.readAllLines(startsPath.toPath).asScala.toSeq
+ }
+ // Each run asks for the latest offset once, after the end of the last
micro-batch, and reads
+ // up to it: a restart neither reads the same rows again nor goes back to
the initial offset.
+ assert(run("first") == Seq("""{"offset":0}"""))
+ assert(run("second") == Seq("""{"offset":5}"""))
+ val written = out.listFiles().filter(_.getName.startsWith("part-"))
+ .flatMap(file => Files.readAllLines(file.toPath).asScala)
+ assert(written.map(_.toLong).sorted.toSeq == (0L until 10L))
+ }
+
+ nativeTest("limit the rows of each micro-batch with admission control") {
+ val session = newSession(defaultPackage)
+ // Runs the query until it read the rows [0, 12), all available from the
start, and returns
+ // the number of rows of each micro-batch and the latest offset that the
last one reported.
+ def run(availableNow: Boolean, options: (String, String)*): (Seq[Long],
String) = {
+ val writer =
session.readStream.format("native_counter").option("max_offset", "12")
+ .options(options.toMap).load()
+ .writeStream
+ .format("noop")
+ .option("checkpointLocation", Utils.createTempDir().getPath)
+ val query = (if (availableNow) writer.trigger(Trigger.AvailableNow())
else writer).start()
+ try {
+ if (availableNow) query.awaitTermination() else
query.processAllAvailable()
+ val batches = query.recentProgress.filter(_.numInputRows > 0).toSeq
+ (batches.map(_.numInputRows), batches.last.sources.head.latestOffset)
+ } finally {
+ query.stop()
+ }
+ }
+ // Each micro-batch reads at most 5 rows, while the progress reports the
latest offset
+ // without the limit.
+ assert(run(availableNow = false, "max_rows_per_trigger" -> "5") ==
+ ((Seq(5L, 5L, 2L), """{"offset":12}""")))
Review Comment:
Done in bcb8f8d: `run` now returns the latest offset of the first
micro-batch, which is `{"offset":12}` while its end is `{"offset":5}`.
--
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]
