HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4474199496
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -969,22 +1228,69 @@ class NativeDataWriter(
/** Creates the `SparkDsTaskContext` of the running task. */
private[ffi] object NativeTask {
+ // The flag of each running task that reads, set once the task failed or
completed by a single
+ // pair of listeners for all its reads, which then also remove it: a task
that reads many
+ // partitions, as after a coalesce, registers them once. The contexts of the
reads keep their
+ // flag, so they still see it once it is removed. A task is the key rather
than its attempt ID,
+ // which is 0 for each context of `TaskContext.empty()`.
+ private val endedTasks = new ConcurrentHashMap[TaskContext, AtomicBoolean]()
Review Comment:
Done in 7b3fbb2: `endedTasks` is a map of Guava's `MapMaker().weakKeys()`,
which compares its keys by identity, so a context that never ends, such as one
that a test or an engine sets on a thread, does not stay. The flags do not
refer to the contexts, and the listeners that do are held by the contexts
themselves.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -886,10 +1018,137 @@ 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 {
+ // 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))
Review Comment:
Done in 7b3fbb2: `wrapEngineStream` catches the
`ExceptionInInitializerError` of the first use of the callbacks and the
`NoClassDefFoundError` of the later ones, and throws an `IllegalStateException`
whose message names the root cause, "Cannot create the callbacks of the streams
of engines: ...", as `load` reports a `LinkageError` of the FFM initialization,
so `NATIVE_DATA_SOURCE_ERROR` shows it. Since Java 18, the cause of such a
`NoClassDefFoundError` keeps the message of the first error, so the later uses
name the root cause too.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -1143,6 +1515,464 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
.mode("append")
.save()
check(writeTasks)
+
+ // A read gets the index of its partition in the scan, also when a task
reads several
+ // partitions, or the partitions of several scans: the task attempt that
read each partition.
+ def readers(dir: File): Map[Int, Long] = dir.listFiles().map { file =>
+ file.getName.toInt -> Files.readString(file.toPath).split(",")(5).toLong
+ }.toMap
+ val (first, second) = (Utils.createTempDir(), Utils.createTempDir())
+ checkRows(
+ range(session, "end" -> "6", "task_path" -> first.getPath).select("id")
+ .union(range(session, "end" -> "6", "task_path" ->
second.getPath).select("id")),
+ (0L until 6L).flatMap(i => Seq(Row(i), Row(i))))
+ assert(readers(first).keySet == Set(0, 1) && readers(second).keySet ==
Set(0, 1))
+ val coalesced = Utils.createTempDir()
+ checkRows(range(session, "end" -> "6", "task_path" ->
coalesced.getPath).select("id")
+ .coalesce(1), (0L until 6L).map(Row(_)))
+ assert(readers(coalesced).keySet == Set(0, 1) &&
readers(coalesced).values.toSet.size == 1)
+ // So does a read of a micro-batch.
+ val (firstStream, secondStream) = (Utils.createTempDir(),
Utils.createTempDir())
+ def counter(dir: File): DataFrame =
session.readStream.format("native_counter")
+ .option("max_offset", "4").option("task_path", dir.getPath).load()
+ val query =
counter(firstStream).union(counter(secondStream)).writeStream.format("noop").start()
+ try query.processAllAvailable() finally query.stop()
+ assert(readers(firstStream).keySet == Set(0) &&
readers(secondStream).keySet == Set(0))
+ }
+
+ nativeTest("open the partitions of batch and micro-batch scans as Arrow C
streams") {
+ val session = newSession(defaultPackage)
+ // The reader factory and the partitions of the batch scan of a query,
which an engine gets
+ // from its plan.
+ def plannedScan(df: DataFrame): (SupportsArrowStreamRead,
Seq[InputPartition]) = {
+ val scan = df.queryExecution.executedPlan.collectFirst { case s:
BatchScanExec => s }.get
+ (scan.readerFactory.asInstanceOf[SupportsArrowStreamRead],
scan.inputPartitions)
+ }
+ // The batch scan of only some of the columns from a library that does not
prune the others
+ // drops them from the batches itself, so its reader factory does not open
streams, which
+ // would have all the columns. With all the columns, it does.
+ val readOnly = newSession(defaultPackage, createPackage(readOnlyLibrary,
dataSources("ro_")))
+ def readsStreams(df: DataFrame): Boolean = {
+ df.queryExecution.executedPlan.collectFirst { case s: BatchScanExec => s
}.get
+ .readerFactory.isInstanceOf[SupportsArrowStreamRead]
+ }
+ val unpruned = readOnly.read.format("ro_native_range").option("end",
"4").load()
+ assert(!readsStreams(unpruned.select("id")) && readsStreams(unpruned))
+ // Runs `f` in a task that only gets the files of the session. It is the
task of the partition
+ // 1, so that the partition of the scan that a read gets is not the
partition of its task.
+ def inTask[T: ClassTag](f: => T): T =
session.artifactManager.withResources {
+ session.sparkContext.parallelize(Seq(0, 1), 2).flatMap(i => if (i == 1)
Seq(f) else Nil)
+ .collect().head
+ }
+ import TestArrowEngine.{callbacks, readIds, readIdsInPlace}
+ // Reads the given partitions in a single task, as an engine does.
+ def readAll(factory: SupportsArrowStreamRead, partitions:
Seq[InputPartition]): Seq[Long] = {
+ inTask {
+ partitions.flatMap { partition =>
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator("engine",
0, Long.MaxValue)
+ try {
+ val stream = ArrowArrayStream.allocateNew(allocator)
+ try factory.openArrowStream(partition, stream.memoryAddress())
catch {
+ case e: Throwable =>
+ stream.close()
+ throw e
+ }
+ readIds(stream, allocator)
+ } finally {
+ allocator.close()
+ }
+ }
+ }
+ }
+ // Checks that `f` opens `delta` more task contexts than it closes. The
count is global to the
+ // JVM, where the tasks of earlier tests may still close their contexts,
so it is read right
+ // before `f`.
+ def checkContexts[T](delta: Int)(f: => T): T = {
+ val before = NativeLibrary.openTaskContexts()
+ val result = f
+ assert(NativeLibrary.openTaskContexts() == before + delta)
+ result
+ }
+ // The size of struct ArrowArrayStream, which the layout test compares
with the header.
+ val streamSize = NativeLibrary.layout().get("ArrowArrayStream").longValue
+
+ // An engine reads both partitions of the scan in a single task.
+ val taskDir = Utils.createTempDir()
+ val (factory, partitions) =
+ plannedScan(range(session, "end" -> "6", "task_path" ->
taskDir.getPath).select("id"))
+ assert(partitions.length == 2)
+ assert(readAll(factory, partitions).sorted == (0L until 6L))
+ // The library got the partition of the scan that it reads, not the one of
the task.
+ val tasks = taskDir.listFiles().map { file =>
+ file.getName.toInt ->
Files.readString(file.toPath).split(",").map(_.toLong).toSeq
+ }.toMap
+ assert(tasks.keySet == Set(0, 1) && tasks.values.map(_(5)).toSet.size ==
1, tasks)
+ assert(NativeLibrary.openEngineStreams() == 0)
+ // The reads of a task share the flag that tells them that the task ended,
with a single pair
+ // of listeners, which remove it once the task ended. The tasks of earlier
tests may still
+ // be ending.
+ eventually(timeout(1.minute))(assert(NativeTask.tasksWithEndedFlags == 0))
+ val flagsInTask = inTask {
+ val ids = partitions.flatMap { partition =>
+ val stream = ArrowArrayStream.allocateNew(ArrowUtils.rootAllocator)
+ try {
+ factory.openArrowStream(partition, stream.memoryAddress())
+ try readIdsInPlace(stream) finally stream.release()
+ } finally {
+ stream.close()
+ }
+ }
+ assert(ids.sorted == (0L until 6L))
+ NativeTask.tasksWithEndedFlags
+ }
+ assert(flagsInTask == 1 && NativeTask.tasksWithEndedFlags == 0)
+
+ // The engine may release the stream after its task ended, and on another
thread: Spark frees
+ // the context of the task once the library released the stream.
+ val released = new File(Utils.createTempDir(), "released")
+ val (laterFactory, laterPartitions) =
+ plannedScan(range(session, "end" -> "6", "released_path" ->
released.getPath).select("id"))
+ // Opens a stream in a task, and returns it unreleased, with what `read`
returned from it in
+ // the task.
+ def openInTask[T: ClassTag](factory: SupportsArrowStreamRead, partition:
InputPartition)(
+ read: ArrowArrayStream => T): (Long, T) = inTask {
+ val stream = Platform.allocateMemory(streamSize)
+ Platform.setMemory(stream, 0.toByte, streamSize)
+ factory.openArrowStream(partition, stream)
+ (stream, read(ArrowArrayStream.wrap(stream)))
+ }
+ // The engine reads the stream to its end in the task, but releases it
later.
+ val (address, readInTask) =
+ checkContexts(1)(openInTask(laterFactory,
laterPartitions.head)(readIdsInPlace))
+ // A copy of the stream, as an engine that moves the stream may leave
behind.
+ val copy = Platform.allocateMemory(streamSize)
+ Utils.tryWithSafeFinally {
+ Platform.copyMemory(null, address, null, copy, streamSize)
+ assert(readInTask == (0L until 3L))
+ assert(NativeLibrary.openEngineStreams() == 1 && !released.exists())
+ checkContexts(-1) {
+ ArrowArrayStream.wrap(address).release()
+ assert(NativeLibrary.openEngineStreams() == 0 && released.exists())
+ }
+ // Once the stream is released, the callbacks of the copy fail with
EINVAL, and its release
+ // only marks it released.
+ checkContexts(0) {
+ val array = ArrowArray.allocateNew(ArrowUtils.rootAllocator)
+ val stale = try
intercept[Exception](ArrowArrayStream.wrap(copy).getNext(array)) finally {
+ array.close()
+ }
+ assert(Utils.exceptionString(stale)
+ .contains("The stream was released, or is not a stream of Spark."))
+ ArrowArrayStream.wrap(copy).release()
+ assert(ArrowArrayStream.wrap(copy).snapshot().release == 0)
+ assert(NativeLibrary.openEngineStreams() == 0)
+ }
+ } {
+ // A failure above may leave the stream open, which would keep the
context of its task.
+ if (ArrowArrayStream.wrap(address).snapshot().release != 0) {
+ ArrowArrayStream.wrap(address).release()
+ }
+ Platform.freeMemory(address)
+ Platform.freeMemory(copy)
+ }
+
+ // Once the task ended, its context tells the library that it was
interrupted, so that the
+ // library stops waiting.
+ val marker = new File(Utils.createTempDir(), "interrupted")
+ val (waitingFactory, waitingPartitions) = plannedScan(
+ range(session, "end" -> "6", "wait_for_interrupt" ->
marker.getPath).select("id"))
+ checkContexts(0) {
+ val (waiting, _) = openInTask(waitingFactory, waitingPartitions.head)(_
=> ())
+ val error = intercept[Exception] {
+ readIds(ArrowArrayStream.wrap(waiting), ArrowUtils.rootAllocator)
+ }
+ assert(Utils.exceptionString(error).contains("the task was killed"))
+ assert(Files.readString(marker.toPath) == "interrupted")
+ assert(NativeLibrary.openEngineStreams() == 0)
+ Platform.freeMemory(waiting)
+ }
+
+ // So does the context of a stream that an engine reads in its task once
the task is killed,
+ // although the task did not end.
+ val killedMarker = new File(Utils.createTempDir(), "interrupted")
+ val (killedFactory, killedPartitions) = plannedScan(range(session, "end"
-> "6",
+ "partitions" -> "1", "wait_for_interrupt" ->
killedMarker.getPath).select("id"))
+ val contextsBeforeKill = NativeLibrary.openTaskContexts()
+ val jobGroup = s"native-data-source-${UUID.randomUUID()}"
+ val killError = new AtomicReference[Throwable]()
+ val engineThread = new Thread(() => {
+ session.sparkContext.setJobGroup(jobGroup, "an engine that waits",
interruptOnCancel = true)
+ try {
+ inTask {
+ // Reads the stream on the thread of the task, which waits in
get_next.
+ val stream = ArrowArrayStream.allocateNew(ArrowUtils.rootAllocator)
+ try {
+ killedFactory.openArrowStream(killedPartitions.head,
stream.memoryAddress())
+ try readIdsInPlace(stream) finally stream.release()
+ } finally {
+ stream.close()
+ }
+ }
+ } catch {
+ case e: Throwable => killError.set(e)
+ }
+ })
+ engineThread.start()
+ eventually(timeout(1.minute)) {
+ assert(new File(killedMarker.getPath + ".waiting").isFile)
+ }
+ session.sparkContext.cancelJobGroup(jobGroup)
+ engineThread.join(TimeUnit.MINUTES.toMillis(1))
+ assert(!engineThread.isAlive && killError.get() != null)
+ eventually(timeout(1.minute)) {
+ assert(killedMarker.isFile && Files.readString(killedMarker.toPath) ==
"interrupted")
+ assert(NativeLibrary.openEngineStreams() == 0)
+ assert(NativeLibrary.openTaskContexts() == contextsBeforeKill)
Review Comment:
Done in 7b3fbb2: `checkContexts` and the waits of both tests of killed tasks
compare the count with `<=`, so a context of an earlier test that closes late
does not fail them. The test of `PipedRDD`, which also checks that a context
stays open with `contexts + 1`, first waits until the contexts of the earlier
tests are closed, so its counts stay exact.
##########
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 = {
+ ColumnarProjection(planned.get.readerFactory, projection)
Review Comment:
Done in 7b3fbb2: for a batch scan of only some of the columns from a library
that does not prune the others, `ColumnarBatchScan` now wraps the native
factory in `ColumnarArrowStreamProjectionReaderFactory`, which projects the
batches as before and implements `SupportsArrowStreamRead`, with the internal
`SupportsProjectedArrowStreamRead` of the native factory. The stream of the
engine then projects the schema of `get_schema` and each array of `get_next` in
place: the struct of the library is moved to memory of Spark, and the struct of
the engine keeps its fields, with only the children of the read schema and a
release of Spark, which releases the moved struct of the library, and so all
its children. The children are not copied, so an engine that moves a projected
child marks it released where the library keeps it, as the C data interface
requires. A projection that fails, such as for a schema that is not a struct,
releases the struct of the library and returns the error. The layout
test now also checks the offsets of `ArrowSchema` and `ArrowArray` that this
uses, the Javadoc and the guide describe it, and "open the partitions of batch
and micro-batch scans as Arrow C streams" reads `select("id")` of the read-only
library through its streams, whose schema only has `id`, and checks that every
projected struct is released.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -222,9 +338,59 @@ abstract class ColumnarOnlyPartitionReaderFactory extends
PartitionReaderFactory
override def supportColumnarReads(partition: InputPartition): Boolean = true
}
-class ColumnarStreamPartitionReaderFactory(reader: ColumnarStreamReader)
+object ColumnarProjection {
Review Comment:
Done in #59931 (e305457), which is merged here:
`ColumnarBatchScan.createReaderFactory` builds the projecting factory itself,
and `ColumnarProjection` is removed. Here, it also picks the factory that
projects the streams of engines, see the comment on L218.
--
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]
