HyukjinKwon commented on code in PR #59942:
URL: https://github.com/apache/spark/pull/59942#discussion_r4418945306
##########
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:
Done in bcb8f8d: `latestOffset` is now two calls, as suggested.
##########
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:
Done in bcb8f8d: `defaultReadLimit` now returns a
`NativeLibrary.ReadLimitInfo` with an `int kind` and a `long maxRows`, which
`latestOffset` also takes.
##########
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:
Done in bcb8f8d: `run` now takes the trigger, and the test checks that
`Trigger.Once()` reads `Seq(12L)`.
--
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]
