andygrove commented on code in PR #7016: URL: https://github.com/apache/datafusion-comet/pull/7016#discussion_r4346551142
########## docs/source/contributor-guide/memory_management.md: ########## @@ -110,12 +110,14 @@ for both at once, and raises the pod's memory request by the same amount. **Comet's JVM-side Arrow allocator is unbounded and accounted by nobody.** `CometArrowAllocator` (`spark/src/main/scala/org/apache/comet/package.scala`) is a single process-wide `new RootAllocator(Long.MaxValue)`. Child allocators are cut from it for FFI stream export -(`CometNativeArrowSource`), broadcast coalescing, and `CometSparkToColumnarExec`. These are real -off-heap bytes in container RSS that neither Spark's `TaskMemoryManager` nor Comet's native memory -pool sees. In practice the volume is modest, a batch at a time per stream, but there is no -ceiling and no backpressure. The executor's memory usage log reports the allocator's total, and -the part of it charged to the import allocator described below, so that the overhead can be sized -for it. +(`CometNativeArrowSource`), broadcast coalescing, the decoded broadcasts that an executor's tasks +share (`DecodedBroadcastCache`), and `CometSparkToColumnarExec`. These are real off-heap bytes in +container RSS that neither Spark's `TaskMemoryManager` nor Comet's native memory pool sees. In +practice the volume is modest, a batch at a time per stream, but there is no ceiling and no +backpressure. A broadcast is the largest case: a task that reads one holds all of it decoded, and +the tasks of an executor that read it at the same time share one such copy. The executor's memory Review Comment: Good catch, and it is right: `DataFusion` counts the batches per plan, so every task's join reserves its whole build side while the JVM holds one copy. I added a paragraph to the allocator section saying that sharing lowers the bytes resident but not the bytes reserved, that spilling behaves as before, and that the untracked figure subtracts all N reservations against one copy in `jvmArrow`, so it undercounts (your 8 tasks and 100 MB example, 700 MB, clamped at zero). The sentence about `reserved()` being a lower bound now carves out tasks that share a broadcast. ########## spark/src/test/scala/org/apache/spark/sql/comet/DecodedBroadcastCacheSuite.scala: ########## @@ -0,0 +1,471 @@ +/* + * 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.comet + +import java.nio.charset.StandardCharsets +import java.util.Properties +import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicInteger + +import scala.concurrent.{Await, Future, Promise} +import scala.concurrent.duration._ +import scala.util.Try + +import org.apache.arrow.vector.{IntVector, VarCharVector} +import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider} +import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, FieldType} +import org.apache.spark.{SparkConf, TaskContextImpl, TaskKilledException} +import org.apache.spark.executor.TaskMetrics +import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.util.io.ChunkedByteBuffer + +import org.apache.comet.{CometArrowAllocator, CometConf} +import org.apache.comet.vector.CometVector + +class DecodedBroadcastCacheSuite extends CometTestBase { + + private val Timeout = 30.seconds + + test("tasks that read a broadcast at the same time decode it once") { + val value = broadcastValue(rows(1000)) + val (cache, decodes) = countingCache() + val before = CometArrowAllocator.getAllocatedMemory + val tasks = (1 to 4).map(i => newTask(i.toLong)) + + assert(readAll(cache.batches(value, tasks.head, "test")) == rows(1000)) + val sharedCopy = CometArrowAllocator.getAllocatedMemory - before + assert(sharedCopy > 0) + tasks.tail.init.foreach { task => + assert(readAll(cache.batches(value, task, "test")) == rows(1000)) + } + assert(decodes.get == 1) + assert(cache.holders(value) == 3) + // Each task closed what it read, the way the native input reader does, so the one shared + // copy is all that the three tasks hold. + assert(CometArrowAllocator.getAllocatedMemory - before == sharedCopy) + + // A consumer that leaves its batches open, as the JVM columnar-to-row path does, has them + // closed when it moves on and when its task completes. + assert(readAll(cache.batches(value, tasks.last, "test"), close = false) == rows(1000)) + assert(decodes.get == 1) + tasks.foreach(_.markTaskCompleted(None)) + assert(cache.size == 0) + assert(CometArrowAllocator.getAllocatedMemory == before) + } + + test("the last task holding a broadcast frees it, and a later task decodes it again") { + val value = broadcastValue(rows(1000)) + val (cache, decodes) = countingCache() + val before = CometArrowAllocator.getAllocatedMemory + val first = newTask(1) + val second = newTask(2) + readAll(cache.batches(value, first, "test")) + readAll(cache.batches(value, second, "test")) + + first.markTaskCompleted(None) + assert(cache.holders(value) == 1) + assert(CometArrowAllocator.getAllocatedMemory > before) + second.markTaskCompleted(None) + assert(cache.size == 0) + assert(CometArrowAllocator.getAllocatedMemory == before) + + val third = newTask(3) + assert(readAll(cache.batches(value, third, "test")) == rows(1000)) + assert(decodes.get == 2) + third.markTaskCompleted(None) + assert(CometArrowAllocator.getAllocatedMemory == before) + } + + test("a task that stops reading part way frees what it read when it completes") { + // Without coalescing, so the broadcast has a batch after the one the task stops at. + val value = (serialize(rows(100)) ++ serialize(rows(200).drop(100))).toArray + val cache = new DecodedBroadcastCache(DecodedBroadcastCache.decode) + val before = CometArrowAllocator.getAllocatedMemory + val task = newTask(1) + val first = cache.batches(value, task, "test").next() + assert(first.numRows() == 100) + task.markTaskCompleted(None) + assert(cache.size == 0) + assert(CometArrowAllocator.getAllocatedMemory == before) + } + + test("a task that comes while another task decodes the broadcast waits for that decode") { + val value = broadcastValue(rows(1000)) + val gated = new GatedCache + val before = CometArrowAllocator.getAllocatedMemory + val loader = newTask(1) + val waiter = newTask(2) + + val loaded = gated.startLoad(value, loader) + val waited = inThread(readAll(gated.cache.batches(value, waiter, "test"))) + gated.awaitHolders(value, 2) + Thread.sleep(200) + assert(!waited.isCompleted) + + gated.finish.countDown() + assert(Await.result(loaded, Timeout) == rows(1000)) + assert(Await.result(waited, Timeout) == rows(1000)) + assert(gated.decodes.get == 1) + + loader.markTaskCompleted(None) + waiter.markTaskCompleted(None) + assert(gated.cache.size == 0) + assert(CometArrowAllocator.getAllocatedMemory == before) + } + + test("a task killed while it waits for another task's decode stops waiting") { + val value = broadcastValue(rows(100)) + val gated = new GatedCache + val before = CometArrowAllocator.getAllocatedMemory + val loader = newTask(1) + val waiter = newTask(2) + + val loaded = gated.startLoad(value, loader) + val waited = inThread(gated.cache.batches(value, waiter, "test").hasNext) + gated.awaitHolders(value, 2) + waiter.markInterrupted("killed by the test") + val killed = intercept[TaskKilledException](Await.result(waited, Timeout)) + assert(killed.reason == "killed by the test") + waiter.markTaskCompleted(Some(killed)) + + gated.finish.countDown() + assert(Await.result(loaded, Timeout) == rows(100)) + loader.markTaskCompleted(None) + assert(gated.cache.size == 0) + assert(CometArrowAllocator.getAllocatedMemory == before) + } + + test("tasks waiting for a decode that fails decode the broadcast for themselves") { + val value = broadcastValue(rows(1000)) + val gated = new GatedCache(failFirst = true) + val before = CometArrowAllocator.getAllocatedMemory + val loader = newTask(1) + val waiter = newTask(2) + + val loaded = gated.startLoad(value, loader) + val waited = inThread(readAll(gated.cache.batches(value, waiter, "test"))) + gated.awaitHolders(value, 2) + gated.finish.countDown() + val failure = intercept[IllegalStateException](Await.result(loaded, Timeout)) + assert(failure.getMessage == "decode failed") + assert(Await.result(waited, Timeout) == rows(1000)) + assert(gated.decodes.get == 1, "the waiting task decoded the broadcast through the cache") + + // The failed entry is no longer handed out, so the next task decodes the broadcast again. + val later = newTask(3) + assert(readAll(gated.cache.batches(value, later, "test")) == rows(1000)) + assert(gated.decodes.get == 2) + + loader.markTaskCompleted(Some(failure)) + waiter.markTaskCompleted(None) + later.markTaskCompleted(None) + assert(gated.cache.size == 0) + assert(CometArrowAllocator.getAllocatedMemory == before) + } + + test("a decode that outlives its task frees what it decoded") { + val value = broadcastValue(rows(100)) + val gated = new GatedCache + val before = CometArrowAllocator.getAllocatedMemory + val task = newTask(1) + + val read = gated.startLoad(value, task) + // The task completes on another thread while its read decodes the broadcast, which gives + // back the task's hold and frees the entry before the decode keeps anything in it. + task.markTaskCompleted(None) + assert(gated.cache.size == 0) + gated.finish.countDown() + assert(Await.result(read, Timeout) == rows(100)) + assert(CometArrowAllocator.getAllocatedMemory == before) + } + + test("each task decodes a broadcast with dictionaries for itself") { + val values = Seq("a", "b", null, "a", "c", "b") + val value = dictionaryBroadcastValue(values) + val before = CometArrowAllocator.getAllocatedMemory + assert(DecodedBroadcastCache.decode(value, "test").isEmpty) + assert(CometArrowAllocator.getAllocatedMemory == before) + + val (cache, decodes) = countingCache() + val tasks = (1 to 2).map(i => newTask(i.toLong)) + tasks.foreach { task => + val strings = cache.batches(value, task, "test").flatMap { batch => + (0 until batch.numRows()).map { i => + if (batch.column(0).isNullAt(i)) null else batch.column(0).getUTF8String(i).toString + } + } + assert(strings.toList == values) + } + assert(decodes.get == 1) + tasks.foreach(_.markTaskCompleted(None)) + assert(cache.size == 0) + assert(CometArrowAllocator.getAllocatedMemory == before) + } + + test("a broadcast of several streams reads every batch of each") { + val batches = Seq(rows(10), rows(300).drop(10), Seq.empty, rows(1000).drop(300)) + // Without coalescing, as for the streams that coalescing leaves as it found them. + val value = batches.flatMap(serialize).toArray + val (cache, decodes) = countingCache() + val before = CometArrowAllocator.getAllocatedMemory + val tasks = (1 to 2).map(i => newTask(i.toLong)) + // The second task leaves each batch open, so moving on to the next one has to close it. + tasks.zip(Seq(true, false)).foreach { case (task, close) => + assert(readAll(cache.batches(value, task, "test"), close) == rows(1000)) + } + assert(decodes.get == 1) + tasks.foreach(_.markTaskCompleted(None)) + assert(cache.size == 0) + assert(CometArrowAllocator.getAllocatedMemory == before) + } + + test("zero-column batches keep their row counts, and an empty broadcast has no batches") { + val empty = (1 to 3).map(_ => new ColumnarBatch(Array.empty[ColumnVector], 5)) + val serialized = Utils.serializeBatches(empty.iterator).map(_._2).toList + val value = Utils.coalesceBroadcastBatches(serialized.iterator)._1 + val cache = new DecodedBroadcastCache(DecodedBroadcastCache.decode) + val task = newTask(1) + assert(cache.batches(value, task, "test").map(_.numRows()).sum == 15) + assert(!cache.batches(Array.empty[ChunkedByteBuffer], task, "test").hasNext) + task.markTaskCompleted(None) + assert(cache.size == 0) + } + + test("CometBatchRDD reads through the shared cache only when asked to") { + val broadcast = spark.sparkContext.broadcast(broadcastValue(rows(100))) + try { + Seq(true, false).foreach { share => + val rdd = new CometBatchRDD(spark.sparkContext, 2, broadcast, share) + val tasks = (1 to 2).map(i => newTask(i.toLong)) + rdd.partitions.zip(tasks).foreach { case (partition, task) => + assert(readAll(rdd.compute(partition, task)) == rows(100)) + } + assert(DecodedBroadcastCache.shared.holders(broadcast.value) == (if (share) 2 else 0)) + tasks.foreach(_.markTaskCompleted(None)) + assert(DecodedBroadcastCache.shared.holders(broadcast.value) == 0) + } + } finally { + broadcast.destroy() + } + } + + test("broadcast joins read the same rows with and without shared decoding") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withParquetTable((0 until 4000).map(i => (i % 400, s"fact-$i")), "fact") { + val dim = (0 until 300).map { i => + (i, if (i % 9 == 0) null else s"dim-$i", (i, s"s$i"), Seq(i, i + 1)) + } + withParquetTable(dim, "dim") { + // The repartition gives the probe side eight tasks, which read the broadcast at once. + val fact = "(SELECT /*+ REPARTITION(8) */ * FROM fact) f" + val hashJoins = Seq( Review Comment: Agreed. The join-type matrix is now `join/broadcast_shared_decode.sql` with a `ConfigMatrix` on the setting, and this test keeps one hash join and the nested loop join over the struct and array columns, with a check that the allocator is back where it started afterwards. I checked that check is not vacuous: with a hold that is never given back it fails. ########## spark/src/test/scala/org/apache/spark/sql/comet/DecodedBroadcastCacheSuite.scala: ########## @@ -0,0 +1,471 @@ +/* + * 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.comet + +import java.nio.charset.StandardCharsets +import java.util.Properties +import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicInteger + +import scala.concurrent.{Await, Future, Promise} +import scala.concurrent.duration._ +import scala.util.Try + +import org.apache.arrow.vector.{IntVector, VarCharVector} +import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider} +import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, FieldType} +import org.apache.spark.{SparkConf, TaskContextImpl, TaskKilledException} +import org.apache.spark.executor.TaskMetrics +import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.util.io.ChunkedByteBuffer + +import org.apache.comet.{CometArrowAllocator, CometConf} +import org.apache.comet.vector.CometVector + +class DecodedBroadcastCacheSuite extends CometTestBase { + + private val Timeout = 30.seconds + + test("tasks that read a broadcast at the same time decode it once") { + val value = broadcastValue(rows(1000)) + val (cache, decodes) = countingCache() + val before = CometArrowAllocator.getAllocatedMemory + val tasks = (1 to 4).map(i => newTask(i.toLong)) + + assert(readAll(cache.batches(value, tasks.head, "test")) == rows(1000)) + val sharedCopy = CometArrowAllocator.getAllocatedMemory - before + assert(sharedCopy > 0) + tasks.tail.init.foreach { task => + assert(readAll(cache.batches(value, task, "test")) == rows(1000)) + } + assert(decodes.get == 1) + assert(cache.holders(value) == 3) + // Each task closed what it read, the way the native input reader does, so the one shared + // copy is all that the three tasks hold. + assert(CometArrowAllocator.getAllocatedMemory - before == sharedCopy) + + // A consumer that leaves its batches open, as the JVM columnar-to-row path does, has them + // closed when it moves on and when its task completes. + assert(readAll(cache.batches(value, tasks.last, "test"), close = false) == rows(1000)) + assert(decodes.get == 1) + tasks.foreach(_.markTaskCompleted(None)) + assert(cache.size == 0) + assert(CometArrowAllocator.getAllocatedMemory == before) + } + + test("the last task holding a broadcast frees it, and a later task decodes it again") { Review Comment: Folded in. The first test now ends with the last task completing, the allocator back to `before`, and a later task decoding again, and the separate test is gone. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometBroadcastValue.scala: ########## @@ -0,0 +1,333 @@ +/* + * 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.comet + +import java.io.DataInputStream +import java.nio.channels.Channels +import java.util.IdentityHashMap +import java.util.concurrent.{CountDownLatch, TimeUnit} + +import scala.collection.mutable.ArrayBuffer +import scala.util.control.NonFatal + +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.{VectorLoader, VectorSchemaRoot, VectorUnloader} +import org.apache.arrow.vector.ipc.ArrowStreamReader +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch +import org.apache.arrow.vector.types.pojo.Schema +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.io.CompressionCodec +import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.io.ChunkedByteBuffer + +import org.apache.comet.CometArrowAllocator +import org.apache.comet.vector.NativeUtil + +/** + * Lets the tasks of an executor that read a Comet broadcast at the same time share one decoded + * copy of it. + * + * A Comet broadcast holds Arrow IPC streams compressed with Spark's codec, and without this cache + * every task that reads it decompresses and decodes all of it. Spark keeps one copy of a + * broadcast's value per executor, so the tasks of an executor get the same value instance, and + * that instance is the key here. The first task to read a value decodes it into buffers that the + * cache owns. Tasks that come while it is decoding wait for it. Each task then reads vectors of + * its own that hold references on those buffers, so a consumer can close what it reads, as + * [[org.apache.spark.sql.comet.execution.arrow.ColumnarBatchArrowReader]] does, without affecting + * another task. + * + * An entry is freed when the last task holding it completes, and a task that reads the broadcast + * after that decodes it again. A native join keeps a build side of one batch, which is what + * coalescing makes of a broadcast, for as long as its task runs, so for such a broadcast the + * cache holds no memory that the tasks would not hold anyway. A native join copies a build side + * of several batches, so the cache keeps one copy of such a broadcast that the tasks would have + * freed once they built their tables. + * + * @param decode + * decodes a broadcast's value; `None` makes each task decode the value for itself + */ +private[comet] class DecodedBroadcastCache( + decode: (Array[ChunkedByteBuffer], String) => Option[DecodedBroadcastCache.Decoded]) { + import DecodedBroadcastCache.{Decoded, KillCheckIntervalMs} + + // The broadcast values that running tasks hold, by identity. Guarded by itself. + private val entries = new IdentityHashMap[Array[ChunkedByteBuffer], Entry]() + + /** + * The batches of `chunks`, the value of a Comet broadcast, for the task of `context`. As with + * [[Utils.decodeBatches]], each batch is valid until the next call to `hasNext` or `next`, or + * until the task completes. + */ + def batches( + chunks: Array[ChunkedByteBuffer], + context: TaskContext, + source: String): Iterator[ColumnarBatch] = + new SharedBatchIterator(chunks, context, source) + + /** The number of broadcast values that running tasks hold. */ + def size: Int = entries.synchronized(entries.size) + + /** The number of holds on the entry for `chunks`, one for each read by a running task. */ + def holders(chunks: Array[ChunkedByteBuffer]): Int = + entries.synchronized(Option(entries.get(chunks)).map(_.holders).getOrElse(0)) + + /** Hands out the entry for `chunks`, creating it if no task holds one. True if it did. */ + private def acquire(chunks: Array[ChunkedByteBuffer]): (Entry, Boolean) = + entries.synchronized { + val existing = entries.get(chunks) + val entry = if (existing != null) { + existing + } else { + val created = new Entry(chunks) + entries.put(chunks, created) + created + } + entry.holders += 1 + (entry, existing == null) + } + + /** Gives back a task's hold on `entry`, and frees the entry once no task holds it. */ + private def release(entry: Entry): Unit = { + val unheld = entries.synchronized { + entry.holders -= 1 + val unheld = entry.holders == 0 + if (unheld) { + forget(entry) + } + unheld + } + if (unheld) { + entry.close() + } + } + + /** Stops handing `entry` out, unless a newer entry has replaced it. Callers hold `entries`. */ + private def forget(entry: Entry): Unit = { + if (entries.get(entry.chunks) eq entry) { + val _ = entries.remove(entry.chunks) + } + } + + private final class Entry(val chunks: Array[ChunkedByteBuffer]) { + // One hold for each read of the broadcast by a running task. Guarded by `entries`. + var holders = 0 + private val done = new CountDownLatch(1) + // Written by the task that loads the entry before `done` opens, and read by others after. + @volatile private var decoded: Option[Decoded] = None + + /** Decodes the broadcast for every task that holds the entry. Runs in the first of them. */ + def load(source: String): Option[Decoded] = { + try { + val result = decode(chunks, source) + // A task whose completion listener runs on another thread can give the entry back + // during the decode. If it was the last holder the entry has been freed already, and + // nothing would free what it kept now. + val kept = entries.synchronized { + if (holders > 0) { + decoded = result + } + holders > 0 + } + if (!kept) { + result.foreach(_.close()) + } + decoded + } catch { + case t: Throwable => + // The task may have failed only because it was killed, so let a later task try again. + // The tasks that are waiting decode the broadcast for themselves. + entries.synchronized(forget(this)) + throw t + } finally { + done.countDown() + } + } + + /** Waits for the task that loads the entry, stopping if this task is killed. */ + def await(context: TaskContext): Option[Decoded] = { + while (!done.await(KillCheckIntervalMs, TimeUnit.MILLISECONDS)) { + context.killTaskIfInterrupted() + } + decoded + } + + /** Frees the decoded batches. Called once no task holds the entry. */ + def close(): Unit = decoded.foreach(_.close()) + } + + /** + * One task's read of a broadcast: views of the shared batches when the entry has them, or the + * task's own decode when it does not. + */ + private final class SharedBatchIterator( + chunks: Array[ChunkedByteBuffer], + context: TaskContext, + source: String) + extends Iterator[ColumnarBatch] { + + // The task takes its hold on the first call to hasNext, which is also when reading the + // broadcast without the cache starts to decode it. + private var started = false + private var failure: Throwable = _ + private var batches: Iterator[ColumnarBatch] = _ + // The root of the view that `next` returned last, closed by the next call to `hasNext`. + private var current: VectorSchemaRoot = _ + + override def hasNext: Boolean = { + if (failure != null) { + throw failure + } + if (!started) { + start() + } + closeCurrent() + batches.hasNext + } + + override def next(): ColumnarBatch = { + if (!hasNext) { + throw new NoSuchElementException + } + batches.next() + } + + private def start(): Unit = { + started = true + val (entry, loading) = acquire(chunks) + context.addTaskCompletionListener[Unit] { _ => + try closeCurrent() + finally release(entry) + } + try { + batches = (if (loading) entry.load(source) else entry.await(context)) match { + case Some(decoded) => + Iterator.tabulate(decoded.numBatches) { i => + current = decoded.view(i) Review Comment: I think it can't in practice, so I took the second option. Spark runs a task's completion listeners on the task's thread, after its body, and native pulls the build side on that thread, so nothing reads once the listener has run. The guard in `load` covered a case the views never covered, so I removed it and its test, and the scaladoc on `batches` now says a read has to stay on the task's thread. Does that match how you read it? -- 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]
