comphead commented on code in PR #7016: URL: https://github.com/apache/datafusion-comet/pull/7016#discussion_r4338820302
########## 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: The cache sits below the join, in `CometBatchRDD.compute`, so all five hash join types run the same cache code. With the setting on by default, `CometJoinSuite`'s broadcast tests already read through the cache on `local[5]`, and the `CometBatchRDD` test above covers the switch. One hash join and the nested loop join with sharing on, over the struct and array columns, plus the `size == 0` check would keep what only this test covers: 2 `checkSparkAnswerAndOperator` calls instead of 12. If the join-type matrix is worth keeping, a `sql-tests/join` fixture with a `ConfigMatrix` on the setting can hold it, and this test keeps one join for the leak check. -- 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]
