dongjoon-hyun commented on code in PR #59934: URL: https://github.com/apache/spark/pull/59934#discussion_r4463375828
########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsArrowStreamRead.java: ########## @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector.read; + +import org.apache.spark.annotation.Evolving; + +/** + * A mix-in interface for a {@link PartitionReaderFactory} that also reads a partition as an + * <a href="https://arrow.apache.org/docs/format/CStreamInterface.html";>Arrow C stream</a>, for an + * engine that processes Arrow data in native code: the engine imports the arrays of the stream + * itself, instead of reading the {@link org.apache.spark.sql.vectorized.ColumnarBatch}es of + * {@link #createColumnarReader(InputPartition)}. The reader factories of the batch and + * micro-batch scans of native data sources implement it, with the stream that the library + * returns. + * <p> + * Each array of the stream is a struct array whose fields are the columns of the read schema of + * the scan. Spark only checks that the source initialized the callbacks of the stream, not its + * arrays, while it checks the arrays of the streams that it reads itself, so the engine checks + * what its import of the arrays requires, such as the schema of the stream against the read + * schema, and the offsets and the lengths of the arrays. The callbacks of the stream can be called + * from any thread, but never concurrently. They call into the JVM, so the engine releases the + * stream before the JVM exits. + * + * @since 4.4.0 + */ +@Evolving +public interface SupportsArrowStreamRead extends PartitionReaderFactory { + + /** + * Initializes an Arrow C stream with the data of a partition. It is called on the thread of the + * running task that reads the partition, and fails otherwise. The caller reads the stream in Review Comment: **5. [Nit] "and fails otherwise" is stricter than the check** `openEngineStream` checks that the thread has the `TaskContext` of a task that is still running (`NativeDataSource.scala` L1010-L1012), not that it is the thread of that task. A thread that got the context with `TaskContext.setTaskContext`, as the stdin writer of `PipedRDD` does, passes too. Could this say something like "It is called in the running task that reads the partition, on a thread that has its `TaskContext`, and fails otherwise"? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -351,17 +378,29 @@ class NativeDataSourceStreamReader( } override def partitions(start: String, end: String): Array[InputPartition] = { - val partitions = call(name, "plan a scan of") { - library.streamPartitions(handle.get, start, end) + if (plannedRange != (start, end)) { + val planned = call(name, "plan a scan of") { + library.streamPartitions(handle.get, start, end) + } + val partitions = NativeInputPartition.from(planned.partitions) + // The executors fetch the state once for the micro-batch, as a broadcast, rather than with + // each of its tasks. Spark only plans the next range once the micro-batch of the previous + // one finished, so the copies of its state 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. + plannedState.foreach(_.unpersist(blocking = false)) Review Comment: **2. [Low] `stop()` keeps the broadcast of the last micro-batch on the executors** Thank you for this `unpersist`, which answers [item 5 of round 6](https://github.com/apache/spark/pull/59934#discussion_r4424626687). It only runs when the next range is planned, though, and `stop()` (L410, outside of the diff) only closes the handle. So once the query stopped, the executors keep the state of its last micro-batch until the ContextCleaner removes it after a GC of the driver, which is the case of item 5 again, for one broadcast per query. Could `stop()` also call `plannedState.foreach(_.unpersist(blocking = false))`? ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java: ########## @@ -903,6 +1158,234 @@ private static boolean isInterrupted(long task) { } } + /** The number of task contexts that are not closed yet, for tests. */ + static int openTaskContexts() { + return tasks.size(); + } + + // --------------------------------------------------------------------------------------------- + // Streams of engines + // --------------------------------------------------------------------------------------------- + + // The errno values that the callbacks of a stream return. + private static final int EINVAL = 22; + private static final int EIO = 5; + // The message of get_last_error for a stream that Spark does not know, which lives as long as + // the JVM. + private static final long UNKNOWN_STREAM_ERROR = + toCString("The stream was released, or is not a stream of Spark."); + + private static final AtomicLong nextEngineStreamId = new AtomicLong(1); + // The streams wrapped for engines that are not released yet, by their private_data. + /** + * The callbacks of the streams that Spark wraps for engines, see {@link EngineStream}, and the + * method handles that call the callbacks of the wrapped streams. They are created when the first + * stream is wrapped, so that a JVM that wraps no stream does not create them, and a failure to + * create them only fails the streams of engines. + */ + private static final class EngineStreamCallbacks { + // The addresses of the C functions that call Java, which live as long as the JVM. + static final long GET_SCHEMA; + static final long GET_NEXT; + static final long GET_LAST_ERROR; + static final long RELEASE; + // The callbacks of struct ArrowArrayStream, called with the address of the function first: + // get_schema and get_next, get_last_error, and release. + static final MethodHandle INT_CALLBACK; + static final MethodHandle POINTER_CALLBACK; + static final MethodHandle VOID_CALLBACK; + + static { + try { + MethodHandles.Lookup lookup = MethodHandles.lookup(); + GET_SCHEMA = Ffm.upcall( + lookup.findStatic(NativeLibrary.class, "engineStreamGetSchema", + MethodType.methodType(int.class, long.class, long.class)), + C.I32, C.PTR, C.PTR); + GET_NEXT = Ffm.upcall( + lookup.findStatic(NativeLibrary.class, "engineStreamGetNext", + MethodType.methodType(int.class, long.class, long.class)), + C.I32, C.PTR, C.PTR); + GET_LAST_ERROR = Ffm.upcall( + lookup.findStatic(NativeLibrary.class, "engineStreamGetLastError", + MethodType.methodType(long.class, long.class)), + C.PTR, C.PTR); + RELEASE = Ffm.upcall( + lookup.findStatic(NativeLibrary.class, "engineStreamRelease", + MethodType.methodType(void.class, long.class)), + C.VOID, C.PTR); + INT_CALLBACK = Ffm.unboundDowncall(C.I32, C.PTR, C.PTR); + POINTER_CALLBACK = Ffm.unboundDowncall(C.PTR, C.PTR); + VOID_CALLBACK = Ffm.unboundDowncall(C.VOID, C.PTR); + } catch (Throwable e) { + // The message names the cause, because the JVM only keeps this message for the later + // uses of the class. + throw new IllegalStateException( + "Cannot create the callbacks of the streams of engines: " + rootCause(e), e); + } + } + } + + private static final ConcurrentHashMap<Long, EngineStream> engineStreams = + new ConcurrentHashMap<>(); + + /** + * The stream of a read that an engine got with {@code SupportsArrowStreamRead}, which ties the + * task context of the read to the stream: the context is closed once the library released the + * stream, whenever and on whichever thread the engine releases it, also after the task ended. + * The stream of the library is moved to memory of Spark, and the stream of the engine gets + * callbacks of Spark, which call the current callbacks of the library's stream, and a + * private_data that identifies the wrapped stream, so that the engine can move the stream too. + */ + private static final class EngineStream { + final long inner; + final Task task; + // The name of the data source, for the errors. + final String name; + // The message of the last error of Spark, rather than of the library, or 0 if the last call + // returned an error of the library, or none. The callbacks never run concurrently. + long sparkError; + + EngineStream(long inner, Task task, String name) { + this.inner = inner; + this.task = task; + this.name = name; + } + + long callback(long offset) { + return Platform.getLong(null, inner + offset); + } + + /** + * Records an error of Spark in the given callback, which get_last_error then returns, and logs + * it. The message is built here, since building it can throw too, as a StackOverflowError on + * the small stack of a thread of an engine can. + */ + int fail(String callback, Throwable e) { + try { + clearError(); + String message = + "The " + callback + " of the stream of the native data source " + name + " failed"; + sparkError = toCString(message + ": " + rootCause(e)); + LOGGER.error(message, e); + } catch (Throwable ignored) { + // An upcall must not throw. + } + return EIO; + } + + void clearError() { + if (sparkError != 0) { + Platform.freeMemory(sparkError); + sparkError = 0; + } + } + } + + /** + * Wraps the stream at the given address, which read initialized with non-NULL callbacks, so that + * releasing it closes the given task context (see {@link EngineStream}). If it fails, the stream + * is left as it is. + * + * @param name the name of the data source, for the errors of the stream + */ + static void wrapEngineStream(long stream, Task task, String name) { + long inner = Platform.allocateMemory(STREAM_SIZE); + long id = nextEngineStreamId.getAndIncrement(); + try { + Platform.copyMemory(null, stream, null, inner, STREAM_SIZE); + engineStreams.put(id, new EngineStream(inner, task, name)); + } catch (Throwable e) { + engineStreams.remove(id); + Platform.freeMemory(inner); + throw e; + } + Platform.putLong(null, stream + STREAM_GET_SCHEMA, EngineStreamCallbacks.GET_SCHEMA); Review Comment: **1. [Low] A failure to create the callbacks leaks the wrapped stream** The first use of `EngineStreamCallbacks` is here, after the `try`, when `engineStreams.put` (L1297) already ran. So if its static initializer fails, the `catch` that removes the entry and frees `inner` does not run, and `releaseOnFailure` (`NativeDataSource.scala` L1033) only releases the stream of the library and closes the task. The class then fails every later use too, with a `NoClassDefFoundError`, so each `openArrowStream` leaves one `EngineStream` in `engineStreams`, with its copy of the stream and a closed task, and `openEngineStreams()` keeps growing. Could this read the four addresses into locals before `allocateMemory`, or move these `putLong` calls into the `try`? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -886,10 +986,132 @@ 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 = { + // A task that failed or was killed is over too, also while its listeners run. + val context = TaskContext.get() + if (context == null || context.isInterrupted() || context.isFailed() || Review Comment: **4. [Nit] A killed task gets "must be called in the running task"** Thank you for these checks, which answer [item 3 of round 6](https://github.com/apache/spark/pull/59934#discussion_r4424626669). A task that was killed is still the running task, though, so for it this message reads as if the engine called `openArrowStream` from the wrong place. Spark reports a kill with a `TaskKilledException` and the reason of the kill, which `context.killTaskIfInterrupted()` throws. Could this call it after the check of `null`, and keep `isFailed()` and `isCompleted()` in this check? ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java: ########## @@ -903,6 +1158,234 @@ private static boolean isInterrupted(long task) { } } + /** The number of task contexts that are not closed yet, for tests. */ + static int openTaskContexts() { + return tasks.size(); + } + + // --------------------------------------------------------------------------------------------- + // Streams of engines + // --------------------------------------------------------------------------------------------- + + // The errno values that the callbacks of a stream return. + private static final int EINVAL = 22; + private static final int EIO = 5; + // The message of get_last_error for a stream that Spark does not know, which lives as long as + // the JVM. + private static final long UNKNOWN_STREAM_ERROR = + toCString("The stream was released, or is not a stream of Spark."); + + private static final AtomicLong nextEngineStreamId = new AtomicLong(1); + // The streams wrapped for engines that are not released yet, by their private_data. + /** + * The callbacks of the streams that Spark wraps for engines, see {@link EngineStream}, and the + * method handles that call the callbacks of the wrapped streams. They are created when the first + * stream is wrapped, so that a JVM that wraps no stream does not create them, and a failure to + * create them only fails the streams of engines. + */ + private static final class EngineStreamCallbacks { + // The addresses of the C functions that call Java, which live as long as the JVM. + static final long GET_SCHEMA; + static final long GET_NEXT; + static final long GET_LAST_ERROR; + static final long RELEASE; + // The callbacks of struct ArrowArrayStream, called with the address of the function first: + // get_schema and get_next, get_last_error, and release. + static final MethodHandle INT_CALLBACK; + static final MethodHandle POINTER_CALLBACK; + static final MethodHandle VOID_CALLBACK; + + static { + try { + MethodHandles.Lookup lookup = MethodHandles.lookup(); + GET_SCHEMA = Ffm.upcall( + lookup.findStatic(NativeLibrary.class, "engineStreamGetSchema", + MethodType.methodType(int.class, long.class, long.class)), + C.I32, C.PTR, C.PTR); + GET_NEXT = Ffm.upcall( + lookup.findStatic(NativeLibrary.class, "engineStreamGetNext", + MethodType.methodType(int.class, long.class, long.class)), + C.I32, C.PTR, C.PTR); + GET_LAST_ERROR = Ffm.upcall( + lookup.findStatic(NativeLibrary.class, "engineStreamGetLastError", + MethodType.methodType(long.class, long.class)), + C.PTR, C.PTR); + RELEASE = Ffm.upcall( + lookup.findStatic(NativeLibrary.class, "engineStreamRelease", + MethodType.methodType(void.class, long.class)), + C.VOID, C.PTR); + INT_CALLBACK = Ffm.unboundDowncall(C.I32, C.PTR, C.PTR); + POINTER_CALLBACK = Ffm.unboundDowncall(C.PTR, C.PTR); + VOID_CALLBACK = Ffm.unboundDowncall(C.VOID, C.PTR); + } catch (Throwable e) { + // The message names the cause, because the JVM only keeps this message for the later + // uses of the class. + throw new IllegalStateException( + "Cannot create the callbacks of the streams of engines: " + rootCause(e), e); + } + } + } + + private static final ConcurrentHashMap<Long, EngineStream> engineStreams = + new ConcurrentHashMap<>(); + + /** + * The stream of a read that an engine got with {@code SupportsArrowStreamRead}, which ties the + * task context of the read to the stream: the context is closed once the library released the + * stream, whenever and on whichever thread the engine releases it, also after the task ended. + * The stream of the library is moved to memory of Spark, and the stream of the engine gets + * callbacks of Spark, which call the current callbacks of the library's stream, and a + * private_data that identifies the wrapped stream, so that the engine can move the stream too. + */ + private static final class EngineStream { + final long inner; + final Task task; + // The name of the data source, for the errors. + final String name; + // The message of the last error of Spark, rather than of the library, or 0 if the last call + // returned an error of the library, or none. The callbacks never run concurrently. + long sparkError; Review Comment: **3. [Nit] `sparkError` is a plain field, although the callbacks may run on different threads** The callbacks never run concurrently, as this comment says, but they may run on different threads one after the other, since `SupportsArrowStreamRead` lets the engine call them from any thread. `fail` writes this field on one thread, and `clearError` may free it and write 0 on the next one, so a thread that saw a stale value would free the message twice, or leak it. In practice, the synchronization with which the engine hands the stream over orders these accesses, but the Java memory model does not promise it for a plain field, and `volatile` costs little here. Could it be `volatile`? ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java: ########## @@ -326,6 +433,9 @@ private static final class Ffm { DOWNCALL = lookup.findVirtual(linker, "downcallHandle", MethodType.methodType(MethodHandle.class, segment, descriptor, options)) .asFixedArity(); + UNBOUND_DOWNCALL = lookup.findVirtual(linker, "downcallHandle", Review Comment: **9. [Nit] The lookup of the unbound `downcallHandle` is still in `Ffm`** Thank you for `EngineStreamCallbacks`, which answers [item 12 of round 6](https://github.com/apache/spark/pull/59934#discussion_r4424626764). Only its initializer uses this handle, through `Ffm.unboundDowncall` (L1217-L1219), but the lookup still runs in the static initializer of `Ffm`, so a failure here still fails all of `Ffm`, unlike what the Javadoc of `EngineStreamCallbacks` says (L1183-L1184). Since `DOWNCALL` looks up the other overload of the same method just before, this can hardly fail alone, so it is mostly about where the lookup lives. Could it move to `EngineStreamCallbacks` too? ########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala: ########## @@ -1143,6 +1426,392 @@ 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) + } + // 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 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))) + } + val (address, copy) = checkContexts(1) { + // The engine reads the stream to its end in the task, but releases it later. + val (address, ids) = openInTask(laterFactory, laterPartitions.head)(readIdsInPlace) + assert(ids == (0L until 3L)) + assert(NativeLibrary.openEngineStreams() == 1 && !released.exists()) + // A copy of the stream, as an engine that moves the stream may leave behind. + val copy = Platform.allocateMemory(streamSize) + Platform.copyMemory(null, address, null, copy, streamSize) + (address, copy) + } + checkContexts(-1) { + ArrowArrayStream.wrap(address).release() + assert(NativeLibrary.openEngineStreams() == 0 && released.exists()) + } + Platform.freeMemory(address) + // 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) + } + 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) + } + + // So does the reader factory of a micro-batch scan, with the state of its micro-batch. + val streamTaskDir = Utils.createTempDir() + val query = session.readStream.format("native_counter").option("max_offset", "4") + .option("task_path", streamTaskDir.getPath).load() + .writeStream.format("noop").start() + val microBatchScan = try { + query.processAllAvailable() + query.asInstanceOf[StreamingQueryWrapper].streamingQuery.lastExecution.executedPlan + .collectFirst { case s: MicroBatchScanExec => s }.get + } finally { + query.stop() + } + // Only the reads of the engine record their tasks from now on. + streamTaskDir.listFiles().foreach(_.delete()) + val streamFactory = microBatchScan.readerFactory.asInstanceOf[SupportsArrowStreamRead] + val streamPartitions = microBatchScan.inputPartitions.toSeq + checkContexts(0) { + assert(readAll(streamFactory, streamPartitions).sorted == (0L until 4L)) + } + assert(NativeLibrary.openEngineStreams() == 0) + // The library got the partition of the scan, the only one of the micro-batch. + assert(streamTaskDir.listFiles().map(_.getName.toInt).toSet == Set(0)) + // The partitions of a micro-batch keep the state that the stream reader returned with them: + // a reader factory that an engine keeps reads them with it, also once Spark planned the next + // micro-batch, and so does the reader of Spark. + def microBatchStream(options: (String, String)*): MicroBatchStream = { + val relation = session.readStream.format("native_counter").options(options.toMap).load() + .queryExecution.analyzed.collectFirst { case relation: StreamingRelationV2 => relation }.get + relation.table.asInstanceOf[SupportsRead].newScanBuilder(relation.extraOptions) + .build().toMicroBatchStream(Utils.createTempDir().getPath) + } + def plan(stream: MicroBatchStream, start: Int, end: Int): Seq[InputPartition] = { + def offset(i: Int): Offset = stream.deserializeOffset(s"""{"offset":$i}""") + stream.planInputPartitions(offset(start), offset(end)).toSeq + } + val stream = microBatchStream("max_offset" -> "4") + try { + val first = plan(stream, 0, 2) + val keptFactory = stream.createReaderFactory() + val second = plan(stream, 2, 4) + // The executors get the state of a micro-batch once, as a broadcast that its partitions + // keep, rather than with the partition of each task. + Seq(first, second).foreach { partitions => + val partition = partitions.head.asInstanceOf[NativeMicroBatchPartition] + assert(partition.state.isDefined && partition.readerState.nonEmpty) + val serialized = new ByteArrayOutputStream() + Utils.tryWithResource(new ObjectOutputStream(serialized))(_.writeObject(partition)) + assert(!serialized.toByteArray.toSeq.containsSlice(partition.readerState.toSeq)) + } + checkContexts(0) { + assert(readAll(keptFactory.asInstanceOf[SupportsArrowStreamRead], first) == Seq(0L, 1L)) + assert(readAll(keptFactory.asInstanceOf[SupportsArrowStreamRead], second) == Seq(2L, 3L)) + val read = inTask { + first.flatMap { partition => + val reader = keptFactory.createColumnarReader(partition) + try { + val ids = ArrayBuffer.empty[Long] + while (reader.next()) { + val batch = reader.get() + (0 until batch.numRows()).foreach(i => ids += batch.column(0).getLong(i)) + } + ids.toSeq + } finally { + reader.close() + } + } + } + assert(read == Seq(0L, 1L)) + } + // A micro-batch without partitions broadcasts no state, although the stream reader + // returned one. + val nextBroadcast = session.sparkContext.broadcast(0).id + 1 Review Comment: **7. [Nit] Correction to my item 7 of round 6: the ID of the next broadcast counts every broadcast** The check that I suggested in [item 7 of round 6](https://github.com/apache/spark/pull/59934#discussion_r4424626708) assumes that nothing else in the `SparkContext` creates a broadcast while `plan(stream, 2, 2)` runs. That is likely in this suite, but a broadcast from another thread of the shared session would make it fail, and the two `broadcast(0)` probes stay until the ContextCleaner removes them. Since the micro-batch has no partition, the test cannot look at a partition either. Could it check the state that the reader planned instead, for example with a `private[ffi]` accessor for tests, or at least `destroy()` the probes? ########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala: ########## @@ -1143,6 +1426,392 @@ 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) + } + // 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 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))) + } + val (address, copy) = checkContexts(1) { Review Comment: **8. [Nit] A failed assertion in this test leaks the stream and its copy** If an assertion from here to L1554 fails, the stream may never be released, and `address` and `copy` are never freed, so the stream of the library and the context of its task stay open for the rest of the run. The later tests count the contexts from their own start, so they still pass, and this only matters once this test already failed. Could the release and the `freeMemory` calls run in a `finally`, for example with `Utils.tryWithSafeFinally`, so that a failure there does not hide the first one? ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java: ########## @@ -903,6 +1158,234 @@ private static boolean isInterrupted(long task) { } } + /** The number of task contexts that are not closed yet, for tests. */ + static int openTaskContexts() { + return tasks.size(); + } + + // --------------------------------------------------------------------------------------------- + // Streams of engines + // --------------------------------------------------------------------------------------------- + + // The errno values that the callbacks of a stream return. + private static final int EINVAL = 22; + private static final int EIO = 5; + // The message of get_last_error for a stream that Spark does not know, which lives as long as + // the JVM. + private static final long UNKNOWN_STREAM_ERROR = + toCString("The stream was released, or is not a stream of Spark."); + + private static final AtomicLong nextEngineStreamId = new AtomicLong(1); + // The streams wrapped for engines that are not released yet, by their private_data. Review Comment: **11. [Nit] The comment of `engineStreams` is now above `EngineStreamCallbacks`** The holder class went in between this comment and `engineStreams` (L1229), so the comment reads as if it described the callbacks, and `engineStreams` has none. Could this comment, and `nextEngineStreamId` with it, move down to `engineStreams`? ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeLibrary.java: ########## @@ -903,6 +1158,234 @@ private static boolean isInterrupted(long task) { } } + /** The number of task contexts that are not closed yet, for tests. */ + static int openTaskContexts() { + return tasks.size(); + } + + // --------------------------------------------------------------------------------------------- + // Streams of engines + // --------------------------------------------------------------------------------------------- + + // The errno values that the callbacks of a stream return. + private static final int EINVAL = 22; + private static final int EIO = 5; + // The message of get_last_error for a stream that Spark does not know, which lives as long as + // the JVM. + private static final long UNKNOWN_STREAM_ERROR = Review Comment: **10. [Nit] Every JVM still allocates `UNKNOWN_STREAM_ERROR`** Only `engineStreamGetLastError` (L1349) returns it, but the static initializer of `NativeLibrary` allocates it, so the driver and each executor that uses a native data source allocate it without an engine. The memory is small, but `EngineStreamCallbacks` exists so that a JVM that wraps no stream does not create what the streams of engines need. Could it move into `EngineStreamCallbacks`, which is always initialized before a callback can run? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala: ########## @@ -970,21 +1192,38 @@ class NativeDataWriter( /** Creates the `SparkDsTaskContext` of the running task. */ private[ffi] object NativeTask { /** - * @param partitionId the partition that the task reads or writes, by default the partition of - * the task + * @param partitionId the `partition_id` of the context: the index of the partition in the scan + * for a read, and by default the partition of the task * @param taskAttemptId the ID of the task attempt, by default the one of the task + * @param read whether the context is the one of a read, which the library may use after the + * task ended: an engine may read the stream of a partition then, and a thread that + * the task does not wait for may be in get_next once it completes, as with + * `PipedRDD`. `is_interrupted` then also reports that the task failed or completed, + * and the context does not keep the task alive. A data writer only runs within its + * task, whose kill `is_interrupted` reports. Outside of a task, it reports nothing. */ def create( partitionId: Option[Int] = None, - taskAttemptId: Option[Long] = None): NativeLibrary.Task = { + taskAttemptId: Option[Long] = None, + read: Boolean = false): NativeLibrary.Task = { val context = Option(TaskContext.get()) + val interrupted: BooleanSupplier = context match { + case Some(taskContext) if read => + val task = new WeakReference(taskContext) + () => Option(task.get()).forall(t => t.isInterrupted() || t.isFailed() || t.isCompleted()) Review Comment: **6. [Nit] Each `is_interrupted` of a read now takes the lock of the task twice** Since [item 2 of round 6](https://github.com/apache/spark/pull/59934#discussion_r4424626654), Spark's own reads get this context too. `isFailed()` and `isCompleted()` of `TaskContextImpl` are `synchronized`, so while the task runs, each call takes its monitor twice, besides the upcall and the lookup in `tasks`, where a read of Spark only did the volatile read of `isInterrupted()` before. The header asks a library that blocks to check it regularly, so a library that checks it for each batch or each row pays this each time. I did not measure it, and next to the upcall it may well be small. Could the context check something cheaper, such as a volatile flag that a completion and a failure listener of the task set, or is this cost fine? -- 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]
