dongjoon-hyun commented on code in PR #59931:
URL: https://github.com/apache/spark/pull/59931#discussion_r4417352691
##########
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:
**3. [Low] `checkpoint_location` is the string of a Hadoop `Path`, not an
escaped URI**
The location is `checkpointDir.toString` from
[`ResolveWriteToStream`](https://github.com/apache/spark/blob/f4a4ed242333c305e6125961d6e894d3e0674a41/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/ResolveWriteToStream.scala#L152-L153),
to which
[`StreamingUtils.getMetadataPath`](https://github.com/apache/spark/blob/f4a4ed242333c305e6125961d6e894d3e0674a41/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/utils/StreamingUtils.scala#L71)
appends `/sources/0`. `Path.toString` has a scheme, but its characters are not
percent-encoded. For a checkpoint at `/data/my ckpt%41`, the library gets
`file:/data/my ckpt%41/sources/0`. A library that parses it as a URI, as this
comment suggests, fails on the space or reads `%41` as `A`, and keeps its
recovery data in another directory.
Could it say that this is a fully qualified Hadoop path, whose characters
are not escaped?
##########
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:
**2. [Medium] Implementing `SupportsTriggerAvailableNow` here would keep
`start` exact for every trigger**
The stream already gets the start through `latestOffset(start, limit)`. If
it also implemented `SupportsTriggerAvailableNow`, Spark would neither wrap it
nor fall back to a single batch. `prepareForTriggerAvailableNow` could only set
a flag, and the first `latestOffset(start, limit)` would ask the library with
the real start and keep that offset as the end of the run. This is the shape of
`PythonMicroBatchStreamWithTriggerAvailableNow` for Python data sources.
The header would then not need an exception for `Trigger.AvailableNow` (item
1), and `start` would mean the same for every trigger, which admission control
will also rely on once it is added. If this is out of scope here, the header
could at least avoid promising the behavior of the wrapper.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -911,6 +946,29 @@ 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()
+ }
+ check("plan a write to", "native_sink", "create_writer") { session =>
+ session.range(2).write.format("native_sink").option("path",
Utils.createTempDir().getPath)
+ .option("states", "2").mode("append").save()
+ }
Review Comment:
**6. [Low] The test does not check that the rejected writer is released**
This only matches the error message. If the `try`/`catch` around
`state(sink, Function.CREATE_WRITER)` in `NativeLibrary.createWriter` were
dropped, the handle would leak, since it is not in a `NativeHandle` yet, and
this test would still pass. The test library only writes a close marker for a
writer that was committed or aborted (`closed_marker`). A marker for a writer
that is closed without either would let this test check the release, and would
also show the case of item 5.
##########
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:
**5. [Low] A writer whose state is rejected is closed without `commit` or
`abort`**
When `create_writer` succeeds but calls `set_state` twice, this closes the
writer. Since `ColumnarBatchWrite` never got a writer, Spark does not call
`abort` either. The rewritten Writes paragraph of the header still says that
the driver calls `commit` or `abort` and then closes the writer. A library that
prepares something in `create_writer`, such as a staging directory (the test
library writes `_CREATED`), then has no call in which to undo it, except
`close_writer`, which cannot fail.
This was already the case when `serialize_writer` returned more than one
value. Now that the state comes from `create_writer` itself, could Spark call
`abort` with no messages before it closes the writer, or the header say that
Spark may close a writer without `commit` or `abort` when it rejects what
`create_writer` returned?
##########
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:
**4. [Low] `stream_partitions`, `commit_offset`, `partitions` and
`initial_offset` still take fixed arguments**
This PR gives `create_stream_reader`, `latest_offset`, `read` and
`create_data_writer` a struct that starts with its size, so that they can get
new inputs within version 1. These four functions keep fixed arguments.
Anything they may need later, such as the ID of the micro-batch for
`stream_partitions` and `commit_offset` if Spark passes it one day, would then
need a new function or version 2.
Is that intended? Since this is the last chance to change their signatures
within version 1, it may be worth giving them an `Info` struct as well, even
one that only has its `size` for now.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java:
##########
@@ -789,6 +904,16 @@ private static byte[] single(Sink sink, Function function)
throws Throwable {
return values.length == 0 ? null : values[0];
}
+ /** Returns the state that a function passed to the set_state sink, or null
if none. */
+ private static byte[] state(Sink sink, Function function) throws Throwable {
Review Comment:
**10. [Nit] `state` repeats `single`, and `Sink.states` repeats `Sink.all`**
They differ only in the list they read and in a word of the message. A
single helper that takes the arrays and that word, such as
`atMostOne(sink.states(), function, "states")`, would keep one copy of the
check, and `Sink` could have one accessor that rethrows `error`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarDataSourceSuite.scala:
##########
@@ -353,6 +388,13 @@ class TestRangeStreamReader(end: Long, schema: StructType)
extends ColumnarStrea
}
}
+object TestRangeStreamReader {
+ // By the option "table" of the read: the checkpoint location of its stream
reader, and the
+ // start offsets that the stream reader got for the latest offset.
+ val checkpointLocations = new ConcurrentHashMap[String, String]()
+ val starts = new ConcurrentHashMap[String, Seq[String]]()
Review Comment:
**9. [Nit] These maps are never cleared**
`starts` keeps the start offsets of earlier runs under the same `table`, and
the new test asserts on the whole list. If the test runs twice in the same JVM,
for example from an sbt shell, the second run appends `{"offset":0}` after the
`{"offset":5}` of the first run, and `starts.tail.forall(...)` fails. Removing
the entries of the table at the start of the test, or using a table name of its
own for each run, would avoid that.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -88,48 +89,56 @@ class NativeDataSource(
if ((declared & NativeLibrary.BATCH_WRITE) != 0) {
capabilities.add(BATCH_WRITE)
}
- if ((declared & NativeLibrary.TRUNCATE) != 0) {
+ if ((declared & NativeLibrary.BATCH_TRUNCATE) != 0) {
capabilities.add(TRUNCATE)
}
capabilities
}
+ // The readers and the writers export the schemas with the time zone of the
session on the
+ // driver, also on the executors, where they pass the same schemas to the
library again.
+
override def batchPlanner(schema: StructType): ColumnarBatchPlanner = {
NativeArrow.checkSupported(schema, name, "plan a scan of")
val executorLocation = NativeDataSourceRegistry.distribute(location)
val library = NativeLibraries.get(location)
+ val timeZoneId = SQLConf.get.sessionLocalTimeZone
Review Comment:
**11. [Nit] The time zone is read in four places, and again by
`ArrowBatchingOptions`**
`batchPlanner`, `streamReader`, `writer` and `streamWriter` each read
`SQLConf.get.sessionLocalTimeZone`. For a write, the schema that
`create_data_writer` gets uses this value, but `ArrowBatchingDataWriter` builds
the arrays that `write` gets with the time zone that
[`ArrowBatchingOptions`](https://github.com/apache/spark/blob/f4a4ed242333c305e6125961d6e894d3e0674a41/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarWrite.scala#L181)
reads separately. They are equal today, because both are read in the same
`createBatchWriterFactory` call, but nothing ties them together. Reading the
time zone once, for example in a field of `NativeDataSource`, or passing the
one of `ArrowBatchingOptions` to the writer, would keep a single source.
##########
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:
**1. [Medium] The sentence about `Trigger.AvailableNow` only holds for a new
query, or with an internal flag**
By default, `ColumnarMicroBatchStream` does not implement
`SupportsTriggerAvailableNow`, so
[`MicroBatchExecution.getTrigger`](https://github.com/apache/spark/blob/f4a4ed242333c305e6125961d6e894d3e0674a41/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L152-L181)
runs `Trigger.AvailableNow` as a single batch (`SingleBatchExecutor`), and
[`constructNextBatch`](https://github.com/apache/spark/blob/f4a4ed242333c305e6125961d6e894d3e0674a41/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L1085-L1091)
passes the real start, from
[`getStartOffset`](https://github.com/apache/spark/blob/f4a4ed242333c305e6125961d6e894d3e0674a41/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala#L1043-L1054).
So Spark asks once, but a query that restarts from its checkpoint gets the end
of its last micro-batch, not the initial offset.
Only the internal `spark.sql.streaming.triggerAvailableNowWrapper.enabled`
passes the initial offset:
[`AvailableNowDataStreamWrapper.prepareForTriggerAvailableNow`](https://github.com/apache/spark/blob/f4a4ed242333c305e6125961d6e894d3e0674a41/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AvailableNowDataStreamWrapper.scala#L59-L66)
calls `latestOffset(initialOffset, ...)`, also after a restart, and returns
that offset for every batch. `start` is then not the offset that the next
micro-batch starts after, and the offset of a library that computes from
`start`, as `native_counter` now does, goes backwards. With the last committed
end at `{"offset":12}` and a step of 5, it returns `{"offset":5}`. Spark plans
the range from 12 to 5, which `stream_partitions` reads as empty, and commits
5, so the next run reads the rows 5 to 12 again.
Since this sentence is frozen with version 1, could it describe the default
instead: once, with the end of the last micro-batch, or the initial offset for
a new query? It could then promise nothing for the wrapper, or the wrapper
could be avoided as in item 2. No test runs `Trigger.AvailableNow` on a native
source, so a test with a restart would pin whichever behavior is chosen.
##########
sql/core/src/test/resources/native-datasource/test_native_datasource.cc:
##########
@@ -1252,20 +1298,24 @@ int createDataWriter(const SparkDsTaskContext* task,
const uint8_t* writer_state
});
}
-int writeBatch(void* handle, struct ArrowArray* array, struct ArrowSchema*
schema) {
+int writeBatch(void* handle, struct ArrowArray* array) {
return guarded([&] {
auto* writer = fromHandle<DataWriter>(handle);
// The batch is released here, so Spark does not release it after the call.
std::unique_ptr<ArrowArray, void (*)(ArrowArray*)> array_owner(
array, [](ArrowArray* a) { a->release(a); });
- std::unique_ptr<ArrowSchema, void (*)(ArrowSchema*)> schema_owner(
- schema, [](ArrowSchema* s) { s->release(s); });
maybeFail(writer->state, "write");
+ const ArrowSchema& schema = writer->schema;
+ if (array->n_children != schema.n_children) {
Review Comment:
**7. [Low] `writeBatch` only compares the number of columns with the schema
of `create_data_writer`**
The comment at the top of the file says that `write` fails unless its
batches have the columns of the schema that `create_data_writer` got, but this
only compares `n_children`. An `ArrowArray` has no names or formats, so the
arrays alone cannot show more. The read side compares the schema of `read` with
the columns that the reader recorded in its state, and the write side could do
the same. `createWriter` could record the names and formats of its schema in
the writer state, and `createDataWriter` could compare them with
`info->schema`. That would catch a schema that differs from the one
`create_writer` got, for example in its time zone (item 11).
##########
sql/core/src/test/resources/native-datasource/test_native_datasource.cc:
##########
@@ -1048,13 +1069,22 @@ int initialOffset(void*, SparkDsSink set_offset, void*
sink) {
return guarded([&] { 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`.
Review Comment:
**8. [Nit] A comment of `NativeDataSourceSuite` still describes the previous
`latest_offset`**
`latest_offset` now returns the start plus `step`, but the test "streaming
read and write" still says that [each call to `latestOffset` makes 5 more rows
available](https://github.com/apache/spark/blob/f4a4ed242333c305e6125961d6e894d3e0674a41/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala#L792).
It could say that each micro-batch makes 5 more rows available after its start.
--
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]
