HyukjinKwon commented on code in PR #59931:
URL: https://github.com/apache/spark/pull/59931#discussion_r4417634705
##########
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:
Done in 84d446a: the test now checks that the rejected writer was aborted
with no messages (`_ABORTED` holds 0) and closed (`_CLOSED`).
##########
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:
Done in 84d446a: `createWriter` records the names and formats of its schema
at every level in the writer state, and `createDataWriter` fails unless its
schema has the same.
##########
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:
Done in 84d446a: it now says that each micro-batch makes 5 more rows
available after its start.
##########
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:
Done in 84d446a: the test removes the entries of its table at its start.
##########
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:
Done in 84d446a: `single` and `state` call one `atMostOne`, and `Sink` has
one accessor that rethrows `error`.
--
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]
