HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4475109964
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,37 +401,117 @@ class NativeDataSourceStreamReader(
offset
}
- override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
- library.streamPartitions(handle.get, start, end)
+ override def partitions(start: String, end: String): Array[InputPartition] =
synchronized {
+ if (plannedRange != (start, end)) {
Review Comment:
Done in #59931 (4a76b34): `ColumnarMicroBatchStream.planInputPartitions`
plans each range once for every columnar stream reader, keeps a failure too,
and drops the range once it is committed, and `NativeDataSourceStreamReader` no
longer caches. The tests of a self-union and of concurrent planning go through
the stream.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -123,53 +132,71 @@ class ColumnarScanBuilder(
override def pruneColumns(requiredSchema: StructType): Unit = {
// Only top-level columns are pruned by the data source. Spark prunes the
nested fields of the
- // columns it reads.
+ // columns it reads. The required columns are only recorded here: a batch
scan asks its planner
+ // to prune the others, and a streaming scan, which has no planner, passes
the required
+ // columns to its stream reader.
val required =
StructType(requiredSchema.fieldNames.flatMap(name => schema.find(_.name
== name)))
- if (required.length < schema.length &&
batchPlanner.pruneColumns(required)) {
- readSchema = required
+ if (required.length < schema.length) {
+ this.requiredSchema = Some(required)
}
}
// The scan takes over the planner, so it is built once.
private lazy val scan = new ColumnarScan(
- shortName, schema, readSchema, dataSource, Option(planner), pushedLimit)
+ shortName, schema, requiredSchema, dataSource, Option(planner),
pushedLimit)
override def build(): Scan = scan
}
/**
+ * A scan of a [[ColumnarDataSource]]. It reads the top-level columns that the
query needs, the
+ * columns of its read schema: a batch scan asks its planner to prune the
other columns when it
+ * plans its partitions, and drops them from the batches of the reader if the
planner does not,
+ * while the stream reader of a streaming scan reads only these columns.
+ *
+ * @param requiredSchema the top-level columns that the query needs, if fewer
than those of the
+ * full schema
* @param pushedDownPlanner the planner that the operations were pushed down
to, if any. Otherwise,
* the scan creates one when its batch is planned.
*/
class ColumnarScan(
shortName: String,
fullSchema: StructType,
- prunedSchema: StructType,
+ requiredSchema: Option[StructType],
dataSource: ColumnarDataSource,
pushedDownPlanner: Option[ColumnarBatchPlanner],
pushedLimit: Option[Int])
extends Scan with SupportsMetadata {
- // Spark may call `toBatch` more than once for the same scan; the partitions
are planned once.
- private lazy val batch = new ColumnarBatchScan(
- pushedDownPlanner.getOrElse(dataSource.batchPlanner(fullSchema)))
+ // Spark may call `toBatch` more than once for the same scan; the partitions
are planned once,
+ // after the planner prunes the columns, if it does. A failure is kept too,
so that the planner
+ // is asked at most once, as the header promises.
+ private lazy val batch: Try[ColumnarBatchScan] = Try {
+ val planner =
pushedDownPlanner.getOrElse(dataSource.batchPlanner(fullSchema))
+ // The ordinals in the full schema of the columns that the scan reads, if
the planner does
+ // not prune the others.
+ val projection = requiredSchema.filterNot(planner.pruneColumns)
Review Comment:
Gone with the projection in #59931 (4a76b34): `pruneColumns` returns nothing
now, and `ColumnarScan` calls it as
`requiredSchema.foreach(planner.pruneColumns)`, inside the memoized `Try`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -886,10 +1047,150 @@ 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.
+ *
+ * @param projection the columns of `schema` that the engine gets, in their
order, if not all of
+ * them: the schema and the arrays of the stream then only
have these children
+ */
+ private[ffi] def openEngineStream[P <: InputPartition: ClassTag](
+ location: NativeLibraryLocation,
+ name: String,
+ kind: Int,
+ partition: InputPartition,
+ schema: StructType,
+ timeZoneId: String,
+ streamAddress: Long,
+ projection: Option[Array[Int]])(
+ 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 {
+ // As a failure of read, a failure to create the callbacks of the
streams of engines, such as
+ // a link error, is an error of the data source.
+ call(name, "read") {
+ NativeLibrary.wrapEngineStream(streamAddress, task, name,
projection.orNull)
+ }
+ } catch {
+ case e: Throwable =>
+ releaseOnFailure(e, streamAddress, task)
+ throw e
+ }
+ }
+
+ /**
+ * Initializes the Arrow C stream at `streamAddress`, whose release callback
is NULL, with the
+ * data of a partition, by calling read of the library in the running task.
Returns the context
+ * of the task, which the library may use until the stream is released: it
is closed then. If
+ * this fails, the stream is released.
+ *
+ * @param partition the partition, whose index in the scan the library gets
as `partition_id`,
+ * rather than the partition of the task: a task may read
several partitions of
+ * a scan, or of several scans, as after a coalesce or a
union
+ */
+ private def openStream(
+ location: NativeLibraryLocation,
+ name: String,
+ kind: Int,
+ readerState: Array[Byte],
+ partition: NativeInputPartition,
+ schema: StructType,
+ timeZoneId: String,
+ streamAddress: Long): NativeLibrary.Task = {
+ val library = NativeLibraries.get(location)
+ val stream = ArrowArrayStream.wrap(streamAddress)
+ val task = NativeTask.create(Some(partition.index), read = true)
+ try {
+ NativeArrow.exportSchema(schema, timeZoneId) { schemaAddress =>
Review Comment:
I kept this one, from a measurement in an earlier round, when the automated
review asked the same: for 1000 columns, converting and exporting the schema
takes 1.83 ms, against 1.74 ms to export an Arrow schema converted once, so the
conversion is about 5% of it, and less for the usual schemas. A cached Arrow
schema would save a few microseconds per partition, while the export, which is
needed each time, dominates.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -222,9 +346,71 @@ abstract class ColumnarOnlyPartitionReaderFactory extends
PartitionReaderFactory
override def supportColumnarReads(partition: InputPartition): Boolean = true
}
-class ColumnarStreamPartitionReaderFactory(reader: ColumnarStreamReader)
+/**
+ * Reads only the columns at `ordinals` of the batches of the readers of
`factory`, for a scan
+ * whose partitions return more columns than the scan reads.
+ */
+class ColumnarProjectionReaderFactory(factory: PartitionReaderFactory,
ordinals: Array[Int])
extends ColumnarOnlyPartitionReaderFactory {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- reader.read(partition)
+ new ColumnarProjectionReader(factory.createColumnarReader(partition),
ordinals)
+ }
+}
+
+/**
+ * A reader factory that opens the Arrow C streams of the partitions for an
engine, see
+ * [[SupportsArrowStreamRead]], and can project them to some of their columns.
+ */
+trait SupportsProjectedArrowStreamRead extends SupportsArrowStreamRead {
+ /**
+ * Opens the stream of `partition`, as `openArrowStream(partition,
streamAddress)` does, but its
+ * schema and its arrays only have the columns at `projection`, in that
order.
+ */
+ def openArrowStream(partition: InputPartition, streamAddress: Long,
projection: Array[Int]): Unit
+}
+
+/**
+ * A [[ColumnarProjectionReaderFactory]] whose Arrow C streams, which an
engine opens, only have the
+ * columns that the scan reads too, as its batches do.
+ */
+class ColumnarArrowStreamProjectionReaderFactory(
+ factory: SupportsProjectedArrowStreamRead,
+ ordinals: Array[Int])
+ extends ColumnarProjectionReaderFactory(factory, ordinals) with
SupportsArrowStreamRead {
+ override def openArrowStream(partition: InputPartition, streamAddress:
Long): Unit = {
+ factory.openArrowStream(partition, streamAddress, ordinals)
}
}
+
+private class ColumnarProjectionReader(
+ reader: PartitionReader[ColumnarBatch],
+ ordinals: Array[Int])
+ extends PartitionReader[ColumnarBatch] {
+
+ // The projection of the current batch of the reader, which `get` returns
until `next`.
+ private var batch: ColumnarBatch = _
+
+ override def next(): Boolean = {
+ val hasNext = reader.next()
+ batch = if (hasNext) new ProjectedColumnarBatch(reader.get(), ordinals)
else null
Review Comment:
Gone with the projection: `ColumnarProjectionReader` and
`ProjectedColumnarBatch` are removed in #59931 (4a76b34).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,37 +401,117 @@ class NativeDataSourceStreamReader(
offset
}
- override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
- library.streamPartitions(handle.get, start, end)
+ override def partitions(start: String, end: String): Array[InputPartition] =
synchronized {
+ if (plannedRange != (start, end)) {
+ plannedPartitions = Try {
+ val planned = call(name, "plan a scan of") {
+ library.streamPartitions(handle.get, start, end)
+ }
+ val partitions = NativeInputPartition.from(planned.partitions)
+ val bytes = bytesOrEmpty(Option(planned.state))
+ // A small state goes with each task, while the executors fetch a
larger one once for the
+ // micro-batch, as a broadcast. Spark only plans the next range once
the micro-batch of the
+ // previous one finished, so the copies of the previous broadcast on
the executors are
+ // removed then. The broadcast itself is destroyed once its
micro-batch is committed,
+ // whose partitions are not read anymore then, also by an engine.
+ uncommittedBroadcasts.lastOption.foreach(_._2.unpersist(blocking =
false))
+ val state = if (bytes.length >
NativeDataSourceStreamReader.maxInlineStateBytes &&
+ partitions.nonEmpty) {
+ val broadcast = SparkSession.active.sparkContext.broadcast(bytes)
Review Comment:
Done in 9711532: the stream reader takes the `SparkContext` when it is
created, on the driver, and broadcasts the states with it, so planning on a
thread without an active session works too.
--
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]
