Skip to content

test: cover DF-promoted CollectLeft broadcast safety gap - #89

Closed
avantgardnerio wants to merge 1 commit into
andygrove:test/1770-unzero-broadcast-thresholdfrom
avantgardnerio:brent/pr1900-static-path-broadcast-safety
Closed

test: cover DF-promoted CollectLeft broadcast safety gap#89
avantgardnerio wants to merge 1 commit into
andygrove:test/1770-unzero-broadcast-thresholdfrom
avantgardnerio:brent/pr1900-static-path-broadcast-safety

Conversation

@avantgardnerio

Copy link
Copy Markdown

Stacked on apache#1900 — adds a single failing test that exposes a gap in the safety guard.

What the test shows

With the new session defaults from apache#1900 (hash_join_single_partition_threshold = 10 MB), DataFusion's own JoinSelection runs during create_physical_plan and can produce HashJoinExec(LEFT, CollectLeft) for small LEFT JOIN big — no swap, because the small side is already on the left. try_collect_left in DF doesn't restrict by join type (it doesn't need to in standalone DF, where there's only one probe task).

Ballista's maybe_promote_to_broadcast short-circuits on partition_mode != Partitioned (planner.rs:294), so the new collect_left_broadcast_safe guard never sees the join. plan_query_stages_internal then lowers it via UnresolvedShuffleExec::new_broadcast at planner.rs:145-184 — re-opening the apache#1055 unsafety that pinning the thresholds to zero used to prevent.

Reproducer output on this branch:

HashJoinExec: mode=CollectLeft, join_type=Left, on=[(k@0, k@0)]
  UnresolvedShuffleExec: broadcast=true, upstream_partitions: 1
  DataSourceExec: partitions=1, partition_sizes=[1]

Why existing tests don't catch it

make_broadcast_test_ctx pins datafusion.optimizer.hash_join_single_partition_threshold = 0 (planner.rs:1494-1497), so DF can never promote on its own in test sessions — only the Ballista promotion path is exercised, which is correctly guarded. The new test mirrors the production session (DF and Ballista thresholds both at 10 MB) so DF's JoinSelection is what does the promotion.

Why AQE is unaffected

DelayJoinSelectionRule runs in plan_preparation_optimizers before DF's JoinSelection and wraps every HashJoinExec in a DynamicJoinSelectionExec, so DF's rule finds nothing to promote. The static path has no equivalent wrapper.

Suggested fix shape (for apache#1900)

In maybe_promote_to_broadcast, when partition_mode == CollectLeft and !collect_left_broadcast_safe(*hash_join.join_type()), demote to Partitioned instead of returning early. Same predicate, same module.

This PR intentionally contains only the test — happy to land the fix here or in apache#1900, whichever you prefer.

Checklist

  • Documentation improved

With the new defaults from this PR, DataFusion's own JoinSelection
runs first under the production session config and can promote a
HashJoinExec to PartitionMode::CollectLeft for `small LEFT JOIN big`
with no swap, because the small side is already on the left.
maybe_promote_to_broadcast then short-circuits on `partition_mode !=
Partitioned`, so the new collect_left_broadcast_safe guard never sees
the join and plan_query_stages_internal lowers it via
UnresolvedShuffleExec::new_broadcast — re-opening the apache#1055 unsafety
that pinning the thresholds to zero used to prevent.

Existing tests miss this because make_broadcast_test_ctx pins DF's
threshold to 0; only the Ballista promotion path is exercised. This
test mirrors the production session (DF and Ballista thresholds both
at 10 MB) and asserts that no LEFT join ends up CollectLeft and no
broadcast reader is produced.

Currently fails on this branch; intended to drive the fix.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@andygrove

Copy link
Copy Markdown
Owner

Confirmed and fixed in apache#1900 (commit b40a670). This was a real gap — DataFusion's own JoinSelection runs during create_physical_plan and, with the non-zero hash_join_single_partition_threshold this PR sets, stamps CollectLeft without restricting by join type, so a pre-promoted unsafe join (e.g. small LEFT JOIN big with the small side already on the left) bypassed the guard that only covered the Partitioned -> CollectLeft promotion path.

The fix moves the broadcast-safety check to the front of maybe_promote_to_broadcast: any CollectLeft hash join whose join type isn't collect_left_broadcast_safe is demoted back to a Partitioned (shuffle) join, re-hash-partitioning both inputs on the join keys. It runs ahead of the threshold check so the correctness guard still fires even if Ballista's own broadcast promotion is disabled while DF's threshold is non-zero.

Your test is pulled into apache#1900 verbatim (just swapped the explanatory comment for a current-state one to match the repo's comment convention) and now passes. Thanks for catching it — I'll leave this PR for you to close.

@avantgardnerio

Copy link
Copy Markdown
Author

This seems like a good case for AQE in DF. If we had that infrastructure, we could port DelayJoinSelectionRule there, and DF would have deferred, and there would have not been another check for Ballista to have to make - IIUC.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants