andygrove opened a new pull request, #7016: URL: https://github.com/apache/datafusion-comet/pull/7016
## Which issue does this PR close? Closes #6806. Part of #6013. ## Rationale for this change Every task that reads a Comet broadcast decompresses and decodes the whole broadcast for itself, in `CometBatchRDD.compute`. Spark keeps one copy of a broadcast's value per executor, but Comet's value is Arrow IPC compressed with Spark's codec, so each probe task of a broadcast hash join repeats the same decode. #6806 profiled a broadcast hash join with a 1M-row build side: about three quarters of each Comet probe task went to that decode, against 6% for building the hash table. ## What changes are included in this PR? - `DecodedBroadcastCache` lets the tasks of an executor that read a broadcast at the same time share one decoded copy of it. The first task to read a broadcast value decodes it into buffers the cache owns. Tasks that arrive during that decode wait for it, and check every 100 ms whether they have been killed. Each task then reads vectors of its own that hold references on the shared buffers, so the native input reader can close what it reads, as it does today, without affecting other tasks. - An entry is freed when the last task holding it completes. The cache therefore never holds memory the tasks would not hold anyway, since without it each of them holds its own decoded copy for as long as it runs. A task that starts after that decodes the broadcast again. - A failed decode is not cached: the tasks waiting for it decode the broadcast for themselves, and the next task to read it tries again. A broadcast whose streams carry dictionaries, which `coalesceBroadcastBatches` leaves as it found them, is still decoded by each task. - `CometBatchRDD.compute` reads through the cache, so it covers every consumer of a Comet broadcast: broadcast hash joins of every join type, broadcast nested loop joins, and JVM consumers. - `spark.comet.exec.broadcastExchange.shareDecoded.enabled`, true by default, turns it off. The tuning guide describes it. - `Utils.decodeBatches`, the cache and the driver-side coalescing read the IPC stream through a channel that reads up to 64 KiB at a time straight into the destination buffer, instead of `Channels.newChannel`, which reads 8 KiB at a time through the JDK's interruptible I/O machinery. A read on an interrupted thread still fails with `ClosedByInterruptException`. Each task still builds its own hash table. Sharing that is the rest of #6013 and needs apache/datafusion#25491, which is in DataFusion 56. The broadcast also stays compressed twice (#2216). Comet's codec layer is now undone once per executor while its tasks overlap, rather than once per task. ## How are these changes tested? `DecodedBroadcastCacheSuite` runs each case through a cache whose decode the test can block or fail: - tasks holding a broadcast at the same time decode it once and between them hold one copy, which is freed when the last of them completes, after which the next task decodes it again - a task that arrives during a decode waits for it, and a task killed while waiting stops - tasks waiting for a decode that fails decode the broadcast themselves, and a later task tries again - a broadcast with dictionaries, one with several streams, zero-column batches and an empty broadcast - a consumer that closes what it reads, one that leaves its batches open, and one that stops part way - `CometBatchRDD` goes through the cache only when the setting is on - inner, left, semi, anti and right broadcast hash joins and a broadcast nested loop join, over string, struct and array build columns and eight probe tasks, give Spark's answer with the setting on and off, and leave nothing in the cache Each of eight mutations fails at least one of these tests: no sharing, no retry after a failed decode, no kill check, a task's last view not closed when it completes, a view not closed when the reader moves on, no dictionary check, `CometBatchRDD` ignoring the setting, and entries never freed. `UtilsSuite` adds a stream larger than one read piece and a decode on an interrupted thread. `CometJoinSuite`, `CometExecSuite`, `CometArrowStreamSuite`, `CometExecIteratorLifecycleSuite`, `CometDppFallbackRepro3949Suite`, `UtilsSuite` and the new suite pass on Spark 4.1 (340 tests). Strict Scala warnings and scalafix pass on Spark 3.5. ### Benchmark A local benchmark that is not part of this PR, in the shape of #6806. `fact` has key, long, double and string columns in one Parquet file per probe task, about 65K rows each. `dim` has unique keys with a long and a string column. The query is `SELECT /*+ BROADCAST(d) */ sum(f.f_long), sum(d.d_long), sum(length(f.f_str)), sum(length(d.d_str)) FROM fact f JOIN dim d ON f.key = d.k`, with AQE off. Each figure is the median of 8 rounds that alternate the setting, after 2 warmup rounds. Run time and CPU are summed over the probe stage's tasks. | Shape | Setting | Wall | Probe tasks' run time | Probe tasks' CPU | Slowest task | | --- | --- | ---: | ---: | ---: | ---: | | 1M-row build, 32 tasks, `local[8]` | off | 285.6 ms | 902.5 ms | 879.0 ms | 32.0 ms | | | on | 244.3 ms | 538.0 ms | 331.0 ms | 30.5 ms | | 1M-row build, 128 tasks, `local[8]` | off | 1321.0 ms | 7749.5 ms | 4503.0 ms | 116.5 ms | | | on | 683.7 ms | 2592.0 ms | 1692.0 ms | 82.0 ms | | 10K-row build, 200 tasks, `local[8]` | off | 312.0 ms | 1149.5 ms | 524.5 ms | 47.5 ms | | | on | 311.8 ms | 1055.0 ms | 490.0 ms | 39.5 ms | | 1M-row build, 32 tasks, `local[1]` | off | 1714.7 ms | 1310.5 ms | 1148.0 ms | 63.5 ms | | | on | 1742.1 ms | 1325.0 ms | 1159.5 ms | 65.0 ms | Under `local[1]` no two tasks overlap, so nothing is shared, and the difference is within run-to-run noise. One decode of the 1M-row broadcast (16.9 MB compressed, 31.8 MB of Arrow IPC), timed on the driver, took 16.4 ms through `Channels.newChannel` and 16.0 ms through the new channel. Environment: Comet main at 66c63bd307 plus this PR, release build, Spark 4.1.3, Scala 2.13, JDK 17.0.10, macOS 26.6.2 on an Apple M3 Max. -- 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]
