Repository navigation
perf: share each broadcast read between the tasks on an executor - #2807
Conversation
github-actions
Bot
added
the
documentation
| &self.metrics, | ||
| partition, | ||
| ))) | ||
| self.fetch(partition, context) |
There was a problem hiding this comment.
nit: can we add else branch here
There was a problem hiding this comment.
Done, it's an if/else now.
| } | ||
| Ok::<_, datafusion::error::DataFusionError>(Arc::new(batches)) | ||
| }) | ||
| .await?; |
There was a problem hiding this comment.
would it make sense to capture last used here, after fetch has finished ?
*entry.last_used.lock().unwrap() = now;
There was a problem hiding this comment.
Yes, it was stamped before the fetch, so a fetch slower than the idle TTL left an entry that looked idle right away and got swept. It's now set in stream once the batches are ready, which covers the fetching task and the ones waiting on it. Added a_slow_fetch_is_not_idle_when_it_finishes.
| }) | ||
| }) | ||
| .clone(); | ||
| *entry.last_used.lock().unwrap() = now; |
There was a problem hiding this comment.
should we move last used to stream ?
There was a problem hiding this comment.
Moved, see the reply above.
| ); | ||
| let reader = self.clone(); | ||
| return Ok( | ||
| cache.stream(key, self.schema.clone(), move || reader.fetch(0, context)) |
There was a problem hiding this comment.
metrics will be updated only once, when reader.fetch is triggered (if no cache entries) do we need to update metrics when broadcast is read from cache ?
There was a problem hiding this comment.
Good catch, cache hits recorded zero output_rows. Tasks served from the cache now record their own output rows, so every task's reader reports what it produced. Fetch metrics stay at zero on a hit since nothing was fetched. Added broadcast_cache_hits_record_output_rows.
milenkovicm
left a comment
milenkovicm
left a comment
There was a problem hiding this comment.
thanks @andygrove in general it looks good, few comments please consider them before merging
The build side of a broadcast hash join is a broadcast ShuffleReaderExec that reads every partition of the upstream stage. Each task runs its own plan, so every task on an executor fetched and decoded the same build side: at 128 partitions a 16-task broadcast stage made 16,384 fetch requests for one 318K-row table. Add an executor-wide BroadcastReadCache. The first task to run a broadcast read fetches it; tasks that ask meanwhile wait for that fetch, and later ones reuse the decoded batches. Entries are keyed by job, stage and the exact locations read, expire after 30 s unused, are not kept above 1 GiB, and are dropped with the job's data. A failed fetch is not cached, so each task reports its own fetch error as before. Controlled by ballista.shuffle.cache_broadcast_reads (default true).
c5982d8 to
25405b7
Compare
Which issue does this PR close?
No issue. Split out of #2792.
Rationale for this change
The build side of a broadcast hash join is a broadcast
ShuffleReaderExecthat reads every output partition of the upstream stage. Each task of the joining stage runs its own copy of the plan, so every task on an executor fetches and decodes the same data. Spark broadcasts a relation once per executor.Each task fetches every upstream output (M map outputs × N partitions), so the cost grows with partition count. In TPC-H SF100 q2 at 128 partitions on two executors, a 16-task broadcast stage made 16,384 fetch requests and decoded the same 318K-row table 16 times. The stage took 3.0 s of task time.
What changes are included in this PR?
BroadcastReadCache(execution_plans/broadcast_cache.rs). The first task on an executor to run a broadcast read fetches it. Tasks that ask while that fetch is running wait for it, and later tasks reuse the decoded batches.ShuffleReaderExec::with_broadcast_cacheattaches the cache.executeuses it only for broadcast readers that have locations, whenballista.shuffle.cache_broadcast_readsis on. The existing fetch body moves into an inherentfetchmethod unchanged.DefaultExecutionEngineowns one cache and attaches it to every broadcast reader in the stages it creates.ExecutionEnginegainsremove_job, with a default that does nothing, which the executor calls when it removes a job's data.ballista.shuffle.cache_broadcast_reads, defaulttrue. Six unit tests cover:remove_jobResults:
Memory. The cached batches are not charged to the memory pool. That matches the fetched batches of a single read, which never were; the pool sees each task's hash table built from them. What changes is how long they live. Without the cache, a task's decoded build side is freed once its hash table is built. With the cache, one decoded copy per broadcast read is kept until it has been unused for 30 s, or until the executor removes the job's data. Idle entries are dropped lazily, on the next cache lookup, and reads above 1 GiB are not kept. So an executor holds at most one extra copy of each recent broadcast read, in exchange for no longer decoding a copy per task. There is no overall cap across entries. Setting
ballista.shuffle.cache_broadcast_reads=falserestores the previous behavior.Are there any user-facing changes?
ballista.shuffle.cache_broadcast_reads, enabled by default.ExecutionEngine::remove_job(with a default implementation),DefaultExecutionEngine::broadcast_cacheandShuffleReaderExec::with_broadcast_cache.