dongjoon-hyun commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4472784521
##########
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)
+ }
}
case class ColumnarStreamOffset(json: String) extends Offset
-class ColumnarMicroBatchStream(reader: ColumnarStreamReader) extends
MicroBatchStream {
+/**
+ * The micro-batch stream of a [[ColumnarDataSource]]. It supports admission
control only so that
+ * the reader gets the offset that each micro-batch starts after, the end of
the last one, for
+ * every trigger: its default read limit reads all the available data. With
+ * `Trigger.AvailableNow`, Spark runs micro-batches up to the first latest
offset that Spark asks
+ * for, rather than wrapping the stream in one that asks for the latest offset
from the initial
+ * offset.
+ *
+ * @param checkpointLocation the location of the checkpoint of the source,
where the stream keeps
+ * the initial offset, under
`ColumnarMicroBatchStream.InitialOffsetLog`
+ */
+class ColumnarMicroBatchStream(reader: ColumnarStreamReader,
checkpointLocation: String)
+ extends MicroBatchStream with SupportsTriggerAvailableNow {
+
+ // The offset that the reader returned last, which the progress of the query
reports.
+ @volatile private var lastLatestOffset: Offset = _
+ // Whether the trigger is Trigger.AvailableNow.
+ @volatile private var availableNow = false
+ // The end of the run with Trigger.AvailableNow: the first latest offset
asked.
+ @volatile private var availableNowEnd: Offset = _
+ // Spark asks for the initial offset for the start of the first latest
offset, and again to plan
+ // the first micro-batch, which both get the same offset. The reader is only
asked when the
+ // checkpoint does not have it yet: a query that restarts before its first
micro-batch is
+ // committed plans it again from the same start, also when the initial
offset of the reader
+ // depends on when it is asked, such as the latest data.
+ private lazy val initial = {
+ val log = new ColumnarInitialOffsetLog(SparkSession.active,
+ new Path(checkpointLocation,
ColumnarMicroBatchStream.InitialOffsetLog).toString)
+ ColumnarStreamOffset(log.get(0).getOrElse {
+ val offset = reader.initialOffset()
+ log.add(0, offset)
+ offset
+ })
+ }
- override def initialOffset(): Offset =
ColumnarStreamOffset(reader.initialOffset())
+ override def initialOffset(): Offset = initial
- override def latestOffset(): Offset =
ColumnarStreamOffset(reader.latestOffset())
+ override def latestOffset(): Offset = {
+ throw SparkException.internalError(
+ "latestOffset(Offset, ReadLimit) should be called instead of
latestOffset().")
+ }
+
+ override def latestOffset(start: Offset, limit: ReadLimit): Offset = {
Review Comment:
This fixes the `Trigger.AvailableNow` end offset at the first `latestOffset`
call instead of in `prepareForTriggerAvailableNow`, which differs from the
contract of `SupportsTriggerAvailableNow`.
When an AvailableNow query restarts from a checkpoint whose last batch is
not committed, `MicroBatchExecution` calls `prepareForTriggerAvailableNow`
while initializing the sources, re-runs that batch first, and only then calls
`latestOffset`. So the data that arrives during the re-run is also processed,
rather than only the data available when the query started.
Could we fix the end offset in `prepareForTriggerAvailableNow`?
##########
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:
When the static initialization of `EngineStreamCallbacks` fails, the first
call throws `ExceptionInInitializerError`, whose `getMessage` is null, so
`NativeCalls.message` falls back to `e.toString`, i.e.
`java.lang.ExceptionInInitializerError`. The later calls get
`NoClassDefFoundError: Could not initialize class ...`. Either way,
`NATIVE_DATA_SOURCE_ERROR` does not show the root cause, including the `Cannot
create the callbacks of the streams of engines: <rootCause>` message that
`NativeLibrary.java` builds for this purpose.
This is also inconsistent with `load()`, which catches a `LinkageError` of
the FFM initialization and reports its root cause. Could we report the root
cause here too?
##########
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))
+ } 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 =>
+ call(name, "read") {
+ library.read(task, kind, schemaAddress, readerState,
partition.bytes, streamAddress)
+ }
+ }
+ // Arrow calls the callbacks of the stream without checking them, and
calling NULL would
+ // crash the JVM.
+ val snapshot = stream.snapshot()
+ if (snapshot.release == 0 || snapshot.get_schema == 0 ||
snapshot.get_next == 0 ||
+ snapshot.get_last_error == 0) {
+ throw error(name, "read", "The library did not initialize the stream,
or a callback of " +
+ "the stream is NULL.")
+ }
+ task
+ } catch {
+ case e: Throwable =>
+ releaseOnFailure(e, streamAddress, task)
+ throw e
+ }
+ }
+
+ /**
+ * After `failure`, releases the stream at `streamAddress`, if the library
initialized it, before
+ * the task that it may use, and marks it released in case its release
callback did not. A
+ * failure to clean up is added to `failure`.
+ */
+ private def releaseOnFailure(
+ failure: Throwable,
+ streamAddress: Long,
+ task: NativeLibrary.Task): Unit = {
+ try {
+ Utils.tryWithSafeFinally {
+ val stream = ArrowArrayStream.wrap(streamAddress)
+ stream.release()
Review Comment:
If `stream.release()` here throws, `markReleased()` is skipped but
`task.close()` still runs. The engine path otherwise uses only FFM, so in a JVM
where Spark's own Arrow reader has never run, the Arrow C Data JNI library may
not be loaded yet, and this JNI call can fail with `UnsatisfiedLinkError`.
In that case, the engine's stream is left with `release != NULL` while the
`SparkDsTaskContext` is already freed. An engine that trusts the contract (the
stream is released on failure) leaks it, and an engine that calls `release`
runs the library's release, which refers to the freed task (use-after-free).
Could we release the stream through FFM like the rest of the engine path, or
keep the task open when the release fails?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -350,37 +382,107 @@ 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 its broadcast on the
executors are removed then.
+ // The driver keeps its own copy, so that the partitions and the
reader factories that an
+ // engine keeps can still fetch it, until the ContextCleaner removes
it once nothing refers
+ // to them.
+ plannedBroadcast.foreach(_.unpersist(blocking = false))
Review Comment:
When the state is larger than 64 KiB, it is broadcast for each micro-batch,
and the previous broadcast only gets `unpersist(blocking = false)` when the
next range is planned. That removes the executor copies, but the driver copy
stays until the `ContextCleaner` removes it after a GC.
For example, with a 1-second trigger and a 1 MB state, the broadcasts can be
promoted to the old generation and survive until
`spark.cleaner.periodicGC.interval` (30 minutes by default). That is up to
about 1.8 GB of broadcast blocks in the driver `BlockManager`, which can evict
cached blocks or put pressure on the driver memory.
Is there a way to bound this? For example, could the broadcast of a
micro-batch be destroyed once the micro-batch is committed, or do engines need
it beyond that?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -123,39 +132,55 @@ 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.
+ private lazy val batch = {
+ 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:
`pruneColumns` is now called lazily while building `ColumnarScan.batch`, but
a failure here is not memoized. `BatchScanExec.batch` is a `@transient lazy
val` that can be evaluated again (for example from `equals`, `hashCode`, or
`inputPartitions`), so the next `toBatch` calls `prune_columns` again:
- With a pushdown (e.g. `select("id").limit(3)`), it is called twice on the
same pushed-down planner, which breaks the "each at most once" contract of the
header.
- Without a pushdown, a new planner and native reader are created each time,
and the previous planner handle leaks until the GC `Cleaner` runs.
This also contradicts the doc of `ColumnarBatchScan` ("A failure is planned
once too"). Could we memoize the result of building the scan, including a
failure?
##########
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:
`endedTasks` is a static map that strongly references its `TaskContext`
keys, and an entry is removed only when the completion or failure listener
runs. So a context that never ends is never removed: for example
`TaskContext.empty()` mentioned above, or a context that an engine or a test
sets with `TaskContext.setTaskContext` before a native read or
`openArrowStream`. Its `TaskContextImpl`, `TaskMetrics`, and registered
listener closures then stay in this map until the JVM exits.
This also defeats holding the task through a `WeakReference`. Could we use
weak keys here?
##########
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:
nit: `ColumnarProjection.apply` has a single caller,
`ColumnarBatchScan.createReaderFactory`. We can inline it there and remove this
object:
```scala
projection.fold(factory)(new ColumnarProjectionReaderFactory(factory, _))
```
##########
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:
For a library that does not prune the columns,
`ColumnarProjectionReaderFactory` wraps the factory and hides
`SupportsArrowStreamRead`, so the engine native path is disabled entirely. The
PR description explains that this is because the arrays of the stream would
still have the other columns. However, `select("id")` on such a library then
sends engines like Gluten back to the Arrow Java import -> `ColumnarToRow` ->
`RowToVelox` path, which is the cost that this PR aims to remove.
Since Spark already intercepts `get_schema` and `get_next` of the engine
streams, how about projecting the children of the struct there instead, so that
the native path is kept for such libraries 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:
`NativeLibrary.openTaskContexts()` is a JVM-wide counter, and this compares
it with `==` to a baseline read before the test. If a context of an earlier
test (e.g. of a killed task or a `PipedRDD` thread) is closed after
`contextsBeforeKill` is read, the counter settles at `contextsBeforeKill - 1`,
and this `eventually` fails after the 1-minute timeout. The same pattern is at
lines 2007, 2036, and 2043, where the comment notes that the tasks of earlier
tests may still be closing.
Could we make these assertions robust to the contexts of earlier tests that
close late, for example by using `<=`?
--
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]
