HyukjinKwon commented on code in PR #59931:
URL: https://github.com/apache/spark/pull/59931#discussion_r4417632386


##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -304,6 +308,31 @@ typedef struct SparkDsCreateInfo {
   int64_t num_options;
 } SparkDsCreateInfo;
 
+/* 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
+   * URI, possibly of a file system that the library cannot access, such as 
s3a or hdfs.
+   */
+  const char* checkpoint_location;
+} SparkDsStreamReaderInfo;
+
+/* 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. With Trigger.AvailableNow, 
Spark asks for the
+   * latest offset only once, when the query starts, with the initial offset.

Review Comment:
   Done in 84d446a: the sentence now says that `start` is the end of the last 
micro-batch, or the initial offset before the first one, for every trigger, and 
that with `Trigger.AvailableNow` Spark asks once and reads up to that offset 
(see 2). A new test runs `Trigger.AvailableNow` twice on the same checkpoint, 
with the wrapper flag on, and checks the starts (`{"offset":0}`, then 
`{"offset":5}`) and that no row is written twice.



##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -189,11 +189,31 @@ 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.
+ */
+class ColumnarMicroBatchStream(reader: ColumnarStreamReader)
+  extends MicroBatchStream with SupportsAdmissionControl {

Review Comment:
   Done in 84d446a: `ColumnarMicroBatchStream` implements 
`SupportsTriggerAvailableNow`. `prepareForTriggerAvailableNow` sets a flag, and 
the first `latestOffset(start, limit)` asks the library with the real start and 
keeps that offset as the end of the run.



##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -304,6 +308,31 @@ typedef struct SparkDsCreateInfo {
   int64_t num_options;
 } SparkDsCreateInfo;
 
+/* 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
+   * URI, possibly of a file system that the library cannot access, such as 
s3a or hdfs.

Review Comment:
   Done in 84d446a: it now says that it is a fully qualified Hadoop path, whose 
characters are not escaped as in a URI.



##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -487,26 +558,22 @@ typedef struct SparkDsLibrary {
   /*
    * Returns the most recent offset available.
    *
+   * info: the offset that the next micro-batch starts after
    * set_offset: called once with the offset, as compact JSON on a single line
    */
-  int (*latest_offset)(void* reader, SparkDsSink set_offset, void* sink);
+  int (*latest_offset)(void* reader, const SparkDsLatestOffsetInfo* info, 
SparkDsSink set_offset,
+                       void* sink);
 
   /*
    * Plans the partitions that read the data after `start` up to and including 
`end`, two offsets
    * as JSON.
    *
    * add_partition: called with each partition, serialized in a format defined 
by the library
+   * set_state: called once with the state of the stream reader that read 
needs on the executors
+   *            to read these partitions, or not at all if there is none
    */
   int (*stream_partitions)(void* reader, const char* start, const char* end,
-                           SparkDsSink add_partition, void* sink);
-
-  /*
-   * Serializes the state of the stream reader that read needs on the 
executors. Spark calls it
-   * after each call to stream_partitions.
-   *
-   * set_state: called once with the state, or not at all if there is none
-   */
-  int (*serialize_stream_reader)(void* reader, SparkDsSink set_state, void* 
sink);
+                           SparkDsSink add_partition, SparkDsSink set_state, 
void* sink);

Review Comment:
   Done in 84d446a: they get `SparkDsPartitionsInfo`, 
`SparkDsInitialOffsetInfo` (both only a size for now), 
`SparkDsStreamPartitionsInfo` (start, end) and `SparkDsCommitOffsetInfo` (end). 
The test library checks their sizes, and the Rust and C++ examples of the docs 
were rebuilt and run.



