HyukjinKwon commented on code in PR #59942:
URL: https://github.com/apache/spark/pull/59942#discussion_r4418944207
##########
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}""")))
+ // A run of Trigger.AvailableNow reads the rows available when it starts,
within the limit.
+ assert(run(availableNow = true, "max_rows_per_trigger" -> "5")._1 ==
Seq(5L, 5L, 2L))
Review Comment:
Done in bcb8f8d: Added the run with `"step" -> "6"`, which reads `Seq(5L,
1L)`.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -584,36 +758,55 @@ typedef struct SparkDsLibrary {
* Creates the writer for one task.
*
* task: the task, with the partition that it writes, valid until the data
writer is closed
- * writer_state: the state returned by serialize_writer, empty if there is
none
- * epoch_id: the ID of the micro-batch for a streaming write, or -1 for a
batch write
+ * info: the schema of the data, the state of the writer, and the ID of the
micro-batch
* data_writer: set to the data writer handle
*/
- int (*create_data_writer)(const SparkDsTaskContext* task, const uint8_t*
writer_state,
- int64_t writer_state_length, int64_t epoch_id,
void** data_writer);
+ int (*create_data_writer)(const SparkDsTaskContext* task, const
SparkDsDataWriterInfo* info,
+ void** data_writer);
/*
* Writes a batch of data.
*
- * array: a struct array whose fields are the columns
- * schema: the schema of the array
+ * array: a struct array whose fields are the columns, with the schema that
create_data_writer
+ * got
*/
- int (*write)(void* data_writer, struct ArrowArray* array, struct
ArrowSchema* schema);
+ int (*write)(void* data_writer, struct ArrowArray* array);
/*
* Commits the data written by the task. If it fails, Spark calls
abort_data_writer. If it
* succeeds, Spark does not call abort_data_writer, even if it then fails
the task, for example
* because set_message was called more than once.
*
+ * info: nothing yet but its size, so that inputs can be added within this
version
* set_message: called once with the message to pass to commit on the
driver, or not at all for
* an empty message
*/
- int (*commit_data_writer)(void* data_writer, SparkDsSink set_message, void*
sink);
+ int (*commit_data_writer)(void* data_writer, const
SparkDsFinishDataWriterInfo* info,
+ SparkDsSink set_message, void* sink);
- /* Aborts the data written by the task. */
- int (*abort_data_writer)(void* data_writer);
+ /*
+ * Aborts the data written by the task.
+ *
+ * info: nothing yet but its size, so that inputs can be added within this
version
+ */
+ int (*abort_data_writer)(void* data_writer, const
SparkDsFinishDataWriterInfo* info);
/* Releases a data writer handle, after its commit or abort. */
void (*close_data_writer)(void* data_writer);
+
+ /*
-------------------------------------------------------------------------------------------
+ * Admission control (driver)
+ *
-------------------------------------------------------------------------------------------
*/
+
+ /*
+ * Returns the limit on the data of each micro-batch of a stream reader,
typically from the
+ * options of the data source, such as a maximum number of rows. Spark asks
once, when the
+ * streaming query starts, for a data source with the capability
SPARK_DS_ADMISSION_CONTROL, and
+ * passes the limit to each latest_offset, except with Trigger.Once.
Review Comment:
Done in bcb8f8d: The header (`limit` and `default_read_limit`) and the guide
now also name `Trigger.AvailableNow` when another source of the query does not
support it and `spark.sql.streaming.triggerAvailableNowWrapper.enabled` is
false.
##########
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;
+} SparkDsReadLimit;
+
+/* 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 without a limit once, before the first
micro-batch that it
+ * plans, and reads up to it: in a single micro-batch, or with a limit, in
micro-batches of
+ * which it asks for each latest offset with that offset as the end.
+ */
+ const char* start;
+ /*
+ * The limit on the data after start, as in SparkDsReadLimit: the default
read limit of a data
+ * source with the capability SPARK_DS_ADMISSION_CONTROL, or
SPARK_DS_READ_ALL_AVAILABLE, as
+ * always for the other data sources and for Trigger.Once. With a limit, the
latest offset is the
+ * end of the data of the next micro-batch, within the limit.
+ */
+ int32_t limit_kind;
+ int64_t max_rows;
Review Comment:
Done in bcb8f8d: `SparkDsLatestOffsetInfo` now has `const SparkDsReadLimit*
limit` instead of `limit_kind` and `max_rows`, which Spark sets to the default
read limit or to a limit of the kind `SPARK_DS_READ_ALL_AVAILABLE`. The struct
is 32 bytes; the bindings, the layout test and the test library are updated.
##########
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;
+} SparkDsReadLimit;
+
+/* 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 without a limit once, before the first
micro-batch that it
+ * plans, and reads up to it: in a single micro-batch, or with a limit, in
micro-batches of
+ * which it asks for each latest offset with that offset as the end.
+ */
+ const char* start;
+ /*
+ * The limit on the data after start, as in SparkDsReadLimit: the default
read limit of a data
+ * source with the capability SPARK_DS_ADMISSION_CONTROL, or
SPARK_DS_READ_ALL_AVAILABLE, as
+ * always for the other data sources and for Trigger.Once. With a limit, the
latest offset is the
+ * end of the data of the next micro-batch, within the limit.
+ */
+ int32_t limit_kind;
+ int64_t max_rows;
+ /*
+ * With Trigger.AvailableNow and a limit, the offset that the run reads up
to, which the latest
+ * offset must not go beyond, so that the run reads the data that was
available when it started;
+ * NULL otherwise.
+ */
+ const char* end;
Review Comment:
Done in bcb8f8d: Renamed to `run_end`, and `runEnd` on the Scala side
(`ColumnarStreamReader.latestOffset`); its doc says that it is not the end of
the next micro-batch.
--
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]
