perf: share one decoded copy of a broadcast among the tasks of an executor

[andygrove]

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, and a task that starts after that decodes the broadcast 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 for a broadcast that coalescing left as several batches the cache keeps one copy that the tasks would have freed once they built their tables.
• 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 decode that finishes after every task holding its entry has completed is freed rather than kept. 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, and the memory management guide lists the cache's allocator.
• Utils.decodeBatches, the cache and the driver-side coalescing read the IPC stream through a channel that fills the destination buffer up to 64 KiB at a time, instead of Channels.newChannel, which reads 8 KiB at a time through the JDK's interruptible I/O machinery. Like that channel, it fails with ClosedByInterruptException before its next piece once its thread is interrupted.

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 decode that outlives its task frees what it decoded
• 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

UtilsSuite adds a stream larger than one read piece, a decode on an interrupted thread, and a read that stops at its next piece once its thread is interrupted.

Each of ten 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, entries never freed, a decode kept after every task holding its entry completed, and an interrupt checked only when a read starts.

CometJoinSuite, CometExecSuite, CometArrowStreamSuite, CometExecIteratorLifecycleSuite, CometDppFallbackRepro3949Suite, UtilsSuite and the new suite pass on Spark 4.1 (342 tests), and strict Scala warnings pass on Spark 3.5. Two local stress runs that are not part of this PR, 4,000 simulated tasks with partial reads and two reads in one task, and 600 with interrupts and kills, on 16 and 12 threads, left nothing allocated or cached.

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 66c63bd plus the first commit of this PR (afa7136), release build, Spark 4.1.3, Scala 2.13, JDK 17.0.10, macOS 26.6.2 on an Apple M3 Max.

[andygrove]

@comphead not sure if this overlaps with any of your PRs? Happy to close this one if it does