fix: preserve broadcast build-side fetch

[shinzoxD]

Summary

• preserve the fetch value when replacing a build-side CoalescePartitionsExec during broadcast insertion
• add an integration regression that compares distributed and single-node results for a limited build side

Testing

• cargo test --features integration --test multi_task_collect_join_repros (8 passed, 2 ignored)
• cargo test distributed_planner::insert_broadcast::tests --lib (7 passed)
• cargo test distributed_planner::normalize_collect_joins::tests --lib (7 passed)
• cargo test --features integration --lib -- --skip protocol::grpc::channel_resolver::tests::fails_establishing_connection (288 passed, 1 ignored, 1 filtered)
• cargo fmt --all -- --check
• cargo check --lib
• cargo clippy --lib -- -D warnings
• git diff --check

Closes #591

Assisted-by: OpenAI Codex

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[shinzoxD]

CI is awaiting maintainer approval for this external-contributor workflow. The regression is reproduced on current main, and the focused integration target plus the 288-test library suite pass locally.

[gabotechs]

Looks good, thanks @shinzoxD!

[gabotechs]

@shinzoxD it seems like the test might be a bit flaky?

https://github.com/datafusion-contrib/datafusion-distributed/actions/runs/31358856004/job/93363560852

[stuhood]

@gabotechs: It looks like this change might violate the conditions that @barbarj was planning to validate:

src/distributed_planner/normalize_collect_joins.rs establishes that a fetch-less CoalescePartitionsExec (a single-node partitioning artifact) should be stripped, but a fetch-bearing CoalescePartitionsExec must be retained on the input below the boundary so the limit is evaluated in the producer stage before distribution.

[gabotechs]

@stuhood 🤔 but is that a problem? without fetch propagation there the query is incorrect.

I see this code handles it gracefully:

datafusion-distributed/src/distributed_planner/normalize_collect_joins.rs

Lines 175 to 186 in e2b6094

	// If the build input is a fetch-less [CoalescePartitionsExec], we can strip it as it's only
	// a remnant of CollectLeft's need to collect input into a single partition. Otherwise, if
	// the build input contains a fetch or is some other node type, we must retain it to ensure
	// correct behavior
	let build_input = if let Some(coalesce) = join.left().downcast_ref::<CoalescePartitionsExec>()
	&& coalesce.fetch().is_none()
	{
	coalesce.input()
	} else {
	join.left()
	};
	let left = Arc::new(RepartitionExec::try_new(

And tests are green.

Although I wonder if it would be better to propagate the fetch to the left RepartitionExec instead.

[barbarj]

@shinzoxD it seems like the test might be a bit flaky?

https://github.com/datafusion-contrib/datafusion-distributed/actions/runs/31358856004/job/93363560852

This looks like the same problem the other fetch-preserving test had. Which row survives is not deterministic, so you need to just assert row counts instead.

[barbarj]

Although I wonder if it would be better to propagate the fetch to the left RepartitionExec instead.

@gabotechs RepartitionExec doesn't support a fetch though, or am I missing something?

[gabotechs]

@barbarj yeah, indeed it doesn't.

[shinzoxD]

Thanks — I don't think this failure is only nondeterministic row selection. This test already selects count(*), so each result set contains one aggregate row; CI failed because the aggregate values were 3750 and 2500.

The fixture has 100 unique build IDs and exactly 50 probe rows for every ID. Any valid LIMIT 50 therefore joins to exactly 2,500 rows regardless of which 50 IDs survive. A value of 3,750 means one execution effectively joined 75 build rows. Comparing only the number of aggregate rows would be vacuous (1 == 1) and hide that mismatch.

This suggests the fetch moved above BroadcastExec is being applied after fan-out / per consumer rather than preserving the original global limit. I'll prepare a focused follow-up that keeps a cardinality assertion and fixes the placement unless you prefer this tracked as an issue first.

[barbarj]

Ah gotcha. @stuhood 's comment is right then. That does sound like it's violating the invariant I made an attempt at validating in #582

[shinzoxD]

Thanks for confirming. The focused follow-up is #624. It retains a fetch-bearing CoalescePartitionsExec as the input to BroadcastExec, while still stripping fetch-less coalesces. The regression passed locally once with the corrected plan and then in five repeated runs; the full upstream CI matrix is now running.