HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4472095051
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -178,30 +204,89 @@ class ColumnarScan(
* The batch of a scan, which plans it once, when Spark first asks for its
partitions or its reader
* factory. A failure is planned once too: the physical plan can ask again,
for example from a lazy
* value, and gets the same error.
+ *
+ * @param projection the ordinals of the columns that the scan reads, if the
partitions return
+ * more columns
*/
-class ColumnarBatchScan(planner: ColumnarBatchPlanner) extends Batch {
+class ColumnarBatchScan(planner: ColumnarBatchPlanner, projection:
Option[Array[Int]] = None)
+ extends Batch {
private lazy val planned: Try[ColumnarPlannedBatch] = Try(planner.plan())
override def planInputPartitions(): Array[InputPartition] =
planned.get.partitions
- override def createReaderFactory(): PartitionReaderFactory =
planned.get.readerFactory
+ override def createReaderFactory(): PartitionReaderFactory = {
Review Comment:
Done in d40d64d: the Javadoc of `SupportsArrowStreamRead`, the guide and the
PR description now say that the reader factories of the micro-batch scans
implement it, and so do those of the batch scans, except for a batch scan of
only some of the columns from a library that does not prune the others, whose
batches Spark projects itself. I did not make the wrapper project the stream:
that would wrap each array of the stream too, to move the fields of the read
schema out of the struct array into a new one and release the others, in
callbacks of Spark. I would leave that to a follow-up if an engine needs it:
such a scan still runs under an engine, which reads its batches from Spark, as
for any scan whose factory does not implement the interface. "open the
partitions of batch and micro-batch scans as Arrow C streams" now uses the
read-only test library, which has no `prune_columns`: the reader factory of
`select("id")` does not implement the interface, and the one that reads all the
col
umns does.
##########
sql/core/src/main/resources/org/apache/spark/sql/datasource/spark_datasource.h:
##########
@@ -304,29 +313,135 @@ typedef struct SparkDsCreateInfo {
int64_t num_options;
} SparkDsCreateInfo;
+/* A batch scan, as create_reader gets it. */
+typedef struct SparkDsReaderInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+} SparkDsReaderInfo;
+
+/* A request for the partitions of a batch scan, as partitions gets it. */
+typedef struct SparkDsPartitionsInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+} SparkDsPartitionsInfo;
+
+/* 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
+ * fully qualified Hadoop path, whose characters are not escaped as in a
URI: a space stays a
+ * space, and % stays %. Its file system may be one that the library cannot
access, such as s3a
+ * or hdfs. Spark keeps what it needs there under the names that start with
_spark, such as the
+ * initial offset of the source, and the library must not use them.
+ */
+ const char* checkpoint_location;
+} SparkDsStreamReaderInfo;
+
+/* A request for the initial offset of a stream reader, as initial_offset gets
it. */
+typedef struct SparkDsInitialOffsetInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+} SparkDsInitialOffsetInfo;
+
+/* 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, for every trigger. With
Trigger.AvailableNow,
+ * Spark asks for the latest offset only once, before the first micro-batch
that it plans, and
+ * reads up to it.
+ */
+ const char* start;
+} SparkDsLatestOffsetInfo;
+
+/* A request for the partitions of a micro-batch, as stream_partitions gets
it. */
+typedef struct SparkDsStreamPartitionsInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /* The offsets, as JSON, that the micro-batch reads the data after, up to
and including end. */
+ const char* start;
+ const char* end;
+} SparkDsStreamPartitionsInfo;
+
+/* The end of the data that Spark processed, as commit_offset gets it. */
+typedef struct SparkDsCommitOffsetInfo {
+ /* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
+ int64_t size;
+ /* The offset, as JSON, up to and including which Spark processed the data.
*/
+ const char* end;
+} SparkDsCommitOffsetInfo;
+
/*
* 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.
+ * that read returns is released, or until the data writer is closed. The
stream of a read may be
+ * released once the task ended, from another thread: an engine that runs a
query in native code
+ * may get the stream from Spark, read it in the task that opened it and
release it later, and
+ * Spark itself may read a stream on a thread that the task does not wait for,
for example the one
+ * of a PipedRDD. The task stays valid until then too, and for a read,
is_interrupted then also
+ * returns true once the task failed or completed, so the library may stop
reading then. Such an
+ * engine may also keep the arrays and the schemas of the stream after it
released the stream, so
+ * the release of an array or a schema must not use the stream or the task.
*/
typedef struct SparkDsTaskContext {
/* The size of this struct in the header that Spark was built with, as in
SparkDsHost. */
int64_t size;
/* The ID of the stage of the task, and the attempt number of the stage. */
int32_t stage_id;
int32_t stage_attempt_number;
- /* The partition that the task reads or writes, and the attempt number of
the task. */
+ /*
+ * The partition: for read, the index of the partition in the list that
partitions or
+ * stream_partitions returned, which is not the partition of the task when
Spark reads several
+ * scans or partitions in one task, for example after a union or a coalesce.
It is unique within
+ * a scan, but a task that reads several scans may get the same index for
each. For
+ * create_data_writer, the partition of the task that writes. Then the
attempt number of the
+ * task.
+ */
int32_t partition_id;
Review Comment:
Thank you. I would keep the current shape: `partition_id` stays the index of
the partition in the scan for a read, as settled in round 3, and
`task_attempt_id`, which the context of a read also has, identifies its task.
With a `partition_index` in `SparkDsReadInfo`, a read would get both, but a
library would then have two numbers to choose from, while the partition of the
task tells a read little, since its task may read several partitions of several
scans. The comment on `partition_id` documents both meanings.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -886,10 +998,135 @@ class NativePartitionReader(
}
}
-/** Writes the batches of a task to the data writer of a native library. */
+object NativePartitionReader {
+ import NativeCalls._
+
+ /**
+ * Initializes the stream at `streamAddress` with the data of a partition
for an engine, see
+ * [[SupportsArrowStreamRead]]: the context of its task is closed once the
stream is released.
+ * It first checks the arguments of `openArrowStream`, before it looks at
the partition, which
+ * must be a `P`, the class of the partitions of the scan of the reader
factory. It then reads
+ * the partition with the state and the partition of the library that
`stateAndPartition`
+ * returns for it.
+ */
+ private[ffi] def openEngineStream[P <: InputPartition: ClassTag](
+ location: NativeLibraryLocation,
+ name: String,
+ kind: Int,
+ partition: InputPartition,
+ schema: StructType,
+ timeZoneId: String,
+ streamAddress: Long)(
+ stateAndPartition: P => (Array[Byte], NativeInputPartition)): Unit = {
+ val context = TaskContext.get()
+ if (context != null) {
+ // A task that was killed gets the error of its kill, as Spark reports
it.
+ context.killTaskIfInterrupted()
+ }
+ // A task that failed is over too, also while its listeners run.
+ if (context == null || context.isFailed() || context.isCompleted()) {
+ throw new IllegalStateException("openArrowStream must be called in the
running task that " +
+ "reads the partition, on a thread that has its TaskContext.")
+ }
+ if (streamAddress == 0 ||
ArrowArrayStream.wrap(streamAddress).snapshot().release != 0) {
+ throw new IllegalArgumentException(
+ "openArrowStream needs the address of a stream whose release callback
is NULL.")
+ }
+ val (readerState, nativePartition) = partition match {
+ case expected: P => stateAndPartition(expected)
+ case _ =>
+ throw new IllegalArgumentException("openArrowStream needs a partition
of the scan of " +
+ s"the reader factory, a ${classTag[P].runtimeClass.getSimpleName},
but got " +
+ s"${Option(partition).map(_.getClass.getName).orNull}.")
+ }
+ val task = openStream(location, name, kind, readerState, nativePartition,
schema, timeZoneId,
+ streamAddress)
+ try {
+ NativeLibrary.wrapEngineStream(streamAddress, task, name)
Review Comment:
Done in d40d64d: `openEngineStream` calls `wrapEngineStream` in `call(name,
"read")`, so a failure to create the callbacks, such as a link error, is a
`NATIVE_DATA_SOURCE_ERROR` that names the data source, like a failure of
`read`. The stream is still released by `releaseOnFailure`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1180,6 +1964,7 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
val marker = new File(Utils.createTempDir(), "released")
val reading = marker.getPath + ".reading"
val df = range(session, "end" -> "10", "partitions" -> "1", "slow_read" ->
marker.getPath)
+ val contexts = NativeLibrary.openTaskContexts()
Review Comment:
Done in d40d64d: "tell a read that its task was killed" reads
`NativeLibrary.openTaskContexts()` before the query, and its last `eventually`
waits until the count is back there, so the test only ends once its task closed
the context of its read.
--
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]
