HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4417463572
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -278,37 +284,54 @@ class NativeBatchPlanner(
override def plan(): ColumnarPlannedBatch = {
try {
- val partitions = call(name, "plan a scan
of")(library.partitions(handle.get))
- val readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeReader(handle.get))))
+ val planned = call(name, "plan a scan
of")(library.partitions(handle.get))
ColumnarPlannedBatch(
- NativeInputPartition.from(partitions),
- new NativePartitionReaderFactory(executorLocation, name, readerState,
readSchema))
+ NativeInputPartition.from(planned.partitions),
+ new NativePartitionReaderFactory(executorLocation, name,
+ bytesOrEmpty(Option(planned.state)), readSchema, timeZoneId))
} finally {
- // The executors only need the serialized state.
+ // The executors only need the state that the reader returned.
handle.close()
}
}
}
-/** Reads the partitions of a planned batch scan of a native data source, on
the executors. */
+/**
+ * Reads the partitions of a planned batch scan of a native data source, on
the executors. An
+ * engine can also get the Arrow C stream of a partition from it, see
[[SupportsArrowStreamRead]].
+ */
class NativePartitionReaderFactory(
location: NativeLibraryLocation,
name: String,
readerState: Array[Byte],
- readSchema: StructType)
- extends ColumnarOnlyPartitionReaderFactory {
+ readSchema: StructType,
+ timeZoneId: String)
+ extends ColumnarOnlyPartitionReaderFactory with SupportsArrowStreamRead {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, readSchema)
+ new NativePartitionReader(location, name, NativeLibrary.READ_BATCH,
readerState,
+ partition.asInstanceOf[NativeInputPartition].bytes, readSchema,
timeZoneId)
+ }
+
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): AutoCloseable = {
+ val task = NativePartitionReader.openStream(location, name,
NativeLibrary.READ_BATCH,
+ readerState, partition.asInstanceOf[NativeInputPartition].bytes,
readSchema, timeZoneId,
+ streamAddress)
+ // In case the caller does not close it, which must release the stream
before then.
+ Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ =>
task.close()))
Review Comment:
Done in 8eae514: there is no completion listener any more. Spark frees the
context once the library released the stream (see 2).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -278,37 +284,54 @@ class NativeBatchPlanner(
override def plan(): ColumnarPlannedBatch = {
try {
- val partitions = call(name, "plan a scan
of")(library.partitions(handle.get))
- val readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeReader(handle.get))))
+ val planned = call(name, "plan a scan
of")(library.partitions(handle.get))
ColumnarPlannedBatch(
- NativeInputPartition.from(partitions),
- new NativePartitionReaderFactory(executorLocation, name, readerState,
readSchema))
+ NativeInputPartition.from(planned.partitions),
+ new NativePartitionReaderFactory(executorLocation, name,
+ bytesOrEmpty(Option(planned.state)), readSchema, timeZoneId))
} finally {
- // The executors only need the serialized state.
+ // The executors only need the state that the reader returned.
handle.close()
}
}
}
-/** Reads the partitions of a planned batch scan of a native data source, on
the executors. */
+/**
+ * Reads the partitions of a planned batch scan of a native data source, on
the executors. An
+ * engine can also get the Arrow C stream of a partition from it, see
[[SupportsArrowStreamRead]].
+ */
class NativePartitionReaderFactory(
location: NativeLibraryLocation,
name: String,
readerState: Array[Byte],
- readSchema: StructType)
- extends ColumnarOnlyPartitionReaderFactory {
+ readSchema: StructType,
+ timeZoneId: String)
+ extends ColumnarOnlyPartitionReaderFactory with SupportsArrowStreamRead {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, readSchema)
+ new NativePartitionReader(location, name, NativeLibrary.READ_BATCH,
readerState,
+ partition.asInstanceOf[NativeInputPartition].bytes, readSchema,
timeZoneId)
+ }
+
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): AutoCloseable = {
Review Comment:
Done in 8eae514: `openArrowStream` returns nothing. Spark wraps the four
callbacks, keyed by a `private_data` of its own, so a moved stream keeps
working and a NULL `private_data` of the library is fine. It frees the context
after the `release` of the library.
##########
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.
+ */
+ const char* start;
+} SparkDsLatestOffsetInfo;
+
/*
* The task that calls read or create_data_writer, on an executor. It stays
valid until the stream
* that read returns is released, or until the data writer is closed.
Review Comment:
Done in 8eae514: the comment now says that the context also stays valid
until the stream is released when an engine releases it after the task ended,
from another thread.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -278,37 +284,54 @@ class NativeBatchPlanner(
override def plan(): ColumnarPlannedBatch = {
try {
- val partitions = call(name, "plan a scan
of")(library.partitions(handle.get))
- val readerState = bytesOrEmpty(
- call(name, "plan a scan
of")(optional(library.serializeReader(handle.get))))
+ val planned = call(name, "plan a scan
of")(library.partitions(handle.get))
ColumnarPlannedBatch(
- NativeInputPartition.from(partitions),
- new NativePartitionReaderFactory(executorLocation, name, readerState,
readSchema))
+ NativeInputPartition.from(planned.partitions),
+ new NativePartitionReaderFactory(executorLocation, name,
+ bytesOrEmpty(Option(planned.state)), readSchema, timeZoneId))
} finally {
- // The executors only need the serialized state.
+ // The executors only need the state that the reader returned.
handle.close()
}
}
}
-/** Reads the partitions of a planned batch scan of a native data source, on
the executors. */
+/**
+ * Reads the partitions of a planned batch scan of a native data source, on
the executors. An
+ * engine can also get the Arrow C stream of a partition from it, see
[[SupportsArrowStreamRead]].
+ */
class NativePartitionReaderFactory(
location: NativeLibraryLocation,
name: String,
readerState: Array[Byte],
- readSchema: StructType)
- extends ColumnarOnlyPartitionReaderFactory {
+ readSchema: StructType,
+ timeZoneId: String)
+ extends ColumnarOnlyPartitionReaderFactory with SupportsArrowStreamRead {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- new NativePartitionReader(
- location, name, readerState,
partition.asInstanceOf[NativeInputPartition].bytes, readSchema)
+ new NativePartitionReader(location, name, NativeLibrary.READ_BATCH,
readerState,
+ partition.asInstanceOf[NativeInputPartition].bytes, readSchema,
timeZoneId)
+ }
+
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): AutoCloseable = {
+ val task = NativePartitionReader.openStream(location, name,
NativeLibrary.READ_BATCH,
+ readerState, partition.asInstanceOf[NativeInputPartition].bytes,
readSchema, timeZoneId,
+ streamAddress)
Review Comment:
Done in 8eae514: the library gets `NativeInputPartition.index`. The test now
reads both partitions of the scan in a single task.
--
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]
