fix: apply broadcast fetch before fan-out - #624
Conversation
| let (broadcast_input, coalesce_fetch) = build_child | ||
| let broadcast_input = build_child | ||
| .downcast_ref::<CoalescePartitionsExec>() | ||
| .filter(|coalesce| coalesce.fetch().is_none()) | ||
| .map_or_else( | ||
| || (Arc::clone(build_child), None), | ||
| |coalesce| (Arc::clone(coalesce.input()), coalesce.fetch()), | ||
| || Arc::clone(build_child), | ||
| |coalesce| Arc::clone(coalesce.input()), |
There was a problem hiding this comment.
🤔 Unless I'm missing something, I think this should be showing up as an intermediate stage that is collapsed to 1 task and immediately broadcasted to several nodes.
Is there any chance we can have a test in this same file that asserts that?
Also, It'd be great to add another test that exercises this in inject_network_boundaries.rs
|
cc @barbarj, if you could give a look to this one that'd be great 🙏 |
| async fn build_side_fetch_is_preserved_by_broadcast() { | ||
| assert_distributed_matches_single_node( | ||
| "SELECT count(*) FROM (SELECT id FROM build_side LIMIT 50) b \ | ||
| JOIN probe_side p ON b.id = p.id", | ||
| true, | ||
| ) | ||
| .await | ||
| .unwrap(); | ||
| } |
There was a problem hiding this comment.
Can we assert the plan snapshot in this test as well? It'd make reasoning about the effects of this change easier. build_side_fetch_is_preserved_by_normalize doesn't do so because of some non-determinism in file partioning, but given this is a different query, I wonder if that might not be a problem here.
What
CoalescePartitionsExecinsideBroadcastExecLIMITrow selectionWhy
This follows up on #618. Moving the fetch to the coalesce above
BroadcastExecchanges a global build-side limit into a per-consumer limit after fan-out. The post-merge unit run demonstrated the resulting correctness failure: single-node and distributedcount(*)returned2500and3750even though every possible set of 50 build IDs has exactly 2,500 probe matches.Keeping the fetch-bearing coalesce as the broadcast input applies the limit once before replication. The resulting distributed plan has
BroadcastExecoverCoalescePartitionsExec: fetch=50.Failure evidence: https://github.com/datafusion-contrib/datafusion-distributed/actions/runs/31358856004/job/93363560852
Validation
cargo fmt --all -- --checkgit diff --checkcargo test --features integration --test multi_task_collect_join_repros build_side_fetch_is_preserved_by_broadcast(passed, then passed 5 repeated runs)