HyukjinKwon commented on code in PR #59932:
URL: https://github.com/apache/spark/pull/59932#discussion_r4418130422
##########
sql/core/src/test/resources/native-datasource/test_native_datasource.cc:
##########
@@ -957,47 +1036,63 @@ int createReader(void* handle, struct ArrowSchema*
schema, void** out) {
for (const auto& column : columns) findColumn(column);
if (returnsZero(source->options, "create_reader")) return;
auto* reader = new Reader();
+ reader->state["kind"] = std::to_string(SPARK_DS_READ_BATCH);
reader->state["columns"] = join(columns);
reader->state["lo"] = "0";
reader->state["hi"] = get(source->options, "end", "10");
reader->state["partitions"] = get(source->options, "partitions", "2");
reader->state["batch_size"] = get(source->options, "batch_size", "1000");
for (const char* option :
- {"fail", "slice", "short", "null_rows", "wait_for_interrupt",
"task_path"}) {
+ {"fail", "slice", "short", "null_rows", "wait_for_interrupt",
"task_path", "states"}) {
reader->state[option] = get(source->options, option, "");
}
*out = reader;
});
}
+// Passes the state to set_state once, or as many times as the option "states"
says.
+void sendState(const Options& options, SparkDsSink set_state, void* sink,
+ const std::string& state) {
+ std::string times = get(options, "states", "");
+ for (int64_t i = 0; i < (times.empty() ? 1 : std::stoll(times)); i++) {
+ sendBytes(set_state, sink, state);
+ }
+}
+
void closeSource(void* handle) { delete fromHandle<DataSource>(handle); }
#ifndef READ_ONLY
-int createStreamReader(void* handle, struct ArrowSchema* schema, void** out) {
+int createStreamReader(void* handle, struct ArrowSchema* schema,
+ const SparkDsStreamReaderInfo* info, void** out) {
return guarded([&] {
auto* source = fromHandle<DataSource>(handle);
if (source->name != name("native_counter")) {
throw std::runtime_error(source->name + " does not support streaming
reads");
}
auto columns = join(importColumns(schema));
+ if (info->checkpoint_location == nullptr) throw std::runtime_error("no
checkpoint location");
+ std::string checkpoint_path = get(source->options, "checkpoint_path", "");
+ if (!checkpoint_path.empty()) writeFile(checkpoint_path,
info->checkpoint_location);
if (returnsZero(source->options, "create_stream_reader")) return;
auto* reader = new Reader();
+ reader->state["kind"] = std::to_string(SPARK_DS_READ_MICRO_BATCH);
reader->state["columns"] = columns;
reader->state["max_offset"] = get(source->options, "max_offset", "10");
reader->state["step"] = get(source->options, "step",
reader->state["max_offset"]);
- reader->state["current"] = "0";
reader->state["batch_size"] = get(source->options, "batch_size", "1000");
- reader->state["commit_path"] = get(source->options, "commit_path", "");
- reader->state["bad_offset"] = get(source->options, "bad_offset", "");
- reader->state["closed_path"] = get(source->options, "closed_path", "");
+ for (const char* option :
+ {"commit_path", "bad_offset", "closed_path", "starts_path",
"states"}) {
+ reader->state[option] = get(source->options, option, "");
+ }
*out = reader;
});
}
int pushLimit(void* handle, int32_t limit, bool* applied) {
- return guarded([&] {
- auto* reader = fromHandle<Reader>(handle);
+ auto* reader = fromHandle<Reader>(handle);
+ if (reader->failed) orderViolation("push_limit after a failed call");
+ return tracked(&reader->failed, [&] {
// Spark calls push_limit at most once.
if (reader->state.count("limit") > 0) throw std::runtime_error("pushLimit
was called twice");
Review Comment:
Done in 0d5c9c2: `latest_offset`, `stream_partitions` and `commit_offset` go
through `tracked`, and record a call after a failed one. A second `push_limit`
is recorded as "push_limit called twice" instead of throwing.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -426,18 +451,26 @@ class NativeDataSourceStreamWriter(
partitionId: Int,
taskId: Long,
epochId: Long): DataWriter[ColumnarBatch] = {
- new NativeDataWriter(location, name, writerState, partitionId, taskId,
epochId)
+ new NativeDataWriter(
+ location, name, schema, timeZoneId, writerState, partitionId, taskId,
epochId)
}
+ // Whether the library committed the micro-batch, after which it is not
aborted, as the header
+ // says, even if Spark fails afterwards.
+ @transient private var committed = false
Review Comment:
Done in 0d5c9c2: `ColumnarBatchWrite` records that the commit returned, and
`ColumnarStreamingWrite` the committed epochs, and neither aborts after it. The
two native writers no longer have the field, and `ColumnarWriter` and
`ColumnarStreamWriter` say that Spark does not abort a committed write or
micro-batch.
##########
docs/sql-native-data-sources.md:
##########
@@ -135,20 +141,24 @@ function reports an error by returning a non-zero value,
with the message return
`last_error`, and Spark reports it as a `NATIVE_DATA_SOURCE_ERROR`. A C++
exception or a Rust panic
must not cross the interface. Each handle that a `create` function returns is
released exactly
once, by the `close` function of its kind. The close functions cannot fail,
and they can run after
-the query, so do the work that can fail, such as flushing data, in `commit`.
Strings are UTF-8.
+the query, so do the work that can fail, such as flushing data, in `commit`.
After a function on
+a handle fails, Spark only closes the handle, except that it aborts a writer
or a data writer whose
+commit or write failed, and it never aborts one whose commit succeeded: the
header lists the order
+of the calls that a library can rely on. Strings are UTF-8.
Review Comment:
Done in 0d5c9c2, as suggested.
##########
sql/core/src/test/resources/native-datasource/test_native_datasource.cc:
##########
@@ -1290,10 +1462,13 @@ int commitDataWriter(void* handle, SparkDsSink
set_message, void* sink) {
int abortDataWriter(void* handle) {
if (usedAfterClose(handle)) return 0;
+ auto* writer = fromHandle<DataWriter>(handle);
+ if (writer->abort_called) orderViolation("abort_data_writer called twice");
+ writer->abort_called = true;
return guarded([&] {
- auto* writer = fromHandle<DataWriter>(handle);
// Spark does not abort a data writer that committed: records it if it
does.
if (writer->committed) {
+ orderViolation("abort_data_writer after a successful commit");
writeFile(writer->state["path"] + "/_ABORTED_COMMITTED_DATA_WRITER", "");
Review Comment:
Done in 0d5c9c2: the `_ABORTED_COMMITTED_DATA_WRITER` marker and its
assertion are gone. `orderViolation` and `starts_path` share `appendLine`.
`orderViolation` still ignores a failure to write, since it also runs outside
of `guarded`, in functions such as the close functions that cannot fail, while
`starts_path` fails the call.
--
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]