##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -569,41 +615,75 @@ void schema(long dataSource, long schemaAddress) throws 
Throwable {
   }
 
   long createReader(long dataSource, long schemaAddress) throws Throwable {
-    return createWithSchema(Function.CREATE_READER, dataSource, schemaAddress);
-  }
-
-  long createStreamReader(long dataSource, long schemaAddress) throws 
Throwable {
-    return createWithSchema(Function.CREATE_STREAM_READER, dataSource, 
schemaAddress);
+    MethodHandle f = function(Function.CREATE_READER);
+    try (Memory memory = new Memory()) {
+      long out = memory.zeroed(8);
+      check((int) f.invokeExact(dataSource, schemaAddress, out), 
Function.CREATE_READER);
+      return readHandle(out, Function.CREATE_READER);
+    }
   }
 
-  private long createWithSchema(Function function, long dataSource, long 
schemaAddress)
+  /**
+   * Creates a stream reader.
+   *
+   * @param checkpointLocation the location of the checkpoint of the source
+   */
+  long createStreamReader(long dataSource, long schemaAddress, String 
checkpointLocation)
       throws Throwable {
-    MethodHandle f = function(function);
+    MethodHandle f = function(Function.CREATE_STREAM_READER);
     try (Memory memory = new Memory()) {
+      long info = memory.zeroed(STREAM_READER_INFO_SIZE);
+      Platform.putLong(null, info, STREAM_READER_INFO_SIZE);
+      Platform.putLong(
+        null, info + STREAM_READER_INFO_CHECKPOINT_LOCATION, 
memory.string(checkpointLocation));
       long out = memory.zeroed(8);
-      check((int) f.invokeExact(dataSource, schemaAddress, out), function);
-      return readHandle(out, function);
+      check((int) f.invokeExact(dataSource, schemaAddress, info, out),
+        Function.CREATE_STREAM_READER);
+      return readHandle(out, Function.CREATE_STREAM_READER);
+    }
+  }
+
+  /** A writer handle, with the state of the writer that the data writers 
need. */
+  static final class WriterHandle {
+    final long handle;
+    /** The state of the writer, or null if it has none. */
+    final byte[] state;
+
+    WriterHandle(long handle, byte[] state) {
+      this.handle = handle;
+      this.state = state;
     }
   }
 
   /**
-   * Creates a writer.
+   * Creates a writer. If the library returns an invalid state, the writer is 
closed.
    *
    * @param mode {@link #WRITE_APPEND} or {@link #WRITE_TRUNCATE}
    * @param queryId the ID of the streaming query of a streaming write, or 
null for a batch write
    */
-  long createWriter(long dataSource, long schemaAddress, int mode, String 
queryId)
+  WriterHandle createWriter(long dataSource, long schemaAddress, int mode, 
String queryId)
       throws Throwable {
     MethodHandle f = function(Function.CREATE_WRITER);
-    try (Memory memory = new Memory()) {
+    try (Memory memory = new Memory(); Sink sink = new Sink()) {
       long info = memory.zeroed(WRITE_INFO_SIZE);
       Platform.putLong(null, info, WRITE_INFO_SIZE);
       Platform.putInt(null, info + WRITE_INFO_MODE, mode);
       Platform.putLong(
         null, info + WRITE_INFO_QUERY_ID, queryId == null ? 0 : 
memory.string(queryId));
       long out = memory.zeroed(8);
-      check((int) f.invokeExact(dataSource, schemaAddress, info, out), 
Function.CREATE_WRITER);
-      return readHandle(out, Function.CREATE_WRITER);
+      check((int) f.invokeExact(dataSource, schemaAddress, info, 
Ffm.STATE_SINK, sink.id, out),
+        Function.CREATE_WRITER);
+      long handle = readHandle(out, Function.CREATE_WRITER);
+      try {
+        return new WriterHandle(handle, state(sink, Function.CREATE_WRITER));
+      } catch (Throwable e) {
+        try {
+          closeWriter(handle);

Review Comment:
   Done in 84d446a: Spark calls `abort` with no messages and an epoch ID of -1, 
then `close_writer`. The header says so in the Writes paragraph and in the doc 
of `abort`.



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