HyukjinKwon commented on code in PR #59931:
URL: https://github.com/apache/spark/pull/59931#discussion_r4423359337
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,17 +359,16 @@ class NativeDataSourceStreamReader(
}
override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
+ val planned = call(name, "plan a scan of") {
library.streamPartitions(handle.get, start, end)
}
- readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeStreamReader(handle.get))))
- NativeInputPartition.from(partitions)
+ readerState = bytesOrEmpty(Option(planned.state))
Review Comment:
Done in bcd6124: `NativeDataSourceStreamReader` now plans each range once.
It keeps the range that it planned last and its partitions in `@transient`
fields. When Spark plans the same range again, as both sides of a union of the
source with itself do, it returns those partitions without calling
`stream_partitions`, so `readerState` stays the state of that call. Only a new
range replaces both. At `stream_partitions`, the header now says that all the
scans of the source in a micro-batch read the partitions of a single call, with
the state that it returned.
The test library now numbers its `stream_partitions` calls in the partitions
and in the state, and `read` fails on a partition of one call with the state of
another. The new test "plan each micro-batch once, also when a query reads a
source twice" unions `native_counter` with itself and checks the rows. With the
new option `ranges_path`, it also checks that `stream_partitions` is called
once for each micro-batch. Without the memo, it fails with "read got the
partition of the call 1 of stream_partitions with the state of the call 2".
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -780,6 +822,7 @@ class NativePartitionReader(
throw error(name, "read", s"The batch has an offset of ${batch.offset},
but Spark only " +
"reads arrays with an offset of 0. Copy sliced arrays before exporting
them.")
}
+ checkBuffers(batch, "The batch")
Review Comment:
Done in bcd6124: the option `short` of the test library now takes
`batch_buffers`, which exports the batches with a NULL pointer to their
buffers. "reject the Arrow data that Spark cannot read" checks that the read
then fails with "The batch has 1 buffers, but a NULL pointer to them."
##########
sql/core/src/test/resources/native-datasource/test_native_datasource.cc:
##########
@@ -1033,58 +1125,86 @@ 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}"); });
+// Sends the offset, or, with the option `bad_option`, no offset for "none", a
pretty-printed one
+// for "multiline", or one that is not JSON for "invalid".
+void sendOffset(const Options& state, const std::string& bad_option,
SparkDsSink set_offset,
+ void* sink, int64_t offset) {
+ std::string bad_offset = get(state, bad_option, "");
+ if (bad_offset == "none") return;
+ if (bad_offset == "multiline") {
+ // Pretty-printed JSON, as serde_json::to_string_pretty writes it.
+ sendBytes(set_offset, sink, "{\n \"offset\": " + std::to_string(offset) +
"\n}");
+ } else if (bad_offset == "invalid") {
+ sendBytes(set_offset, sink, "offset " + std::to_string(offset));
+ } else {
+ sendBytes(set_offset, sink, "{\"offset\":" + std::to_string(offset) + "}");
+ }
}
-// Each call makes `step` more rows available, up to `max_offset`.
-int latestOffset(void* handle, SparkDsSink set_offset, void* sink) {
+int initialOffset(void* handle, const SparkDsInitialOffsetInfo* 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);
- std::string bad_offset = get(state, "bad_offset", "");
- if (bad_offset == "none") return;
- if (bad_offset == "multiline") {
- // Pretty-printed JSON, as serde_json::to_string_pretty writes it.
- sendBytes(set_offset, sink, "{\n \"offset\": " + std::to_string(latest)
+ "\n}");
- } else if (bad_offset == "invalid") {
- sendBytes(set_offset, sink, "offset " + std::to_string(latest));
- } else {
- sendBytes(set_offset, sink, "{\"offset\":" + std::to_string(latest) +
"}");
+ checkSize("SparkDsInitialOffsetInfo", info->size,
sizeof(SparkDsInitialOffsetInfo));
+ const Options& state = fromHandle<Reader>(handle)->state;
+ int64_t offset = 0;
+ std::string path = get(state, "initial_offsets_path", "");
+ if (!path.empty()) {
+ FILE* file = std::fopen(path.c_str(), "a+");
Review Comment:
Done in bcd6124: `initial_offset` of the test library now rewinds the stream
before it counts the lines, so the count no longer depends on where the C
library starts reading an `"a+"` stream. The writes still append.
--
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]
