dongjoon-hyun commented on code in PR #59931:
URL: https://github.com/apache/spark/pull/59931#discussion_r4417868873


##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -189,11 +189,51 @@ 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]]. It supports admission 
control only so that
+ * the data source gets the offset that each micro-batch starts after: its 
default read limit reads
+ * all the available data, which Spark then always asks for.
+ */
+/**
+ * The micro-batch stream of a [[ColumnarDataSource]]. The reader gets the 
start of each
+ * micro-batch, the end of the last one, for every trigger: with 
`Trigger.AvailableNow`, Spark runs
+ * micro-batches up to the latest offset that the first micro-batch got, 
rather than wrapping the
+ * stream in one that asks for the latest offset from the initial offset.
+ */
+class ColumnarMicroBatchStream(reader: ColumnarStreamReader)
+  extends MicroBatchStream with SupportsTriggerAvailableNow {
+
+  // The offset that the reader returned last, which the progress of the query 
reports.
+  @volatile private var lastLatestOffset: Offset = _
+  // With Trigger.AvailableNow, whether the latest offset of the run is the 
first one asked.
+  @volatile private var availableNow = false
+  @volatile private var availableNowEnd: Offset = _
 
   override def initialOffset(): Offset = 
ColumnarStreamOffset(reader.initialOffset())

Review Comment:
   **1. [Medium] `initial_offset` is now called twice for the first 
micro-batch, so `latest_offset` and `stream_partitions` may get different 
starts**
   
   Since this stream is now a `SupportsAdmissionControl`, through 
`SupportsTriggerAvailableNow`, `constructNextBatch` takes [the 
`SupportsAdmissionControl` 
branch](https://github.com/apache/spark/blob/84d446a51f5b45da4adcaf57e42f8372e5a3f739/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L1085-L1091),
 whose start comes from 
[`getStartOffset`](https://github.com/apache/spark/blob/84d446a51f5b45da4adcaf57e42f8372e5a3f739/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L1043-L1054):
 it calls `initialOffset()` while the query has no start offset yet. Then 
[`runBatch`](https://github.com/apache/spark/blob/84d446a51f5b45da4adcaf57e42f8372e5a3f739/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L1224-L1232)
 calls `initialOffset` again to plan the same micro-batch. Before this PR, the 
plain `MicroBatchStream` branch only called `latestOffse
 t()`, so a run called `initial_offset` once.
   
   `initialOffset` is not kept here, so `SparkDsLatestOffsetInfo.start` and 
`SparkDsStreamPartitionsInfo.start` of the first micro-batch come from two 
calls of the library. The header does not ask `initial_offset` to return the 
same offset each time. A library whose initial offset is the latest data at the 
time of the call, as `startingOffsets=latest` is for Kafka, could return 100 to 
the first call and compute a latest offset of 105 from it, then return 110 to 
the second call: `stream_partitions` then gets the range from 110 to 105. With 
`Trigger.AvailableNow`, the end of the run also comes from the first value.
   
   Could the stream keep the initial offset, for example with `private lazy val 
initial = ColumnarStreamOffset(reader.initialOffset())`? Independently, the 
header could say that Spark may call `initial_offset` more than once, also when 
the query restarts before its first micro-batch is committed, and that it must 
return the same offset each time. A library whose initial offset is not 
deterministic could keep it under `checkpoint_location`, as the Kafka source 
keeps its initial offsets in the checkpoint.



##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -911,6 +979,34 @@ class NativeDataSourceSuite extends QueryTest with 
SharedSparkSession {
         "msg" -> "create failed with the error code 1."))
   }
 
+  nativeTest("reject a state returned more than once") {
+    def check(action: String, name: String, function: String)(f: SparkSession 
=> Unit): Unit = {
+      val e = intercept[Exception](f(newSession(defaultPackage)))
+      checkError(
+        exception = findError(e, "NATIVE_DATA_SOURCE_ERROR"),
+        condition = "NATIVE_DATA_SOURCE_ERROR",
+        parameters = Map("action" -> action, "name" -> name,
+          "msg" -> s"$function returned 2 states instead of one."))
+    }
+    check("plan a scan of", "native_range", "partitions") { session =>
+      range(session, "states" -> "2").collect()
+    }
+    check("plan a scan of", "native_counter", "stream_partitions") { session =>
+      val query = session.readStream.format("native_counter").option("states", 
"2").load()
+        .writeStream.format("noop").start()
+      try query.processAllAvailable() finally query.stop()
+    }
+    val dir = Utils.createTempDir()
+    check("plan a write to", "native_sink", "create_writer") { session =>
+      session.range(2).write.format("native_sink").option("path", dir.getPath)
+        .option("states", "2").mode("append").save()
+    }
+    // The rejected writer is aborted, with no messages, so that the library 
can undo what
+    // create_writer prepared, and released.
+    assert(Files.readString(new File(dir, "_ABORTED").toPath) == "0")

Review Comment:
   **4. [Low] `_ABORTED` holds the number of NULL messages, so this does not 
check that the abort got no messages**
   
   `abortWrite` of the test library writes [the number of messages that are 
NULL](https://github.com/apache/spark/blob/84d446a51f5b45da4adcaf57e42f8372e5a3f739/sql/core/src/test/resources/native-datasource/test_native_datasource.cc#L1270-L1276)
 to `_ABORTED`, so `"0"` also holds if Spark passes messages that are not NULL. 
Could the test library also record `num_messages`, for example as `"0,0"`?
   
   Also, only the batch write is covered. The rejected writer of a streaming 
write gets an epoch ID of -1 although it has a query ID, so a check with 
`writeStream` and the option `states` set to 2 would pin that `_ABORTED`, not 
`_ABORTED_0`, is written and that the writer is closed.



##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -189,11 +189,51 @@ 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]]. It supports admission 
control only so that
+ * the data source gets the offset that each micro-batch starts after: its 
default read limit reads
+ * all the available data, which Spark then always asks for.
+ */

Review Comment:
   **5. [Nit] Two Scaladoc comments precede `ColumnarMicroBatchStream`**
   
   The first one is left from the first commit, and only the second one 
applies. Also, the comment of `availableNow`, "whether the latest offset of the 
run is the first one asked", rather describes `availableNowEnd`, which has no 
comment. Maybe "Whether the trigger is Trigger.AvailableNow" for the first, and 
"The end of the run with Trigger.AvailableNow: the first latest offset asked" 
for the second.



##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -185,28 +193,32 @@ class NativeDataSource(
     }
   }
 
-  /** Creates a writer for a batch write, or for a streaming write with the ID 
of its query. */
+  /**
+   * Creates a writer for a batch write, or for a streaming write with the ID 
of its query, with
+   * the state that its data writers need.
+   */
   private def createWriter(
       library: NativeLibrary,
       schema: StructType,
+      timeZoneId: String,

Review Comment:
   **8. [Nit] `createWriter` takes the time zone that is now a field**
   
   `timeZoneId` is now a field of `NativeDataSource`, and both callers pass 
that field to this parameter, which shadows it. Without the parameter, the time 
zone of the schema that `create_writer` gets would have a single source.



##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -414,19 +519,23 @@ typedef struct SparkDsLibrary {
    * Creates a reader for a micro-batch streaming scan.
    *
    * schema: the schema to read: the one the user specified, or the one 
returned by schema
+   * info: the location of the checkpoint of the source
    * reader: set to the stream reader handle
    */
-  int (*create_stream_reader)(void* source, struct ArrowSchema* schema, void** 
reader);
+  int (*create_stream_reader)(void* source, struct ArrowSchema* schema,
+                              const SparkDsStreamReaderInfo* info, void** 
reader);

Review Comment:
   **3. [Low] `create_reader` still takes no struct that can grow, unlike 
`create_stream_reader`**
   
   `create_stream_reader` now gets a `SparkDsStreamReaderInfo`, and 
`partitions` and `initial_offset` get a struct with only its size, so that 
inputs can be added within version 1. 
[`create_reader`](https://github.com/apache/spark/blob/84d446a51f5b45da4adcaf57e42f8372e5a3f739/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L510-L516),
 right above, stays `(source, schema, reader)`, so a batch reader cannot get 
anything more when it is created, such as an ID of the scan as 
`SparkDsWriteInfo` has the ID of the query, without a new version.
   
   The same holds for `commit` and `abort`, and for `commit_data_writer` and 
`abort_data_writer`. This PR already gives a second meaning to the `epoch_id` 
of `abort`, -1 for a writer whose `create_writer` Spark rejected, which a field 
of an `Info` struct could carry instead. Since these signatures are frozen with 
version 1, could the functions that may need more inputs take an `Info` struct 
too?



##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -304,6 +309,62 @@ typedef struct SparkDsCreateInfo {
   int64_t num_options;
 } SparkDsCreateInfo;
 
+/* 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. */
+typedef struct SparkDsInitialOffsetInfo {
+  /* The size of this struct in the header that Spark was built with, as in 
SparkDsHost. */
+  int64_t size;
+} SparkDsInitialOffsetInfo;
+
+/* 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 only once, when the query starts, and 
reads up to it.

Review Comment:
   **2. [Low] With `Trigger.AvailableNow`, the latest offset is asked before 
the first micro-batch that Spark plans, not always when the query starts**
   
   `prepareForTriggerAvailableNow` only sets a flag, and the end of the run is 
the first offset that `latestOffset(start, limit)` returns. When the offset log 
of the checkpoint ends with a micro-batch that was not committed, 
[`populateStartOffsets`](https://github.com/apache/spark/blob/84d446a51f5b45da4adcaf57e42f8372e5a3f739/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L855-L860)
 marks it as constructed, and Spark runs it again before it asks for any latest 
offset. The first `latest_offset` then comes after that micro-batch, with its 
end as `start`, and the end of the run includes the data that arrived while it 
ran.
   
   Could this say "only once, before the first micro-batch that it plans", or 
similar? [The streaming reads of the 
lifecycle](https://github.com/apache/spark/blob/84d446a51f5b45da4adcaf57e42f8372e5a3f739/sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h#L65-L67)
 also say that Spark asks for the `latest_offset` for each micro-batch, which 
could mention this exception.



##########
sql/core/src/test/resources/native-datasource/test_native_datasource.cc:
##########
@@ -1033,28 +1067,38 @@ int planPartitions(void* handle, SparkDsSink 
add_partition, void* sink) {
       int64_t stop = std::min(hi, start + step);
       sendBytes(add_partition, sink, std::to_string(start) + "," + 
std::to_string(stop));
     }
+    sendState(reader->state, set_state, sink, serialize(reader->state));
   });
 }
 
-int serializeReader(void* handle, SparkDsSink set_state, void* sink) {
-  return guarded([&] { sendBytes(set_state, sink, 
serialize(fromHandle<Reader>(handle)->state)); });
-}
-
 void closeReader(void* handle) { delete fromHandle<Reader>(handle); }
 
 #ifndef READ_ONLY
 
-int initialOffset(void*, SparkDsSink set_offset, void* sink) {
-  return guarded([&] { sendBytes(set_offset, sink, "{\"offset\":0}"); });
+int initialOffset(void*, const SparkDsInitialOffsetInfo* info, SparkDsSink 
set_offset,
+                  void* sink) {
+  return guarded([&] {
+    checkSize("SparkDsInitialOffsetInfo", info->size, 
sizeof(SparkDsInitialOffsetInfo));
+    sendBytes(set_offset, sink, "{\"offset\":0}");
+  });
 }
 
-// Each call makes `step` more rows available, up to `max_offset`.
-int latestOffset(void* handle, SparkDsSink set_offset, void* sink) {
+// Makes `step` more rows available after the start of the micro-batch, up to 
`max_offset`.
+int latestOffset(void* handle, const SparkDsLatestOffsetInfo* info, 
SparkDsSink set_offset,
+                 void* sink) {
   return guarded([&] {
-    Options& state = fromHandle<Reader>(handle)->state;
-    int64_t latest = std::min(getLong(state, "max_offset", 10),
-                              getLong(state, "current", 0) + getLong(state, 
"step", 10));
-    state["current"] = std::to_string(latest);
+    const Options& state = fromHandle<Reader>(handle)->state;
+    if (info->start == nullptr) throw std::runtime_error("no start offset");

Review Comment:
   **6. [Nit] The test library does not check the size of 
`SparkDsLatestOffsetInfo`, `SparkDsStreamReaderInfo` or `SparkDsWriteInfo`**
   
   `latestOffset` reads `info->start` here, 
[`createStreamReader`](https://github.com/apache/spark/blob/84d446a51f5b45da4adcaf57e42f8372e5a3f739/sql/core/src/test/resources/native-datasource/test_native_datasource.cc#L1005-L1007)
 reads `info->checkpoint_location` and 
[`createWriter`](https://github.com/apache/spark/blob/84d446a51f5b45da4adcaf57e42f8372e5a3f739/sql/core/src/test/resources/native-datasource/test_native_datasource.cc#L1203-L1204)
 reads `info->mode`, without the `checkSize` that the other functions call. 
`readPartition` and `createDataWriter` check the size with `offsetof` inline 
instead. With `checkSize` everywhere, a call in `NativeLibrary` that forgets to 
set the size, which `memory.zeroed` leaves at 0, would fail for every struct.



##########
docs/sql-native-data-sources.md:
##########
@@ -101,32 +101,36 @@ Spark calls the functions of the table as follows:
   that the data source does not support before it computes the data to write. 
When the user does
   not specify a schema, Spark asks for it with `schema`.
 - **Batch reads**: on the driver, Spark creates a reader with `create_reader`, 
pushes operations
-  down with `push_limit` and `prune_columns`, plans the partitions with 
`partitions`, and
-  serializes the state of the reader with `serialize_reader`. Then it calls 
`read` on the
-  executors, in one task for each partition, and `read` exports an 
`ArrowArrayStream`. Its
-  callbacks can be called from any thread, but never concurrently.
-- **Streaming reads**: on the driver, Spark creates a stream reader with 
`create_stream_reader`.
-  For each micro-batch, it asks for the `latest_offset`, plans the partitions 
with
-  `stream_partitions`, and 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, and rejects an offset that is not valid JSON or contains a line 
break, such as
-  pretty-printed JSON.
+  down with `push_limit` and `prune_columns`, and plans the partitions with 
`partitions`, which
+  also returns the state of the reader. Then it calls `read` on the executors, 
in one task for each
+  partition, with a `SparkDsReadInfo` that has the partition, the schema to 
read and the state of
+  the reader, and `read` exports an `ArrowArrayStream`. Its callbacks can be 
called from any
+  thread, but never concurrently.
+- **Streaming reads**: on the driver, Spark creates a stream reader with 
`create_stream_reader`,
+  which gets the location of the checkpoint of the source in a 
`SparkDsStreamReaderInfo`. For each
+  micro-batch, 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,
+  and 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, and rejects
+  an offset that is not valid JSON or contains a line break, such as 
pretty-printed JSON.
 - **Writes**: on the driver, right before the tasks of the write (or of a 
micro-batch) start, Spark
   creates a writer with `create_writer`, which gets a `SparkDsWriteInfo` with 
the mode of the write,
   `SPARK_DS_WRITE_APPEND` or `SPARK_DS_WRITE_TRUNCATE`, and the ID of the 
query of a streaming
-  write. With adaptive query execution, it first runs the shuffle stages of 
the data to write, and
-  creates no writer if they fail, but the last stage runs in the tasks of the 
write, so a writer can
-  still be aborted because computing the data failed: do irreversible work, 
such as removing the
-  existing data for an overwrite, in `commit`. On the executors, each task 
creates a data writer
-  with `create_data_writer`, writes the batches of its data with `write`, and 
calls
-  `commit_data_writer`, or `abort_data_writer` if the task or the commit 
fails. Then the driver
-  calls `commit` with the messages of the tasks, or `abort` if the write 
failed. A streaming write
-  commits each micro-batch with its epoch ID, and the ID of the query stays 
the same when the query
-  restarts from its checkpoint: together, they identify the output of a 
micro-batch, so that a
-  library can skip an epoch that it already committed. A micro-batch streaming 
query creates a new
-  stream writer for each micro-batch, and Spark closes the previous one when 
it creates the next
-  one. Because Spark is not told when a streaming query stops, the last stream 
writer of a query is
-  closed once it is no longer referenced, or when the query is restarted.
+  write, and returns the state of the writer. With adaptive query execution, 
it first runs the
+  shuffle stages of the data to write, and creates no writer if they fail, but 
the last stage runs
+  in the tasks of the write, so a writer can still be aborted because 
computing the data failed:
+  do irreversible work, such as removing the existing data for an overwrite, 
in `commit`. On the
+  executors, each task creates a data writer with `create_data_writer`, which 
gets the schema of
+  the data and the state of the writer in a `SparkDsDataWriterInfo`, writes 
the batches of its data
+  with `write`, and calls `commit_data_writer`, or `abort_data_writer` if the 
task or the commit
+  fails. Then the driver calls `commit` with the messages of the tasks, or 
`abort` if the write
+  failed. A streaming write commits each micro-batch with its epoch ID, and 
the ID of the query

Review Comment:
   **7. [Nit] The guide does not say that a rejected writer is aborted**
   
   The header now says that when Spark rejects what `create_writer` returned, 
it calls `abort` with no messages and an epoch ID of -1, also for a streaming 
writer, before `close_writer`. This paragraph still says that `abort` only 
follows a failed write. Since a library that only reads the guide may take an 
epoch ID of -1 for a batch write, could it mention this case too?



-- 
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]

Reply via email to