dongjoon-hyun commented on code in PR #59942: URL: https://github.com/apache/spark/pull/59942#discussion_r4418390289
########## 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: **2. [Low] The header says that each `latest_offset` gets the limit, but some calls get `READ_ALL_AVAILABLE`** With a limit, Spark also calls `latest_offset` with `SPARK_DS_READ_ALL_AVAILABLE`: in each trigger but `Trigger.AvailableNow`, for the progress ([`ColumnarTable.scala` L240](https://github.com/apache/spark/blob/51de9259ebcffbc0183b76a063e6b97912390582/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala#L240)), and once at the start of a run of `Trigger.AvailableNow`, for the end of the run ([L230](https://github.com/apache/spark/blob/51de9259ebcffbc0183b76a063e6b97912390582/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala#L230)). Two sentences of the header then do not hold: - "passes it to each latest_offset", here and in the guide (`sql-native-data-sources.md` L118); - "For each micro-batch, or only before the first one with Trigger.AvailableNow, it asks for the latest_offset" (L64-66 of this file): with a limit, Spark asks before each micro-batch of `Trigger.AvailableNow` too. The comment of `latest_offset` itself (L679) still says "Returns the most recent offset available", without the limit and the end in its info. A library that counts its calls as triggers, for example a rate source that makes more rows available at each call, or one that reserves data in `latest_offset`, would move twice as fast with a limit. Could the header say that 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? ########## 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: **4. [Low] The assertion on the reported latest offset passes without the latest offset without the limit** The last micro-batch reads the rows after 10, so its limited latest offset is `{"offset":12}` too, and this passes if the progress reports the limited offset, as before this PR. The first micro-batch tells them apart: its `endOffset` is `{"offset":5}`, while its `latestOffset` must be `{"offset":12}`. Could `run` return the latest offset of the first micro-batch instead (L901)? ```scala (batches.map(_.numInputRows), batches.head.sources.head.latestOffset) ``` ########## 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: **1. [Medium] Without `Trigger.AvailableNow`, a limited stream calls `latest_offset` twice in each trigger** With a limit and any trigger but `Trigger.AvailableNow`, this branch asks the library for the latest offset twice: once with `READ_ALL_AVAILABLE` for the progress, and once with the limit. Spark calls `latestOffset` in every iteration of its loop, also when there is no new data, and then only sleeps for [`spark.sql.streaming.pollingDelay`](https://github.com/apache/spark/blob/51de9259ebcffbc0183b76a063e6b97912390582/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4839-L4845), 10 ms by default. So an idle query with a limit makes about 200 calls a second instead of 100, and every call of a library that lists remote files or asks a broker is doubled. The two calls also see the data at two different times. If rows arrive between them, the limited offset can be beyond the reported latest offset, so the progress shows a `latestOffset` behind its `endOffset`, and a lag computed from the two becomes negative. Kafka computes both from a single fetch ([`KafkaMicroBatchStream.latestOffset`](https://github.com/apache/spark/blob/51de9259ebcffbc0183b76a063e6b97912390582/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala#L136-L183)), and a Python data source reports its own latest offset with [`reportLatestOffset`](https://github.com/apache/spark/blob/51de9259ebcffbc0183b76a063e6b97912390582/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonMicroBatchStream.scala#L141-L148). Since version 1 is not released yet, could one call return both? For example, `SparkDsLatestOffsetInfo` could get an optional `set_latest` sink that the library calls with the latest offset without the limit, and Spark would only ask again when the library does not call it. ########## 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: **5. [Low] The test of `Trigger.AvailableNow` passes without the end of the run** `step` defaults to `max_offset`, so the rows `[0, 12)` are all available from the start: the latest offset without the limit is 12 at every call, and the end of the run never lowers a limited offset. The run reads 5, 5 and 2 rows also if Spark passes no `end`, or if the library ignores it. With fewer rows available when the run starts, the end matters, for example with a `step` of 6: ```scala // Only the rows [0, 6) are available when the run starts. assert(run(availableNow = true, "max_rows_per_trigger" -> "5", "step" -> "6")._1 == Seq(5L, 1L)) ``` Without the end, the run would read 5, 5 and 2 rows instead, since `native_counter` makes 6 more rows available after each start. ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java: ########## @@ -630,86 +775,167 @@ boolean pruneColumns(long reader, String[] columnNames) throws Throwable { } } - byte[][] partitions(long reader) throws Throwable { - MethodHandle f = function(Function.PARTITIONS); - try (Sink sink = new Sink()) { - check((int) f.invokeExact(reader, Ffm.SINK, sink.id), Function.PARTITIONS); - return sink.all(); + /** The partitions that a reader planned, with the state of the reader that reads them. */ + static final class PlannedPartitions { + final byte[][] partitions; + /** The state of the reader, or null if it has none. */ + final byte[] state; + + PlannedPartitions(byte[][] partitions, byte[] state) { + this.partitions = partitions; + this.state = state; } } - byte[] serializeReader(long reader) throws Throwable { - return serialize(Function.SERIALIZE_READER, reader); + PlannedPartitions partitions(long reader) throws Throwable { + MethodHandle f = function(Function.PARTITIONS); + try (Memory memory = new Memory(); Sink sink = new Sink()) { + long info = memory.zeroed(PARTITIONS_INFO_SIZE); + Platform.putLong(null, info, PARTITIONS_INFO_SIZE); + check((int) f.invokeExact(reader, info, Ffm.SINK, Ffm.STATE_SINK, sink.id), + Function.PARTITIONS); + return new PlannedPartitions(sink.values(), state(sink, Function.PARTITIONS)); + } } void closeReader(long reader) throws Throwable { function(Function.CLOSE_READER).invokeExact(reader); } String initialOffset(long streamReader) throws Throwable { - return offset(Function.INITIAL_OFFSET, streamReader); + MethodHandle f = function(Function.INITIAL_OFFSET); + try (Memory memory = new Memory(); Sink sink = new Sink()) { + long info = memory.zeroed(INITIAL_OFFSET_INFO_SIZE); + Platform.putLong(null, info, INITIAL_OFFSET_INFO_SIZE); + check((int) f.invokeExact(streamReader, info, Ffm.SINK, sink.id), Function.INITIAL_OFFSET); + return offset(Function.INITIAL_OFFSET, sink); + } } - String latestOffset(long streamReader) throws Throwable { - return offset(Function.LATEST_OFFSET, streamReader); + /** + * Returns the latest offset. + * + * @param start the offset that the next micro-batch starts after + * @param limitKind {@link #READ_ALL_AVAILABLE} or {@link #READ_MAX_ROWS} + * @param maxRows the maximum number of rows with {@link #READ_MAX_ROWS} + * @param end the offset that a run of Trigger.AvailableNow reads up to, or null + */ + String latestOffset(long streamReader, String start, int limitKind, long maxRows, String end) + throws Throwable { + MethodHandle f = function(Function.LATEST_OFFSET); + try (Memory memory = new Memory(); Sink sink = new Sink()) { + long info = memory.zeroed(LATEST_OFFSET_INFO_SIZE); + Platform.putLong(null, info, LATEST_OFFSET_INFO_SIZE); + Platform.putLong(null, info + LATEST_OFFSET_INFO_START, memory.string(start)); + Platform.putInt(null, info + LATEST_OFFSET_INFO_LIMIT_KIND, limitKind); + Platform.putLong(null, info + LATEST_OFFSET_INFO_MAX_ROWS, maxRows); + Platform.putLong(null, info + LATEST_OFFSET_INFO_END, end == null ? 0 : memory.string(end)); + check((int) f.invokeExact(streamReader, info, Ffm.SINK, sink.id), Function.LATEST_OFFSET); + return offset(Function.LATEST_OFFSET, sink); + } + } + + /** + * Returns the default read limit of a stream reader, as {kind, maximum number of rows}, of + * which the kind is {@link #READ_ALL_AVAILABLE} if the library sets nothing. + */ + long[] defaultReadLimit(long streamReader) throws Throwable { Review Comment: **10. [Nit] `defaultReadLimit` returns the kind and the maximum in a `long[]`** The kind is an `int32_t`, read with `getInt`, widened into a `long[]`, and narrowed again with `kind.toInt` in `NativeDataSourceStreamReader`, which also depends on the order of the two values. A field added to `SparkDsReadLimit` later would change the length of the array and the pattern `val Array(kind, maxRows)`. Could it return a small class with an `int kind` and a `long maxRows`, as `streamPartitions` returns a `PlannedPartitions`? ########## 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") { Review Comment: **11. [Nit] No test checks that `Trigger.Once` reads all the available data** The header, the guide and the description say that `Trigger.Once` ignores the limit, but this test covers only the default trigger and `Trigger.AvailableNow`, and `TestRangeStreamReader` in `ColumnarDataSourceSuite` ignores the limit. A change of `ColumnarMicroBatchStream` that asked with its default read limit rather than the limit that Spark passes would read 5 of the 12 rows with `Trigger.Once`, and the tests would still pass. Could `run` take the trigger, and the test check `Seq(12L)` with `Trigger.Once()`, which is deprecated but still supported? ########## 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: **6. [Low] `Trigger.AvailableNow` also ignores the limit when Spark falls back to a single micro-batch** The header says that Spark passes the limit "except with Trigger.Once", here and at L388-389, and the guide says that "`Trigger.Once` ignores the limit". With `Trigger.AvailableNow`, when another source of the query does not support it and [`spark.sql.streaming.triggerAvailableNowWrapper.enabled`](https://github.com/apache/spark/blob/51de9259ebcffbc0183b76a063e6b97912390582/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4342-L4352) is false, as by default, Spark [falls back to a single micro-batch](https://github.com/apache/spark/blob/51de9259ebcffbc0183b76a063e6b97912390582/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L152-L181), whose sources [get `ReadLimit.allAvailable()`](https://github.com/apache/spark/blob/51de9259ebcffbc0183b76a063e6b97912390582/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L386-L397). Spark does not call `prepareForTriggerAv ailableNow` then either, so a native source reads all the available data in one micro-batch, the large first micro-batch that this PR avoids, while the warning only says that the limit "is ignored when Trigger.Once is used". Could the header and the guide say "except with Trigger.Once, and with Trigger.AvailableNow when another source of the query does not support it"? ########## 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: **7. [Low] `limit_kind` and `max_rows` repeat the fields of `SparkDsReadLimit`** The limit has two layouts now: `SparkDsReadLimit`, which `default_read_limit` fills, and these two fields, which `NativeLibrary.latestOffset` writes from a `ReadLimit`. A kind added later within version 1, such as a maximum number of bytes for `ReadLimit.maxBytes`, would need a field in both structs, two offsets in the bindings and in the layout test, and a library would read the same limit from two places. Since the layout is not frozen yet, could the info point to the limit instead, for example with `const SparkDsReadLimit* limit`, which Spark sets to the default read limit, or to a limit of the kind `SPARK_DS_READ_ALL_AVAILABLE`? ########## 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: **3. [Low] The header does not say what to do when the next unit of data is larger than the limit** `max_rows` is "the maximum number of rows of a micro-batch", and the latest offset is "within the limit" (L389-390), while Spark's [`ReadMaxRows`](https://github.com/apache/spark/blob/51de9259ebcffbc0183b76a063e6b97912390582/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/streaming/ReadMaxRows.java#L23-L24) is only approximate. A library whose offsets cannot split a file or a group of records, with a unit larger than `max_rows`, can only follow the header by returning `start`. Spark then sees no new data: a query with a processing-time trigger never reads further, and a run of `Trigger.AvailableNow` [terminates](https://github.com/apache/spark/blob/51de9259ebcffbc0183b76a063e6b97912390582/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L794-L797) without the rest of the data, and without an error. Could the header say that, when there is data after `start`, the latest offset must be after `start`, even if the next unit is larger than the limit? The file source does the same with `maxBytesPerTrigger`: it always [takes the first file](https://github.com/apache/spark/blob/51de9259ebcffbc0183b76a063e6b97912390582/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/FileStreamSource.scala#L146-L163), whatever its size. ########## 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: **8. [Nit] `end` is the end of a micro-batch in the other structs** In `SparkDsStreamPartitionsInfo` (L408) and `SparkDsCommitOffsetInfo` (L416), `end` is the end of a micro-batch, while here it is the end of a whole run of `Trigger.AvailableNow`, which a library could take for the end of the next micro-batch. The Scala side has the same overlap between `ColumnarStreamReader.latestOffset(start, limit, end)` and `partitions(start, end)`, and `TestRangeStreamReader` had to name it `runEnd`, next to its own `end`. Since the names of the fields are fixed with 4.4.0, could it be `run_end`, or `available_now_end`? ########## 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) + } Review Comment: **9. [Nit] The three branches of `latestOffset` can be two calls** `ask(allAvailable, None)` is in three places, and each branch sets `lastLatestOffset`. The latest offset without the limit is the same in each branch: `availableNowEnd`, asked once, with `Trigger.AvailableNow`, and a new one otherwise. Since `availableNowEnd` stays null without `Trigger.AvailableNow`, this does the same, with the calls in the same order: ```scala 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 latest = if (!availableNow) { ask(ReadLimit.allAvailable(), None) } else { if (availableNowEnd == null) availableNowEnd = ask(ReadLimit.allAvailable(), None) availableNowEnd } lastLatestOffset = latest if (limit == ReadLimit.allAvailable()) latest else ask(limit, Option(availableNowEnd)) } ``` -- 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]
